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

import mean_field;
import test_helpers;

using namespace mean_field;
namespace prepared_test   = gravity_prepared_test_utils;
namespace gravity_context = operators::context::gravity_field;

TEST_CASE(
    "Gravity Field Linearization Context Applies Selective Invalidation",
    tags::integration &tags::gravity &tags::contexts
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    gravity_context::GravityFieldLinearizationContext context(
        f, *f.domainMapperStateless
    );

    mfem::Vector density = prepared_test::make_deterministic_vector(
        f.densityFes->GetTrueVSize(), 0.11
    );
    mfem::Vector displacement     = prepared_test::make_displacement(f, 0.0);
    mfem::Vector gravity_gradient = prepared_test::make_deterministic_vector(
        f.gravityFluxFes->GetTrueVSize(), 0.37
    );
    mfem::Vector gravity_potential = prepared_test::make_deterministic_vector(
        f.gravityPotentialFes->GetTrueVSize(), 0.63
    );

    gravity_context::GravityFieldRevisions revisions;

    auto make_state = [&]() {
        return gravity_context::GravityFieldStateView{
            .density           = density,
            .displacement      = displacement,
            .gravity_gradient  = gravity_gradient,
            .gravity_potential = gravity_potential
        };
    };

    REQUIRE_FALSE(context.IsPrepared());

    const gravity_context::GravityFieldPreparationReport initial_report =
        context.Prepare(make_state(), revisions);

    REQUIRE(context.IsPrepared());
    CHECK(initial_report.geometry.reconstructed_operators);
    CHECK(initial_report.geometry.rebuilt_mass_operator);
    CHECK(initial_report.geometry.rebuilt_source_operator);
    CHECK(initial_report.geometry.refreshed_variation_state);
    CHECK(initial_report.updated_density);
    CHECK(initial_report.updated_gravity_gradient);
    CHECK(initial_report.DidAnyWork());

    const auto initial_mass_preparations =
        context.GetGeometryContext().GetMassOperator().GetPreparationCount();
    const auto initial_source_preparations =
        context.GetGeometryContext().GetSourceOperator().GetPreparationCount();

    const gravity_context::GravityFieldPreparationReport repeated_report =
        context.Prepare(make_state(), revisions);

    CHECK_FALSE(repeated_report.DidAnyWork());
    CHECK(
        context.GetGeometryContext().GetMassOperator().GetPreparationCount() ==
        initial_mass_preparations
    );
    CHECK(
        context.GetGeometryContext()
            .GetSourceOperator()
            .GetPreparationCount() == initial_source_preparations
    );

    gravity_potential(0) += 0.25;
    ++revisions.gravity_potential.value;

    const gravity_context::GravityFieldPreparationReport potential_report =
        context.Prepare(make_state(), revisions);

    CHECK_FALSE(potential_report.DidAnyWork());
    CHECK(
        context.GetRevisions().gravity_potential == revisions.gravity_potential
    );

    density(0) += 0.5;
    ++revisions.density.value;

    const gravity_context::GravityFieldPreparationReport density_report =
        context.Prepare(make_state(), revisions);

    CHECK(density_report.updated_density);
    CHECK_FALSE(density_report.updated_gravity_gradient);
    CHECK_FALSE(density_report.geometry.DidAnyWork());
    CHECK(context.GetDensity()(0) == density(0));

    gravity_gradient(0) -= 0.4;
    ++revisions.gravity_gradient.value;

    const gravity_context::GravityFieldPreparationReport gradient_report =
        context.Prepare(make_state(), revisions);

    CHECK_FALSE(gradient_report.updated_density);
    CHECK(gradient_report.updated_gravity_gradient);
    CHECK_FALSE(gradient_report.geometry.DidAnyWork());
    CHECK(context.GetGravityGradient()(0) == gravity_gradient(0));

    displacement = prepared_test::make_displacement(f, 1.0);
    ++revisions.displacement.value;

    const gravity_context::GravityFieldPreparationReport displacement_report =
        context.Prepare(make_state(), revisions);

    CHECK_FALSE(displacement_report.geometry.reconstructed_operators);
    CHECK(displacement_report.geometry.rebuilt_mass_operator);
    CHECK(displacement_report.geometry.rebuilt_source_operator);
    CHECK(displacement_report.geometry.refreshed_variation_state);
    CHECK_FALSE(displacement_report.updated_density);
    CHECK_FALSE(displacement_report.updated_gravity_gradient);
    CHECK(
        context.GetGeometryContext().GetMassOperator().GetPreparationCount() ==
        initial_mass_preparations + 1
    );
    CHECK(
        context.GetGeometryContext()
            .GetSourceOperator()
            .GetPreparationCount() == initial_source_preparations + 1
    );

    ++revisions.discretization.value;

    const gravity_context::GravityFieldPreparationReport discretization_report =
        context.Prepare(make_state(), revisions);

    CHECK(discretization_report.geometry.reconstructed_operators);
    CHECK(discretization_report.geometry.rebuilt_mass_operator);
    CHECK(discretization_report.geometry.rebuilt_source_operator);
    CHECK(discretization_report.updated_density);
    CHECK(discretization_report.updated_gravity_gradient);
    CHECK(
        context.GetGeometryContext().GetMassOperator().GetPreparationCount() ==
        1
    );
    CHECK(
        context.GetGeometryContext()
            .GetSourceOperator()
            .GetPreparationCount() == 1
    );
}

