#include <algorithm>
#include <array>
#include <cmath>
#include <cstdint>
#include <limits>
#include <type_traits>

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

import mean_field;
import test_helpers;

namespace mass_normalization_test_utils {
    using CoupledForm = mean_field::utils::blocks::barotropic_equilibrium_form;

    constexpr auto densityValue =
        mean_field::utils::blocks::get_value_block<CoupledForm>(mean_field::utils::blocks::density_field.mass_term);

    constexpr auto displacementValue = mean_field::utils::blocks::get_value_block<CoupledForm>(
        mean_field::utils::blocks::displacement_field.geometry_term
    );

    constexpr auto massResidual = mean_field::utils::blocks::get_residual_block<CoupledForm>(
        mean_field::utils::blocks::barotropic_constant_field.mass_normalization_term
    );

    [[nodiscard]] mean_field::operators::MassNormalizationLayout make_layout(const mean_field::fem::FEM &f) {
        const std::array<int, CoupledForm::value_block_count> valueSizes{
            f.densityFes->GetTrueVSize(),          f.displacementFes->GetTrueVSize(), f.gravityFluxFes->GetTrueVSize(),
            f.gravityPotentialFes->GetTrueVSize(), f.enthalpyFes->GetTrueVSize(),     1
        };

        const std::array<int, CoupledForm::residual_block_count> residualSizes{
            f.gravityFluxFes->GetTrueVSize(),  f.gravityPotentialFes->GetTrueVSize(), f.densityFes->GetTrueVSize(),
            f.displacementFes->GetTrueVSize(), f.enthalpyFes->GetTrueVSize(),         1
        };

        return {valueSizes, residualSizes};
    }

