#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <type_traits>

#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 gravity_context = mean_field::operators::context::gravity_field;
    namespace preconditioning = mean_field::preconditioning;

    using FixedAMG            = backend::HypreBoomerAMG<backend::FixedCycles>;
    using AdaptiveAMG         = backend::HypreBoomerAMG<backend::SolveToTolerance>;
    using FixedGravityLDU =
        preconditioning::GravityFieldBlock<backend::Diagonal, FixedAMG, preconditioning::GravityApproximateLDU>;
    using ChebyshevGravityLDU = preconditioning::
        GravityFieldBlock<backend::MatrixFreeChebyshev, FixedAMG, preconditioning::GravityApproximateLDU>;
    using AdaptiveGravityLDU =
        preconditioning::GravityFieldBlock<backend::Diagonal, AdaptiveAMG, preconditioning::GravityApproximateLDU>;

    using DensityIdentity =
        preconditioning::IdentityBlock<blocks::density::mass::value, blocks::density::mass::residual>;
    using SurfaceIdentity = preconditioning::IdentityBlock<
        blocks::surface_deformation::parameters::value,
        blocks::surface_deformation::shape_equilibrium::residual>;
    using EnthalpyIdentity =
        preconditioning::IdentityBlock<blocks::enthalpy::specific::value, blocks::enthalpy::specific::residual>;
    using MassIdentity = preconditioning::IdentityBlock<
        blocks::fixed_total_mass::mass_normalization::value,
        blocks::fixed_total_mass::mass_normalization::residual>;
    using FixedGravityPlan = preconditioning::
        PreconditionerPlan<DensityIdentity, SurfaceIdentity, FixedGravityLDU, EnthalpyIdentity, MassIdentity>;
    using AdaptiveGravityPlan = preconditioning::
        PreconditionerPlan<DensityIdentity, SurfaceIdentity, AdaptiveGravityLDU, EnthalpyIdentity, MassIdentity>;

    template <typename Policy>
    mfem::Vector applyKnownFactorization(
        Policy policy,
        const mfem::Vector &rightHandSide,
        preconditioning::GravityFactorizationStatistics *statistics = nullptr
    ) {
        mfem::Vector massDiagonal(2);
        massDiagonal     = 1.0;
        auto massInverse = backend::prepare(backend::Diagonal{}, massDiagonal);

        mfem::DenseMatrix schurMatrix(1);
        schurMatrix(0, 0) = 5.0;
        auto schurInverse = backend::prepare(backend::DenseDirect{}, schurMatrix);

        mfem::DenseMatrix divergence(1, 2);
        divergence(0, 0) = 2.0;
        divergence(0, 1) = -1.0;

        preconditioning::GravityFactorizationOperator<Policy> factorization(
            policy, massInverse, schurInverse, divergence
        );
        mfem::Vector action(factorization.Height());
        action = std::numeric_limits<double>::quiet_NaN();
        factorization.Mult(rightHandSide, action);
        if (statistics != nullptr) {
            *statistics = factorization.GetStatistics();
        }
        return action;
    }

    void checkVector(
        const mfem::Vector &computed,
        const std::array<
            double,
            3> &expected
    ) {
        REQUIRE(computed.Size() == static_cast<int>(expected.size()));
        for (int index = 0; index < computed.Size(); ++index) {
            CHECK(computed(index) == Catch::Approx(expected[static_cast<std::size_t>(index)]).margin(2.0e-14));
        }
    }

    template <typename Policy> void checkExactDenseRecovery(Policy policy) {
        mfem::DenseMatrix mass(2);
        mass(0, 0)       = 2.0;
        mass(0, 1)       = 0.5;
        mass(1, 0)       = 0.5;
        mass(1, 1)       = 1.5;
        auto massInverse = backend::prepare(backend::DenseDirect{}, mass);

        mfem::DenseMatrix divergence(1, 2);
        divergence(0, 0) = 1.0;
        divergence(0, 1) = -2.0;

        mfem::Vector divergenceTranspose(2);
        divergenceTranspose(0) = 1.0;
        divergenceTranspose(1) = -2.0;
        mfem::Vector massInverseDivergenceTranspose(2);
        massInverse.Mult(divergenceTranspose, massInverseDivergenceTranspose);

        mfem::DenseMatrix schur(1);
        schur(0, 0)       = divergenceTranspose * massInverseDivergenceTranspose;
        auto schurInverse = backend::prepare(backend::DenseDirect{}, schur);

        preconditioning::GravityFactorizationOperator<Policy> factorization(
            policy, massInverse, schurInverse, divergence
        );

        mfem::Vector exact(3);
        exact(0) = 0.7;
        exact(1) = -1.2;
        exact(2) = 0.4;

        mfem::Vector rightHandSide(3);
        mfem::Vector exactGradient(exact.GetData(), 2);
        mfem::Vector gradientRightHandSide(rightHandSide.GetData(), 2);
        mass.Mult(exactGradient, gradientRightHandSide);
        gradientRightHandSide(0) += divergence(0, 0) * exact(2);
        gradientRightHandSide(1) += divergence(0, 1) * exact(2);
        rightHandSide(2) = divergence(0, 0) * exact(0) + divergence(0, 1) * exact(1);

        mfem::Vector action(3);
        action                            = 0.0;
        const double *const actionStorage = action.GetData();
        factorization.Mult(rightHandSide, action);

        CHECK(action.GetData() == actionStorage);
        for (int index = 0; index < action.Size(); ++index) {
            CHECK(action(index) == Catch::Approx(exact(index)).margin(2.0e-13));
        }

        mfem::Vector repeated(3);
        repeated = 0.0;
        factorization.Mult(rightHandSide, repeated);
        for (int index = 0; index < repeated.Size(); ++index) {
            CHECK(repeated(index) == action(index));
        }
    }

    struct PreparedGeometry final {
        mean_field::fem::FEM finiteElements;
        gravity_context::GravityFieldGeometryContext context;

        explicit PreparedGeometry(const mean_field::utils::Args &arguments)
            : finiteElements(
                  mean_field::fem::setup_fem(
                      arguments.mesh_file,
                      arguments,
                      0
                  )
              ),
              context(
                  finiteElements,
                  *finiteElements.domainMapperStateless
              ) {
            mfem::Vector displacementTrue(finiteElements.displacementFes->GetTrueVSize());
            displacementTrue                = 0.0;
            const mfem::Vector displacement = context.GetDisplacementMap().gather(displacementTrue);
            context.PreparePrimal(displacement, {.value = 1}, {.value = 1});
        }
    };
} // namespace

