#include <algorithm>
#include <array>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <numbers>
#include <type_traits>
#include <utility>

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>
#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace {
    namespace backend         = mean_field::preconditioning::backend;
    namespace blocks          = mean_field::utils::blocks;
    namespace preconditioning = mean_field::preconditioning;

    struct MaterialValue final : blocks::value_block_base { };
    struct GravityValue final : blocks::value_block_base { };
    struct MaterialResidual final : blocks::residual_block_base { };
    struct GravityResidual final : blocks::residual_block_base { };

    using MockForm = blocks::block_form<
        blocks::type_list<MaterialValue, GravityValue>,
        blocks::type_list<MaterialResidual, GravityResidual>>;
    using MockJacobian = blocks::type_list<
        blocks::block_row<MaterialResidual, MaterialValue, GravityValue>,
        blocks::block_row<GravityResidual, MaterialValue, GravityValue>>;
    using MockMaterialComponent = preconditioning::ComponentDeclaration<
        blocks::type_list<MaterialValue>,
        blocks::type_list<MaterialResidual>,
        blocks::type_list<preconditioning::Coupling<MaterialResidual, MaterialValue>>,
        preconditioning::IdentityOperatorCharacteristics,
        backend::Identity,
        preconditioning::NoPreparationDependencies>;
    using MockGravityComponent = preconditioning::ComponentDeclaration<
        blocks::type_list<GravityValue>,
        blocks::type_list<GravityResidual>,
        blocks::type_list<preconditioning::Coupling<GravityResidual, GravityValue>>,
        preconditioning::IdentityOperatorCharacteristics,
        backend::Identity,
        preconditioning::NoPreparationDependencies>;
    using MockStructure = preconditioning::StellarStructureBlock<
        MockMaterialComponent,
        MockGravityComponent,
        MockForm,
        MockJacobian,
        preconditioning::ApproximateStellarBlockLDU>;

    template <typename MaterialComponent, typename GravityComponent>
    concept MockComponentsCanCompose = requires {
        typename preconditioning::StellarStructureBlock<
            MaterialComponent, GravityComponent, MockForm, MockJacobian, preconditioning::IndependentStellarSubsystems>;
    };

    class KnownCrossCouplings final {
    public:
        [[nodiscard]] constexpr int MaterialSize() const noexcept {
            return 1;
        }

        [[nodiscard]] constexpr int GravitySize() const noexcept {
            return 1;
        }

        void ApplyMaterialToGravity(
            const mfem::Vector &materialDirection,
            mfem::Vector &gravityAction
        ) const {
            gravityAction(0) = 3.0 * materialDirection(0);
        }

        void ApplyGravityToMaterial(
            const mfem::Vector &gravityDirection,
            mfem::Vector &materialAction
        ) const {
            materialAction(0) = 7.0 * gravityDirection(0);
        }
    };

    template <typename Policy>
    [[nodiscard]] mfem::Vector applyKnownFactorization(
        Policy policy,
        preconditioning::StellarStructureFactorizationStatistics *statistics = nullptr
    ) {
        mfem::Vector materialDiagonal(1);
        mfem::Vector gravityDiagonal(1);
        materialDiagonal(0)  = 2.0;
        gravityDiagonal(0)   = 5.0;
        auto materialInverse = backend::prepare(backend::Diagonal{}, materialDiagonal);
        auto gravityInverse  = backend::prepare(backend::Diagonal{}, gravityDiagonal);
        const KnownCrossCouplings couplings;
        preconditioning::StellarStructureFactorizationOperator<Policy, KnownCrossCouplings> factorization(
            policy, materialInverse, gravityInverse, couplings
        );
        mfem::Vector rightHandSide(2);
        mfem::Vector action(2);
        rightHandSide(0) = 11.0;
        rightHandSide(1) = 13.0;
        factorization.Mult(rightHandSide, action);
        if (statistics != nullptr) {
            *statistics = factorization.GetStatistics();
        }
        return action;
    }

    using PolytropicModel   = mean_field::model::StellarModel<mean_field::models::SpecificationSet<
        mean_field::eos::Polytrope,
        mean_field::surface::Isobaric,
        mean_field::integral::FixedTotalMass,
        mean_field::constraint::FixedCentralDensity>>;
    using PolytropicProblem = mean_field::equilibrium::StellarEquilibriumProblem<PolytropicModel>;
    using PolytropicMaterialSurfaceDescriptor = preconditioning::MaterialSurfaceDescriptorFor<PolytropicProblem>;
    using MaterialComponent =
        decltype(preconditioning::materialSurfaceBlock(std::declval<const PolytropicProblem &>()));
    using FixedAMG = backend::HypreBoomerAMG<backend::FixedCycles>;
    using GravityComponent =
        preconditioning::GravityFieldBlock<backend::Diagonal, FixedAMG, preconditioning::GravityApproximateLDU>;
    using PolytropicStructure = decltype(preconditioning::stellarStructureBlock(
        std::declval<const PolytropicProblem &>(),
        std::declval<MaterialComponent>(),
        std::declval<GravityComponent>(),
        preconditioning::IndependentStellarSubsystems{}
    ));

    template <typename Problem>
    concept CanPrepareDefaultStellarStructure = requires(const Problem &problem) {
        preconditioning::prepare(problem, preconditioning::stellarStructureBlock(problem));
    };

    struct DistinctPhysicalCore final { };

    [[nodiscard]] mean_field::operators::StellarEquilibriumDependencies
    makeDependencies(const std::uint64_t revision = 1) {
        return {
            .discretization     = {.identity = 9101, .revision = 1},
            .density            = {.identity = 9103, .revision = revision},
            .surfaceDeformation = {.identity = 9107, .revision = revision},
            .gravityGradient    = {.identity = 9111, .revision = revision},
            .gravityPotential   = {.identity = 9117, .revision = revision},
            .enthalpy           = {.identity = 9123, .revision = revision},
            .bernoulliConstant  = {.identity = 9129, .revision = revision},
            .rotation           = {.identity = 9131, .revision = revision},
            .targetMass         = {.identity = 9137, .revision = 1}
        };
    }

    [[nodiscard]] mean_field::physics::RigidRotation zeroRotation() {
        mfem::Vector angularVelocity(3);
        mfem::Vector center(3);
        angularVelocity = 0.0;
        center          = 0.0;
        return {angularVelocity, center};
    }

    [[nodiscard]] double relativeError(
        const mfem::Vector &left,
        const mfem::Vector &right
    ) {
        mfem::Vector difference(left);
        difference -= right;
        return difference.Norml2() / std::max({1.0, left.Norml2(), right.Norml2()});
    }
} // namespace

