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

import mean_field;
import test_helpers;

namespace prepared_hydrostatic_jacobian_test_utils {
    mean_field::operators::context::hydrostatic::HydrostaticEquilibriumDependencies make_dependencies() {
        return {
            .discretization    = {.identity = 307, .revision = 2},
            .enthalpy          = {.identity = 311, .revision = 3},
            .gravityPotential  = {.identity = 313, .revision = 5},
            .displacement      = {.identity = 317, .revision = 7},
            .rotation          = {.identity = 331, .revision = 11},
            .bernoulliConstant = {.identity = 337, .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.23
    ) {
        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.41
    ) {
        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.13 * scale;
        angularVelocity(1) = -0.19 * scale;
        angularVelocity(2) = 0.47 * scale;

        mfem::Vector center(3);

        center(0) = 0.031;
        center(1) = -0.023;
        center(2) = 0.017;

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

    void centered_residual_difference(
        const mean_field::fem::FEM &f,
        const mean_field::physics::RigidRotation &rotation,
        const mfem::Vector &enthalpyPlus,
        const mfem::Vector &enthalpyMinus,
        const mfem::Vector &gravityPotentialPlus,
        const mfem::Vector &gravityPotentialMinus,
        const mfem::Vector &displacement,
        const double bernoulliConstantPlus,
        const double bernoulliConstantMinus,
        mfem::Vector &difference
    ) {
        mfem::Vector residualPlus;
        mfem::Vector residualMinus;

        field_dof_test_utils::apply_hydrostatic_reference(
            f, rotation, enthalpyPlus, gravityPotentialPlus, displacement, bernoulliConstantPlus, residualPlus
        );

        field_dof_test_utils::apply_hydrostatic_reference(
            f, rotation, enthalpyMinus, gravityPotentialMinus, displacement, bernoulliConstantMinus, residualMinus
        );

        // The two states are separated by one complete variation:
        // x_+ = x + 0.5 dx and x_- = x - 0.5 dx.
        difference = residualPlus;
        difference -= residualMinus;
    }

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

TEST_CASE(
    "Prepared Hydrostatic Algebraic Jacobian Matches Centered Residual "
    "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_jacobian_test_utils::make_enthalpy(f);

    const mfem::Vector gravityPotential = prepared_hydrostatic_jacobian_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_jacobian_test_utils::make_rotation();

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

    const mfem::Vector enthalpyVariation = prepared_hydrostatic_jacobian_test_utils::make_enthalpy(f, 0.71);

    const mfem::Vector gravityPotentialVariation =
        prepared_hydrostatic_jacobian_test_utils::make_gravity_potential(f, 0.83);

    constexpr double bernoulliConstantVariation = -0.31;

    mfem::Vector enthalpyAction;
    mfem::Vector gravityPotentialAction;
    mfem::Vector bernoulliConstantAction;
    mfem::Vector combinedAction;

    preparedOperator.ApplyEnthalpyJacobianAction(enthalpyVariation, enthalpyAction);

    preparedOperator.ApplyGravityPotentialJacobianAction(gravityPotentialVariation, gravityPotentialAction);

    preparedOperator.ApplyBernoulliConstantJacobianAction(bernoulliConstantVariation, bernoulliConstantAction);

    preparedOperator.ApplyAlgebraicJacobianAction(
        enthalpyVariation, gravityPotentialVariation, bernoulliConstantVariation, combinedAction
    );

    mfem::Vector enthalpyPlus(enthalpy);
    mfem::Vector enthalpyMinus(enthalpy);
    mfem::Vector gravityPotentialPlus(gravityPotential);
    mfem::Vector gravityPotentialMinus(gravityPotential);

    enthalpyPlus.Add(0.5, enthalpyVariation);
    enthalpyMinus.Add(-0.5, enthalpyVariation);

    mfem::Vector enthalpyReference;

    prepared_hydrostatic_jacobian_test_utils::centered_residual_difference(
        f, rotation, enthalpyPlus, enthalpyMinus, gravityPotential, gravityPotential, displacement, bernoulliConstant,
        bernoulliConstant, enthalpyReference
    );

    enthalpyPlus  = enthalpy;
    enthalpyMinus = enthalpy;

    gravityPotentialPlus.Add(0.5, gravityPotentialVariation);

    gravityPotentialMinus.Add(-0.5, gravityPotentialVariation);

    mfem::Vector gravityPotentialReference;

    prepared_hydrostatic_jacobian_test_utils::centered_residual_difference(
        f, rotation, enthalpy, enthalpy, gravityPotentialPlus, gravityPotentialMinus, displacement, bernoulliConstant,
        bernoulliConstant, gravityPotentialReference
    );

    gravityPotentialPlus  = gravityPotential;
    gravityPotentialMinus = gravityPotential;

    mfem::Vector bernoulliConstantReference;

    prepared_hydrostatic_jacobian_test_utils::centered_residual_difference(
        f, rotation, enthalpy, enthalpy, gravityPotential, gravityPotential, displacement,
        bernoulliConstant + 0.5 * bernoulliConstantVariation, bernoulliConstant - 0.5 * bernoulliConstantVariation,
        bernoulliConstantReference
    );

    enthalpyPlus.Add(0.5, enthalpyVariation);
    enthalpyMinus.Add(-0.5, enthalpyVariation);

    gravityPotentialPlus.Add(0.5, gravityPotentialVariation);

    gravityPotentialMinus.Add(-0.5, gravityPotentialVariation);

    mfem::Vector combinedReference;

    prepared_hydrostatic_jacobian_test_utils::centered_residual_difference(
        f, rotation, enthalpyPlus, enthalpyMinus, gravityPotentialPlus, gravityPotentialMinus, displacement,
        bernoulliConstant + 0.5 * bernoulliConstantVariation, bernoulliConstant - 0.5 * bernoulliConstantVariation,
        combinedReference
    );

    mfem::Vector sumOfBlocks(enthalpyAction);
    sumOfBlocks += gravityPotentialAction;
    sumOfBlocks += bernoulliConstantAction;

    const double enthalpyError =
        prepared_hydrostatic_jacobian_test_utils::relative_error(enthalpyAction, enthalpyReference, f.mesh->GetComm());

    const double gravityPotentialError = prepared_hydrostatic_jacobian_test_utils::relative_error(
        gravityPotentialAction, gravityPotentialReference, f.mesh->GetComm()
    );

    const double bernoulliConstantError = prepared_hydrostatic_jacobian_test_utils::relative_error(
        bernoulliConstantAction, bernoulliConstantReference, f.mesh->GetComm()
    );

    const double combinedError =
        prepared_hydrostatic_jacobian_test_utils::relative_error(combinedAction, combinedReference, f.mesh->GetComm());

    const double blockSumError =
        prepared_hydrostatic_jacobian_test_utils::relative_error(combinedAction, sumOfBlocks, f.mesh->GetComm());

    INFO("Enthalpy block error = " << enthalpyError);
    INFO("Gravity-potential block error = " << gravityPotentialError);
    INFO("Bernoulli-constant block error = " << bernoulliConstantError);
    INFO("Combined algebraic action error = " << combinedError);
    INFO("Combined-versus-summed-block error = " << blockSumError);

    const auto &statistics = preparedOperator.GetAlgebraicJacobianStatistics();

    CHECK(report.preparedAlgebraicJacobianBlocks);
    CHECK(statistics.preparations == 1);
    CHECK(statistics.enthalpyApplications == 1);
    CHECK(statistics.gravityPotentialApplications == 1);
    CHECK(statistics.bernoulliConstantApplications == 1);
    CHECK(statistics.combinedApplications == 1);

    CHECK(enthalpyError < 5.0e-12);
    CHECK(gravityPotentialError < 5.0e-12);
    CHECK(bernoulliConstantError < 5.0e-12);
    CHECK(combinedError < 5.0e-12);
    CHECK(blockSumError < 5.0e-13);
}

TEST_CASE(
    "Prepared Hydrostatic Algebraic Jacobian Reuses And Rebuilds Only With "
    "Geometry",
    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);

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

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

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

    double bernoulliConstant                    = 0.37;

    mean_field::physics::RigidRotation rotation = prepared_hydrostatic_jacobian_test_utils::make_rotation(0.8);

    auto dependencies                           = prepared_hydrostatic_jacobian_test_utils::make_dependencies();

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

    const mfem::Vector enthalpyVariation = prepared_hydrostatic_jacobian_test_utils::make_enthalpy(f, 0.67);

    const mfem::Vector gravityPotentialVariation =
        prepared_hydrostatic_jacobian_test_utils::make_gravity_potential(f, 0.79);

    constexpr double bernoulliConstantVariation = 0.28;

    mfem::Vector initialAction;

    preparedOperator.ApplyAlgebraicJacobianAction(
        enthalpyVariation, gravityPotentialVariation, bernoulliConstantVariation, initialAction
    );

    enthalpy          = prepared_hydrostatic_jacobian_test_utils::make_enthalpy(f, 1.13);

    gravityPotential  = prepared_hydrostatic_jacobian_test_utils::make_gravity_potential(f, 1.31);

    bernoulliConstant = 0.62;
    rotation          = prepared_hydrostatic_jacobian_test_utils::make_rotation(1.4);

    ++dependencies.enthalpy.revision;
    ++dependencies.gravityPotential.revision;
    ++dependencies.bernoulliConstant.revision;
    ++dependencies.rotation.revision;

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

    mfem::Vector baseStateChangedAction;

    preparedOperator.ApplyAlgebraicJacobianAction(
        enthalpyVariation, gravityPotentialVariation, bernoulliConstantVariation, baseStateChangedAction
    );

    CHECK_FALSE(baseStateReport.contextReport.preparedGeometryState);
    CHECK(baseStateReport.contextReport.preparedRotationDependencies);
    CHECK(baseStateReport.contextReport.preparedBaseState);
    CHECK_FALSE(baseStateReport.preparedAlgebraicJacobianBlocks);
    CHECK(preparedOperator.GetAlgebraicJacobianStatistics().preparations == 1);
    CHECK(
        prepared_hydrostatic_jacobian_test_utils::relative_error(
            baseStateChangedAction, initialAction, f.mesh->GetComm()
        ) == 0.0
    );

    const mfem::Vector secondEnthalpyVariation = prepared_hydrostatic_jacobian_test_utils::make_enthalpy(f, 1.57);

    const mfem::Vector secondGravityPotentialVariation =
        prepared_hydrostatic_jacobian_test_utils::make_gravity_potential(f, 1.73);

    mfem::Vector secondDirectionAction;

    preparedOperator.ApplyAlgebraicJacobianAction(
        secondEnthalpyVariation, secondGravityPotentialVariation, -0.19, secondDirectionAction
    );

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

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

    ++dependencies.displacement.revision;

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

    mfem::Vector geometryChangedAction;

    preparedOperator.ApplyAlgebraicJacobianAction(
        enthalpyVariation, gravityPotentialVariation, bernoulliConstantVariation, geometryChangedAction
    );

    mfem::Vector enthalpyPlus(enthalpy);
    mfem::Vector enthalpyMinus(enthalpy);
    mfem::Vector gravityPotentialPlus(gravityPotential);
    mfem::Vector gravityPotentialMinus(gravityPotential);

    enthalpyPlus.Add(0.5, enthalpyVariation);
    enthalpyMinus.Add(-0.5, enthalpyVariation);

    gravityPotentialPlus.Add(0.5, gravityPotentialVariation);

    gravityPotentialMinus.Add(-0.5, gravityPotentialVariation);

    mfem::Vector geometryReference;

    prepared_hydrostatic_jacobian_test_utils::centered_residual_difference(
        f, rotation, enthalpyPlus, enthalpyMinus, gravityPotentialPlus, gravityPotentialMinus, displacement,
        bernoulliConstant + 0.5 * bernoulliConstantVariation, bernoulliConstant - 0.5 * bernoulliConstantVariation,
        geometryReference
    );

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

    const double geometryEffect = prepared_hydrostatic_jacobian_test_utils::relative_error(
        geometryChangedAction, initialAction, f.mesh->GetComm()
    );

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

    INFO("Algebraic Jacobian change after deformation update = " << geometryEffect);

    const auto &statistics = preparedOperator.GetAlgebraicJacobianStatistics();

    CHECK(geometryReport.contextReport.preparedGeometryState);
    CHECK(geometryReport.preparedAlgebraicJacobianBlocks);
    CHECK(statistics.preparations == 2);
    CHECK(statistics.enthalpyApplications == 0);
    CHECK(statistics.gravityPotentialApplications == 0);
    CHECK(statistics.bernoulliConstantApplications == 0);
    CHECK(statistics.combinedApplications == 4);

    CHECK(geometryReferenceError < 5.0e-12);
    CHECK(geometryEffect > 1.0e-8);
}