#include <catch2/catch_test_macros.hpp>
#include <cmath>
#include <cstdint>
#include <mfem.hpp>
#include <type_traits>

import mean_field;
import test_helpers;

namespace barotropic_closure_context_test_utils {
    namespace field   = mean_field::field;
    namespace domain  = mean_field::utils::domain;
    namespace context = mean_field::operators::context::barotropic;

    using Schema      = domain::CoreEnvelopeVacuumDomainSchema;

    struct Maps final {
        field::FieldDofMap density;
        field::FieldDofMap enthalpy;
        field::FieldDofMap displacement;

        explicit Maps(const mean_field::fem::FEM &f)
            : density(
                  field::make_field_dof_map<
                      field::Density,
                      Schema>(*f.densityFes)
              ),
              enthalpy(
                  field::make_field_dof_map<
                      field::Enthalpy,
                      Schema>(*f.enthalpyFes)
              ),
              displacement(
                  field::make_field_dof_map<
                      field::Displacement,
                      Schema>(*f.displacementFes)
              ) {
        }
    };

    [[nodiscard]] context::BarotropicClosureDependencies make_dependencies() {
        return {
            .discretization = {.identity = 101, .revision = 2},
            .density        = {.identity = 103, .revision = 3},
            .enthalpy       = {.identity = 107, .revision = 5},
            .displacement   = {.identity = 109, .revision = 7}
        };
    }

    [[nodiscard]] context::BarotropicClosureStateView make_state(
        const mfem::Vector &density,
        const mfem::Vector &enthalpy,
        const mfem::Vector &displacement
    ) {
        return {.density = density, .enthalpy = enthalpy, .displacement = displacement};
    }

    [[nodiscard]] mfem::Vector reduce(
        const field::FieldDofMap &map,
        const mfem::Vector &full
    ) {
        return map.gather(full);
    }

    [[nodiscard]] mfem::Vector make_density(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction value(f.densityFes.get());
        mfem::FunctionCoefficient coefficient([phase](const mfem::Vector &position) {
            return 0.71 + 0.05 * std::sin(0.63 * position(0) + phase) + 0.02 * position(1);
        });
        value.ProjectCoefficient(coefficient);
        mfem::Vector result;
        value.GetTrueDofs(result);
        return result;
    }

    [[nodiscard]] mfem::Vector make_enthalpy(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction value(f.enthalpyFes.get());
        mfem::FunctionCoefficient coefficient([phase](const mfem::Vector &position) {
            return 0.93 + 0.04 * std::cos(0.57 * position(1) - phase) + 0.015 * position(2);
        });
        value.ProjectCoefficient(coefficient);
        mfem::Vector result;
        value.GetTrueDofs(result);
        return result;
    }

    [[nodiscard]] double relative_error(
        const mfem::Vector &left,
        const mfem::Vector &right,
        const MPI_Comm communicator
    ) {
        return gravity_prepared_test_utils::relative_error(left, right, communicator);
    }
} // namespace barotropic_closure_context_test_utils

TEST_CASE(
    "Prepared Barotropic Closure Owns Its Linearization Context",
    tags::barotrope &tags::closure &tags::contexts &tags::prepared &tags::field &tags::unit
) {
    using Operator = mean_field::operators::PreparedBarotropicClosureOperator;
    using Context  = mean_field::operators::context::barotropic::BarotropicClosureLinearizationContext;

    STATIC_REQUIRE_FALSE(std::is_copy_constructible_v<Context>);
    STATIC_REQUIRE_FALSE(std::is_copy_assignable_v<Context>);
    STATIC_REQUIRE_FALSE(std::is_move_constructible_v<Context>);
    STATIC_REQUIRE_FALSE(std::is_move_assignable_v<Context>);

    auto args              = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());

    const mean_field::eos::Polytrope equationOfState(3.0, 1.5);
    Operator preparedOperator(f, *f.domainMapperStateless, equationOfState);

    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK_FALSE(preparedOperator.GetContext().IsPrepared());
    CHECK(&preparedOperator.GetContext() == &preparedOperator.GetContext());

    const auto statistics = preparedOperator.GetContextPreparationStatistics();
    CHECK(statistics.staticPreparations == 0);
    CHECK(statistics.geometryPreparations == 0);
    CHECK(statistics.baseStatePreparations == 0);
}