TEST_CASE(
    "Stellar Structure Composition Derives Both Cross-Subsystem Graphs",
    "[preconditioning][stellar_structure][unit][type_contract]"
) {
    using ExpectedMaterialToGravity = blocks::type_list<
        preconditioning::Coupling<blocks::gravity::gradient::residual, blocks::surface_deformation::parameters::value>,
        preconditioning::Coupling<blocks::gravity::poisson::residual, blocks::density::mass::value>,
        preconditioning::Coupling<blocks::gravity::poisson::residual, blocks::surface_deformation::parameters::value>>;
    using ExpectedGravityToMaterial = blocks::type_list<
        preconditioning::Coupling<
            blocks::surface_deformation::shape_equilibrium::residual, blocks::gravity::gradient::value>,
        preconditioning::Coupling<blocks::enthalpy::specific::residual, blocks::gravity::poisson::value>>;

    STATIC_CHECK(preconditioning::PreconditionerComponent<PolytropicStructure>);
    STATIC_CHECK(preconditioning::ExecutableStellarStructureRuntimeFor<
                 mean_field::operators::PreparedStellarEquilibriumOperator>);
    STATIC_CHECK_FALSE(preconditioning::ExecutableStellarStructureRuntimeFor<DistinctPhysicalCore>);
    STATIC_CHECK(preconditioning::StellarStructureRuntimeFor<
                 PolytropicMaterialSurfaceDescriptor,
                 mean_field::operators::PreparedStellarEquilibriumOperator>);
    STATIC_CHECK_FALSE(preconditioning::StellarStructureRuntimeFor<
                       PolytropicMaterialSurfaceDescriptor,
                       DistinctPhysicalCore>);
    STATIC_CHECK(preconditioning::StellarStructurePreconditionerProblem<PolytropicProblem>);
    STATIC_CHECK(preconditioning::StellarStructurePreparableFor<
                 PolytropicProblem,
                 MaterialComponent,
                 GravityComponent>);
    STATIC_CHECK(CanPrepareDefaultStellarStructure<PolytropicProblem>);
    STATIC_CHECK(std::same_as<typename PolytropicStructure::MaterialToGravityCouplings, ExpectedMaterialToGravity>);
    STATIC_CHECK(std::same_as<typename PolytropicStructure::GravityToMaterialCouplings, ExpectedGravityToMaterial>);
    STATIC_CHECK(PolytropicStructure::MaterialToGravityCouplings::size == 3);
    STATIC_CHECK(PolytropicStructure::GravityToMaterialCouplings::size == 2);
    STATIC_CHECK(PolytropicStructure::RequiredCouplings::size == 17);
    STATIC_CHECK(preconditioning::backend::ArnoldiAdmissible<typename PolytropicStructure::BackendType>);

    using FixedMassIdentity = preconditioning::IdentityBlock<
        blocks::fixed_total_mass::mass_normalization::value, blocks::fixed_total_mass::mass_normalization::residual>;
    using FixedCentralDensityIdentity = preconditioning::IdentityBlock<
        blocks::fixed_central_density::central_value::value, blocks::fixed_central_density::central_value::residual>;
    using CompletePlan =
        preconditioning::PreconditionerPlan<PolytropicStructure, FixedMassIdentity, FixedCentralDensityIdentity>;
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<CompletePlan, typename PolytropicProblem::FormType>);
    STATIC_CHECK(
        preconditioning::CompatiblePreconditionerFor<
            CompletePlan, typename PolytropicProblem::FormType, typename PolytropicProblem::JacobianFormType>
    );

    STATIC_CHECK(MockComponentsCanCompose<MockMaterialComponent, MockGravityComponent>);
    STATIC_CHECK_FALSE(MockComponentsCanCompose<MockMaterialComponent, MockMaterialComponent>);
    STATIC_CHECK(preconditioning::PreconditionerComponent<MockStructure>);
    STATIC_CHECK(MockStructure::MaterialToGravityCouplings::size == 1);
    STATIC_CHECK(MockStructure::GravityToMaterialCouplings::size == 1);
}

