#include <algorithm>
#include <array>
#include <cmath>
#include <numbers>

#include <catch2/catch_test_macros.hpp>

#include <mfem.hpp>

import mean_field;
import test_helpers;

namespace gravity_displacement_force_analytic_test_utils {
    struct AffineCase {
        const char *name;
        std::array<double, 3> scales;
    };

    [[nodiscard]] double analytic_sphere_volume(const double radius) {
        return (4.0 / 3.0) * std::numbers::pi * radius * radius * radius;
    }

    [[nodiscard]] double determinant(
        const std::array<
            double,
            3> &scales
    ) {
        return scales[0] * scales[1] * scales[2];
    }

    [[nodiscard]] double relative_scalar_error(
        const double computed,
        const double expected
    ) {
        return std::abs(computed - expected) / std::max(std::abs(expected), 1.0e-30);
    }

    [[nodiscard]] mfem::Vector make_constant_density(
        const mean_field::fem::FEM &f,
        const double densityValue
    ) {
        mfem::ParGridFunction densityField(f.densityFes.get());
        mfem::ConstantCoefficient densityCoefficient(densityValue);
        densityField.ProjectCoefficient(densityCoefficient);

        mfem::Vector densityTrue;
        densityField.GetTrueDofs(densityTrue);
        return densityTrue;
    }

    [[nodiscard]] mfem::Vector make_reference_gravity(
        const mean_field::fem::FEM &f,
        const std::array<
            double,
            3> &referenceGravity
    ) {
        mfem::ParGridFunction gravityField(f.gravityFluxFes.get());

        mfem::VectorFunctionCoefficient gravityCoefficient(
            f.mesh->Dimension(), [referenceGravity](const mfem::Vector &, mfem::Vector &value) {
                value.SetSize(3);

                for (int component = 0; component < 3; ++component) {
                    value(component) = referenceGravity[static_cast<std::size_t>(component)];
                }
            }
        );

        gravityField.ProjectCoefficient(gravityCoefficient);

        mfem::Vector gravityTrue;
        gravityField.GetTrueDofs(gravityTrue);
        return gravityTrue;
    }

    [[nodiscard]] mfem::Vector make_radial_gravity(
        const mean_field::fem::FEM &f,
        const double radialCoefficient
    ) {
        mfem::ParGridFunction gravityField(f.gravityFluxFes.get());

        mfem::VectorFunctionCoefficient gravityCoefficient(
            f.mesh->Dimension(), [radialCoefficient](const mfem::Vector &position, mfem::Vector &value) {
                value.SetSize(position.Size());

                for (int component = 0; component < position.Size(); ++component) {
                    value(component) = radialCoefficient * position(component);
                }
            }
        );

        gravityField.ProjectCoefficient(gravityCoefficient);

        mfem::Vector gravityTrue;
        gravityField.GetTrueDofs(gravityTrue);
        return gravityTrue;
    }

    [[nodiscard]] mfem::Vector make_affine_displacement(
        const mean_field::fem::FEM &f,
        const std::array<
            double,
            3> &scales
    ) {
        mfem::ParGridFunction displacementField(f.displacementFes.get());

        mfem::VectorFunctionCoefficient displacementCoefficient(
            f.mesh->Dimension(), [scales](const mfem::Vector &position, mfem::Vector &value) {
                value.SetSize(position.Size());

                for (int component = 0; component < position.Size(); ++component) {
                    value(component) = (scales[static_cast<std::size_t>(component)] - 1.0) * position(component);
                }
            }
        );

        displacementField.ProjectCoefficient(displacementCoefficient);

        mfem::Vector displacementTrue;
        displacementField.GetTrueDofs(displacementTrue);
        return displacementTrue;
    }

    [[nodiscard]] mfem::Vector make_constant_test_direction(
        const mean_field::fem::FEM &f,
        const int selectedComponent
    ) {
        mfem::ParGridFunction testField(f.displacementFes.get());

        mfem::VectorFunctionCoefficient testCoefficient(
            f.mesh->Dimension(), [selectedComponent](const mfem::Vector &position, mfem::Vector &value) {
                value.SetSize(position.Size());
                value                    = 0.0;
                value(selectedComponent) = 1.0;
            }
        );

        testField.ProjectCoefficient(testCoefficient);

        mfem::Vector testTrue;
        testField.GetTrueDofs(testTrue);
        return testTrue;
    }