TEST_CASE(
    "Gravity Field Blocks Expose Complete Compile-Time Ownership And Backend Contracts",
    tags::preconditioning_gravity_unit
) {
    using Form         = blocks::surface_deformed_stellar_equilibrium_form;
    using JacobianForm = blocks::surface_deformed_stellar_equilibrium_jacobian_form;

    STATIC_CHECK(preconditioning::PreconditionerComponent<FixedGravityLDU>);
    STATIC_CHECK(preconditioning::PreconditionerComponent<ChebyshevGravityLDU>);
    STATIC_CHECK(preconditioning::PreconditionerComponent<AdaptiveGravityLDU>);
    STATIC_CHECK(preconditioning::backend::ArnoldiAdmissible<typename FixedGravityLDU::BackendType>);
    STATIC_CHECK(preconditioning::backend::ArnoldiAdmissible<typename ChebyshevGravityLDU::BackendType>);
    STATIC_CHECK_FALSE(preconditioning::backend::ArnoldiAdmissible<typename AdaptiveGravityLDU::BackendType>);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<FixedGravityPlan, Form>);
    STATIC_CHECK(preconditioning::CompatiblePreconditionerFor<FixedGravityPlan, Form, JacobianForm>);
    STATIC_CHECK(preconditioning::StationaryLinearPreconditionerPlan<FixedGravityPlan>);
    STATIC_CHECK(preconditioning::CompletePreconditionerFor<AdaptiveGravityPlan, Form>);
    STATIC_CHECK(preconditioning::CompatiblePreconditionerFor<AdaptiveGravityPlan, Form, JacobianForm>);
    STATIC_CHECK_FALSE(preconditioning::StationaryLinearPreconditionerPlan<AdaptiveGravityPlan>);
    STATIC_CHECK(FixedGravityLDU::RequiredCouplings::size == 3);
}