    [[nodiscard]] mfem::Vector make_density(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction field(f.densityFes.get());
        mfem::FunctionCoefficient coefficient([phase](const mfem::Vector &position) {
            return 0.91 + 0.07 * std::sin(0.83 * position(0) + phase) + 0.05 * std::cos(0.61 * position(1) - phase) +
                   0.03 * position(2) * position(2);
        });
        field.ProjectCoefficient(coefficient);

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

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

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

    [[nodiscard]] mfem::Vector make_density_direction(
        const mean_field::fem::FEM &f,
        const double phase
    ) {
        mfem::ParGridFunction field(f.densityFes.get());
        mfem::FunctionCoefficient coefficient([phase](const mfem::Vector &position) {
            return 0.19 * std::sin(0.71 * position(0) + phase) - 0.13 * std::cos(0.89 * position(1) - phase) +
                   0.08 * position(2);
        });
        field.ProjectCoefficient(coefficient);

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

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

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

    [[nodiscard]] mfem::Vector make_displacement_direction(
        const mean_field::fem::FEM &f,
        const double scale
    ) {
        mfem::ParGridFunction field(f.displacementFes.get());
        mfem::VectorFunctionCoefficient coefficient(
            f.mesh->Dimension(), [scale](const mfem::Vector &position, mfem::Vector &value) {
                value.SetSize(3);
                value(0) = scale * (0.07 * position(0) + 0.018 * position(1) * position(2));
                value(1) = scale * (-0.05 * position(1) + 0.013 * position(0) * position(2));
                value(2) = scale * (0.04 * position(2) - 0.011 * position(0) * position(1));
            }
        );
        field.ProjectCoefficient(coefficient);

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

    [[nodiscard]] mean_field::operators::MassNormalizationDependencies make_dependencies() {
        return {
            .discretization = {.identity = 701, .revision = 3},
            .density        = {.identity = 709, .revision = 5},
            .displacement   = {.identity = 719, .revision = 7},
            .targetMass     = {.identity = 727, .revision = 11}
        };
    }

    [[nodiscard]] mean_field::operators::context::gravity_field::GravityFieldRevisions make_gravity_revisions(
        const mean_field::operators::MassNormalizationDependencies &dependencies,
        const std::uint64_t gravityGradientRevision  = 13,
        const std::uint64_t gravityPotentialRevision = 17
    ) {
        return {
            .discretization    = {.value = dependencies.discretization.revision},
            .displacement      = {.value = dependencies.displacement.revision},
            .density           = {.value = dependencies.density.revision},
            .gravity_gradient  = {.value = gravityGradientRevision},
            .gravity_potential = {.value = gravityPotentialRevision}
        };
    }

    void prepare_gravity_context(
        mean_field::operators::context::gravity_field::GravityFieldLinearizationContext &context,
        const mean_field::fem::FEM &f,
        const mfem::Vector &density,
        const mfem::Vector &displacement,
        const mean_field::operators::MassNormalizationDependencies &dependencies,
        const std::uint64_t gravityGradientRevision  = 13,
        const std::uint64_t gravityPotentialRevision = 17
    ) {
        mfem::Vector gravityGradient(f.gravityFluxFes->GetTrueVSize());
        gravityGradient = 0.0;

        mfem::Vector gravityPotential(f.gravityPotentialFes->GetTrueVSize());
        gravityPotential = 0.0;

        context.Prepare(
            {.density           = density,
             .displacement      = displacement,
             .gravity_gradient  = gravityGradient,
             .gravity_potential = gravityPotential},
            make_gravity_revisions(dependencies, gravityGradientRevision, gravityPotentialRevision)
        );
    }

    [[nodiscard]] double residual_value(const mean_field::operators::PreparedMassNormalizationOperator &massOperator) {
        mfem::Vector residual;
        massOperator.BuildResidual(residual);
        REQUIRE(residual.Size() == 1);
        return residual(0);
    }

    [[nodiscard]] double relative_error(
        const double computed,
        const double reference
    ) {
        return std::abs(computed - reference) /
               std::max(std::abs(reference), 100.0 * std::numeric_limits<double>::epsilon());
    }
} // namespace mass_normalization_test_utils

TEST_CASE(
    "Prepared Mass Normalization Has The Analytic Affine Volume Scaling",
    tags::barotrope &tags::prepared &tags::analytic_comparison
) {
    using Operator = mean_field::operators::PreparedMassNormalizationOperator;

    STATIC_REQUIRE_FALSE(std::is_copy_constructible_v<Operator>);
    STATIC_REQUIRE_FALSE(std::is_copy_assignable_v<Operator>);
    STATIC_REQUIRE_FALSE(std::is_move_constructible_v<Operator>);
    STATIC_REQUIRE_FALSE(std::is_move_assignable_v<Operator>);

    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());

    const double densityValue  = 1.37;
    const double targetMass    = 0.73;
    const double affineScale   = 0.086;

    const mfem::Vector density = mass_normalization_test_utils::make_constant_density(f, densityValue);
    mfem::Vector displacement(f.displacementFes->GetTrueVSize());
    displacement      = 0.0;

    auto dependencies = mass_normalization_test_utils::make_dependencies();

    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );

    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, density, displacement, dependencies);

    Operator massOperator(f, *f.domainMapperStateless, gravityContext);

    const auto initialReport = massOperator.Prepare({.targetMass = targetMass}, dependencies);

    CHECK(initialReport.rebuiltStaticPlan);
    CHECK(initialReport.refreshedGeometry);
    CHECK(initialReport.refreshedDensity);
    CHECK(initialReport.updatedTargetMass);
    CHECK(initialReport.assembledResidual);

    const double undeformedMass = massOperator.GetCurrentMass();
    const mean_field::mapping::COORDINATE_SPACE volumeCoordinates =
        f.has_mapping() ? mean_field::mapping::COORDINATE_SPACE::PHYSICAL
                        : mean_field::mapping::COORDINATE_SPACE::REFERENCE;

    const double independentlyIntegratedMass =
        densityValue * mean_field::analysis::get_mesh_volume(f, volumeCoordinates, mean_field::utils::DOMAINS::STELLAR);

