#include <cstdint>

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

import mean_field;
import test_helpers;

namespace barotropic_closure_context_test_utils {
    mfem::Vector project_field(
        mfem::ParFiniteElementSpace &finiteElementSpace,
        mfem::Coefficient &coefficient
    ) {
        mfem::ParGridFunction field(&finiteElementSpace);

        field.ProjectCoefficient(coefficient);

        mfem::Vector trueVector;
        field.GetTrueDofs(trueVector);

        return trueVector;
    }

    mfem::Vector make_density(const mean_field::fem::FEM &f) {
        mfem::FunctionCoefficient coefficient([](const mfem::Vector &position) {
            return 0.55 + 0.025 * position(0) - 0.010 * position(1);
        });

        return project_field(*f.densityFes, coefficient);
    }

    mfem::Vector make_enthalpy(const mean_field::fem::FEM &f) {
        mfem::FunctionCoefficient coefficient([](const mfem::Vector &position) {
            return 0.90 + 0.020 * position(0) - 0.010 * position(1) +
                   0.005 * position(2);
        });

        return project_field(*f.enthalpyFes, coefficient);
    }
} // namespace barotropic_closure_context_test_utils

TEST_CASE(
    "Barotropic Closure Context Tracks Independent Revisions",
    tags::barotrope &tags::closure &tags::hydro &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);

    const mean_field::physics::PolytropicBarotrope barotrope(3.0, 1.5);

    mean_field::operators::context::barotropic::
        BarotropicClosureLinearizationContext context(
            f, *f.domainMapperStateless, barotrope
        );

    mfem::Vector density =
        barotropic_closure_context_test_utils::make_density(f);

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

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

    mean_field::operators::context::barotropic::BarotropicClosureRevisions
        revisions{.density = 3, .enthalpy = 5, .displacement = 7};

    CHECK_FALSE(context.IsPrepared());
    CHECK(context.GetPreparationCount() == 0);

    context.Prepare(density, enthalpy, displacement, revisions);

    REQUIRE(context.IsPrepared());

    CHECK(context.MatchesRevisions(revisions));
    CHECK(context.GetRevisions() == revisions);
    CHECK(context.GetPreparationCount() == 1);

    CHECK(context.GetOperator().GetPreparationCount() == 1);

    context.Prepare(density, enthalpy, displacement, revisions);

    CHECK(context.GetPreparationCount() == 1);

    const double frozenDensityValue = context.GetBaseDensityTrue()(0);

    density(0) += 0.125;

    CHECK(context.GetBaseDensityTrue()(0) == frozenDensityValue);

    context.Prepare(density, enthalpy, displacement, revisions);

    CHECK(context.GetPreparationCount() == 1);
    CHECK(context.GetBaseDensityTrue()(0) == frozenDensityValue);

    ++revisions.density;

    context.Prepare(density, enthalpy, displacement, revisions);

    CHECK(context.GetPreparationCount() == 2);
    CHECK(context.GetBaseDensityTrue()(0) == density(0));

    enthalpy(0) += 0.050;
    ++revisions.enthalpy;

    context.Prepare(density, enthalpy, displacement, revisions);

    CHECK(context.GetPreparationCount() == 3);
    CHECK(context.GetBaseEnthalpyTrue()(0) == enthalpy(0));

    displacement = gravity_prepared_test_utils::make_displacement(f, 1.0);

    ++revisions.displacement;

    context.Prepare(density, enthalpy, displacement, revisions);

    CHECK(context.GetPreparationCount() == 4);
    CHECK(context.GetRevisions() == revisions);
    CHECK(context.MatchesRevisions(revisions));

    CHECK(context.GetOperator().GetPreparationCount() == 4);
}