TEST_CASE(
    "Gravity Factorization Policies Preserve Their Signed Block Algebra",
    tags::preconditioning_gravity_unit
) {
    mfem::Vector rightHandSide(3);
    rightHandSide(0) = 3.0;
    rightHandSide(1) = 4.0;
    rightHandSide(2) = 7.0;

    checkVector(applyKnownFactorization(preconditioning::GravityBlockDiagonal{}, rightHandSide), {3.0, 4.0, 1.4});
    checkVector(applyKnownFactorization(preconditioning::GravityLowerTriangular{}, rightHandSide), {3.0, 4.0, -1.0});
    checkVector(applyKnownFactorization(preconditioning::GravityUpperTriangular{}, rightHandSide), {5.8, 2.6, -1.4});

    preconditioning::GravityFactorizationStatistics statistics;
    checkVector(
        applyKnownFactorization(preconditioning::GravityApproximateLDU{}, rightHandSide, &statistics), {5.0, 3.0, -1.0}
    );
    CHECK(statistics.applications == 1);
    CHECK(statistics.massInverseApplications == 2);
    CHECK(statistics.potentialSchurApplications == 1);
    CHECK(statistics.divergenceApplications == 1);
    CHECK(statistics.transposeDivergenceApplications == 1);
}

TEST_CASE(
    "Exact Gravity LDU Recovers A Dense Coupled Saddle-Point System",
    tags::preconditioning_gravity_unit
) {
    checkExactDenseRecovery(preconditioning::GravityApproximateLDU{});
}

TEST_CASE(
    "Assembled Gravity Divergence Matches The Prepared Matrix-Free Couplings",
    tags::preconditioning_gravity_integration
) {
    const auto arguments = test_utils::setup_args();
    PreparedGeometry geometry(arguments);

    const auto assembledDivergence = preconditioning::assembleGravityDivergenceSurrogate(geometry.finiteElements);
    const mfem::Operator &preparedDivergence = geometry.context.GetDivergenceOperator();

    const mfem::Vector flux                  = gravity_prepared_test_utils::make_deterministic_vector(
        geometry.finiteElements.gravityFluxFes->GetTrueVSize(), 0.31
    );
    mfem::Vector assembledForward(assembledDivergence->Height());
    mfem::Vector preparedForward(preparedDivergence.Height());
    assembledDivergence->Mult(flux, assembledForward);
    preparedDivergence.Mult(flux, preparedForward);

    const mfem::Vector potential = gravity_prepared_test_utils::make_deterministic_vector(
        geometry.finiteElements.gravityPotentialFes->GetTrueVSize(), 0.73
    );
    mfem::Vector assembledTranspose(assembledDivergence->Width());
    mfem::Vector preparedTranspose(preparedDivergence.Width());
    assembledDivergence->MultTranspose(potential, assembledTranspose);
    preparedDivergence.MultTranspose(potential, preparedTranspose);

    const MPI_Comm communicator = geometry.finiteElements.mesh->GetComm();
    CHECK(gravity_prepared_test_utils::relative_error(assembledForward, preparedForward, communicator) <= 2.0e-12);
    CHECK(gravity_prepared_test_utils::relative_error(assembledTranspose, preparedTranspose, communicator) <= 2.0e-12);

    const auto &gradientMap  = geometry.context.GetMassOperator().GetFluxMap();
    const auto &potentialMap = geometry.context.GetSourceOperator().GetPotentialMap();
    preconditioning::ReducedGravityDivergenceOperator reducedDivergence(preparedDivergence, gradientMap, potentialMap);
    const mfem::Vector reducedFlux =
        gravity_prepared_test_utils::make_deterministic_vector(gradientMap.reduced_size(), 0.47);
    mfem::Vector reducedAction(reducedDivergence.Height());
    reducedDivergence.Mult(reducedFlux, reducedAction);

    const mfem::Vector trueFlux = gradientMap.scatter(reducedFlux);
    mfem::Vector trueAction(potentialMap.full_size());
    preparedDivergence.Mult(trueFlux, trueAction);
    const mfem::Vector expectedReducedAction = potentialMap.gather(trueAction);
    CHECK(gravity_prepared_test_utils::relative_error(reducedAction, expectedReducedAction, communicator) <= 2.0e-14);
}

