#include <algorithm>
#include <cmath>
#include <limits>

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

import mean_field;
import test_helpers;

namespace pressure_force_context_test_utils {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

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

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

    [[nodiscard]]
    mfem::Vector make_enthalpy_true(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::Vector enthalpy(f.enthalpyFes->GetTrueVSize());

        for (int index = 0; index < enthalpy.Size(); ++index) {
            const double position = static_cast<double>(index + 1);

            enthalpy(index) =
                0.95 + 0.08 * std::sin(0.17 * position + phase) + 0.03 * std::cos(0.11 * position - 0.5 * phase);
        }

        return enthalpy;
    }

    [[nodiscard]]
    double relative_difference(
        const mfem::Vector &left,
        const mfem::Vector &right,
        const MPI_Comm communicator
    ) {
        MFEM_VERIFY(
            left.Size() == right.Size(), "Cannot compare pressure-force context vectors with "
                                         "different sizes."
        );

        mfem::Vector difference(left);

        difference -= right;

        const double scale = std::max(
            {gravity_prepared_test_utils::global_norm(left, communicator),
             gravity_prepared_test_utils::global_norm(right, communicator),
             100.0 * std::numeric_limits<double>::epsilon()}
        );

        return gravity_prepared_test_utils::global_norm(difference, communicator) / scale;
    }

    [[nodiscard]]
    mean_field::operators::context::pressure_force::PressureForceDependencies make_dependencies() {
        return {
            .discretization = {.identity = 101, .revision = 7},
            .enthalpy       = {.identity = 103, .revision = 11},
            .displacement   = {.identity = 107, .revision = 13}
        };
    }
} // namespace pressure_force_context_test_utils

TEST_CASE(
    "Pressure Force Context Applies Selective Invalidation In FieldDof Coordinates",
    tags::barotrope &tags::pressure &tags::prepared &tags::field &tags::unit
) {
    mean_field::utils::Args args = test_utils::setup_args();

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

    REQUIRE(f.okay());

    const pressure_force_context_test_utils::Maps maps(f);

    mean_field::operators::context::pressure_force::PressureForceLinearizationContext context(
        f, *f.domainMapperStateless, maps.enthalpy, maps.displacement
    );

    mfem::Vector enthalpy = maps.enthalpy.gather(pressure_force_context_test_utils::make_enthalpy_true(f, 0.23));

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

    const mfem::Vector changedDisplacement =
        maps.displacement.gather(gravity_prepared_test_utils::make_displacement(f, 0.79));

    mfem::Vector displacement(initialDisplacement);

    auto dependencies = pressure_force_context_test_utils::make_dependencies();

    const mean_field::operators::context::pressure_force::PressureForceStateView state{
        .enthalpy = enthalpy, .displacement = displacement
    };

    CHECK_FALSE(context.IsPrepared());

    const auto initialReport = context.Prepare(state, dependencies);

    REQUIRE(context.IsPrepared());

    CHECK(context.MatchesDependencies(dependencies));

    CHECK(initialReport.DidAnyWork());

    CHECK(initialReport.preparedStaticDependencies);

    CHECK(initialReport.preparedGeometryState);

    CHECK(initialReport.preparedMaterialState);

    CHECK(initialReport.updatedEnthalpy);

    CHECK(initialReport.updatedDisplacement);

    REQUIRE(maps.enthalpy.reduced_size() < maps.enthalpy.full_size());

    CHECK(maps.displacement.is_identity());

    CHECK(context.GetBaseEnthalpy().Size() == maps.enthalpy.reduced_size());

    CHECK(context.GetDisplacement().Size() == maps.displacement.reduced_size());

    const auto unchangedReport = context.Prepare(state, dependencies);

    CHECK_FALSE(unchangedReport.DidAnyWork());

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

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

    enthalpy(0) += 0.125;

    displacement               = changedDisplacement;

    const auto unstampedReport = context.Prepare(state, dependencies);

    CHECK_FALSE(unstampedReport.DidAnyWork());

    CHECK(
        pressure_force_context_test_utils::relative_difference(
            context.GetBaseEnthalpy(), frozenEnthalpy, f.mesh->GetComm()
        ) == 0.0
    );

    CHECK(
        pressure_force_context_test_utils::relative_difference(
            context.GetDisplacement(), frozenDisplacement, f.mesh->GetComm()
        ) == 0.0
    );

    ++dependencies.enthalpy.revision;

    const auto enthalpyReport = context.Prepare(state, dependencies);

    CHECK(enthalpyReport.DidAnyWork());

    CHECK_FALSE(enthalpyReport.preparedStaticDependencies);

    CHECK_FALSE(enthalpyReport.preparedGeometryState);

    CHECK(enthalpyReport.preparedMaterialState);

    CHECK(enthalpyReport.updatedEnthalpy);

    CHECK_FALSE(enthalpyReport.updatedDisplacement);

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

    CHECK(
        pressure_force_context_test_utils::relative_difference(
            context.GetDisplacement(), frozenDisplacement, f.mesh->GetComm()
        ) == 0.0
    );

    ++dependencies.displacement.revision;

    const auto displacementReport = context.Prepare(state, dependencies);

    CHECK(displacementReport.DidAnyWork());

    CHECK_FALSE(displacementReport.preparedStaticDependencies);

    CHECK(displacementReport.preparedGeometryState);

    CHECK(displacementReport.preparedMaterialState);

    CHECK_FALSE(displacementReport.updatedEnthalpy);

    CHECK(displacementReport.updatedDisplacement);

    CHECK(
        pressure_force_context_test_utils::relative_difference(
            context.GetDisplacement(), changedDisplacement, f.mesh->GetComm()
        ) == 0.0
    );

    ++dependencies.discretization.revision;

    const auto discretizationReport = context.Prepare(state, dependencies);

    CHECK(discretizationReport.preparedStaticDependencies);

    CHECK(discretizationReport.preparedGeometryState);

    CHECK(discretizationReport.preparedMaterialState);

    CHECK(discretizationReport.updatedEnthalpy);

    CHECK(discretizationReport.updatedDisplacement);

    const auto &statistics = context.GetPreparationStatistics();

    CHECK(statistics.staticPreparations == 2);

    CHECK(statistics.geometryPreparations == 3);

    CHECK(statistics.materialPreparations == 4);
}

