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

import mean_field;
import test_helpers;

namespace prepared_hydrostatic_complete_test_utils {
    mean_field::operators::context::hydrostatic::
        HydrostaticEquilibriumDependencies
        make_dependencies() {
        return {
            .discretization    = {.identity = 503, .revision = 2},
            .enthalpy          = {.identity = 509, .revision = 3},
            .gravityPotential  = {.identity = 521, .revision = 5},
            .displacement      = {.identity = 523, .revision = 7},
            .rotation          = {.identity = 541, .revision = 11},
            .bernoulliConstant = {.identity = 547, .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
        };
    }

    mean_field::physics::RigidRotation make_rotation() {
        mfem::Vector angularVelocity(3);

        angularVelocity(0) = 0.18;
        angularVelocity(1) = -0.13;
        angularVelocity(2) = 0.49;

        mfem::Vector center(3);

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

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

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

        const mfem::Vector secondField =
            gravity_prepared_test_utils::make_displacement(f, secondPhase);

        direction -= secondField;
        return direction;
    }

    void centered_complete_difference(
        const mean_field::fem::FEM &f,
        const mean_field::physics::RigidRotation &rotation,
        const mfem::Vector &baseEnthalpy,
        const mfem::Vector &baseGravityPotential,
        const mfem::Vector &baseDisplacement,
        const double baseBernoulliConstant,
        const mfem::Vector &enthalpyVariation,
        const mfem::Vector &gravityPotentialVariation,
        const mfem::Vector &displacementVariation,
        const double bernoulliConstantVariation,
        const double step,
        mfem::Vector &difference
    ) {
        mfem::Vector enthalpyPlus(baseEnthalpy);
        mfem::Vector enthalpyMinus(baseEnthalpy);
        mfem::Vector gravityPotentialPlus(baseGravityPotential);
        mfem::Vector gravityPotentialMinus(baseGravityPotential);
        mfem::Vector displacementPlus(baseDisplacement);
        mfem::Vector displacementMinus(baseDisplacement);

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

        gravityPotentialPlus.Add(step, gravityPotentialVariation);

        gravityPotentialMinus.Add(-step, gravityPotentialVariation);

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

        const double bernoulliConstantPlus =
            baseBernoulliConstant + step * bernoulliConstantVariation;

        const double bernoulliConstantMinus =
            baseBernoulliConstant - step * bernoulliConstantVariation;

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

        mean_field::operators::kernels::apply_hydrostatic_equilibrium(
            f, *f.domainMapperStateless, rotation, enthalpyPlus,
            gravityPotentialPlus, displacementPlus, bernoulliConstantPlus,
            residualPlus
        );

        mean_field::operators::kernels::apply_hydrostatic_equilibrium(
            f, *f.domainMapperStateless, rotation, enthalpyMinus,
            gravityPotentialMinus, displacementMinus, bernoulliConstantMinus,
            residualMinus
        );

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

    void set_block(
        mfem::Vector &packedDirection,
        const mean_field::operators::HydrostaticJacobianBlockLayout &layout,
        const mean_field::operators::HydrostaticJacobianInputBlock block,
        const mfem::Vector &values
    ) {
        REQUIRE(layout.Size(block) == values.Size());

        const int offset = layout.Offset(block);

        for (int entry = 0; entry < values.Size(); ++entry) {
            packedDirection(offset + entry) = values(entry);
        }
    }
} // namespace prepared_hydrostatic_complete_test_utils

TEST_CASE(
    "Prepared Hydrostatic Complete Jacobian Matches Sum And Centered "
    "Differences",
    tags::barotrope &tags::hydro &tags::integration &tags::jacobian
        &tags::prepared &tags::self_consistency
) {
    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 =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.enthalpyFes->GetTrueVSize(), 0.34
        );

    const mfem::Vector gravityPotential =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.gravityPotentialFes->GetTrueVSize(), 0.57
        );

    const mfem::Vector displacement =
        gravity_prepared_test_utils::make_displacement(f, 0.68);

    constexpr double bernoulliConstant = 0.43;

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

