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

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);
}

TEST_CASE(
    "Prepared Hydrostatic Equilibrium Reports Candidate Mapping And Arithmetic Failures Without Unwinding",
    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);
    REQUIRE(f.okay());

    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.73);
    displacement *= 1.0e200;

    auto dependencies           = prepared_hydrostatic_test_utils::make_dependencies();
    const auto rotation         = prepared_hydrostatic_test_utils::make_rotation();

    const auto mappingRejection = preparedOperator.TryPrepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, 0.41), dependencies,
        rotation
    );
    REQUIRE_FALSE(mappingRejection.has_value());
    CHECK(
        (mappingRejection.error().reason ==
             mean_field::operators::HydrostaticEquilibriumPreparationRejectionReason::inverted_geometry ||
         mappingRejection.error().reason ==
             mean_field::operators::HydrostaticEquilibriumPreparationRejectionReason::non_finite_geometry)
    );
    CHECK(mappingRejection.error().mappingStatus != mean_field::mapping::MappingStatus::valid);
    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK_THROWS_AS(
        preparedOperator.Prepare(
            prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, 0.41), dependencies,
            rotation
        ),
        std::domain_error
    );

    displacement = 0.0;
    ++dependencies.displacement.revision;
    ++dependencies.rotation.revision;

    const auto enormousRotation    = prepared_hydrostatic_test_utils::make_rotation(1.0e200);
    const auto arithmeticRejection = preparedOperator.TryPrepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, 0.41), dependencies,
        enormousRotation
    );
    REQUIRE_FALSE(arithmeticRejection.has_value());
    CHECK(
        arithmeticRejection.error().reason ==
        mean_field::operators::HydrostaticEquilibriumPreparationRejectionReason::non_finite_residual
    );
    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK_THROWS_AS(
        preparedOperator.Prepare(
            prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, 0.41), dependencies,
            enormousRotation
        ),
        std::domain_error
    );

    ++dependencies.rotation.revision;
    const auto recovered = preparedOperator.TryPrepare(
        prepared_hydrostatic_test_utils::make_state(enthalpy, gravityPotential, displacement, 0.41), dependencies,
        rotation
    );
    REQUIRE(recovered.has_value());
    CHECK(recovered->preparedResidual);
    CHECK(preparedOperator.IsPrepared());
}

TEST_CASE(
    "Hydrostatic Tensor Preparation Matches Pointwise Preparation And Rebuilds On Mode Changes",
    tags::barotrope_hydrostatic_prepared_residual &tags::unit
) {
    const auto args = test_utils::setup_args();
    auto f = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());
    using namespace prepared_hydrostatic_test_utils;
    mean_field::operators::PreparedHydrostaticEquilibriumOperator fast(f, *f.domainMapperStateless);
    mean_field::operators::PreparedHydrostaticEquilibriumOperator reference(f, *f.domainMapperStateless);
    auto h = make_enthalpy(f), phi = make_gravity_potential(f);
    const auto dh = make_enthalpy(f, 0.53), dp = make_gravity_potential(f, 0.71);
    const auto du = field_dof_test_utils::make_supported_displacement(f, 0.31);
    const auto rotation = make_rotation();
    auto dependencies = make_dependencies();

    for (double scale : {0.0, 0.42, 0.73}) {
        CAPTURE(scale);
        const auto displacement = field_dof_test_utils::make_supported_displacement(f, scale);
        h *= 1.01; phi *= 0.99;
        ++dependencies.displacement.revision;
        ++dependencies.enthalpy.revision;
        ++dependencies.gravityPotential.revision;
        const auto state = make_state(h, phi, displacement, 0.39);
        REQUIRE(fast.Prepare(state, dependencies, rotation).DidAnyWork());
        REQUIRE(reference.PrepareReference(state, dependencies, rotation).DidAnyWork());
        CHECK(fast.GetTensorDisplacementElementCount() > 0);
        const auto compare = [&](auto &&apply) {
            mfem::Vector a, b;
            apply(fast, a); apply(reference, b);
            CHECK(gravity_prepared_test_utils::relative_error(a, b, f.mesh->GetComm()) < 2.0e-11);
        };
        compare([](auto &op, auto &out) { op.BuildResidual(out); });
        compare([&](auto &op, auto &out) { op.ApplyEnthalpyJacobianAction(dh, out); });
        compare([&](auto &op, auto &out) { op.ApplyGravityPotentialJacobianAction(dp, out); });
        compare([](auto &op, auto &out) { op.ApplyBernoulliConstantJacobianAction(0.13, out); });
        compare([&](auto &op, auto &out) { op.ApplyDisplacementJacobianAction(du, out); });
        compare([&](auto &op, auto &out) { op.ApplyCompleteJacobianAction(dh, dp, 0.13, du, out); });
        compare([](auto &op, auto &out) { op.AssembleEnthalpyJacobianDiagonal(out); });
        CHECK_FALSE(fast.Prepare(state, dependencies, rotation).DidAnyWork());
        CHECK_FALSE(reference.PrepareReference(state, dependencies, rotation).DidAnyWork());

        CHECK(fast.PrepareReference(state, dependencies, rotation).DidAnyWork());
        compare([](auto &op, auto &out) { op.BuildResidual(out); });
        CHECK_FALSE(fast.PrepareReference(state, dependencies, rotation).DidAnyWork());
        CHECK(fast.Prepare(state, dependencies, rotation).DidAnyWork());
        compare([&](auto &op, auto &out) { op.ApplyCompleteJacobianAction(dh, dp, 0.13, du, out); });
        CHECK_FALSE(fast.Prepare(state, dependencies, rotation).DidAnyWork());
    }

    mfem::ParGridFunction foldedField(f.displacementFes.get());
    mfem::VectorFunctionCoefficient fold(3, [](const mfem::Vector &x, mfem::Vector &u) {
        u = x; u *= -2.0;
    });
    foldedField.ProjectCoefficient(fold);
    mfem::Vector foldedTrue; foldedField.GetTrueDofs(foldedTrue);
    const auto folded = fast.GetDisplacementMap().gather(foldedTrue);
    ++dependencies.displacement.revision;
    const auto rejectedState = make_state(h, phi, folded, 0.39);
    const auto rejected = fast.TryPrepare(rejectedState, dependencies, rotation);
    const auto referenceRejected = reference.TryPrepareReference(rejectedState, dependencies, rotation);
    REQUIRE_FALSE(rejected.has_value());
    REQUIRE_FALSE(referenceRejected.has_value());
    CHECK(rejected.error().reason == referenceRejected.error().reason);
    CHECK(rejected.error().mappingStatus == mean_field::mapping::MappingStatus::non_positive_determinant);
    CHECK_FALSE(fast.IsPrepared());
    CHECK_FALSE(reference.IsPrepared());

    const auto recovered = make_state(h, phi, du, 0.39);
    ++dependencies.displacement.revision;
    REQUIRE(fast.TryPrepare(recovered, dependencies, rotation).has_value());
    REQUIRE(reference.TryPrepareReference(recovered, dependencies, rotation).has_value());
    mfem::Vector a, b;
    fast.BuildResidual(a); reference.BuildResidual(b);
    CHECK(gravity_prepared_test_utils::relative_error(a, b, f.mesh->GetComm()) < 2.0e-11);
}