    [[nodiscard]] mfem::Vector make_dilation_test_direction(const mean_field::fem::FEM &f) {
        mfem::ParGridFunction testField(f.displacementFes.get());

        mfem::VectorFunctionCoefficient testCoefficient(
            f.mesh->Dimension(), [](const mfem::Vector &position, mfem::Vector &value) { value = position; }
        );

        testField.ProjectCoefficient(testCoefficient);

        mfem::Vector testTrue;
        testField.GetTrueDofs(testTrue);
        return testTrue;
    }

    void set_mass_normalized_density(
        mean_field::fem::FEM &f,
        const double targetMass,
        mfem::ParGridFunction &densityField
    ) {
        const mfem::Vector stellarDensityTrue = gravity_prepared_test_utils::make_domain_supported_density(f, true);

        densityField.SetFromTrueDofs(stellarDensityTrue);

        const double unnormalizedMass =
            mean_field::analysis::domain_integrate_grid_function(f, densityField, mean_field::utils::DOMAINS::STELLAR);

        MFEM_VERIFY(unnormalizedMass > 0.0, "The analytic gravity-force test obtained non-positive mass.");

        densityField *= targetMass / unnormalizedMass;
    }
} // namespace gravity_displacement_force_analytic_test_utils

TEST_CASE(
    "Gravity Displacement Force Matches Analytic Affine Resultants",
    tags::gravity &tags::accuracy &tags::analytic_comparison &tags::integration
) {
    mean_field::utils::Args args = test_utils::setup_args();

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

    REQUIRE(f.okay());
    REQUIRE(f.domainMapperStateless != nullptr);
    REQUIRE(f.domainMapperStateless != nullptr);

    constexpr double densityValue = 1.37;

    constexpr std::array<double, 3> physicalGravity{0.31, -0.47, 0.22};

    constexpr std::array<gravity_displacement_force_analytic_test_utils::AffineCase, 3> affineCases{
        {{.name = "identity geometry", .scales = {1.0, 1.0, 1.0}},
         {.name = "volume-preserving affine geometry", .scales = {1.14, 0.93, 1.0 / (1.14 * 0.93)}},
         {.name = "volume-changing affine geometry", .scales = {1.11, 0.96, 1.07}}}
    };

    const mfem::Vector density = gravity_displacement_force_analytic_test_utils::make_constant_density(f, densityValue);

    const double referenceVolume =
        gravity_displacement_force_analytic_test_utils::analytic_sphere_volume(mean_field::utils::RADIUS);

    constexpr double relativeTolerance = 5.0e-6;

    for (const gravity_displacement_force_analytic_test_utils::AffineCase &affineCase : affineCases) {
        DYNAMIC_SECTION(affineCase.name) {
            const double mapDeterminant =
                gravity_displacement_force_analytic_test_utils::determinant(affineCase.scales);

            REQUIRE(mapDeterminant > 0.0);

            std::array<double, 3> referenceGravity{};

            /*
             * For x = A X, the H(div) Piola relation is
             *
             *     g_phys = A g_ref / det(A).
             *
             * Prescribe the RT pullback that represents the requested
             * constant physical gravity field exactly.
             */
            for (int component = 0; component < 3; ++component) {
                referenceGravity[static_cast<std::size_t>(component)] =
                    mapDeterminant * physicalGravity[static_cast<std::size_t>(component)] /
                    affineCase.scales[static_cast<std::size_t>(component)];
            }

            const mfem::Vector gravityGradient =
                gravity_displacement_force_analytic_test_utils::make_reference_gravity(f, referenceGravity);

            const mfem::Vector displacement =
                gravity_displacement_force_analytic_test_utils::make_affine_displacement(f, affineCase.scales);

            mfem::Vector residual;

            mean_field::operators::kernels::apply_gravity_displacement_force_residual(
                f, *f.domainMapperStateless, density, gravityGradient, displacement, residual
            );

            for (int component = 0; component < 3; ++component) {
                const mfem::Vector testDirection =
                    gravity_displacement_force_analytic_test_utils::make_constant_test_direction(f, component);

                const double computedResultant =
                    gravity_prepared_test_utils::global_dot(residual, testDirection, f.mesh->GetComm());

                const double expectedResultant = densityValue * physicalGravity[static_cast<std::size_t>(component)] *
                                                 mapDeterminant * referenceVolume;

                const double relativeError     = gravity_displacement_force_analytic_test_utils::relative_scalar_error(
                    computedResultant, expectedResultant
                );

                CAPTURE(component);
                INFO("Map determinant = " << mapDeterminant);
                INFO("Computed resultant = " << computedResultant);
                INFO("Analytic resultant = " << expectedResultant);
                INFO("Relative resultant error = " << relativeError);

                CHECK(relativeError < relativeTolerance);
            }
        }
    }
}