    preparedOperator.Prepare(
        prepared_hydrostatic_complete_test_utils::make_state(
            enthalpy, gravityPotential, displacement, bernoulliConstant
        ),
        prepared_hydrostatic_complete_test_utils::make_dependencies(), rotation
    );

    const mfem::Vector enthalpyVariation =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.enthalpyFes->GetTrueVSize(), 1.07
        );

    const mfem::Vector gravityPotentialVariation =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.gravityPotentialFes->GetTrueVSize(), 1.31
        );

    const mfem::Vector displacementVariation =
        prepared_hydrostatic_complete_test_utils::make_displacement_direction(
            f, 1.19, 0.38
        );

    constexpr double bernoulliConstantVariation = -0.37;

    mfem::Vector enthalpyAction;
    mfem::Vector gravityPotentialAction;
    mfem::Vector bernoulliConstantAction;
    mfem::Vector displacementAction;
    mfem::Vector completeAction;

    preparedOperator.ApplyEnthalpyJacobianAction(
        enthalpyVariation, enthalpyAction
    );

    preparedOperator.ApplyGravityPotentialJacobianAction(
        gravityPotentialVariation, gravityPotentialAction
    );

    preparedOperator.ApplyBernoulliConstantJacobianAction(
        bernoulliConstantVariation, bernoulliConstantAction
    );

    preparedOperator.ApplyDisplacementJacobianAction(
        displacementVariation, displacementAction
    );

    preparedOperator.ApplyCompleteJacobianAction(
        enthalpyVariation, gravityPotentialVariation,
        bernoulliConstantVariation, displacementVariation, completeAction
    );

    mfem::Vector summedAction(enthalpyAction);
    summedAction += gravityPotentialAction;
    summedAction += bernoulliConstantAction;
    summedAction += displacementAction;

    constexpr double finiteDifferenceStep = 1.0e-5;

    mfem::Vector centeredDifference;

    prepared_hydrostatic_complete_test_utils::centered_complete_difference(
        f, rotation, enthalpy, gravityPotential, displacement,
        bernoulliConstant, enthalpyVariation, gravityPotentialVariation,
        displacementVariation, bernoulliConstantVariation, finiteDifferenceStep,
        centeredDifference
    );

    const double summationError = gravity_prepared_test_utils::relative_error(
        completeAction, summedAction, f.mesh->GetComm()
    );

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

    INFO("Complete-action summation error = " << summationError);

    INFO(
        "Complete-action centered-difference error = "
        << centeredDifferenceError
    );

    CHECK(preparedOperator.GetCompleteJacobianStatistics().applications == 1);

    CHECK(summationError < 5.0e-12);
    CHECK(centeredDifferenceError < 2.0e-7);
}

