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

import mean_field;
import test_helpers;

using namespace mean_field;
using Catch::Matchers::WithinAbs;
namespace prepared_test = gravity_prepared_test_utils;

namespace {
    [[nodiscard]] mfem::Vector make_folding_displacement(const mean_field::fem::FEM &f) {
        mfem::ParGridFunction field(f.displacementFes.get());
        mfem::VectorFunctionCoefficient coefficient(
            f.mesh->Dimension(), [](const mfem::Vector &position, mfem::Vector &value) {
                value.SetSize(position.Size());
                for (int dimension = 0; dimension < position.Size(); ++dimension) {
                    value(dimension) = -2.0 * position(dimension);
                }
            }
        );
        field.ProjectCoefficient(coefficient);

        mfem::Vector displacementTrue;
        field.GetTrueDofs(displacementTrue);
        return displacementTrue;
    }
} // namespace

TEST_CASE(
    "Prepared Mapped Gravity Source Reports Invalid Candidate Geometry Without Unwinding",
    tags::gravity_prepared_unit &tags::geometry
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());

    operators::PreparedMappedGravitySourceOperator preparedOperator(f, *f.domainMapperStateless);
    mfem::Vector displacement = preparedOperator.GetDisplacementMap().gather(make_folding_displacement(f));

    const auto rejected       = preparedOperator.TryPrepare(displacement);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().reason == operators::GravitySourcePreparationRejectionReason::invalid_mapping);
    CHECK(rejected.error().mappingStatus == mapping::MappingStatus::non_positive_determinant);
    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK_FALSE(preparedOperator.HasVariationData());
    CHECK(preparedOperator.GetPreparationCount() == 0);

    CHECK_THROWS_AS(preparedOperator.Prepare(displacement), std::domain_error);
    CHECK_FALSE(preparedOperator.IsPrepared());
    CHECK(preparedOperator.GetPreparationCount() == 0);

    displacement         = 0.0;
    const auto recovered = preparedOperator.TryPrepare(displacement);
    REQUIRE(recovered.has_value());
    CHECK(preparedOperator.IsPrepared());
    CHECK(preparedOperator.HasVariationData());
    CHECK(preparedOperator.GetPreparationCount() == 1);
}

TEST_CASE(
    "Prepared Mapped Gravity Source Reuses Tables Across Preparation Modes And Rejection",
    tags::gravity_prepared_unit &tags::geometry
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(f.okay());

    operators::PreparedMappedGravitySourceOperator operation(f, *f.domainMapperStateless);
    const mfem::Vector densityTrue      = prepared_test::make_deterministic_vector(f.densityFes->GetTrueVSize(), 0.41);
    const mfem::Vector density          = operation.GetDensityMap().gather(densityTrue);
    const mfem::Vector displacementTrue = prepared_test::make_displacement(f, 0.4);
    const mfem::Vector displacement     = operation.GetDisplacementMap().gather(displacementTrue);
    const mfem::Vector directionTrue    = prepared_test::make_displacement(f, 0.7);
    const MPI_Comm communicator         = f.mesh->GetComm();

    operation.Prepare(displacement);
    mfem::Vector baselineAction;
    mfem::Vector baselineVariation;
    operation.Mult(density, baselineAction);
    operation.MultDisplacementVariationTrue(densityTrue, directionTrue, baselineVariation);

    const mfem::Vector primalDisplacementTrue = prepared_test::make_displacement(f, 1.0);
    operation.PreparePrimal(operation.GetDisplacementMap().gather(primalDisplacementTrue));
    REQUIRE(operation.IsPrepared());
    CHECK_FALSE(operation.HasVariationData());
    mfem::Vector primalAction;
    mfem::Vector referenceActionTrue;
    operation.Mult(density, primalAction);
    operators::kernels::apply_mapped_source(
        f, *f.domainMapperStateless, densityTrue, primalDisplacementTrue, referenceActionTrue
    );
    CHECK_THAT(
        prepared_test::relative_error(
            primalAction, operation.GetPotentialMap().gather(referenceActionTrue), communicator
        ),
        WithinAbs(0.0, 2.0e-11)
    );

    operation.Prepare(displacement);
    REQUIRE(operation.HasVariationData());
    mfem::Vector repeatedAction;
    mfem::Vector repeatedVariation;
    operation.Mult(density, repeatedAction);
    operation.MultDisplacementVariationTrue(densityTrue, directionTrue, repeatedVariation);
    CHECK(prepared_test::relative_error(repeatedAction, baselineAction, communicator) < 2.0e-14);
    CHECK(prepared_test::relative_error(repeatedVariation, baselineVariation, communicator) < 2.0e-14);

    const auto rejected = operation.TryPrepare(operation.GetDisplacementMap().gather(make_folding_displacement(f)));
    REQUIRE_FALSE(rejected.has_value());
    CHECK_FALSE(operation.IsPrepared());
    CHECK_FALSE(operation.HasVariationData());
    REQUIRE(operation.TryPrepare(displacement).has_value());
    REQUIRE(operation.HasVariationData());
    operation.Mult(density, repeatedAction);
    operation.MultDisplacementVariationTrue(densityTrue, directionTrue, repeatedVariation);
    CHECK(prepared_test::relative_error(repeatedAction, baselineAction, communicator) < 2.0e-14);
    CHECK(prepared_test::relative_error(repeatedVariation, baselineVariation, communicator) < 2.0e-14);
    CHECK(operation.GetPreparationCount() == 4);
}