TEST_CASE(
    "Gravity Field Linearization Context Owns Frozen Base Fields",
    tags::integration &tags::gravity &tags::contexts
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    gravity_context::GravityFieldLinearizationContext context(
        f, *f.domainMapperStateless
    );

    mfem::Vector density = prepared_test::make_deterministic_vector(
        f.densityFes->GetTrueVSize(), 0.13
    );
    mfem::Vector displacement     = prepared_test::make_displacement(f, 0.4);
    mfem::Vector gravity_gradient = prepared_test::make_deterministic_vector(
        f.gravityFluxFes->GetTrueVSize(), 0.47
    );
    mfem::Vector gravity_potential = prepared_test::make_deterministic_vector(
        f.gravityPotentialFes->GetTrueVSize(), 0.71
    );

    gravity_context::GravityFieldRevisions revisions;

    context.Prepare(
        {.density           = density,
         .displacement      = displacement,
         .gravity_gradient  = gravity_gradient,
         .gravity_potential = gravity_potential},
        revisions
    );

    const mfem::Vector frozen_density = context.GetDensity();
    const mfem::Vector frozen_displacement =
        context.GetGeometryContext().GetDisplacement();
    const mfem::Vector frozen_gravity_gradient = context.GetGravityGradient();

    density                                    = 0.0;
    displacement                               = 0.0;
    gravity_gradient                           = 0.0;
    gravity_potential                          = 0.0;

    CHECK(
        prepared_test::relative_error(
            context.GetDensity(), frozen_density, f.mesh->GetComm()
        ) == 0.0
    );
    CHECK(
        prepared_test::relative_error(
            context.GetGeometryContext().GetDisplacement(), frozen_displacement,
            f.displacementFes->GetComm()
        ) == 0.0
    );
    CHECK(
        prepared_test::relative_error(
            context.GetGravityGradient(), frozen_gravity_gradient,
            f.gravityFluxFes->GetComm()
        ) == 0.0
    );

    const gravity_context::GravityFieldPreparationReport
        unchanged_revision_report = context.Prepare(
            {.density           = density,
             .displacement      = displacement,
             .gravity_gradient  = gravity_gradient,
             .gravity_potential = gravity_potential},
            revisions
        );

    CHECK_FALSE(unchanged_revision_report.DidAnyWork());
    CHECK(
        prepared_test::relative_error(
            context.GetDensity(), frozen_density, f.mesh->GetComm()
        ) == 0.0
    );
    CHECK(
        prepared_test::relative_error(
            context.GetGeometryContext().GetDisplacement(), frozen_displacement,
            f.displacementFes->GetComm()
        ) == 0.0
    );
    CHECK(
        prepared_test::relative_error(
            context.GetGravityGradient(), frozen_gravity_gradient,
            f.gravityFluxFes->GetComm()
        ) == 0.0
    );
}

TEST_CASE(
    "Gravity Field Geometry Contexts Have Independent Prepared State",
    tags::integration &tags::gravity &tags::contexts
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    gravity_context::GravityFieldGeometryContext first_context(
        f, *f.domainMapperStateless
    );
    gravity_context::GravityFieldGeometryContext second_context(
        f, *f.domainMapperStateless
    );

    const mfem::Vector first_displacement =
        prepared_test::make_displacement(f, 0.0);
    const mfem::Vector second_displacement =
        prepared_test::make_displacement(f, 1.0);
    const mfem::Vector gravity_gradient =
        prepared_test::make_deterministic_vector(
            f.gravityFluxFes->GetTrueVSize(), 0.35
        );

    first_context.Prepare(first_displacement, {.value = 0}, {.value = 0});
    second_context.Prepare(second_displacement, {.value = 0}, {.value = 0});

    mfem::Vector first_action;
    mfem::Vector second_action_before;
    mfem::Vector second_action_after;

    first_context.GetMassOperator().Mult(gravity_gradient, first_action);
    second_context.GetMassOperator().Mult(
        gravity_gradient, second_action_before
    );

    const mfem::Vector updated_first_displacement =
        prepared_test::make_displacement(f, 0.6);
    first_context.Prepare(
        updated_first_displacement, {.value = 0}, {.value = 1}
    );

    second_context.GetMassOperator().Mult(
        gravity_gradient, second_action_after
    );

    const MPI_Comm communicator            = f.gravityFluxFes->GetComm();
    const double independent_context_error = prepared_test::relative_error(
        second_action_after, second_action_before, communicator
    );
    const double distinct_geometry_difference = prepared_test::relative_error(
        first_action, second_action_before, communicator
    );

    INFO(
        "Second-context change after preparing first context = "
        << independent_context_error
    );
    INFO(
        "Difference between independently prepared geometries = "
        << distinct_geometry_difference
    );

    CHECK(independent_context_error < 2.0e-14);
    CHECK(distinct_geometry_difference > 1.0e-5);
}