#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>
#include <cmath>
#include <memory>
#include <mfem.hpp>
#include <stdexcept>

import mean_field;
import test_helpers;

using namespace mean_field;
using Catch::Matchers::WithinAbs;
namespace prepared_test = gravity_prepared_test_utils;

namespace {
    // Deliberately not a KelvinCompactification dynamic type: the prepared
    // operator must keep arbitrary exterior maps on the general derivative.
    class DelegatingExteriorMap final : public mapping::compactification::ExteriorDomainMap {
    public:
        explicit DelegatingExteriorMap(mapping::compactification::options::KelvinCompactificationOptions options)
            : kelvin(options) {}
        mapping::MappingStatus Evaluate(const mapping::compactification::ExteriorMapInput &input,
                                       mapping::compactification::ExteriorMapResult &result) const override {
            return kelvin.Evaluate(input, result);
        }
        mapping::MappingStatus EvaluateVariation(const mapping::compactification::ExteriorMapInput &input,
                                                const mapping::compactification::ExteriorMapResult &result,
                                                const mapping::compactification::ExteriorMapDirection &direction,
                                                mapping::compactification::ExteriorMapVariation &variation) const override {
            return kelvin.EvaluateVariation(input, result, direction, variation);
        }
        std::string_view GetName() const noexcept override { return "DelegatingExteriorMap"; }
    private:
        mapping::compactification::KelvinCompactification kelvin;
    };

    [[nodiscard]] mfem::Vector make_folding_displacement(const mean_field::fem::FEM &f) {
        mfem::ParGridFunction field(f.displacementFes.get());
        mfem::VectorFunctionCoefficient coefficient(
            f.mesh->Dimension(), [](const mfem::Vector &position, mfem::Vector &value) {
                value.SetSize(position.Size());
                for (int dimension = 0; dimension < position.Size(); ++dimension) {
                    value(dimension) = -2.0 * position(dimension);
                }
            }
        );
        field.ProjectCoefficient(coefficient);

        mfem::Vector displacementTrue;
        field.GetTrueDofs(displacementTrue);
        return displacementTrue;
    }
} // namespace

TEST_CASE(
    "Prepared Hdiv Geometry Variation Preserves Reference Piola Contractions",
    tags::gravity_prepared_jacobian_accuracy
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());

    operators::PreparedMappedHDivMassOperator preparedOperator(f, *f.domainMapperStateless);
    const MPI_Comm communicator  = f.gravityFluxFes->GetComm();
    const mfem::Vector first     = prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(), 0.31);
    const mfem::Vector second    = prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(), 0.79);
    const mfem::Vector direction = prepared_test::make_displacement(f, 0.43);
    mfem::Vector firstAction;

    bool hasStellar = false;
    bool hasVacuum  = false;
    for (int element = 0; element < f.mesh->GetNE(); ++element) {
        const bool vacuum = f.domainMapperStateless->IsCompactifiedElement(*f.mesh->GetElementTransformation(element));
        hasVacuum         = hasVacuum || vacuum;
        hasStellar        = hasStellar || !vacuum;
    }
    const int localDomains[2]{hasStellar ? 1 : 0, hasVacuum ? 1 : 0};
    int globalDomains[2]{};
    REQUIRE(MPI_Allreduce(localDomains, globalDomains, 2, MPI_INT, MPI_MAX, communicator) == MPI_SUCCESS);
    REQUIRE(globalDomains[0] != 0);
    REQUIRE(globalDomains[1] != 0);

    // The stateless path still constructs each physically mapped RT basis;
    // compare both an undeformed and a changed prepared geometry, including
    // Kelvin exterior elements, against the compact forward/dual contractions.
    for (const double scale : {0.0, 0.7}) {
        const mfem::Vector displacementTrue = prepared_test::make_displacement(f, scale);
        const mfem::Vector displacement     = preparedOperator.GetDisplacementMap().gather(displacementTrue);
        preparedOperator.Prepare(displacement);
        CHECK(preparedOperator.GetFastVariationElementCount() == f.mesh->GetNE());
        CHECK(preparedOperator.GetTensorVariationElementCount() == f.mesh->GetNE());
        preparedOperator.MultDisplacementVariationTrue(first, direction, firstAction);

        mfem::Vector referenceAction;
        preparedOperator.MultDisplacementVariationTrueCachedReference(first, direction, referenceAction);
        CHECK(prepared_test::relative_error(firstAction, referenceAction, communicator) < 2.0e-11);
        preparedOperator.MultDisplacementVariationTrueReference(first, direction, referenceAction);
        CHECK(prepared_test::relative_error(firstAction, referenceAction, communicator) < 2.0e-11);
        operators::kernels::apply_mapped_hdiv_mass_variation(
            f, *f.domainMapperStateless, first, displacementTrue, direction, referenceAction
        );
        const double error = prepared_test::relative_error(firstAction, referenceAction, communicator);
        INFO("Deformation scale = " << scale);
        INFO("Prepared/stateless geometry-variation relative error = " << error);
        CHECK(error < 2.0e-11);
    }

    mfem::Vector secondAction;
    preparedOperator.MultDisplacementVariationTrue(second, direction, secondAction);
    const double firstSecond = prepared_test::global_dot(first, secondAction, communicator);
    const double secondFirst = prepared_test::global_dot(second, firstAction, communicator);
    CHECK(prepared_test::relative_scalar_error(firstSecond, secondFirst) < 2.0e-11);
    CHECK(preparedOperator.GetPreparationCount() == 2);
}