TEST_CASE(
    "Stellar Structure Factorizations Preserve Independent Triangular And Approximate LDU Algebra",
    "[preconditioning][stellar_structure][unit][factorization]"
) {
    const auto check = [](const mfem::Vector &value, const std::array<double, 2> expected) {
        REQUIRE(value.Size() == 2);
        CHECK(value(0) == Catch::Approx(expected[0]).margin(2.0e-14));
        CHECK(value(1) == Catch::Approx(expected[1]).margin(2.0e-14));
        mfem::Vector expectedVector(2);
        expectedVector(0) = expected[0];
        expectedVector(1) = expected[1];
        CHECK(relativeError(value, expectedVector) <= 2.0e-14);
    };

    check(applyKnownFactorization(preconditioning::IndependentStellarSubsystems{}), {5.5, 2.6});
    check(applyKnownFactorization(preconditioning::MaterialThenGravityTriangular{}), {5.5, -0.7});
    check(applyKnownFactorization(preconditioning::GravityThenMaterialTriangular{}), {-3.6, 2.6});

    preconditioning::StellarStructureFactorizationStatistics statistics;
    check(applyKnownFactorization(preconditioning::ApproximateStellarBlockLDU{}, &statistics), {7.95, -0.7});
    CHECK(statistics.applications == 1);
    CHECK(statistics.materialSurfaceInverseApplications == 2);
    CHECK(statistics.gravityInverseApplications == 1);
    CHECK(statistics.materialToGravityApplications == 1);
    CHECK(statistics.gravityToMaterialApplications == 1);
}