    CHECK(mass_normalization_test_utils::relative_error(undeformedMass, independentlyIntegratedMass) < 1.0e-12);

    CHECK(
        mass_normalization_test_utils::relative_error(
            mass_normalization_test_utils::residual_value(massOperator), undeformedMass - targetMass
        ) < 2.0e-15
    );

    displacement = mass_normalization_test_utils::make_affine_displacement(f, affineScale);
    ++dependencies.displacement.revision;

    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, density, displacement, dependencies);

    const auto deformedReport = massOperator.Prepare({.targetMass = targetMass}, dependencies);

    CHECK_FALSE(deformedReport.rebuiltStaticPlan);
    CHECK(deformedReport.refreshedGeometry);
    CHECK_FALSE(deformedReport.refreshedDensity);

    const double expectedScale = std::pow(1.0 + affineScale, 3);
    const double measuredScale = massOperator.GetCurrentMass() / undeformedMass;

    INFO("Expected affine mass scale = " << expectedScale);
    INFO("Measured affine mass scale = " << measuredScale);
    CHECK(mass_normalization_test_utils::relative_error(measuredScale, expectedScale) < 5e-7);
}

TEST_CASE(
    "Prepared Mass Normalization Density Jacobian Matches Centered Difference",
    tags::barotrope &tags::prepared &tags::jacobian &tags::accuracy
) {
    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());

    mfem::Vector density                = mass_normalization_test_utils::make_density(f, 0.31);
    const mfem::Vector densityDirection = mass_normalization_test_utils::make_density_direction(f, 0.67);
    const mfem::Vector displacement     = mass_normalization_test_utils::make_displacement_direction(f, 0.43);

    auto dependencies                   = mass_normalization_test_utils::make_dependencies();

    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );
    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, density, displacement, dependencies);

    mean_field::operators::PreparedMassNormalizationOperator massOperator(f, *f.domainMapperStateless, gravityContext);
    massOperator.Prepare({.targetMass = 1.23}, dependencies);

    mfem::Vector analyticAction;
    massOperator.ApplyDensityJacobianAction(densityDirection, analyticAction);

    constexpr double epsilon = 1.0e-3;
    mfem::Vector densityPlus(density);
    densityPlus.Add(epsilon, densityDirection);
    ++dependencies.density.revision;
    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, densityPlus, displacement, dependencies);
    massOperator.Prepare({.targetMass = 1.23}, dependencies);
    const double residualPlus = mass_normalization_test_utils::residual_value(massOperator);

    mfem::Vector densityMinus(density);
    densityMinus.Add(-epsilon, densityDirection);
    ++dependencies.density.revision;
    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, densityMinus, displacement, dependencies);
    massOperator.Prepare({.targetMass = 1.23}, dependencies);
    const double residualMinus    = mass_normalization_test_utils::residual_value(massOperator);

    const double finiteDifference = (residualPlus - residualMinus) / (2.0 * epsilon);

    INFO("Density action = " << analyticAction(0));
    INFO("Density centered difference = " << finiteDifference);
    CHECK(mass_normalization_test_utils::relative_error(analyticAction(0), finiteDifference) < 3.0e-8);
}

