#include <catch2/catch_test_macros.hpp>
#include <cstdint>
#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::gravity_context
) {
    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(context.GetDensityMap().reduced_size(), 0.11);
    mfem::Vector displacement = context.GetDisplacementMap().gather(prepared_test::make_displacement(f, 0.0));
    mfem::Vector gravity_gradient =
        prepared_test::make_deterministic_vector(context.GetGravityGradientMap().reduced_size(), 0.37);
    mfem::Vector gravity_potential =
        prepared_test::make_deterministic_vector(context.GetGravityPotentialMap().reduced_size(), 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.rebuilt_divergence_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 &geometry_context = context.GetGeometryContext();
    CHECK(geometry_context.GetMassOperator().HasVariationData());
    CHECK(geometry_context.GetSourceOperator().HasVariationData());
    CHECK(geometry_context.GetDivergenceOperator().Width() == f.gravityFluxFes->GetTrueVSize());
    CHECK(geometry_context.GetDivergenceOperator().Height() == f.gravityPotentialFes->GetTrueVSize());
    CHECK(geometry_context.GetTransposeDivergenceOperator().Width() == f.gravityPotentialFes->GetTrueVSize());
    CHECK(geometry_context.GetTransposeDivergenceOperator().Height() == f.gravityFluxFes->GetTrueVSize());

    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.GetDensityTrue()(context.GetDensityMap().true_dof(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.GetGravityGradientTrue()(context.GetGravityGradientMap().true_dof(0)) == gravity_gradient(0));

    displacement = context.GetDisplacementMap().gather(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_FALSE(displacement_report.geometry.rebuilt_divergence_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.geometry.rebuilt_divergence_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 Geometry Context Distinguishes Primal And Linearization Preparation",
    tags::gravity_context
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    gravity_context::GravityFieldGeometryContext context(f, *f.domainMapperStateless);
    const mfem::Vector displacement = context.GetDisplacementMap().gather(prepared_test::make_displacement(f, 0.25));

    const gravity_context::GravityFieldGeometryPreparation primal_report =
        context.PreparePrimal(displacement, {.value = 0}, {.value = 0});

    REQUIRE(context.IsPrepared());
    CHECK(primal_report.reconstructed_operators);
    CHECK(primal_report.rebuilt_mass_operator);
    CHECK(primal_report.rebuilt_source_operator);
    CHECK(primal_report.rebuilt_divergence_operator);
    CHECK_FALSE(primal_report.refreshed_variation_state);
    CHECK_FALSE(context.GetMassOperator().HasVariationData());
    CHECK_FALSE(context.GetSourceOperator().HasVariationData());

    const std::uint64_t primal_mass_preparations   = context.GetMassOperator().GetPreparationCount();
    const std::uint64_t primal_source_preparations = context.GetSourceOperator().GetPreparationCount();
    const gravity_context::GravityFieldGeometryPreparation repeated_primal_report =
        context.PreparePrimal(displacement, {.value = 0}, {.value = 0});

    CHECK_FALSE(repeated_primal_report.DidAnyWork());
    CHECK(context.GetMassOperator().GetPreparationCount() == primal_mass_preparations);
    CHECK(context.GetSourceOperator().GetPreparationCount() == primal_source_preparations);

    const gravity_context::GravityFieldGeometryPreparation upgrade_report =
        context.Prepare(displacement, {.value = 0}, {.value = 0});

    CHECK_FALSE(upgrade_report.reconstructed_operators);
    CHECK(upgrade_report.rebuilt_mass_operator);
    CHECK(upgrade_report.rebuilt_source_operator);
    CHECK_FALSE(upgrade_report.rebuilt_divergence_operator);
    CHECK(upgrade_report.refreshed_variation_state);
    CHECK(context.GetMassOperator().HasVariationData());
    CHECK(context.GetSourceOperator().HasVariationData());
    CHECK(context.GetMassOperator().GetPreparationCount() == primal_mass_preparations + 1);
    CHECK(context.GetSourceOperator().GetPreparationCount() == primal_source_preparations + 1);
}

TEST_CASE(
    "Gravity Field Linearization Context Owns Frozen Base Fields",
    tags::gravity_context
) {
    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(context.GetDensityMap().reduced_size(), 0.13);
    mfem::Vector displacement = context.GetDisplacementMap().gather(prepared_test::make_displacement(f, 0.4));
    mfem::Vector gravity_gradient =
        prepared_test::make_deterministic_vector(context.GetGravityGradientMap().reduced_size(), 0.47);
    mfem::Vector gravity_potential =
        prepared_test::make_deterministic_vector(context.GetGravityPotentialMap().reduced_size(), 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.GetDensityTrue();
    const mfem::Vector frozen_displacement     = context.GetGeometryContext().GetDisplacementTrue();
    const mfem::Vector frozen_gravity_gradient = context.GetGravityGradientTrue();

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

    CHECK(prepared_test::relative_error(context.GetDensityTrue(), frozen_density, f.mesh->GetComm()) == 0.0);
    CHECK(
        prepared_test::relative_error(
            context.GetGeometryContext().GetDisplacementTrue(), frozen_displacement, f.displacementFes->GetComm()
        ) == 0.0
    );
    CHECK(
        prepared_test::relative_error(
            context.GetGravityGradientTrue(), 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.GetDensityTrue(), frozen_density, f.mesh->GetComm()) == 0.0);
    CHECK(
        prepared_test::relative_error(
            context.GetGeometryContext().GetDisplacementTrue(), frozen_displacement, f.displacementFes->GetComm()
        ) == 0.0
    );
    CHECK(
        prepared_test::relative_error(
            context.GetGravityGradientTrue(), frozen_gravity_gradient, f.gravityFluxFes->GetComm()
        ) == 0.0
    );
}

TEST_CASE(
    "Gravity Field Geometry Contexts Have Independent Prepared State",
    tags::gravity_context
) {
    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 =
        first_context.GetDisplacementMap().gather(prepared_test::make_displacement(f, 0.0));
    const mfem::Vector second_displacement =
        second_context.GetDisplacementMap().gather(prepared_test::make_displacement(f, 1.0));
    const mfem::Vector gravity_gradient = prepared_test::make_field_map<field::Gravity>(f).gather(
        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 =
        first_context.GetDisplacementMap().gather(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);
}