TEST_CASE(
    "Barotropic Closure Context Applies Selective Invalidation In Reduced Field Coordinates",
    tags::barotrope &tags::closure &tags::contexts &tags::prepared &tags::field &tags::unit
) {
    using namespace barotropic_closure_context_test_utils;

    auto args              = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.okay());

    const Maps maps(f);

    const mean_field::eos::Polytrope equationOfState(3.0, 1.5);

    mean_field::operators::PreparedBarotropicClosureOperator preparedOperator(
        f, *f.domainMapperStateless, equationOfState
    );

    mfem::Vector density  = reduce(maps.density, make_density(f, 0.17));

    mfem::Vector enthalpy = reduce(maps.enthalpy, make_enthalpy(f, 0.29));

    const mfem::Vector initialDisplacement =
        reduce(maps.displacement, gravity_prepared_test_utils::make_displacement(f, 0.35));

    /*
     * A second smooth, orientation-preserving geometry.
     *
     * Do not manufacture a new geometry by perturbing an arbitrary H1
     * coefficient. A modest change in one high-order displacement DOF can
     * correspond to a very large local displacement gradient and can invert
     * an element.
     */
    const mfem::Vector changedDisplacement =
        reduce(maps.displacement, gravity_prepared_test_utils::make_displacement(f, 0.85));

    mfem::Vector displacement(initialDisplacement);

    auto dependencies        = make_dependencies();

    const auto initialReport = preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);

    const auto &context      = preparedOperator.GetContext();

    REQUIRE(preparedOperator.IsPrepared());

    REQUIRE(context.IsPrepared());

    CHECK(context.MatchesDependencies(dependencies));

    CHECK(context.GetDependencies() == dependencies);

    CHECK(initialReport.contextReport.preparedStaticDependencies);

    CHECK(initialReport.contextReport.preparedGeometryState);

    CHECK(initialReport.contextReport.preparedBaseState);

    CHECK(initialReport.contextReport.updatedDensity);

    CHECK(initialReport.contextReport.updatedEnthalpy);

    CHECK(initialReport.contextReport.updatedDisplacement);

    CHECK(initialReport.preparedElementData);

    CHECK(initialReport.DidAnyWork());

    CHECK(preparedOperator.GetPreparationCount() == 1);

    const mfem::Vector frozenDensity      = context.GetBaseDensity();

    const mfem::Vector frozenEnthalpy     = context.GetBaseEnthalpy();

    const mfem::Vector frozenDisplacement = context.GetDisplacement();

    /*
     * Modify all three candidate states without updating their dependency
     * stamps.
     *
     * The context must continue exposing the previously frozen state.
     */
    density(0) += 0.25;
    enthalpy(0) -= 0.18;
    displacement              = changedDisplacement;

    const auto repeatedReport = preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);

    CHECK_FALSE(repeatedReport.DidAnyWork());

    CHECK_FALSE(repeatedReport.contextReport.updatedDensity);

    CHECK_FALSE(repeatedReport.contextReport.updatedEnthalpy);

    CHECK_FALSE(repeatedReport.contextReport.updatedDisplacement);

    CHECK_FALSE(repeatedReport.preparedElementData);

    CHECK(preparedOperator.GetPreparationCount() == 1);

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

    CHECK(relative_error(context.GetBaseDensity(), frozenDensity, communicator) == 0.0);

    CHECK(relative_error(context.GetBaseEnthalpy(), frozenEnthalpy, communicator) == 0.0);

    CHECK(relative_error(context.GetDisplacement(), frozenDisplacement, communicator) == 0.0);

    /*
     * Density invalidation.
     *
     * Geometry remains frozen because the displacement dependency did not
     * change.
     */
    ++dependencies.density.revision;

    const auto densityReport = preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);

    CHECK_FALSE(densityReport.contextReport.preparedStaticDependencies);

    CHECK_FALSE(densityReport.contextReport.preparedGeometryState);

    CHECK(densityReport.contextReport.preparedBaseState);

    CHECK(densityReport.contextReport.updatedDensity);

    CHECK_FALSE(densityReport.contextReport.updatedEnthalpy);

    CHECK_FALSE(densityReport.contextReport.updatedDisplacement);

    CHECK(densityReport.preparedElementData);

    CHECK(context.GetBaseDensity()(0) == density(0));

    /*
     * The candidate displacement has changed, but because its revision has
     * not changed the context must still retain the original geometry.
     */
    CHECK(relative_error(context.GetDisplacement(), frozenDisplacement, communicator) == 0.0);

    /*
     * Enthalpy invalidation.
     */
    ++dependencies.enthalpy.revision;

    const auto enthalpyReport = preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);

    CHECK_FALSE(enthalpyReport.contextReport.preparedStaticDependencies);

    CHECK_FALSE(enthalpyReport.contextReport.preparedGeometryState);

    CHECK(enthalpyReport.contextReport.preparedBaseState);

    CHECK_FALSE(enthalpyReport.contextReport.updatedDensity);

    CHECK(enthalpyReport.contextReport.updatedEnthalpy);

    CHECK_FALSE(enthalpyReport.contextReport.updatedDisplacement);

    CHECK(enthalpyReport.preparedElementData);

    CHECK(context.GetBaseEnthalpy()(0) == enthalpy(0));

    CHECK(relative_error(context.GetDisplacement(), frozenDisplacement, communicator) == 0.0);

    /*
     * Displacement invalidation.
     *
     * The changed geometry is now intentionally accepted. Because it came
     * from the smooth test displacement projection rather than an arbitrary
     * single H1 coefficient mutation, it remains a valid mapping.
     */
    ++dependencies.displacement.revision;

    const auto displacementReport = preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);

    CHECK_FALSE(displacementReport.contextReport.preparedStaticDependencies);

    CHECK(displacementReport.contextReport.preparedGeometryState);

    CHECK(displacementReport.contextReport.preparedBaseState);

    CHECK_FALSE(displacementReport.contextReport.updatedDensity);

    CHECK_FALSE(displacementReport.contextReport.updatedEnthalpy);

    CHECK(displacementReport.contextReport.updatedDisplacement);

    CHECK(displacementReport.preparedElementData);

    CHECK(relative_error(context.GetDisplacement(), changedDisplacement, communicator) == 0.0);

    /*
     * Discretization invalidates everything.
     */
    ++dependencies.discretization.revision;

    const auto discretizationReport =
        preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);

    CHECK(discretizationReport.contextReport.preparedStaticDependencies);

    CHECK(discretizationReport.contextReport.preparedGeometryState);

    CHECK(discretizationReport.contextReport.preparedBaseState);

    CHECK(discretizationReport.contextReport.updatedDensity);

    CHECK(discretizationReport.contextReport.updatedEnthalpy);

    CHECK(discretizationReport.contextReport.updatedDisplacement);

    CHECK(discretizationReport.preparedElementData);

    const auto finalStatistics = preparedOperator.GetContextPreparationStatistics();

    CHECK(finalStatistics.staticPreparations == 2);

    CHECK(finalStatistics.geometryPreparations == 3);

    CHECK(finalStatistics.baseStatePreparations == 5);

    CHECK(preparedOperator.GetPreparationCount() == 5);
}