TEST_CASE(
    "Prepared Mass Normalization Geometry Jacobian Matches Centered Difference",
    tags::barotrope &tags::prepared &tags::jacobian &tags::geometry
) {
    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());

    const mfem::Vector density               = mass_normalization_test_utils::make_density(f, 0.37);
    mfem::Vector displacement                = mass_normalization_test_utils::make_displacement_direction(f, 0.51);
    const mfem::Vector displacementDirection = mass_normalization_test_utils::make_displacement_direction(f, -0.79);

    auto dependencies                        = mass_normalization_test_utils::make_dependencies();

    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );
    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, density, displacement, dependencies);

    mean_field::operators::PreparedMassNormalizationOperator massOperator(f, *f.domainMapperStateless, gravityContext);
    massOperator.Prepare({.targetMass = 1.11}, dependencies);

    mfem::Vector analyticAction;
    massOperator.ApplyDisplacementJacobianAction(displacementDirection, analyticAction);

    constexpr double epsilon = 1.0e-6;
    mfem::Vector displacementPlus(displacement);
    displacementPlus.Add(epsilon, displacementDirection);
    ++dependencies.displacement.revision;
    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, density, displacementPlus, dependencies);
    massOperator.Prepare({.targetMass = 1.11}, dependencies);
    const double residualPlus = mass_normalization_test_utils::residual_value(massOperator);

    mfem::Vector displacementMinus(displacement);
    displacementMinus.Add(-epsilon, displacementDirection);
    ++dependencies.displacement.revision;
    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, density, displacementMinus, dependencies);
    massOperator.Prepare({.targetMass = 1.11}, dependencies);
    const double residualMinus    = mass_normalization_test_utils::residual_value(massOperator);

    const double finiteDifference = (residualPlus - residualMinus) / (2.0 * epsilon);

    INFO("Geometry action = " << analyticAction(0));
    INFO("Geometry centered difference = " << finiteDifference);
    CHECK(mass_normalization_test_utils::relative_error(analyticAction(0), finiteDifference) < 2.0e-7);
}

TEST_CASE(
    "Prepared Mass Normalization Selectively Refreshes Its Cached State",
    tags::barotrope &tags::prepared &tags::contexts &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());

    mfem::Vector density                   = mass_normalization_test_utils::make_density(f, 0.29);
    mfem::Vector displacement              = mass_normalization_test_utils::make_displacement_direction(f, 0.41);
    auto dependencies                      = mass_normalization_test_utils::make_dependencies();

    std::uint64_t gravityPotentialRevision = 17;

    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );
    mass_normalization_test_utils::prepare_gravity_context(
        gravityContext, f, density, displacement, dependencies, 13, gravityPotentialRevision
    );

    mean_field::operators::PreparedMassNormalizationOperator massOperator(f, *f.domainMapperStateless, gravityContext);

    massOperator.Prepare({.targetMass = 1.0}, dependencies);
    const std::uint64_t preparationCount = massOperator.GetPreparationCount();

    const auto repeated                  = massOperator.Prepare({.targetMass = 1.0}, dependencies);
    CHECK_FALSE(repeated.DidAnyWork());
    CHECK(massOperator.GetPreparationCount() == preparationCount);

    ++gravityPotentialRevision;
    mass_normalization_test_utils::prepare_gravity_context(
        gravityContext, f, density, displacement, dependencies, 13, gravityPotentialRevision
    );

    const auto potentialOnly = massOperator.Prepare({.targetMass = 1.0}, dependencies);
    CHECK_FALSE(potentialOnly.DidAnyWork());

    const double residualBeforeTargetChange = mass_normalization_test_utils::residual_value(massOperator);
    const double massBeforeTargetChange     = massOperator.GetCurrentMass();

    ++dependencies.targetMass.revision;
    const auto targetOnly = massOperator.Prepare({.targetMass = 1.4}, dependencies);
    CHECK(targetOnly.updatedTargetMass);
    CHECK(targetOnly.assembledResidual);
    CHECK_FALSE(targetOnly.rebuiltStaticPlan);
    CHECK_FALSE(targetOnly.refreshedGeometry);
    CHECK_FALSE(targetOnly.refreshedDensity);
    CHECK(massOperator.GetCurrentMass() == massBeforeTargetChange);
    CHECK(
        mass_normalization_test_utils::relative_error(
            mass_normalization_test_utils::residual_value(massOperator) - residualBeforeTargetChange, -0.4
        ) < 2.0e-15
    );

    const double massBeforeDensityChange = massOperator.GetCurrentMass();
    density                              = mass_normalization_test_utils::make_density(f, 0.83);
    ++dependencies.density.revision;
    mass_normalization_test_utils::prepare_gravity_context(
        gravityContext, f, density, displacement, dependencies, 13, gravityPotentialRevision
    );

    const auto densityOnly = massOperator.Prepare({.targetMass = 1.4}, dependencies);
    CHECK(densityOnly.refreshedDensity);
    CHECK(densityOnly.assembledResidual);
    CHECK_FALSE(densityOnly.refreshedGeometry);
    CHECK(massOperator.GetCurrentMass() != massBeforeDensityChange);

    displacement = mass_normalization_test_utils::make_displacement_direction(f, 0.87);
    ++dependencies.displacement.revision;
    mass_normalization_test_utils::prepare_gravity_context(
        gravityContext, f, density, displacement, dependencies, 13, gravityPotentialRevision
    );

    const auto geometryOnly = massOperator.Prepare({.targetMass = 1.4}, dependencies);
    CHECK(geometryOnly.refreshedGeometry);
    CHECK(geometryOnly.assembledResidual);
    CHECK_FALSE(geometryOnly.refreshedDensity);
}