TEST_CASE(
    "Gravity Displacement Force Reproduces Analytic Homogeneous Sphere Work",
    tags::gravity &tags::accuracy &tags::analytic_comparison &tags::integration
) {
    mean_field::utils::Args args = test_utils::setup_args();

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

    REQUIRE(f.okay());
    REQUIRE(f.domainMapperStateless != nullptr);

    const double radius       = mean_field::utils::RADIUS;
    const double mass         = mean_field::utils::MASS;
    const double volume       = gravity_displacement_force_analytic_test_utils::analytic_sphere_volume(radius);

    const double densityValue = mass / volume;
    const double radialGravityCoefficient = mean_field::utils::G * mass / (radius * radius * radius);

    const mfem::Vector density = gravity_displacement_force_analytic_test_utils::make_constant_density(f, densityValue);

    const mfem::Vector gravityGradient =
        gravity_displacement_force_analytic_test_utils::make_radial_gravity(f, radialGravityCoefficient);

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

    mfem::Vector residual;

    mean_field::operators::kernels::apply_gravity_displacement_force_residual(
        f, *f.domainMapperStateless, density, gravityGradient, displacement, residual
    );

    const mfem::Vector dilationDirection =
        gravity_displacement_force_analytic_test_utils::make_dilation_test_direction(f);

    const double computedWork = gravity_prepared_test_utils::global_dot(residual, dilationDirection, f.mesh->GetComm());

    const double analyticWork = (3.0 / 5.0) * mean_field::utils::G * mass * mass / radius;

    const double relativeError =
        gravity_displacement_force_analytic_test_utils::relative_scalar_error(computedWork, analyticWork);

    INFO("Computed positive gravity work = " << computedWork);
    INFO("Analytic positive gravity work = " << analyticWork);
    INFO("Computed gravitational virial = " << -computedWork);
    INFO("Analytic binding energy = " << -analyticWork);
    INFO("Relative analytic work error = " << relativeError);

    REQUIRE(computedWork > 0.0);
    CHECK(relativeError < 1.0e-5);
}

TEST_CASE(
    "Solved Homogeneous Sphere Gravity Force Matches Analytic Virial",
    tags::gravity &tags::accuracy &tags::analytic_comparison &tags::integration &tags::initialization
) {
    mean_field::utils::Args args = test_utils::setup_args();
    args.p.rtol                  = 1.0e-13;
    args.p.max_iters             = std::max(args.p.max_iters, 1000);

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

    REQUIRE(f.okay());
    REQUIRE(f.domainMapperStateless != nullptr);

    mfem::ParGridFunction displacementField(f.displacementFes.get());
    displacementField = 0.0;

    REQUIRE(f.domainMapperStateless != nullptr);
    *f.displacement     = 0.0;

    const double radius = mean_field::utils::RADIUS;
    const double mass   = mean_field::utils::MASS;

    mfem::ParGridFunction densityField(f.densityFes.get());

    gravity_displacement_force_analytic_test_utils::set_mass_normalized_density(f, mass, densityField);

    const mean_field::physics::GravitySolution gravitySolution =
        mean_field::physics::solve_gravity_field(f, args, densityField, displacementField);

    mfem::Vector densityTrue;
    mfem::Vector gravityGradientTrue;
    mfem::Vector displacementTrue;

    densityField.GetTrueDofs(densityTrue);
    gravitySolution.gradPhi.GetTrueDofs(gravityGradientTrue);
    displacementField.GetTrueDofs(displacementTrue);

    mfem::Vector residual;

    mean_field::operators::kernels::apply_gravity_displacement_force_residual(
        f, *f.domainMapperStateless, densityTrue, gravityGradientTrue, displacementTrue, residual
    );

    const mfem::Vector dilationDirection =
        gravity_displacement_force_analytic_test_utils::make_dilation_test_direction(f);

    const double computedWork = gravity_prepared_test_utils::global_dot(residual, dilationDirection, f.mesh->GetComm());

    const double analyticWork = (3.0 / 5.0) * mean_field::utils::G * mass * mass / radius;

    const double relativeError =
        gravity_displacement_force_analytic_test_utils::relative_scalar_error(computedWork, analyticWork);

    INFO("Solved-field positive gravity work = " << computedWork);
    INFO("Analytic positive gravity work = " << analyticWork);
    INFO("Solved-field gravitational virial = " << -computedWork);
    INFO("Analytic homogeneous-sphere binding energy = " << -analyticWork);
    INFO("Relative solved-field virial error = " << relativeError);

    REQUIRE(computedWork > 0.0);
    CHECK(relativeError < 1.0e-5);
}