#include <memory>

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

import mean_field;
import test_helpers;

namespace {
    mfem::Vector make_zero_displacement(const mean_field::fem::FEM &f) {
        mfem::Vector displacement(f.displacementFes->GetTrueVSize());
        displacement = 0.0;
        return displacement;
    }

    mfem::Vector project_constant(
        mfem::ParFiniteElementSpace &finiteElementSpace,
        const double value
    ) {
        mfem::ConstantCoefficient coefficient(value);
        mfem::ParGridFunction field(&finiteElementSpace);

        field.ProjectCoefficient(coefficient);

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

    namespace barotropic_closure_geometry_test_utils {
        mfem::Vector project_scalar_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_base_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) + 0.006 * position(2);
            });

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

        mfem::Vector make_base_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_scalar_field(*f.enthalpyFes, coefficient);
        }
    } // namespace barotropic_closure_geometry_test_utils
} // namespace

TEST_CASE(
    "Barotropic Closure Vanishes For A Representable Constant State",
    tags::hydro &tags::residuals &tags::unit &tags::closure &tags::kernels &tags::barotrope
) {
    auto args              = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    const mean_field::eos::Polytrope barotrope(3.0, 1.5);

    constexpr double enthalpyValue  = 0.8;

    const double densityValue       = barotrope.density_from_enthalpy(enthalpyValue);

    const mfem::Vector enthalpy     = project_constant(*f.enthalpyFes, enthalpyValue);

    const mfem::Vector density      = project_constant(*f.densityFes, densityValue);

    const mfem::Vector displacement = make_zero_displacement(f);

    mfem::Vector residual;
    mfem::Vector scale;

    mean_field::operators::kernels::apply_barotropic_closure(
        f, *f.domainMapperStateless, barotrope, density, enthalpy, displacement, residual
    );

    mean_field::operators::kernels::apply_barotropic_closure_density_action(
        f, *f.domainMapperStateless, barotrope, density, displacement, scale
    );

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

    const double relativeResidual = gravity_prepared_test_utils::global_norm(residual, communicator) /
                                    gravity_prepared_test_utils::global_norm(scale, communicator);

    INFO("Relative constant-state closure residual = " << relativeResidual);

    CHECK(relativeResidual < 5.0e-12);
}

TEST_CASE(
    "Barotropic Closure Density Action Matches The Stellar Mass Matrix",
    tags::hydro &tags::jacobian &tags::unit &tags::closure &tags::kernels &tags::barotrope
) {
    auto args              = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    const mean_field::eos::Polytrope barotrope(3.0, 1.5);

    const mfem::Vector displacement = make_zero_displacement(f);

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

    mfem::Vector kernelAction;

    mean_field::operators::kernels::apply_barotropic_closure_density_action(
        f, *f.domainMapperStateless, barotrope, densityVariation, displacement, kernelAction
    );

    using Schema  = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    using Stellar = mean_field::utils::domain::Stellar;

    mfem::Array<int> stellarMarker(f.mesh->attributes.Max());
    stellarMarker = 0;

    for (int attributeIndex = 0; attributeIndex < f.mesh->attributes.Size(); ++attributeIndex) {
        const int attribute = f.mesh->attributes[attributeIndex];

        if (Schema::template attribute_belongs_to<Stellar>(attribute)) {
            stellarMarker[attribute - 1] = 1;
        }
    }

    mfem::ParBilinearForm massForm(f.densityFes.get());

    massForm.AddDomainIntegrator(new mfem::MassIntegrator(), stellarMarker);

    massForm.Assemble();
    massForm.Finalize();

    std::unique_ptr<mfem::HypreParMatrix> massMatrix(massForm.ParallelAssemble());

    REQUIRE(massMatrix != nullptr);
    REQUIRE(massMatrix->Width() == densityVariation.Size());

    mfem::Vector referenceAction(massMatrix->Height());
    referenceAction = 0.0;

    massMatrix->Mult(densityVariation, referenceAction);
    const double relativeError =
        gravity_prepared_test_utils::relative_error(kernelAction, referenceAction, f.mesh->GetComm());

    INFO("Density-action mass-matrix error = " << relativeError);

    CHECK(relativeError < 5.0e-12);
}