TEST_CASE(
    "Stellar Structure Cross Actions Are Exact Restricted Jacobian Actions And Compose Prepared Blocks",
    "[preconditioning][stellar_structure][integration]"
) {
    using namespace mean_field;
    const utils::Args arguments = test_utils::setup_args();
    fem::FEM finiteElements     = fem::setup_fem(arguments.mesh_file, arguments, 0);
    REQUIRE(finiteElements.okay());

    constexpr double radius         = utils::RADIUS;
    constexpr double mass           = utils::MASS;
    const double polytropicConstant = 2.0 * utils::G * radius * radius / std::numbers::pi_v<double>;
    const double centralDensity     = std::numbers::pi_v<double> * mass / (4.0 * radius * radius * radius);
    const auto stellarModel         = model::StellarModel(
        eos::Polytrope({.n = 1.0, .K = polytropicConstant}),
        surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{mass}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{centralDensity}})
    );
    auto problem   = equilibrium::discretize(stellarModel, finiteElements);
    auto projected = seed::makeProjectedEquilibriumState(problem, seed::LaneEmden({.radialSampleCount = 512}));
    problem.Prepare(projected.values, makeDependencies(), zeroRotation());
    const auto &physical = problem.GetPreparedOperator().GetPhysicalOperator();

    preconditioning::StellarStructureCrossJacobianOperator cross(physical);
    mfem::Vector direction(cross.Width());
    for (int index = 0; index < direction.Size(); ++index) {
        direction(index) = 0.01 * std::sin(0.29 * static_cast<double>(index + 1));
    }
    mfem::Vector crossAction(cross.Height());
    cross.Mult(direction, crossAction);

    const mfem::Vector materialDirection(direction.GetData(), cross.MaterialSize());
    const mfem::Vector gravityDirection(direction.GetData() + cross.MaterialSize(), cross.GravitySize());
    const auto &materialOffsets = cross.GetMaterialOffsets();
    const auto &gravityOffsets  = cross.GetGravityOffsets();

    mfem::Vector materialOnlyDirection(physical.Width());
    materialOnlyDirection         = 0.0;
    auto materialOnlyView         = physical.GetRootManifest().stateView(materialOnlyDirection);
    mfem::Vector materialDensity  = materialOnlyView.block(blocks::density_field.mass_term);
    mfem::Vector materialSurface  = materialOnlyView.block(blocks::surface_deformation_field.parameters_term);
    mfem::Vector materialEnthalpy = materialOnlyView.block(blocks::enthalpy_field.specific_term);
    const mfem::Vector sourceDensity(
        const_cast<mfem::real_t *>(materialDirection.GetData()) + materialOffsets[0],
        materialOffsets[1] - materialOffsets[0]
    );
    const mfem::Vector sourceSurface(
        const_cast<mfem::real_t *>(materialDirection.GetData()) + materialOffsets[1],
        materialOffsets[2] - materialOffsets[1]
    );
    const mfem::Vector sourceEnthalpy(
        const_cast<mfem::real_t *>(materialDirection.GetData()) + materialOffsets[2],
        materialOffsets[3] - materialOffsets[2]
    );
    materialDensity  = sourceDensity;
    materialSurface  = sourceSurface;
    materialEnthalpy = sourceEnthalpy;
    mfem::Vector materialOnlyAction;
    physical.Mult(materialOnlyDirection, materialOnlyAction);
    const auto materialOnlyActionView = physical.GetRootManifest().residualView(materialOnlyAction);

    mfem::Vector gravityOnlyDirection(physical.Width());
    gravityOnlyDirection          = 0.0;
    auto gravityOnlyView          = physical.GetRootManifest().stateView(gravityOnlyDirection);
    mfem::Vector gravityGradient  = gravityOnlyView.block(blocks::gravity_field.gradient_term);
    mfem::Vector gravityPotential = gravityOnlyView.block(blocks::gravity_field.poisson_term);
    const mfem::Vector sourceGravityGradient(
        const_cast<mfem::real_t *>(gravityDirection.GetData()) + gravityOffsets[0],
        gravityOffsets[1] - gravityOffsets[0]
    );
    const mfem::Vector sourceGravityPotential(
        const_cast<mfem::real_t *>(gravityDirection.GetData()) + gravityOffsets[1],
        gravityOffsets[2] - gravityOffsets[1]
    );
    gravityGradient  = sourceGravityGradient;
    gravityPotential = sourceGravityPotential;
    mfem::Vector gravityOnlyAction;
    physical.Mult(gravityOnlyDirection, gravityOnlyAction);
    const auto gravityOnlyActionView = physical.GetRootManifest().residualView(gravityOnlyAction);

    mfem::Vector expected(cross.Height());
    expected = 0.0;
    expected.SetVector(gravityOnlyActionView.block(blocks::density_field.mass_term), materialOffsets[0]);
    expected.SetVector(
        gravityOnlyActionView.block(blocks::surface_deformation_field.shape_equilibrium_term), materialOffsets[1]
    );
    expected.SetVector(gravityOnlyActionView.block(blocks::enthalpy_field.specific_term), materialOffsets[2]);
    expected.SetVector(
        materialOnlyActionView.block(blocks::gravity_field.gradient_term), cross.MaterialSize() + gravityOffsets[0]
    );
    expected.SetVector(
        materialOnlyActionView.block(blocks::gravity_field.poisson_term), cross.MaterialSize() + gravityOffsets[1]
    );

    const mfem::Vector expectedMaterial(expected.GetData(), cross.MaterialSize());
    const mfem::Vector expectedGravity(expected.GetData() + cross.MaterialSize(), cross.GravitySize());
    const mfem::Vector expectedDensity(
        expectedMaterial.GetData() + materialOffsets[0], materialOffsets[1] - materialOffsets[0]
    );
    const mfem::Vector expectedSurface(
        expectedMaterial.GetData() + materialOffsets[1], materialOffsets[2] - materialOffsets[1]
    );
    const mfem::Vector expectedEnthalpy(
        expectedMaterial.GetData() + materialOffsets[2], materialOffsets[3] - materialOffsets[2]
    );
    const mfem::Vector expectedGravityGradient(
        expectedGravity.GetData() + gravityOffsets[0], gravityOffsets[1] - gravityOffsets[0]
    );
    const mfem::Vector expectedGravityPotential(
        expectedGravity.GetData() + gravityOffsets[1], gravityOffsets[2] - gravityOffsets[1]
    );

    const mfem::Vector crossMaterial(crossAction.GetData(), cross.MaterialSize());
    const mfem::Vector crossGravity(crossAction.GetData() + cross.MaterialSize(), cross.GravitySize());
    const mfem::Vector crossDensity(
        crossMaterial.GetData() + materialOffsets[0], materialOffsets[1] - materialOffsets[0]
    );
    const mfem::Vector crossSurface(
        crossMaterial.GetData() + materialOffsets[1], materialOffsets[2] - materialOffsets[1]
    );
    const mfem::Vector crossEnthalpy(
        crossMaterial.GetData() + materialOffsets[2], materialOffsets[3] - materialOffsets[2]
    );
    const mfem::Vector crossGravityGradient(
        crossGravity.GetData() + gravityOffsets[0], gravityOffsets[1] - gravityOffsets[0]
    );
    const mfem::Vector crossGravityPotential(
        crossGravity.GetData() + gravityOffsets[1], gravityOffsets[2] - gravityOffsets[1]
    );

    INFO(
        "gravity-to-material density-row error = " << relativeError(crossDensity, expectedDensity)
                                                   << ", actual norm = " << crossDensity.Norml2()
                                                   << ", expected norm = " << expectedDensity.Norml2()
    );
    INFO(
        "gravity-to-material surface-row error = " << relativeError(crossSurface, expectedSurface)
                                                   << ", actual norm = " << crossSurface.Norml2()
                                                   << ", expected norm = " << expectedSurface.Norml2()
    );
    INFO(
        "gravity-to-material enthalpy-row error = " << relativeError(crossEnthalpy, expectedEnthalpy)
                                                    << ", actual norm = " << crossEnthalpy.Norml2()
                                                    << ", expected norm = " << expectedEnthalpy.Norml2()
    );
    INFO(
        "material-to-gravity gradient-row error = " << relativeError(crossGravityGradient, expectedGravityGradient)
                                                    << ", actual norm = " << crossGravityGradient.Norml2()
                                                    << ", expected norm = " << expectedGravityGradient.Norml2()
    );
    INFO(
        "material-to-gravity Poisson-row error = " << relativeError(crossGravityPotential, expectedGravityPotential)
                                                   << ", actual norm = " << crossGravityPotential.Norml2()
                                                   << ", expected norm = " << expectedGravityPotential.Norml2()
    );
    CHECK(relativeError(crossDensity, expectedDensity) <= 2.0e-12);
    CHECK(relativeError(crossSurface, expectedSurface) <= 2.0e-12);
    CHECK(relativeError(crossEnthalpy, expectedEnthalpy) <= 2.0e-12);
    CHECK(relativeError(crossGravityGradient, expectedGravityGradient) <= 2.0e-12);
    CHECK(relativeError(crossGravityPotential, expectedGravityPotential) <= 2.0e-12);
    CHECK(relativeError(crossAction, expected) <= 2.0e-12);

    auto materialBlock = preconditioning::materialSurfaceBlock(problem);
    auto gravityBlock  = preconditioning::GravityFieldBlock(
        backend::Diagonal{}, FixedAMG{backend::FixedCycles{.cycles = 1}}, preconditioning::GravityApproximateLDU{}
    );
    auto structure = preconditioning::stellarStructureBlock(
        problem, materialBlock, gravityBlock, preconditioning::ApproximateStellarBlockLDU{}
    );
    auto prepared = preconditioning::prepare(problem, structure);
    mfem::Vector rightHandSide(prepared.Width());
    mfem::Vector correction(prepared.Height());
    for (int index = 0; index < rightHandSide.Size(); ++index) {
        rightHandSide(index) = std::cos(0.17 * static_cast<double>(index + 1));
    }
    prepared.Mult(rightHandSide, correction);
    for (int index = 0; index < correction.Size(); ++index) {
        REQUIRE(std::isfinite(correction(index)));
    }
    const auto &statistics = prepared.GetFactorization().GetStatistics();
    CHECK(statistics.applications == 1);
    CHECK(statistics.materialSurfaceInverseApplications == 2);
    CHECK(statistics.gravityInverseApplications == 1);
    CHECK(statistics.materialToGravityApplications == 1);
    CHECK(statistics.gravityToMaterialApplications == 1);

    const auto unchanged = prepared.Refresh();
    CHECK_FALSE(unchanged.DidAnyWork());
    CHECK(prepared.IsCurrent());
}