TEST_CASE(
    "Pressure Force Context Uses Identity And Revision In Every Dependency",
    tags::barotrope &tags::pressure &tags::prepared &tags::field &tags::unit
) {
    mean_field::utils::Args args = test_utils::setup_args();

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

    REQUIRE(f.okay());

    const pressure_force_context_test_utils::Maps maps(f);

    mean_field::operators::context::pressure_force::PressureForceLinearizationContext context(
        f, *f.domainMapperStateless, maps.enthalpy, maps.displacement
    );

    const mfem::Vector enthalpy = maps.enthalpy.gather(pressure_force_context_test_utils::make_enthalpy_true(f, 0.61));

    const mfem::Vector displacement = maps.displacement.gather(gravity_prepared_test_utils::make_displacement(f, 0.73));

    const mean_field::operators::context::pressure_force::PressureForceStateView state{
        .enthalpy = enthalpy, .displacement = displacement
    };

    auto dependencies = pressure_force_context_test_utils::make_dependencies();

    context.Prepare(state, dependencies);

    ++dependencies.enthalpy.identity;

    dependencies.enthalpy.revision    = 0;

    const auto enthalpyIdentityReport = context.Prepare(state, dependencies);

    CHECK_FALSE(enthalpyIdentityReport.preparedStaticDependencies);

    CHECK_FALSE(enthalpyIdentityReport.preparedGeometryState);

    CHECK(enthalpyIdentityReport.preparedMaterialState);

    CHECK(enthalpyIdentityReport.updatedEnthalpy);

    CHECK_FALSE(enthalpyIdentityReport.updatedDisplacement);

    ++dependencies.displacement.identity;

    dependencies.displacement.revision    = 0;

    const auto displacementIdentityReport = context.Prepare(state, dependencies);

    CHECK_FALSE(displacementIdentityReport.preparedStaticDependencies);

    CHECK(displacementIdentityReport.preparedGeometryState);

    CHECK(displacementIdentityReport.preparedMaterialState);

    CHECK_FALSE(displacementIdentityReport.updatedEnthalpy);

    CHECK(displacementIdentityReport.updatedDisplacement);

    ++dependencies.discretization.identity;

    dependencies.discretization.revision    = 0;

    const auto discretizationIdentityReport = context.Prepare(state, dependencies);

    CHECK(discretizationIdentityReport.preparedStaticDependencies);

    CHECK(discretizationIdentityReport.preparedGeometryState);

    CHECK(discretizationIdentityReport.preparedMaterialState);

    CHECK(discretizationIdentityReport.updatedEnthalpy);

    CHECK(discretizationIdentityReport.updatedDisplacement);

    CHECK(context.MatchesDependencies(dependencies));

    const auto &statistics = context.GetPreparationStatistics();

    CHECK(statistics.staticPreparations == 2);

    CHECK(statistics.geometryPreparations == 3);

    CHECK(statistics.materialPreparations == 4);
}