TEST_CASE(
    "Barotropic Closure Jacobian Matches A Combined Centered Difference",
    tags::hydro &tags::jacobian &tags::unit &tags::closure &tags::kernels &tags::barotrope
) {
    auto args              = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    const mean_field::eos::Polytrope barotrope(3.0, 1.5);

    mfem::FunctionCoefficient densityCoefficient([](const mfem::Vector &position) {
        return 0.4 + 0.03 * position(0) - 0.01 * position(1);
    });

    mfem::FunctionCoefficient enthalpyCoefficient([](const mfem::Vector &position) {
        return 0.9 + 0.02 * position(0) - 0.01 * position(1);
    });

    mfem::FunctionCoefficient enthalpyVariationCoefficient([](const mfem::Vector &position) {
        return 0.07 + 0.015 * position(0) + 0.008 * position(2);
    });

    mfem::ParGridFunction densityField(f.densityFes.get());
    mfem::ParGridFunction enthalpyField(f.enthalpyFes.get());
    mfem::ParGridFunction enthalpyVariationField(f.enthalpyFes.get());

    densityField.ProjectCoefficient(densityCoefficient);
    enthalpyField.ProjectCoefficient(enthalpyCoefficient);
    enthalpyVariationField.ProjectCoefficient(enthalpyVariationCoefficient);

    mfem::Vector density;
    mfem::Vector enthalpy;
    mfem::Vector enthalpyVariation;

    densityField.GetTrueDofs(density);
    enthalpyField.GetTrueDofs(enthalpy);
    enthalpyVariationField.GetTrueDofs(enthalpyVariation);

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

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

    constexpr double differenceStep = 1.0e-6;

    const mfem::Vector plusDensity =
        gravity_prepared_test_utils::linear_combination(density, 1.0, densityVariation, differenceStep);

    const mfem::Vector minusDensity =
        gravity_prepared_test_utils::linear_combination(density, 1.0, densityVariation, -differenceStep);

    const mfem::Vector plusEnthalpy =
        gravity_prepared_test_utils::linear_combination(enthalpy, 1.0, enthalpyVariation, differenceStep);

    const mfem::Vector minusEnthalpy =
        gravity_prepared_test_utils::linear_combination(enthalpy, 1.0, enthalpyVariation, -differenceStep);

    mfem::Vector plusResidual;
    mfem::Vector minusResidual;

    mean_field::operators::kernels::apply_barotropic_closure(
        f, *f.domainMapperStateless, barotrope, plusDensity, plusEnthalpy, displacement, plusResidual
    );

    mean_field::operators::kernels::apply_barotropic_closure(
        f, *f.domainMapperStateless, barotrope, minusDensity, minusEnthalpy, displacement, minusResidual
    );

    mfem::Vector finiteDifference(plusResidual);
    finiteDifference -= minusResidual;
    finiteDifference *= 1.0 / (2.0 * differenceStep);

    mfem::Vector densityAction;
    mfem::Vector enthalpyAction;

    mean_field::operators::kernels::apply_barotropic_closure_density_action(
        f, *f.domainMapperStateless, barotrope, densityVariation, displacement, densityAction
    );

    mean_field::operators::kernels::apply_barotropic_closure_enthalpy_action(
        f, *f.domainMapperStateless, barotrope, enthalpy, enthalpyVariation, displacement, enthalpyAction
    );

    mfem::Vector analyticAction(densityAction);
    analyticAction += enthalpyAction;

    const double relativeError =
        gravity_prepared_test_utils::relative_error(analyticAction, finiteDifference, f.mesh->GetComm());

    INFO("Combined EOS Jacobian error = " << relativeError);

    CHECK(relativeError < 2.0e-8);
}

TEST_CASE(
    "Barotropic Closure Density Action Excludes Vacuum And Uses Mapped Volume",
    tags::hydro &tags::mapping &tags::unit &tags::closure &tags::barotrope &tags::kernels
) {
    auto args              = test_utils::setup_args();
    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    const mean_field::eos::Polytrope barotrope(3.0, 1.5);

    const mfem::Vector stellarDensity       = gravity_prepared_test_utils::make_domain_supported_density(f, true);

    const mfem::Vector vacuumDensity        = gravity_prepared_test_utils::make_domain_supported_density(f, false);

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

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

    mfem::Vector stellarAction;
    mfem::Vector vacuumAction;
    mfem::Vector deformedAction;

    mean_field::operators::kernels::apply_barotropic_closure_density_action(
        f, *f.domainMapperStateless, barotrope, stellarDensity, identityDisplacement, stellarAction
    );

    mean_field::operators::kernels::apply_barotropic_closure_density_action(
        f, *f.domainMapperStateless, barotrope, vacuumDensity, identityDisplacement, vacuumAction
    );

    mean_field::operators::kernels::apply_barotropic_closure_density_action(
        f, *f.domainMapperStateless, barotrope, stellarDensity, deformedDisplacement, deformedAction
    );

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

    const double stellarNorm    = gravity_prepared_test_utils::global_norm(stellarAction, communicator);

    const double vacuumNorm     = gravity_prepared_test_utils::global_norm(vacuumAction, communicator);

    const double geometryChange =
        gravity_prepared_test_utils::relative_error(deformedAction, stellarAction, communicator);

    INFO("Stellar action norm = " << stellarNorm);
    INFO("Vacuum action norm = " << vacuumNorm);
    INFO("Relative mapped-volume change = " << geometryChange);

    CHECK(stellarNorm > 0.0);
    CHECK(vacuumNorm <= 1.0e-13 * stellarNorm);
    CHECK(geometryChange > 1.0e-5);
}