TEST_CASE(
    "Prepared Mapped Gravity Source Matches Stateless Kernel",
    tags::gravity_prepared
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    operators::PreparedMappedGravitySourceOperator prepared_operator(f, *f.domainMapperStateless);
    REQUIRE(prepared_operator.Width() == prepared_operator.GetDensityMap().reduced_size());
    REQUIRE(prepared_operator.Height() == prepared_operator.GetPotentialMap().reduced_size());

    const mfem::Vector density_true = prepared_test::make_deterministic_vector(f.densityFes->GetTrueVSize(), 0.41);
    const mfem::Vector density      = prepared_operator.GetDensityMap().gather(density_true);
    const MPI_Comm communicator     = f.mesh->GetComm();

    mfem::Vector identity_action;
    mfem::Vector deformed_action;

    for (const double deformation_scale : {0.0, 1.0}) {
        const mfem::Vector displacement_true = prepared_test::make_displacement(f, deformation_scale);
        const mfem::Vector displacement      = prepared_operator.GetDisplacementMap().gather(displacement_true);

        prepared_operator.Prepare(displacement);

        mfem::Vector prepared_action;

        prepared_operator.Mult(density, prepared_action);
        mfem::Vector reference_action_true;
        operators::kernels::apply_mapped_source(
            f, *f.domainMapperStateless, density_true, displacement_true, reference_action_true
        );
        const mfem::Vector reference_action = prepared_operator.GetPotentialMap().gather(reference_action_true);

        const double relative_error = prepared_test::relative_error(prepared_action, reference_action, communicator);

        INFO("Deformation scale = " << deformation_scale);
        INFO("Prepared source norm = " << prepared_test::global_norm(prepared_action, communicator));
        INFO("Reference source norm = " << prepared_test::global_norm(reference_action, communicator));
        INFO("Relative prepared-source error = " << relative_error);

        REQUIRE(prepared_operator.IsPrepared());
        CHECK_THAT(relative_error, WithinAbs(0.0, 2.0e-11));

        if (deformation_scale == 0.0) {
            identity_action = prepared_action;
        } else {
            deformed_action = prepared_action;
        }
    }

    const double geometry_change = prepared_test::relative_error(deformed_action, identity_action, communicator);

    INFO("Relative source change under deformation = " << geometry_change);

    CHECK(prepared_operator.GetPreparationCount() == 2);
    CHECK(geometry_change > 1.0e-5);
}