TEST_CASE("Prepared Hdiv Tangent Retains General Exterior Map Fallback",
          tags::gravity_prepared_jacobian_accuracy) {
    auto args = test_utils::setup_args();
    auto f = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());
    mapping::DomainMapper mapper(args.domain_mapper_options, std::make_unique<DelegatingExteriorMap>(args.kelvin_options));
    operators::PreparedMappedHDivMassOperator prepared(f, mapper);
    const auto base = prepared_test::make_displacement(f, 0.7);
    const auto direction = prepared_test::make_displacement(f, 0.43);
    const auto flux = prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(), 0.31);
    prepared.Prepare(prepared.GetDisplacementMap().gather(base));
    int stellarElements = 0;
    for (int element = 0; element < f.mesh->GetNE(); ++element)
        stellarElements += !mapper.IsCompactifiedElement(*f.mesh->GetElementTransformation(element));
    CHECK(prepared.GetFastVariationElementCount() == stellarElements);
    mfem::Vector actual, expected;
    prepared.MultDisplacementVariationTrue(flux, direction, actual);
    operators::kernels::apply_mapped_hdiv_mass_variation(f, mapper, flux, base, direction, expected);
    CHECK(prepared_test::relative_error(actual, expected, f.mesh->GetComm()) < 2.0e-11);
    prepared.PreparePrimal(prepared.GetDisplacementMap().gather(base));
    CHECK_FALSE(prepared.HasVariationData());
    CHECK(prepared.GetFastVariationElementCount() == 0);
    CHECK(prepared.GetFastMassElementCount() == f.mesh->GetNE());
    mfem::Vector primal, primalReference;
    const auto reducedFlux = prepared.GetFluxMap().gather(flux);
    prepared.Mult(reducedFlux, primal);
    prepared.MultReference(reducedFlux, primalReference);
    CHECK(prepared_test::relative_error(primal, primalReference, f.mesh->GetComm()) < 2.0e-11);
    prepared.Prepare(prepared.GetDisplacementMap().gather(base));
    prepared.MultDisplacementVariationTrue(flux, direction, actual);
    CHECK(prepared_test::relative_error(actual, expected, f.mesh->GetComm()) < 2.0e-11);
}