TEST_CASE(
    "Prepared Gravity Block Diagonal Is Legacy Equivalent And Allocation Stable",
    tags::preconditioning_gravity_integration
) {
    const auto arguments = test_utils::setup_args();
    PreparedGeometry geometry(arguments);

    mean_field::operators::ReducedGravityFieldPreconditioner legacy(geometry.finiteElements, geometry.context);
    const auto block = preconditioning::GravityFieldBlock(
        backend::Diagonal{}, FixedAMG{backend::FixedCycles{.cycles = 1}}, preconditioning::GravityBlockDiagonal{}
    );
    auto prepared                    = preconditioning::prepare(geometry.finiteElements, geometry.context, block);

    const mfem::Vector rightHandSide = gravity_prepared_test_utils::make_deterministic_vector(prepared.Width(), 0.59);
    mfem::Vector legacyAction(prepared.Height());
    mfem::Vector preparedAction(prepared.Height());
    legacyAction                           = 0.0;
    preparedAction                         = 0.0;
    double *const preparedStorage          = preparedAction.GetData();

    const std::uint64_t massPreparations   = geometry.context.GetMassOperator().GetPreparationCount();
    const std::uint64_t sourcePreparations = geometry.context.GetSourceOperator().GetPreparationCount();
    legacy.Mult(rightHandSide, legacyAction);
    prepared.Mult(rightHandSide, preparedAction);

    CHECK(preparedAction.GetData() == preparedStorage);
    CHECK(geometry.context.GetMassOperator().GetPreparationCount() == massPreparations);
    CHECK(geometry.context.GetSourceOperator().GetPreparationCount() == sourcePreparations);
    CHECK(
        gravity_prepared_test_utils::relative_error(
            preparedAction, legacyAction, geometry.finiteElements.mesh->GetComm()
        ) <= 2.0e-12
    );
}