TEST_CASE(
    "Barotropic Closure Context Reprepares A Consistent Frozen State",
    tags::barotrope &tags::closure &tags::hydro &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);

    const mean_field::physics::PolytropicBarotrope barotrope(3.0, 1.5);

    mean_field::operators::context::barotropic::
        BarotropicClosureLinearizationContext context(
            f, *f.domainMapperStateless, barotrope
        );

    const mfem::Vector density =
        barotropic_closure_context_test_utils::make_density(f);

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

    const mfem::Vector identityDisplacement =
        gravity_prepared_test_utils::make_displacement(f, 0.0);

    const mfem::Vector densityVariation =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.densityFes->GetTrueVSize(), 0.37
        );

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

    const mfem::Vector displacementVariation =
        gravity_prepared_test_utils::make_deterministic_vector(
            f.displacementFes->GetTrueVSize(), 0.37
        );

    mean_field::operators::context::barotropic::BarotropicClosureRevisions
        revisions{.density = 11, .enthalpy = 13, .displacement = 17};

    context.Prepare(density, enthalpy, identityDisplacement, revisions);

    mfem::Vector initialResidual;
    mfem::Vector initialAction;

    context.BuildResidual(initialResidual);

    context.GetOperator().Mult(
        densityVariation, enthalpyVariation, displacementVariation,
        initialAction
    );

    mfem::Vector changedDensity(density);
    changedDensity.Add(0.025, densityVariation);

    mfem::Vector changedEnthalpy(enthalpy);
    changedEnthalpy.Add(0.015, enthalpyVariation);

    const mfem::Vector deformedDisplacement =
        gravity_prepared_test_utils::make_displacement(f, 1.0);

    context.Prepare(
        changedDensity, changedEnthalpy, deformedDisplacement, revisions
    );

    mfem::Vector unchangedResidual;
    mfem::Vector unchangedAction;

    context.BuildResidual(unchangedResidual);

    context.GetOperator().Mult(
        densityVariation, enthalpyVariation, displacementVariation,
        unchangedAction
    );

    const MPI_Comm communicator = f.mesh->GetComm();

    CHECK(context.GetPreparationCount() == 1);

    CHECK(
        gravity_prepared_test_utils::relative_error(
            unchangedResidual, initialResidual, communicator
        ) < 5.0e-15
    );

    CHECK(
        gravity_prepared_test_utils::relative_error(
            unchangedAction, initialAction, communicator
        ) < 5.0e-15
    );

    ++revisions.density;
    ++revisions.enthalpy;
    ++revisions.displacement;

    context.Prepare(
        changedDensity, changedEnthalpy, deformedDisplacement, revisions
    );

    mfem::Vector preparedResidual;
    mfem::Vector preparedAction;

    context.BuildResidual(preparedResidual);

    context.GetOperator().Mult(
        densityVariation, enthalpyVariation, displacementVariation,
        preparedAction
    );

    mfem::Vector referenceResidual;
    mfem::Vector referenceDensityAction;
    mfem::Vector referenceEnthalpyAction;
    mfem::Vector referenceDisplacementAction;

    mean_field::operators::kernels::apply_barotropic_closure(
        f, *f.domainMapperStateless, barotrope, changedDensity, changedEnthalpy,
        deformedDisplacement, referenceResidual
    );

    mean_field::operators::kernels::apply_barotropic_closure_density_action(
        f, *f.domainMapperStateless, barotrope, densityVariation,
        deformedDisplacement, referenceDensityAction
    );

    mean_field::operators::kernels::apply_barotropic_closure_enthalpy_action(
        f, *f.domainMapperStateless, barotrope, changedEnthalpy,
        enthalpyVariation, deformedDisplacement, referenceEnthalpyAction
    );

    mean_field::operators::kernels::
        apply_barotropic_closure_displacement_action(
            f, *f.domainMapperStateless, barotrope, changedDensity,
            changedEnthalpy, deformedDisplacement, displacementVariation,
            referenceDisplacementAction
        );

    mfem::Vector referenceAction(referenceDensityAction);
    referenceAction += referenceEnthalpyAction;
    referenceAction += referenceDisplacementAction;
    const double residualError = gravity_prepared_test_utils::relative_error(
        preparedResidual, referenceResidual, communicator
    );

    const double actionError = gravity_prepared_test_utils::relative_error(
        preparedAction, referenceAction, communicator
    );

    const double residualChange = gravity_prepared_test_utils::relative_error(
        preparedResidual, initialResidual, communicator
    );

    const double actionChange = gravity_prepared_test_utils::relative_error(
        preparedAction, initialAction, communicator
    );

    INFO("Prepared-context residual error = " << residualError);

    INFO("Prepared-context Jacobian error = " << actionError);

    INFO("Residual change after valid revision = " << residualChange);

    INFO("Jacobian change after valid revision = " << actionChange);

    CHECK(context.GetPreparationCount() == 2);
    CHECK(residualError < 5.0e-12);
    CHECK(actionError < 5.0e-12);
    CHECK(residualChange > 1.0e-6);
    CHECK(actionChange > 1.0e-6);
}