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

import mean_field;
import test_helpers;

namespace hydrostatic_context_test_utils {
    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies make_dependencies() {
        return {
            .discretization    = {.identity = 101, .revision = 2},
            .enthalpy          = {.identity = 103, .revision = 3},
            .gravityPotential  = {.identity = 107, .revision = 5},
            .displacement      = {.identity = 109, .revision = 7},
            .rotation          = {.identity = 113, .revision = 11},
            .bernoulliConstant = {.identity = 127, .revision = 13}
        };
    }

    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumStateView make_state(
        const mfem::Vector &enthalpy,
        const mfem::Vector &gravityPotential,
        const mfem::Vector &displacement,
        const double bernoulliConstant
    ) {
        return {
            .enthalpy          = enthalpy,
            .gravityPotential  = gravityPotential,
            .displacement      = displacement,
            .bernoulliConstant = bernoulliConstant
        };
    }

    void check_base_only(const mean_field::operators::context::hydrostatic::HydrostaticPreparationReport &report) {
        CHECK_FALSE(report.preparedStaticDependencies);
        CHECK_FALSE(report.preparedGeometryState);
        CHECK_FALSE(report.preparedRotationDependencies);
        CHECK(report.preparedBaseState);
    }
} // namespace hydrostatic_context_test_utils

TEST_CASE(
    "Hydrostatic Context Applies Selective Invalidation",
    tags::barotrope_hydrostatic_context &tags::unit
) {
    auto args              = test_utils::setup_args();

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

    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumContext context(f, *f.domainMapperStateless);

    mfem::Vector enthalpy =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Enthalpy>(*f.enthalpyFes, 0.17);

    mfem::Vector gravityPotential =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Gravity>(
            *f.gravityPotentialFes, 0.31
        );

    mfem::Vector displacement = field_dof_test_utils::make_supported_displacement(f, 0.35);

    double bernoulliConstant  = 0.73;

    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies dependencies =
        hydrostatic_context_test_utils::make_dependencies();

    CHECK_FALSE(context.IsPrepared());
    CHECK_FALSE(context.MatchesDependencies(dependencies));

    const auto initialStatistics = context.GetPreparationStatistics();

    CHECK(initialStatistics.staticPreparations == 0);
    CHECK(initialStatistics.geometryPreparations == 0);
    CHECK(initialStatistics.rotationPreparations == 0);
    CHECK(initialStatistics.baseStatePreparations == 0);

    const auto initialReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    REQUIRE(context.IsPrepared());
    CHECK(context.MatchesDependencies(dependencies));
    CHECK(context.GetDependencies() == dependencies);

    CHECK(initialReport.preparedStaticDependencies);
    CHECK(initialReport.preparedGeometryState);
    CHECK(initialReport.preparedRotationDependencies);
    CHECK(initialReport.preparedBaseState);
    CHECK(initialReport.updatedEnthalpy);
    CHECK(initialReport.updatedGravityPotential);
    CHECK(initialReport.updatedDisplacement);
    CHECK(initialReport.updatedBernoulliConstant);
    CHECK(initialReport.DidAnyWork());

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

    const mfem::Vector frozenGravityPotential = context.GetBaseGravityPotentialTrue();

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

    const double frozenBernoulliConstant      = context.GetBernoulliConstant();

    enthalpy(0) += 0.125;
    gravityPotential(0) -= 0.075;
    displacement(0) += 0.050;
    bernoulliConstant += 0.20;

    const auto repeatedReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    CHECK_FALSE(repeatedReport.DidAnyWork());
    CHECK_FALSE(repeatedReport.updatedEnthalpy);
    CHECK_FALSE(repeatedReport.updatedGravityPotential);
    CHECK_FALSE(repeatedReport.updatedDisplacement);
    CHECK_FALSE(repeatedReport.updatedBernoulliConstant);

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

    CHECK(
        gravity_prepared_test_utils::relative_error(context.GetBaseEnthalpyTrue(), frozenEnthalpy, communicator) == 0.0
    );

    CHECK(
        gravity_prepared_test_utils::relative_error(
            context.GetBaseGravityPotentialTrue(), frozenGravityPotential, communicator
        ) == 0.0
    );

    CHECK(
        gravity_prepared_test_utils::relative_error(context.GetDisplacementTrue(), frozenDisplacement, communicator) ==
        0.0
    );

    CHECK(context.GetBernoulliConstant() == frozenBernoulliConstant);

    ++dependencies.enthalpy.revision;

    const auto enthalpyReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    hydrostatic_context_test_utils::check_base_only(enthalpyReport);

    CHECK(enthalpyReport.updatedEnthalpy);
    CHECK_FALSE(enthalpyReport.updatedGravityPotential);
    CHECK_FALSE(enthalpyReport.updatedDisplacement);
    CHECK_FALSE(enthalpyReport.updatedBernoulliConstant);
    CHECK(context.GetBaseEnthalpyTrue()(context.GetEnthalpyMap().true_dof(0)) == enthalpy(0));

    ++dependencies.gravityPotential.revision;

    const auto gravityPotentialReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    hydrostatic_context_test_utils::check_base_only(gravityPotentialReport);

    CHECK_FALSE(gravityPotentialReport.updatedEnthalpy);
    CHECK(gravityPotentialReport.updatedGravityPotential);
    CHECK_FALSE(gravityPotentialReport.updatedDisplacement);
    CHECK_FALSE(gravityPotentialReport.updatedBernoulliConstant);

    CHECK(context.GetBaseGravityPotentialTrue()(context.GetGravityPotentialMap().true_dof(0)) == gravityPotential(0));

    ++dependencies.bernoulliConstant.revision;

    const auto bernoulliReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    hydrostatic_context_test_utils::check_base_only(bernoulliReport);

    CHECK_FALSE(bernoulliReport.updatedEnthalpy);
    CHECK_FALSE(bernoulliReport.updatedGravityPotential);
    CHECK_FALSE(bernoulliReport.updatedDisplacement);
    CHECK(bernoulliReport.updatedBernoulliConstant);

    CHECK(context.GetBernoulliConstant() == bernoulliConstant);

    ++dependencies.rotation.revision;

    const auto rotationReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    CHECK_FALSE(rotationReport.preparedStaticDependencies);
    CHECK_FALSE(rotationReport.preparedGeometryState);
    CHECK(rotationReport.preparedRotationDependencies);
    CHECK(rotationReport.preparedBaseState);
    CHECK_FALSE(rotationReport.updatedEnthalpy);
    CHECK_FALSE(rotationReport.updatedGravityPotential);
    CHECK_FALSE(rotationReport.updatedDisplacement);
    CHECK_FALSE(rotationReport.updatedBernoulliConstant);

    ++dependencies.displacement.revision;

    const auto displacementReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    CHECK_FALSE(displacementReport.preparedStaticDependencies);

    CHECK(displacementReport.preparedGeometryState);
    CHECK(displacementReport.preparedRotationDependencies);
    CHECK(displacementReport.preparedBaseState);
    CHECK_FALSE(displacementReport.updatedEnthalpy);
    CHECK_FALSE(displacementReport.updatedGravityPotential);
    CHECK(displacementReport.updatedDisplacement);
    CHECK_FALSE(displacementReport.updatedBernoulliConstant);

    CHECK(context.GetDisplacementTrue()(context.GetDisplacementMap().true_dof(0)) == displacement(0));

    ++dependencies.discretization.revision;

    const auto discretizationReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    CHECK(discretizationReport.preparedStaticDependencies);
    CHECK(discretizationReport.preparedGeometryState);
    CHECK(discretizationReport.preparedRotationDependencies);
    CHECK(discretizationReport.preparedBaseState);
    CHECK(discretizationReport.updatedEnthalpy);
    CHECK(discretizationReport.updatedGravityPotential);
    CHECK(discretizationReport.updatedDisplacement);
    CHECK(discretizationReport.updatedBernoulliConstant);

    const auto finalStatistics = context.GetPreparationStatistics();

    CHECK(finalStatistics.staticPreparations == 2);
    CHECK(finalStatistics.geometryPreparations == 3);
    CHECK(finalStatistics.rotationPreparations == 4);
    CHECK(finalStatistics.baseStatePreparations == 7);
}