TEST_CASE(
    "Barotropic Closure Context Uses Identity And Revision For Every Dependency",
    tags::barotrope &tags::closure &tags::contexts &tags::prepared &tags::field &tags::unit
) {
    using namespace barotropic_closure_context_test_utils;

    auto args              = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());

    const Maps maps(f);
    const mean_field::eos::Polytrope equationOfState(3.0, 1.5);
    mean_field::operators::PreparedBarotropicClosureOperator preparedOperator(
        f, *f.domainMapperStateless, equationOfState
    );

    mfem::Vector density      = reduce(maps.density, make_density(f, 0.41));
    mfem::Vector enthalpy     = reduce(maps.enthalpy, make_enthalpy(f, 0.53));
    mfem::Vector displacement = reduce(maps.displacement, gravity_prepared_test_utils::make_displacement(f, 0.60));

    auto dependencies         = make_dependencies();
    preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);

    const mfem::Vector frozenDensity = preparedOperator.GetContext().GetBaseDensity();
    density(0) += 0.19;

    const auto sameStampReport = preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);
    CHECK_FALSE(sameStampReport.DidAnyWork());
    CHECK(preparedOperator.GetContext().GetBaseDensity()(0) == frozenDensity(0));

    ++dependencies.density.identity;
    ++dependencies.density.revision;

    const auto newIdentityReport = preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);
    CHECK_FALSE(newIdentityReport.contextReport.preparedStaticDependencies);
    CHECK_FALSE(newIdentityReport.contextReport.preparedGeometryState);
    CHECK(newIdentityReport.contextReport.preparedBaseState);
    CHECK(newIdentityReport.contextReport.updatedDensity);
    CHECK(preparedOperator.GetContext().GetBaseDensity()(0) == density(0));
    CHECK(preparedOperator.GetContext().MatchesDependencies(dependencies));

    ++dependencies.displacement.identity;
    ++dependencies.displacement.revision;

    const auto displacementIdentityReport =
        preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);
    CHECK_FALSE(displacementIdentityReport.contextReport.preparedStaticDependencies);
    CHECK(displacementIdentityReport.contextReport.preparedGeometryState);
    CHECK(displacementIdentityReport.contextReport.preparedBaseState);
    CHECK(displacementIdentityReport.contextReport.updatedDisplacement);

    ++dependencies.discretization.identity;
    ++dependencies.discretization.revision;

    const auto discretizationIdentityReport =
        preparedOperator.Prepare(make_state(density, enthalpy, displacement), dependencies);
    CHECK(discretizationIdentityReport.contextReport.preparedStaticDependencies);
    CHECK(discretizationIdentityReport.contextReport.preparedGeometryState);
    CHECK(discretizationIdentityReport.contextReport.preparedBaseState);
}