TEST_CASE(
    "Prepared Hydrostatic MFEM Adapter Uses Four Block Layout And Reuses "
    "Preparation",
    tags::barotrope &tags::hydro &tags::integration &tags::jacobian
        &tags::mfem_operators &tags::prepared &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 =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.enthalpyFes->GetTrueVSize(), 0.41
        );

    const mfem::Vector gravityPotential =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.gravityPotentialFes->GetTrueVSize(), 0.63
        );

    const mfem::Vector displacement =
        gravity_prepared_test_utils::make_displacement(f, 0.74);

    constexpr double bernoulliConstant = 0.38;

    preparedOperator.Prepare(
        prepared_hydrostatic_complete_test_utils::make_state(
            enthalpy, gravityPotential, displacement, bernoulliConstant
        ),
        prepared_hydrostatic_complete_test_utils::make_dependencies(),
        prepared_hydrostatic_complete_test_utils::make_rotation()
    );

    mean_field::operators::PreparedHydrostaticEquilibriumJacobianOperator
        adapter(f, preparedOperator);

    const auto &layout = adapter.GetLayout();

    CHECK(
        layout.Offset(
            mean_field::operators::HydrostaticJacobianInputBlock::enthalpy
        ) == 0
    );

    CHECK(
        layout.Offset(
            mean_field::operators::HydrostaticJacobianInputBlock::
                gravityPotential
        ) == f.enthalpyFes->GetTrueVSize()
    );

    CHECK(
        layout.Size(
            mean_field::operators::HydrostaticJacobianInputBlock::
                bernoulliConstant
        ) == 1
    );

    CHECK(adapter.Width() == layout.GetTotalSize());
    CHECK(adapter.Height() == layout.GetResidualSize());
    CHECK(adapter.Height() == f.enthalpyFes->GetTrueVSize());

    const mfem::Vector enthalpyVariation =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.enthalpyFes->GetTrueVSize(), 1.12
        );

    const mfem::Vector gravityPotentialVariation =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.gravityPotentialFes->GetTrueVSize(), 1.39
        );

    const mfem::Vector displacementVariation =
        prepared_hydrostatic_complete_test_utils::make_displacement_direction(
            f, 1.28, 0.49
        );

    constexpr double bernoulliConstantVariation = 0.29;

    mfem::Vector packedDirection(adapter.Width());
    packedDirection = 0.0;

    prepared_hydrostatic_complete_test_utils::set_block(
        packedDirection, layout,
        mean_field::operators::HydrostaticJacobianInputBlock::enthalpy,
        enthalpyVariation
    );

    prepared_hydrostatic_complete_test_utils::set_block(
        packedDirection, layout,
        mean_field::operators::HydrostaticJacobianInputBlock::gravityPotential,
        gravityPotentialVariation
    );

    prepared_hydrostatic_complete_test_utils::set_block(
        packedDirection, layout,
        mean_field::operators::HydrostaticJacobianInputBlock::displacement,
        displacementVariation
    );

    packedDirection(layout.Offset(
        mean_field::operators::HydrostaticJacobianInputBlock::bernoulliConstant
    )) = bernoulliConstantVariation;

    mfem::Vector directAction;
    mfem::Vector adapterAction;

    preparedOperator.ApplyCompleteJacobianAction(
        enthalpyVariation, gravityPotentialVariation,
        bernoulliConstantVariation, displacementVariation, directAction
    );

    adapter.Mult(packedDirection, adapterAction);

    const double adapterError = gravity_prepared_test_utils::relative_error(
        adapterAction, directAction, f.mesh->GetComm()
    );

    const auto contextStatisticsBefore =
        preparedOperator.GetContextPreparationStatistics();

    const auto algebraicStatisticsBefore =
        preparedOperator.GetAlgebraicJacobianStatistics();

    const auto displacementStatisticsBefore =
        preparedOperator.GetDisplacementJacobianStatistics();

    mfem::Vector secondPackedDirection(packedDirection);
    secondPackedDirection *= -0.61;

    mfem::Vector secondAdapterAction;
    adapter.Mult(secondPackedDirection, secondAdapterAction);

    mfem::Vector expectedSecondAction(adapterAction);
    expectedSecondAction *= -0.61;

    const double adapterLinearityError =
        gravity_prepared_test_utils::relative_error(
            secondAdapterAction, expectedSecondAction, f.mesh->GetComm()
        );

    const auto &contextStatisticsAfter =
        preparedOperator.GetContextPreparationStatistics();

    const auto &algebraicStatisticsAfter =
        preparedOperator.GetAlgebraicJacobianStatistics();

    const auto &displacementStatisticsAfter =
        preparedOperator.GetDisplacementJacobianStatistics();

    INFO("MFEM adapter/direct-action error = " << adapterError);

    INFO("MFEM adapter linearity error = " << adapterLinearityError);

    CHECK(adapterError < 5.0e-12);
    CHECK(adapterLinearityError < 5.0e-12);
    CHECK(secondAdapterAction.Size() == adapter.Height());

    CHECK(contextStatisticsAfter == contextStatisticsBefore);

    CHECK(
        algebraicStatisticsAfter.preparations ==
        algebraicStatisticsBefore.preparations
    );

    CHECK(
        displacementStatisticsAfter.preparations ==
        displacementStatisticsBefore.preparations
    );

    CHECK(preparedOperator.GetCompleteJacobianStatistics().applications == 3);
}