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

import mean_field;
import test_helpers;

namespace prepared_hydrostatic_test_utils {
    static mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies make_dependencies() {
        return {
            .discretization    = {.identity = 211, .revision = 2},
            .enthalpy          = {.identity = 223, .revision = 3},
            .gravityPotential  = {.identity = 227, .revision = 5},
            .displacement      = {.identity = 229, .revision = 7},
            .rotation          = {.identity = 233, .revision = 11},
            .bernoulliConstant = {.identity = 239, .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
        };
    }

    mfem::Vector make_enthalpy(
        const mean_field::fem::FEM &f,
        const double phase = 0.19
    ) {
        return field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Enthalpy>(
            *f.enthalpyFes, phase
        );
    }

    mfem::Vector make_gravity_potential(
        const mean_field::fem::FEM &f,
        const double phase = 0.37
    ) {
        return field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Gravity>(
            *f.gravityPotentialFes, phase
        );
    }

    mean_field::physics::RigidRotation make_rotation(const double scale = 1.0) {
        mfem::Vector angularVelocity(3);

        angularVelocity(0) = 0.17 * scale;
        angularVelocity(1) = -0.11 * scale;
        angularVelocity(2) = 0.43 * scale;

        mfem::Vector center(3);

        center(0) = 0.037;
        center(1) = -0.029;
        center(2) = 0.021;

        return mean_field::physics::RigidRotation(angularVelocity, center);
    }
} // namespace prepared_hydrostatic_test_utils

TEST_CASE(
    "Prepared Hydrostatic Residual Matches Stateless Kernel",
    tags::barotrope_hydrostatic_prepared_residual &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::PreparedHydrostaticEquilibriumOperator preparedOperator(f, *f.domainMapperStateless);

    const mfem::Vector enthalpy                       = prepared_hydrostatic_test_utils::make_enthalpy(f);

    const mfem::Vector gravityPotential               = prepared_hydrostatic_test_utils::make_gravity_potential(f);

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

    constexpr double bernoulliConstant                = 0.41;

    const mean_field::physics::RigidRotation rotation = prepared_hydrostatic_test_utils::make_rotation();

    const auto dependencies                           = prepared_hydrostatic_test_utils::make_dependencies();

    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK(preparedOperator.GetResidualPreparationCount() == 0);
    CHECK(preparedOperator.GetResidualApplicationCount() == 0);

    const auto report = preparedOperator.Prepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies, rotation
    );

    mfem::Vector preparedResidual;
    mfem::Vector referenceResidual;

    preparedOperator.BuildResidual(preparedResidual);

    field_dof_test_utils::apply_hydrostatic_reference(
        f, rotation, enthalpy, gravityPotential, displacement, bernoulliConstant, referenceResidual
    );

    const double relativeError =
        gravity_prepared_test_utils::relative_error(preparedResidual, referenceResidual, f.mesh->GetComm());

    INFO("Prepared hydrostatic residual relative error = " << relativeError);

    const auto &statistics = preparedOperator.GetContextPreparationStatistics();

    CHECK(preparedOperator.IsPrepared());
    CHECK(report.contextReport.preparedStaticDependencies);
    CHECK(report.contextReport.preparedGeometryState);
    CHECK(report.contextReport.preparedRotationDependencies);
    CHECK(report.contextReport.preparedBaseState);
    CHECK(report.updatedRotation);
    CHECK(report.preparedResidual);
    CHECK(report.DidAnyWork());

    CHECK(preparedOperator.GetStellarElementCount() > 0);
    CHECK(statistics.staticPreparations == 1);
    CHECK(statistics.geometryPreparations == 1);
    CHECK(statistics.rotationPreparations == 1);
    CHECK(statistics.baseStatePreparations == 1);
    CHECK(preparedOperator.GetResidualPreparationCount() == 1);
    CHECK(preparedOperator.GetResidualApplicationCount() == 1);

    CHECK(relativeError < 2.0e-12);
}