TEST_CASE(
    "Prepared Mapped Hdiv Mass Reports Invalid Candidate Geometry Without Unwinding",
    tags::gravity_prepared_unit &tags::geometry
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());

    operators::PreparedMappedHDivMassOperator preparedOperator(f, *f.domainMapperStateless);
    mfem::Vector displacement = preparedOperator.GetDisplacementMap().gather(make_folding_displacement(f));

    const auto rejected       = preparedOperator.TryPrepare(displacement);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().reason == operators::HDivMassPreparationRejectionReason::invalid_mapping);
    CHECK(rejected.error().mappingStatus == mapping::MappingStatus::non_positive_determinant);
    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK_FALSE(preparedOperator.HasVariationData());
    CHECK(preparedOperator.GetPreparationCount() == 0);

    CHECK_THROWS_AS(preparedOperator.Prepare(displacement), std::domain_error);
    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK(preparedOperator.GetPreparationCount() == 0);

    displacement         = 0.0;
    const auto recovered = preparedOperator.TryPrepare(displacement);
    REQUIRE(recovered.has_value());
    CHECK(preparedOperator.IsPrepared());
    CHECK(preparedOperator.HasVariationData());
    CHECK(preparedOperator.GetPreparationCount() == 1);
}

TEST_CASE(
    "Prepared Mapped Hdiv Mass Matches Stateless Kernel",
    tags::gravity_prepared
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    operators::PreparedMappedHDivMassOperator prepared_operator(f, *f.domainMapperStateless);
    REQUIRE(prepared_operator.Width() == prepared_operator.GetFluxMap().reduced_size());
    REQUIRE(prepared_operator.Height() == prepared_operator.GetFluxMap().reduced_size());

    const mfem::Vector gravity_gradient_true =
        prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(), 0.21);
    const mfem::Vector gravity_gradient = prepared_operator.GetFluxMap().gather(gravity_gradient_true);
    const MPI_Comm communicator         = f.gravityFluxFes->GetComm();

    mfem::Vector identity_action;
    mfem::Vector deformed_action;

    for (const double deformation_scale : {0.0, 1.0}) {
        const mfem::Vector displacement_true = prepared_test::make_displacement(f, deformation_scale);
        const mfem::Vector displacement      = prepared_operator.GetDisplacementMap().gather(displacement_true);

        prepared_operator.Prepare(displacement);

        mfem::Vector prepared_action;

        prepared_operator.Mult(gravity_gradient, prepared_action);
        CHECK(prepared_operator.GetFastMassElementCount() == f.mesh->GetNE());
        mfem::Vector mfemAction;
        prepared_operator.MultReference(gravity_gradient, mfemAction);
        CHECK(prepared_test::relative_error(prepared_action, mfemAction, communicator) < 2.0e-11);
        // Reference calls must restore the optimized mode and leave PA data intact.
        prepared_operator.Mult(gravity_gradient, mfemAction);
        CHECK(prepared_test::relative_error(prepared_action, mfemAction, communicator) < 2.0e-14);
        mfem::Vector reference_action_true;
        operators::kernels::apply_mapped_hdiv_mass(
            f, *f.domainMapperStateless, gravity_gradient_true, displacement_true, reference_action_true
        );
        const mfem::Vector reference_action = prepared_operator.GetFluxMap().gather(reference_action_true);

        const double relative_error = prepared_test::relative_error(prepared_action, reference_action, communicator);

        INFO("Deformation scale = " << deformation_scale);
        INFO("Prepared action norm = " << prepared_test::global_norm(prepared_action, communicator));
        INFO("Reference action norm = " << prepared_test::global_norm(reference_action, communicator));
        INFO("Relative prepared-operator error = " << relative_error);

        REQUIRE(prepared_operator.IsPrepared());
        CHECK_THAT(relative_error, WithinAbs(0.0, 2.0e-11));

        if (deformation_scale == 0.0) {
            identity_action = prepared_action;
        } else {
            deformed_action = prepared_action;
        }
    }

    const double geometry_change = prepared_test::relative_error(deformed_action, identity_action, communicator);

    INFO("Relative action change under deformation = " << geometry_change);

    CHECK(prepared_operator.GetPreparationCount() == 2);
    CHECK(geometry_change > 1.0e-5);
}

