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

import mean_field;
import test_helpers;

namespace prepared_hydrostatic_displacement_test_utils {
    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies make_dependencies() {
        return {
            .discretization    = {.identity = 401, .revision = 2},
            .enthalpy          = {.identity = 409, .revision = 3},
            .gravityPotential  = {.identity = 419, .revision = 5},
            .displacement      = {.identity = 421, .revision = 7},
            .rotation          = {.identity = 431, .revision = 11},
            .bernoulliConstant = {.identity = 433, .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.29
    ) {
        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.47
    ) {
        return field_dof_test_utils::make_deterministic_supported_vector<mean_field::field::Gravity>(
            *f.gravityPotentialFes, phase
        );
    }

    mfem::Vector make_displacement_direction(
        const mean_field::fem::FEM &f,
        const double firstPhase,
        const double secondPhase
    ) {
        mfem::Vector direction         = field_dof_test_utils::make_supported_displacement(f, firstPhase);

        const mfem::Vector secondField = field_dof_test_utils::make_supported_displacement(f, secondPhase);

        direction -= secondField;
        return direction;
    }

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

        angularVelocity(0) = 0.16 * scale;
        angularVelocity(1) = -0.14 * scale;
        angularVelocity(2) = 0.46 * scale;

        mfem::Vector center(3);

        center(0) = 0.034;
        center(1) = -0.026;
        center(2) = 0.019;

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

    void centered_displacement_difference(
        const mean_field::fem::FEM &f,
        const mean_field::physics::RigidRotation &rotation,
        const mfem::Vector &enthalpy,
        const mfem::Vector &gravityPotential,
        const mfem::Vector &baseDisplacement,
        const mfem::Vector &displacementVariation,
        const double bernoulliConstant,
        const double step,
        mfem::Vector &difference
    ) {
        mfem::Vector displacementPlus(baseDisplacement);
        mfem::Vector displacementMinus(baseDisplacement);

        displacementPlus.Add(step, displacementVariation);
        displacementMinus.Add(-step, displacementVariation);

        mfem::Vector residualPlus;
        mfem::Vector residualMinus;

        field_dof_test_utils::apply_hydrostatic_reference(
            f, rotation, enthalpy, gravityPotential, displacementPlus, bernoulliConstant, residualPlus
        );

        field_dof_test_utils::apply_hydrostatic_reference(
            f, rotation, enthalpy, gravityPotential, displacementMinus, bernoulliConstant, residualMinus
        );

        difference = residualPlus;
        difference -= residualMinus;
        difference /= 2.0 * step;
    }

    double relative_error(
        const mfem::Vector &actual,
        const mfem::Vector &expected,
        const MPI_Comm communicator
    ) {
        return gravity_prepared_test_utils::relative_error(actual, expected, communicator);
    }
} // namespace prepared_hydrostatic_displacement_test_utils

TEST_CASE(
    "Prepared Hydrostatic Displacement Jacobian Matches Centered Differences",
    tags::barotrope_hydrostatic_prepared_jacobian &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_displacement_test_utils::make_enthalpy(f);

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

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

    constexpr double bernoulliConstant  = 0.39;

    const mean_field::physics::RigidRotation rotation =
        prepared_hydrostatic_displacement_test_utils::make_rotation(0.9);

    const auto dependencies = prepared_hydrostatic_displacement_test_utils::make_dependencies();

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

    const mfem::Vector firstVariation =
        prepared_hydrostatic_displacement_test_utils::make_displacement_direction(f, 1.17, 0.31);

    const mfem::Vector secondVariation =
        prepared_hydrostatic_displacement_test_utils::make_displacement_direction(f, 1.43, 0.58);

    mfem::Vector combinedVariation(firstVariation);
    combinedVariation += secondVariation;

    mfem::Vector firstAction;
    mfem::Vector secondAction;
    mfem::Vector combinedAction;

    preparedOperator.ApplyDisplacementJacobianAction(firstVariation, firstAction);

    preparedOperator.ApplyDisplacementJacobianAction(secondVariation, secondAction);

    preparedOperator.ApplyDisplacementJacobianAction(combinedVariation, combinedAction);

    constexpr double finiteDifferenceStep = 1.0e-5;

    mfem::Vector centeredDifference;

    prepared_hydrostatic_displacement_test_utils::centered_displacement_difference(
        f, rotation, enthalpy, gravityPotential, displacement, firstVariation, bernoulliConstant, finiteDifferenceStep,
        centeredDifference
    );

    mfem::Vector sumOfActions(firstAction);
    sumOfActions += secondAction;

    const double centeredDifferenceError = prepared_hydrostatic_displacement_test_utils::relative_error(
        firstAction, centeredDifference, f.mesh->GetComm()
    );

    const double linearityError =
        prepared_hydrostatic_displacement_test_utils::relative_error(combinedAction, sumOfActions, f.mesh->GetComm());

    INFO("Prepared displacement centered-difference error = " << centeredDifferenceError);

    INFO("Prepared displacement linearity error = " << linearityError);

    const auto &statistics = preparedOperator.GetDisplacementJacobianStatistics();

    CHECK(report.preparedDisplacementJacobianData);
    CHECK(statistics.preparations == 1);
    CHECK(statistics.applications == 3);
    CHECK(centeredDifferenceError < 2.0e-7);
    CHECK(linearityError < 5.0e-12);
}

TEST_CASE(
    "Prepared Hydrostatic Displacement Jacobian Reuses And Refreshes Frozen "
    "Data",
    tags::barotrope_hydrostatic_prepared_jacobian &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_displacement_test_utils::make_enthalpy(f, 0.37);

    const mfem::Vector gravityPotential = prepared_hydrostatic_displacement_test_utils::make_gravity_potential(f, 0.53);

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

    constexpr double bernoulliConstant  = 0.36;

    mean_field::physics::RigidRotation rotation = prepared_hydrostatic_displacement_test_utils::make_rotation(0.75);

    auto dependencies                           = prepared_hydrostatic_displacement_test_utils::make_dependencies();

    preparedOperator.Prepare(
        prepared_hydrostatic_displacement_test_utils::make_state(
            enthalpy, gravityPotential, displacement, bernoulliConstant
        ),
        dependencies, rotation
    );

    const mfem::Vector displacementVariation =
        prepared_hydrostatic_displacement_test_utils::make_displacement_direction(f, 1.09, 0.27);

    const mfem::Vector secondVariation =
        prepared_hydrostatic_displacement_test_utils::make_displacement_direction(f, 1.36, 0.64);

    mfem::Vector initialAction;
    mfem::Vector secondDirectionAction;

    preparedOperator.ApplyDisplacementJacobianAction(displacementVariation, initialAction);

    preparedOperator.ApplyDisplacementJacobianAction(secondVariation, secondDirectionAction);

    CHECK(preparedOperator.GetDisplacementJacobianStatistics().preparations == 1);

    rotation = prepared_hydrostatic_displacement_test_utils::make_rotation(1.45);

    ++dependencies.rotation.revision;

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

    mfem::Vector rotationUpdatedAction;

    preparedOperator.ApplyDisplacementJacobianAction(displacementVariation, rotationUpdatedAction);

    constexpr double finiteDifferenceStep = 1.0e-5;

    mfem::Vector rotationReference;

    prepared_hydrostatic_displacement_test_utils::centered_displacement_difference(
        f, rotation, enthalpy, gravityPotential, displacement, displacementVariation, bernoulliConstant,
        finiteDifferenceStep, rotationReference
    );

    const double rotationReferenceError = prepared_hydrostatic_displacement_test_utils::relative_error(
        rotationUpdatedAction, rotationReference, f.mesh->GetComm()
    );

    const double rotationEffect = prepared_hydrostatic_displacement_test_utils::relative_error(
        rotationUpdatedAction, initialAction, f.mesh->GetComm()
    );

    CHECK_FALSE(rotationReport.contextReport.preparedGeometryState);
    CHECK(rotationReport.contextReport.preparedRotationDependencies);
    CHECK(rotationReport.contextReport.preparedBaseState);
    CHECK(rotationReport.updatedRotation);
    CHECK(rotationReport.preparedDisplacementJacobianData);
    CHECK_FALSE(rotationReport.preparedAlgebraicJacobianBlocks);
    CHECK(rotationReferenceError < 2.0e-7);
    CHECK(rotationEffect > 1.0e-8);

    displacement = field_dof_test_utils::make_supported_displacement(f, 0.91);

    ++dependencies.displacement.revision;

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

    mfem::Vector geometryUpdatedAction;

    preparedOperator.ApplyDisplacementJacobianAction(displacementVariation, geometryUpdatedAction);

    mfem::Vector geometryReference;

    prepared_hydrostatic_displacement_test_utils::centered_displacement_difference(
        f, rotation, enthalpy, gravityPotential, displacement, displacementVariation, bernoulliConstant,
        finiteDifferenceStep, geometryReference
    );

    const double geometryReferenceError = prepared_hydrostatic_displacement_test_utils::relative_error(
        geometryUpdatedAction, geometryReference, f.mesh->GetComm()
    );

    const double geometryEffect = prepared_hydrostatic_displacement_test_utils::relative_error(
        geometryUpdatedAction, rotationUpdatedAction, f.mesh->GetComm()
    );

    INFO("Rotation-updated displacement Jacobian error = " << rotationReferenceError);

    INFO("Displacement Jacobian change after rotation update = " << rotationEffect);

    INFO("Geometry-updated displacement Jacobian error = " << geometryReferenceError);

    INFO("Displacement Jacobian change after geometry update = " << geometryEffect);

    const auto &contextStatistics      = preparedOperator.GetContextPreparationStatistics();

    const auto &displacementStatistics = preparedOperator.GetDisplacementJacobianStatistics();

    CHECK(geometryReport.contextReport.preparedGeometryState);
    CHECK(geometryReport.contextReport.preparedRotationDependencies);
    CHECK(geometryReport.contextReport.preparedBaseState);
    CHECK_FALSE(geometryReport.updatedRotation);
    CHECK(geometryReport.preparedAlgebraicJacobianBlocks);
    CHECK(geometryReport.preparedDisplacementJacobianData);

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

    CHECK(displacementStatistics.preparations == 3);
    CHECK(displacementStatistics.applications == 4);
    CHECK(geometryReferenceError < 2.0e-7);
    CHECK(geometryEffect > 1.0e-8);
}