TEST_CASE(
    "Prepared Mapped Gravity Source Preserves Linearity And Excludes Vacuum",
    tags::gravity_prepared
) {
    auto args  = test_utils::setup_args();
    fem::FEM f = fem::setup_fem(args.mesh_file, args, 0);

    operators::PreparedMappedGravitySourceOperator prepared_operator(f, *f.domainMapperStateless);
    REQUIRE(prepared_operator.Width() == prepared_operator.GetDensityMap().reduced_size());
    REQUIRE(prepared_operator.Height() == prepared_operator.GetPotentialMap().reduced_size());
    const mfem::Vector displacement =
        prepared_operator.GetDisplacementMap().gather(prepared_test::make_displacement(f, 1.0));
    prepared_operator.Prepare(displacement);

    const mfem::Vector first = prepared_operator.GetDensityMap().gather(
        prepared_test::make_deterministic_vector(f.densityFes->GetTrueVSize(), 0.27)
    );
    const mfem::Vector second = prepared_operator.GetDensityMap().gather(
        prepared_test::make_deterministic_vector(f.densityFes->GetTrueVSize(), 0.79)
    );
    const mfem::Vector combination = prepared_test::linear_combination(first, 1.3, second, -0.6);
    const mfem::Vector stellar_density =
        prepared_operator.GetDensityMap().gather(prepared_test::make_domain_supported_density(f, true));
    const mfem::Vector vacuum_density =
        prepared_operator.GetDensityMap().gather(prepared_test::make_domain_supported_density(f, false));

    mfem::Vector first_action;
    mfem::Vector second_action;
    mfem::Vector combination_action;
    mfem::Vector stellar_action;
    mfem::Vector vacuum_action;
    mfem::Vector transpose_action;

    prepared_operator.Mult(first, first_action);
    prepared_operator.Mult(second, second_action);
    prepared_operator.Mult(combination, combination_action);
    prepared_operator.Mult(stellar_density, stellar_action);
    prepared_operator.Mult(vacuum_density, vacuum_action);

    const mfem::Vector potential =
        prepared_test::make_deterministic_vector(prepared_operator.GetPotentialMap().reduced_size(), 1.17);
    prepared_operator.MultTranspose(potential, transpose_action);

    const mfem::Vector expected_combination = prepared_test::linear_combination(first_action, 1.3, second_action, -0.6);
    const MPI_Comm communicator             = f.mesh->GetComm();

    const double linearity_error =
        prepared_test::relative_error(combination_action, expected_combination, communicator);
    const double stellar_norm             = prepared_test::global_norm(stellar_action, communicator);
    const double vacuum_norm              = prepared_test::global_norm(vacuum_action, communicator);
    const double forward_pairing          = prepared_test::global_dot(first_action, potential, communicator);
    const double transpose_pairing        = prepared_test::global_dot(first, transpose_action, communicator);
    const double transpose_error          = prepared_test::relative_scalar_error(forward_pairing, transpose_pairing);
    const std::uint64_t preparation_count = prepared_operator.GetPreparationCount();

    mfem::Vector repeated_action;
    prepared_operator.Mult(first, repeated_action);

    INFO("Relative source linearity error = " << linearity_error);
    INFO("Stellar source norm = " << stellar_norm);
    INFO("Vacuum-only source norm = " << vacuum_norm);
    INFO("Relative source-transpose pairing error = " << transpose_error);

    CHECK_THAT(linearity_error, WithinAbs(0.0, 2.0e-12));
    CHECK(stellar_norm > 0.0);
    CHECK(vacuum_norm <= 1.0e-13 * stellar_norm);
    CHECK_THAT(transpose_error, WithinAbs(0.0, 2.0e-12));
    CHECK(prepared_test::relative_error(repeated_action, first_action, communicator) < 2.0e-14);
    CHECK(prepared_operator.GetPreparationCount() == preparation_count);
}