TEST_CASE(
    "Prepared Mapped Hdiv Mass Preserves Operator Identities",
    tags::gravity_prepared
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    operators::PreparedMappedHDivMassOperator prepared_operator(f, *f.domainMapperStateless);
    REQUIRE(prepared_operator.Width() == prepared_operator.GetFluxMap().reduced_size());
    REQUIRE(prepared_operator.Height() == prepared_operator.GetFluxMap().reduced_size());
    const mfem::Vector displacement =
        prepared_operator.GetDisplacementMap().gather(prepared_test::make_displacement(f, 1.0));
    prepared_operator.Prepare(displacement);

    const mfem::Vector first = prepared_operator.GetFluxMap().gather(
        prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(), 0.17)
    );
    const mfem::Vector second = prepared_operator.GetFluxMap().gather(
        prepared_test::make_deterministic_vector(f.gravityFluxFes->GetTrueVSize(), 0.83)
    );
    const mfem::Vector combination = prepared_test::linear_combination(first, 1.7, second, -0.4);

    mfem::Vector first_action;
    mfem::Vector second_action;
    mfem::Vector combination_action;
    mfem::Vector zero_action;

    prepared_operator.Mult(first, first_action);
    prepared_operator.Mult(second, second_action);
    prepared_operator.Mult(combination, combination_action);

    mfem::Vector expected_combination = prepared_test::linear_combination(first_action, 1.7, second_action, -0.4);

    mfem::Vector zero(first.Size());
    zero = 0.0;
    prepared_operator.Mult(zero, zero_action);

    const MPI_Comm communicator       = f.gravityFluxFes->GetComm();

    const double first_second_product = prepared_test::global_dot(first, second_action, communicator);
    const double second_first_product = prepared_test::global_dot(second, first_action, communicator);
    const double symmetry_error = prepared_test::relative_scalar_error(first_second_product, second_first_product);
    const double linearity_error =
        prepared_test::relative_error(combination_action, expected_combination, communicator);
    const double first_energy             = prepared_test::global_dot(first, first_action, communicator);
    const double second_energy            = prepared_test::global_dot(second, second_action, communicator);
    const std::uint64_t preparation_count = prepared_operator.GetPreparationCount();

    mfem::Vector repeated_action;
    prepared_operator.Mult(first, repeated_action);

    INFO("u^T M v = " << first_second_product);
    INFO("v^T M u = " << second_first_product);
    INFO("Relative symmetry error = " << symmetry_error);
    INFO("Relative linearity error = " << linearity_error);
    INFO("u^T M u = " << first_energy);
    INFO("v^T M v = " << second_energy);

    CHECK_THAT(symmetry_error, WithinAbs(0.0, 2.0e-12));
    CHECK_THAT(linearity_error, WithinAbs(0.0, 2.0e-12));
    CHECK_THAT(prepared_test::global_norm(zero_action, communicator), WithinAbs(0.0, 1.0e-14));
    CHECK(first_energy > 0.0);
    CHECK(second_energy > 0.0);
    CHECK(prepared_test::relative_error(repeated_action, first_action, communicator) < 2.0e-14);
    CHECK(prepared_operator.GetPreparationCount() == preparation_count);
}

TEST_CASE(
    "Prepared Mapped Hdiv Mass Diagonal Is Positive Across Both Domains",
    tags::gravity_prepared
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    operators::PreparedMappedHDivMassOperator prepared_operator(f, *f.domainMapperStateless);
    const mfem::Vector displacement =
        prepared_operator.GetDisplacementMap().gather(prepared_test::make_displacement(f, 1.0));
    prepared_operator.Prepare(displacement);

    mfem::Vector diagonal;
    mfem::Vector true_diagonal;
    prepared_operator.AssembleDiagonal(diagonal);
    prepared_operator.AssembleTrueDiagonal(true_diagonal);

    REQUIRE(diagonal.Size() == prepared_operator.Height());
    REQUIRE(true_diagonal.Size() == prepared_operator.GetFluxMap().full_size());

    const mfem::Vector gathered_true_diagonal = prepared_operator.GetFluxMap().gather(true_diagonal);

    for (int i = 0; i < diagonal.Size(); ++i) {
        REQUIRE(std::isfinite(diagonal(i)));
        CHECK(diagonal(i) > 0.0);
        CHECK_THAT(diagonal(i), WithinAbs(gathered_true_diagonal(i), 1.0e-14 * std::abs(diagonal(i))));
    }
}