TEST_CASE(
    "Barotropic Closure Displacement Action Matches Centered Differences",
    tags::barotrope &tags::closure &tags::hydro &tags::integration &tags::jacobian &tags::mapping &tags::physics
        &tags::kernels
) {
    auto args              = test_utils::setup_args();

    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.domainMapperStateless != nullptr);

    const mean_field::eos::Polytrope barotrope(3.0, 1.5);

    const mfem::Vector baseDensity           = barotropic_closure_geometry_test_utils::make_base_density(f);

    const mfem::Vector baseEnthalpy          = barotropic_closure_geometry_test_utils::make_base_enthalpy(f);

    const mfem::Vector displacementVariation = gravity_prepared_test_utils::make_displacement(f, 0.65);

    constexpr double differenceStep          = 1.0e-5;

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

    for (const double deformationScale : {0.0, 1.0}) {
        DYNAMIC_SECTION("Base deformation scale = " << deformationScale) {
            const mfem::Vector baseDisplacement = gravity_prepared_test_utils::make_displacement(f, deformationScale);

            mfem::Vector plusDisplacement(baseDisplacement);

            mfem::Vector minusDisplacement(baseDisplacement);

            plusDisplacement.Add(differenceStep, displacementVariation);

            minusDisplacement.Add(-differenceStep, displacementVariation);

            mfem::Vector plusResidual;
            mfem::Vector minusResidual;
            mfem::Vector analyticAction;

            mean_field::operators::kernels::apply_barotropic_closure(
                f, *f.domainMapperStateless, barotrope, baseDensity, baseEnthalpy, plusDisplacement, plusResidual
            );

            mean_field::operators::kernels::apply_barotropic_closure(
                f, *f.domainMapperStateless, barotrope, baseDensity, baseEnthalpy, minusDisplacement, minusResidual
            );

            mean_field::operators::kernels::apply_barotropic_closure_displacement_action(
                f, *f.domainMapperStateless, barotrope, baseDensity, baseEnthalpy, baseDisplacement,
                displacementVariation, analyticAction
            );

            mfem::Vector finiteDifference(plusResidual);

            finiteDifference -= minusResidual;
            finiteDifference *= 1.0 / (2.0 * differenceStep);

            const double analyticNorm = gravity_prepared_test_utils::global_norm(analyticAction, communicator);

            const double finiteDifferenceNorm =
                gravity_prepared_test_utils::global_norm(finiteDifference, communicator);

            const double relativeError =
                gravity_prepared_test_utils::relative_error(analyticAction, finiteDifference, communicator);

            INFO("Base deformation scale = " << deformationScale);

            INFO("Analytic geometry-action norm = " << analyticNorm);

            INFO("Finite-difference geometry-action norm = " << finiteDifferenceNorm);

            INFO("Geometry-action relative error = " << relativeError);

            REQUIRE(analyticNorm > 1.0e-12);
            REQUIRE(finiteDifferenceNorm > 1.0e-12);

            CHECK(relativeError < 5.0e-8);
        }
    }
}