TEST_CASE(
    "Prepared Mass Normalization Complete Action And Coupled Routing Are Exact",
    tags::barotrope &tags::prepared &tags::jacobian &tags::mfem_operators
) {
    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());

    const mfem::Vector density               = mass_normalization_test_utils::make_density(f, 0.47);
    const mfem::Vector displacement          = mass_normalization_test_utils::make_displacement_direction(f, 0.57);
    const mfem::Vector densityDirection      = mass_normalization_test_utils::make_density_direction(f, 0.71);
    const mfem::Vector displacementDirection = mass_normalization_test_utils::make_displacement_direction(f, -0.63);

    const auto dependencies                  = mass_normalization_test_utils::make_dependencies();

    mean_field::operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        f, *f.domainMapperStateless
    );
    mass_normalization_test_utils::prepare_gravity_context(gravityContext, f, density, displacement, dependencies);

    mean_field::operators::PreparedMassNormalizationOperator massOperator(f, *f.domainMapperStateless, gravityContext);
    massOperator.Prepare({.targetMass = 1.19}, dependencies);

    mfem::Vector densityAction;
    mfem::Vector displacementAction;
    mfem::Vector completeAction;

    massOperator.ApplyDensityJacobianAction(densityDirection, densityAction);
    massOperator.ApplyDisplacementJacobianAction(displacementDirection, displacementAction);
    massOperator.ApplyCompleteJacobianAction(densityDirection, displacementDirection, completeAction);

    CHECK(
        mass_normalization_test_utils::relative_error(completeAction(0), densityAction(0) + displacementAction(0)) <
        2.0e-15
    );

    const auto layout = mass_normalization_test_utils::make_layout(f);
    mean_field::operators::PreparedMassNormalizationJacobianOperator adapter(layout, massOperator);

    mfem::Vector direction(layout.value_offsets().Last());
    direction = 0.0;

    for (int entry = 0; entry < densityDirection.Size(); ++entry) {
        direction(layout.offset(mass_normalization_test_utils::densityValue) + entry) = densityDirection(entry);
    }

    for (int entry = 0; entry < displacementDirection.Size(); ++entry) {
        direction(layout.offset(mass_normalization_test_utils::displacementValue) + entry) =
            displacementDirection(entry);
    }

    mfem::Vector coupledAction;
    adapter.Mult(direction, coupledAction);

    const int massOffset = layout.offset(mass_normalization_test_utils::massResidual);

    REQUIRE(coupledAction.Size() == layout.residual_offsets().Last());
    CHECK(coupledAction(massOffset) == completeAction(0));

    for (int entry = 0; entry < coupledAction.Size(); ++entry) {
        if (entry != massOffset) {
            CHECK(coupledAction(entry) == 0.0);
        }
    }

    CHECK(&massOperator.GetFEM() == &f);
    CHECK(&massOperator.GetGravityContext() == &gravityContext);
    CHECK(adapter.GetLayout().residual_offsets().Last() == layout.residual_offsets().Last());
}