TEST_CASE(
    "Prepared Gravity Blocks Refresh Explicitly Without Repreparing Geometry",
    tags::preconditioning_gravity_integration
) {
    const auto arguments = test_utils::setup_args();
    PreparedGeometry geometry(arguments);

    const auto block = preconditioning::GravityFieldBlock(
        backend::Diagonal{}, FixedAMG{backend::FixedCycles{.cycles = 1}}, preconditioning::GravityApproximateLDU{}
    );
    auto prepared             = preconditioning::prepare(geometry.finiteElements, geometry.context, block);
    const auto chebyshevBlock = preconditioning::GravityFieldBlock(
        backend::MatrixFreeChebyshev{.order = 2, .powerIterations = 10}, FixedAMG{backend::FixedCycles{.cycles = 1}},
        preconditioning::GravityApproximateLDU{}
    );
    auto chebyshevPrepared        = preconditioning::prepare(geometry.finiteElements, geometry.context, chebyshevBlock);

    const auto unchanged          = prepared.Refresh(geometry.finiteElements, geometry.context);
    const auto unchangedChebyshev = chebyshevPrepared.Refresh(geometry.finiteElements, geometry.context);
    CHECK_FALSE(unchanged.DidAnyWork());
    CHECK_FALSE(unchangedChebyshev.DidAnyWork());
    CHECK(prepared.GetStatistics().refreshChecks == 1);
    CHECK(prepared.GetStatistics().noOpRefreshes == 1);

    const mfem::Vector displacementTrue = gravity_prepared_test_utils::make_displacement(geometry.finiteElements, 0.4);
    const mfem::Vector displacement     = geometry.context.GetDisplacementMap().gather(displacementTrue);
    geometry.context.PreparePrimal(displacement, {.value = 1}, {.value = 2});
    CHECK_FALSE(prepared.IsCurrent());
    CHECK_FALSE(chebyshevPrepared.IsCurrent());

    mfem::Vector rightHandSide(prepared.Width());
    mfem::Vector action(prepared.Height());
    rightHandSide = 1.0;
    action        = 0.0;
    CHECK_THROWS_AS(prepared.Mult(rightHandSide, action), std::logic_error);
    CHECK_THROWS_AS(chebyshevPrepared.Mult(rightHandSide, action), std::logic_error);

    const std::uint64_t massPreparations   = geometry.context.GetMassOperator().GetPreparationCount();
    const std::uint64_t sourcePreparations = geometry.context.GetSourceOperator().GetPreparationCount();
    const auto changed                     = prepared.Refresh(geometry.finiteElements, geometry.context);
    const auto changedChebyshev            = chebyshevPrepared.Refresh(geometry.finiteElements, geometry.context);

    CHECK(changed.geometryChanged);
    CHECK_FALSE(changed.discretizationChanged);
    CHECK(changed.rebuiltMassInverse);
    CHECK_FALSE(changed.rebuiltDivergenceBinding);
    CHECK(changed.rebuiltPotentialSchur);
    CHECK(prepared.IsCurrent());
    CHECK(changedChebyshev.geometryChanged);
    CHECK(changedChebyshev.rebuiltMassInverse);
    CHECK(changedChebyshev.rebuiltPotentialSchur);
    CHECK(chebyshevPrepared.IsCurrent());
    CHECK(chebyshevPrepared.GetMassInverse().GetStatistics().setups == 2);
    CHECK(geometry.context.GetMassOperator().GetPreparationCount() == massPreparations);
    CHECK(geometry.context.GetSourceOperator().GetPreparationCount() == sourcePreparations);
    CHECK(prepared.GetStatistics().refreshes == 1);

    mfem::Vector refreshedAction(chebyshevPrepared.Height());
    refreshedAction = 0.0;
    chebyshevPrepared.Mult(rightHandSide, refreshedAction);
    for (int index = 0; index < refreshedAction.Size(); ++index) {
        CHECK(std::isfinite(refreshedAction(index)));
    }

    // A discretization revision reconstructs the matrix-free mass operator.  The
    // owning gravity block must reject every route to its now-stale inverse until
    // refresh has rebound and rebuilt the Chebyshev smoother.
    geometry.context.PreparePrimal(displacement, {.value = 2}, {.value = 2});
    CHECK_FALSE(chebyshevPrepared.IsCurrent());
    CHECK_THROWS_AS(chebyshevPrepared.Mult(rightHandSide, action), std::logic_error);
    CHECK_THROWS_AS(chebyshevPrepared.GetMassInverse(), std::logic_error);

    const auto reconstructed = chebyshevPrepared.Refresh(geometry.finiteElements, geometry.context);
    CHECK(reconstructed.discretizationChanged);
    CHECK_FALSE(reconstructed.geometryChanged);
    CHECK(reconstructed.rebuiltMassInverse);
    CHECK(reconstructed.rebuiltDivergenceBinding);
    CHECK(reconstructed.rebuiltPotentialSchur);
    CHECK(chebyshevPrepared.IsCurrent());
    CHECK(chebyshevPrepared.GetMassInverse().GetStatistics().setups == 3);

    mfem::Vector firstReconstructedAction(chebyshevPrepared.Height());
    mfem::Vector secondReconstructedAction(chebyshevPrepared.Height());
    firstReconstructedAction  = 0.0;
    secondReconstructedAction = 0.0;
    chebyshevPrepared.Mult(rightHandSide, firstReconstructedAction);
    chebyshevPrepared.Mult(rightHandSide, secondReconstructedAction);
    mfem::Vector repeatabilityError(firstReconstructedAction);
    repeatabilityError -= secondReconstructedAction;
    CHECK(repeatabilityError.Norml2() <= 2.0e-14 * std::max(1.0, firstReconstructedAction.Norml2()));
    for (int index = 0; index < firstReconstructedAction.Size(); ++index) {
        CHECK(std::isfinite(firstReconstructedAction(index)));
    }
}