TEST_CASE(
    "Barotropic Closure Displacement Action Is Linear In Its Direction",
    tags::barotrope &tags::closure &tags::hydro &tags::jacobian &tags::mapping &tags::physics &tags::unit &tags::kernels
) {
    auto args              = test_utils::setup_args();

    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.domainMapperStateless != nullptr);

    const mean_field::eos::Polytrope barotrope(3.0, 1.5);

    const mfem::Vector baseDensity      = barotropic_closure_geometry_test_utils::make_base_density(f);

    const mfem::Vector baseEnthalpy     = barotropic_closure_geometry_test_utils::make_base_enthalpy(f);

    const mfem::Vector baseDisplacement = gravity_prepared_test_utils::make_displacement(f, 0.8);

    const mfem::Vector firstDirection   = gravity_prepared_test_utils::make_displacement(f, 0.4);

    mfem::Vector secondDirection =
        gravity_prepared_test_utils::make_deterministic_vector(f.displacementFes->GetTrueVSize(), 0.91);

    secondDirection *= 0.01;

    constexpr double firstScale  = 1.7;
    constexpr double secondScale = -0.43;

    mfem::Vector combinedDirection(firstDirection);

    combinedDirection *= firstScale;

    combinedDirection.Add(secondScale, secondDirection);

    mfem::Vector zeroDirection(f.displacementFes->GetTrueVSize());
    zeroDirection = 0.0;

    mfem::Vector firstAction;
    mfem::Vector secondAction;
    mfem::Vector combinedAction;
    mfem::Vector zeroAction;

    mean_field::operators::kernels::apply_barotropic_closure_displacement_action(
        f, *f.domainMapperStateless, barotrope, baseDensity, baseEnthalpy, baseDisplacement, firstDirection, firstAction
    );

    mean_field::operators::kernels::apply_barotropic_closure_displacement_action(
        f, *f.domainMapperStateless, barotrope, baseDensity, baseEnthalpy, baseDisplacement, secondDirection,
        secondAction
    );

    mean_field::operators::kernels::apply_barotropic_closure_displacement_action(
        f, *f.domainMapperStateless, barotrope, baseDensity, baseEnthalpy, baseDisplacement, combinedDirection,
        combinedAction
    );

    mean_field::operators::kernels::apply_barotropic_closure_displacement_action(
        f, *f.domainMapperStateless, barotrope, baseDensity, baseEnthalpy, baseDisplacement, zeroDirection, zeroAction
    );

    mfem::Vector expectedAction(firstAction);

    expectedAction *= firstScale;

    expectedAction.Add(secondScale, secondAction);

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

    const double expectedNorm   = gravity_prepared_test_utils::global_norm(expectedAction, communicator);

    const double linearityError =
        gravity_prepared_test_utils::relative_error(combinedAction, expectedAction, communicator);

    const double zeroActionNorm = gravity_prepared_test_utils::global_norm(zeroAction, communicator);

    INFO("Expected combined-action norm = " << expectedNorm);

    INFO("Directional-linearity error = " << linearityError);

    INFO("Zero-direction action norm = " << zeroActionNorm);

    REQUIRE(expectedNorm > 1.0e-12);

    CHECK(linearityError < 5.0e-12);

    CHECK(zeroActionNorm <= 5.0e-14 * expectedNorm);
}

TEST_CASE(
    "Barotropic Closure Displacement Action Excludes Vacuum",
    tags::barotrope &tags::closure &tags::hydro &tags::mapping &tags::physics &tags::unit &tags::kernels
) {
    auto args              = test_utils::setup_args();

    mean_field::fem::FEM f = mean_field::fem::setup_fem(args.mesh_file, args, 0);

    REQUIRE(f.domainMapperStateless != nullptr);

    const mean_field::eos::Polytrope barotrope(3.0, 1.5);

    const mfem::Vector stellarDensity = gravity_prepared_test_utils::make_domain_supported_density(f, true);

    const mfem::Vector vacuumDensity  = gravity_prepared_test_utils::make_domain_supported_density(f, false);

    mfem::Vector zeroEnthalpy(f.enthalpyFes->GetTrueVSize());
    zeroEnthalpy                             = 0.0;

    const mfem::Vector baseDisplacement      = gravity_prepared_test_utils::make_displacement(f, 0.7);

    const mfem::Vector displacementVariation = gravity_prepared_test_utils::make_displacement(f, 0.5);

    mfem::Vector stellarAction;
    mfem::Vector vacuumAction;

    mean_field::operators::kernels::apply_barotropic_closure_displacement_action(
        f, *f.domainMapperStateless, barotrope, stellarDensity, zeroEnthalpy, baseDisplacement, displacementVariation,
        stellarAction
    );

    mean_field::operators::kernels::apply_barotropic_closure_displacement_action(
        f, *f.domainMapperStateless, barotrope, vacuumDensity, zeroEnthalpy, baseDisplacement, displacementVariation,
        vacuumAction
    );

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

    const double stellarNorm    = gravity_prepared_test_utils::global_norm(stellarAction, communicator);

    const double vacuumNorm     = gravity_prepared_test_utils::global_norm(vacuumAction, communicator);

    INFO("Stellar geometry-action norm = " << stellarNorm);

    INFO("Vacuum geometry-action norm = " << vacuumNorm);

    REQUIRE(stellarNorm > 1.0e-12);

    CHECK(vacuumNorm <= 1.0e-13 * stellarNorm);
}