TEST_CASE(
    "Hydrostatic Context Uses Identity In Every Dependency",
    tags::barotrope_hydrostatic_context &tags::unit
) {
    auto args              = test_utils::setup_args();

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

    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumContext context(f, *f.domainMapperStateless);

    mfem::Vector enthalpy =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Enthalpy>(*f.enthalpyFes, 0.23);

    const mfem::Vector gravityPotential =
        field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Gravity>(
            *f.gravityPotentialFes, 0.41
        );

    const mfem::Vector displacement    = field_dof_test_utils::make_supported_displacement(f, 0.60);

    constexpr double bernoulliConstant = 0.81;

    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies dependencies =
        hydrostatic_context_test_utils::make_dependencies();

    const auto preparedEnthalpyDependency = dependencies.enthalpy;

    auto olderSameIdentity                = preparedEnthalpyDependency;
    --olderSameIdentity.revision;

    auto resetNewIdentity = preparedEnthalpyDependency;
    ++resetNewIdentity.identity;
    resetNewIdentity.revision = 0;

    CHECK_FALSE(olderSameIdentity.CanFollow(preparedEnthalpyDependency));

    CHECK(resetNewIdentity.CanFollow(preparedEnthalpyDependency));

    context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    const mfem::Vector firstFrozenEnthalpy = context.GetBaseEnthalpyTrue();

    enthalpy(0) += 0.33;

    const auto sameStampReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    CHECK_FALSE(sameStampReport.DidAnyWork());
    CHECK(
        context.GetBaseEnthalpyTrue()(context.GetEnthalpyMap().true_dof(0)) ==
        firstFrozenEnthalpy(context.GetEnthalpyMap().true_dof(0))
    );

    ++dependencies.enthalpy.identity;
    dependencies.enthalpy.revision = 0;

    const auto newIdentityReport   = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    hydrostatic_context_test_utils::check_base_only(newIdentityReport);

    CHECK(newIdentityReport.updatedEnthalpy);
    CHECK(context.GetBaseEnthalpyTrue()(context.GetEnthalpyMap().true_dof(0)) == enthalpy(0));
    CHECK(context.MatchesDependencies(dependencies));

    ++dependencies.rotation.identity;
    dependencies.rotation.revision       = 0;

    const auto newRotationIdentityReport = context.Prepare(
        hydrostatic_context_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies
    );

    CHECK_FALSE(newRotationIdentityReport.preparedStaticDependencies);

    CHECK_FALSE(newRotationIdentityReport.preparedGeometryState);

    CHECK(newRotationIdentityReport.preparedRotationDependencies);

    CHECK(newRotationIdentityReport.preparedBaseState);

    const auto statistics = context.GetPreparationStatistics();

    CHECK(statistics.staticPreparations == 1);
    CHECK(statistics.geometryPreparations == 1);
    CHECK(statistics.rotationPreparations == 2);
    CHECK(statistics.baseStatePreparations == 3);
}