TEST_CASE(
    "Prepared Hydrostatic Residual Reuses And Selectively Rebuilds Data",
    tags::barotrope_hydrostatic_prepared_residual &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::PreparedHydrostaticEquilibriumOperator preparedOperator(f, *f.domainMapperStateless);

    mfem::Vector enthalpy         = prepared_hydrostatic_test_utils::make_enthalpy(f);

    mfem::Vector gravityPotential = prepared_hydrostatic_test_utils::make_gravity_potential(f);

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

    double bernoulliConstant      = 0.36;

    const mfem::Vector initialEnthalpy(enthalpy);
    const mfem::Vector initialGravityPotential(gravityPotential);
    const mfem::Vector initialDisplacement(displacement);
    const double initialBernoulliConstant                    = bernoulliConstant;

    const mean_field::physics::RigidRotation initialRotation = prepared_hydrostatic_test_utils::make_rotation(0.80);

    const mean_field::physics::RigidRotation changedRotation = prepared_hydrostatic_test_utils::make_rotation(1.25);

    auto dependencies                                        = prepared_hydrostatic_test_utils::make_dependencies();

    preparedOperator.Prepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies, initialRotation
    );

    mfem::Vector initialResidual;
    preparedOperator.BuildResidual(initialResidual);

    enthalpy(0) += 0.29;
    gravityPotential(0) -= 0.17;
    displacement = field_dof_test_utils::make_supported_displacement(f, 0.73);
    bernoulliConstant += 0.23;

    const auto unchangedReport = preparedOperator.Prepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies, changedRotation
    );

    mfem::Vector unchangedResidual;
    preparedOperator.BuildResidual(unchangedResidual);

    CHECK_FALSE(unchangedReport.DidAnyWork());
    CHECK_FALSE(unchangedReport.updatedRotation);
    CHECK_FALSE(unchangedReport.preparedResidual);
    CHECK(gravity_prepared_test_utils::relative_error(unchangedResidual, initialResidual, f.mesh->GetComm()) == 0.0);
    CHECK(preparedOperator.GetResidualPreparationCount() == 1);

    // Only the enthalpy stamp changes. The altered potential,
    // displacement, constant, and rotation remain intentionally frozen.
    ++dependencies.enthalpy.revision;

    const auto enthalpyReport = preparedOperator.Prepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies, changedRotation
    );

    mfem::Vector enthalpyResidual;
    mfem::Vector enthalpyReference;

    preparedOperator.BuildResidual(enthalpyResidual);

    field_dof_test_utils::apply_hydrostatic_reference(
        f, initialRotation, enthalpy, initialGravityPotential, initialDisplacement, initialBernoulliConstant,
        enthalpyReference
    );

    CHECK_FALSE(enthalpyReport.contextReport.preparedStaticDependencies);
    CHECK_FALSE(enthalpyReport.contextReport.preparedGeometryState);
    CHECK_FALSE(enthalpyReport.contextReport.preparedRotationDependencies);
    CHECK(enthalpyReport.contextReport.preparedBaseState);
    CHECK_FALSE(enthalpyReport.updatedRotation);
    CHECK(enthalpyReport.preparedResidual);
    CHECK(
        gravity_prepared_test_utils::relative_error(enthalpyResidual, enthalpyReference, f.mesh->GetComm()) < 2.0e-12
    );

    ++dependencies.rotation.revision;

    const auto rotationReport = preparedOperator.Prepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies, changedRotation
    );

    mfem::Vector rotationResidual;
    mfem::Vector rotationReference;

    preparedOperator.BuildResidual(rotationResidual);

    field_dof_test_utils::apply_hydrostatic_reference(
        f, changedRotation, enthalpy, initialGravityPotential, initialDisplacement, initialBernoulliConstant,
        rotationReference
    );

    CHECK_FALSE(rotationReport.contextReport.preparedStaticDependencies);
    CHECK_FALSE(rotationReport.contextReport.preparedGeometryState);
    CHECK(rotationReport.contextReport.preparedRotationDependencies);
    CHECK(rotationReport.contextReport.preparedBaseState);
    CHECK(rotationReport.updatedRotation);
    CHECK(rotationReport.preparedResidual);
    CHECK(
        gravity_prepared_test_utils::relative_error(rotationResidual, rotationReference, f.mesh->GetComm()) < 2.0e-12
    );

    ++dependencies.displacement.revision;

    const auto displacementReport = preparedOperator.Prepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, bernoulliConstant),
        dependencies, changedRotation
    );

    mfem::Vector displacementResidual;
    mfem::Vector displacementReference;

    preparedOperator.BuildResidual(displacementResidual);

    field_dof_test_utils::apply_hydrostatic_reference(
        f, changedRotation, enthalpy, initialGravityPotential, displacement, initialBernoulliConstant,
        displacementReference
    );

    CHECK_FALSE(displacementReport.contextReport.preparedStaticDependencies);
    CHECK(displacementReport.contextReport.preparedGeometryState);
    CHECK(displacementReport.contextReport.preparedRotationDependencies);
    CHECK(displacementReport.contextReport.preparedBaseState);
    CHECK_FALSE(displacementReport.updatedRotation);
    CHECK(displacementReport.preparedResidual);
    CHECK(
        gravity_prepared_test_utils::relative_error(displacementResidual, displacementReference, f.mesh->GetComm()) <
        2.0e-12
    );

    const auto &statistics = preparedOperator.GetContextPreparationStatistics();

    CHECK(statistics.staticPreparations == 1);
    CHECK(statistics.geometryPreparations == 2);
    CHECK(statistics.rotationPreparations == 3);
    CHECK(statistics.baseStatePreparations == 4);
    CHECK(preparedOperator.GetResidualPreparationCount() == 4);
    CHECK(preparedOperator.GetResidualApplicationCount() == 5);

    const double displacementEffect =
        gravity_prepared_test_utils::relative_error(displacementResidual, rotationResidual, f.mesh->GetComm());

    INFO("Residual change after displacement update = " << displacementEffect);

    CHECK(displacementEffect > 1.0e-8);
}