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

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

import mean_field;
import test_helpers;

namespace angular_momentum_test_utils {
    [[nodiscard]] mean_field::operators::AngularMomentumDependencies makeDependencies() {
        return {
            .discretization = {.identity = 15013, .revision = 3},
            .density = {.identity = 15017, .revision = 5},
            .displacement = {.identity = 15031, .revision = 7},
            .rotation = {.identity = 15053, .revision = 11}
        };
    }

    [[nodiscard]] mean_field::operators::context::gravity_field::GravityFieldRevisions makeGravityRevisions(
        const mean_field::operators::AngularMomentumDependencies &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 prepareGravityContext(
        mean_field::operators::context::gravity_field::GravityFieldLinearizationContext &context,
        const mean_field::fem::FEM &finiteElements,
        const mfem::Vector &density,
        const mfem::Vector &displacement,
        const mean_field::operators::AngularMomentumDependencies &dependencies,
        const std::uint64_t gravityGradientRevision = 13,
        const std::uint64_t gravityPotentialRevision = 17
    ) {
        mfem::Vector gravityGradient(finiteElements.gravityFluxFes->GetTrueVSize());
        mfem::Vector gravityPotential(finiteElements.gravityPotentialFes->GetTrueVSize());
        gravityGradient = 0.0;
        gravityPotential = 0.0;
        context.Prepare(
            {.density = context.GetDensityMap().gather(density),
             .displacement = context.GetDisplacementMap().gather(displacement),
             .gravity_gradient = context.GetGravityGradientMap().gather(gravityGradient),
             .gravity_potential = context.GetGravityPotentialMap().gather(gravityPotential)},
            makeGravityRevisions(dependencies, gravityGradientRevision, gravityPotentialRevision)
        );
    }

    [[nodiscard]] mfem::Vector projectDensity(
        const mean_field::fem::FEM &finiteElements,
        const double phase
    ) {
        mfem::ParGridFunction field(finiteElements.densityFes.get());
        mfem::FunctionCoefficient coefficient([phase](const mfem::Vector &position) {
            return 0.94 + 0.08 * std::sin(0.71 * position(0) + phase) +
                   0.05 * std::cos(0.63 * position(1) - phase) + 0.03 * position(2) * position(2);
        });
        field.ProjectCoefficient(coefficient);
        mfem::Vector result;
        field.GetTrueDofs(result);
        return result;
    }

    [[nodiscard]] mfem::Vector projectDensityDirection(
        const mean_field::fem::FEM &finiteElements,
        const double phase
    ) {
        mfem::ParGridFunction field(finiteElements.densityFes.get());
        mfem::FunctionCoefficient coefficient([phase](const mfem::Vector &position) {
            return 0.17 * std::sin(0.83 * position(0) + phase) -
                   0.12 * std::cos(0.79 * position(1) - phase) + 0.06 * position(2);
        });
        field.ProjectCoefficient(coefficient);
        mfem::Vector result;
        field.GetTrueDofs(result);
        return result;
    }

    [[nodiscard]] mfem::Vector projectConstantDensity(
        const mean_field::fem::FEM &finiteElements,
        const double value,
        mfem::ParGridFunction *fieldOutput = nullptr
    ) {
        mfem::ParGridFunction field(finiteElements.densityFes.get());
        mfem::ConstantCoefficient coefficient(value);
        field.ProjectCoefficient(coefficient);
        if (fieldOutput != nullptr) {
            *fieldOutput = field;
        }
        mfem::Vector result;
        field.GetTrueDofs(result);
        return result;
    }

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

    [[nodiscard]] mfem::Vector projectDisplacementDirection(
        const mean_field::fem::FEM &finiteElements,
        const double scale
    ) {
        mfem::ParGridFunction field(finiteElements.displacementFes.get());
        mfem::VectorFunctionCoefficient coefficient(
            finiteElements.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]] double residual(const mean_field::operators::PreparedAngularMomentumOperator &operation) {
        mfem::Vector value;
        operation.BuildResidual(value);
        REQUIRE(value.Size() == 1);
        return value(0);
    }

    [[nodiscard]] double relativeError(const double actual, const double expected) {
        return std::abs(actual - expected) /
               std::max({std::abs(actual), std::abs(expected), 100.0 * std::numeric_limits<double>::epsilon()});
    }
} // namespace angular_momentum_test_utils

TEST_CASE(
    "Prepared Angular Momentum Satisfies Moment Scaling And The Parallel Axis Theorem",
    "[fixed-angular-momentum][physics][analytic]"
) {
    using namespace mean_field;
    using Catch::Approx;
    using Operator = operators::PreparedAngularMomentumOperator;

    STATIC_CHECK_FALSE(std::is_copy_constructible_v<Operator>);
    STATIC_CHECK_FALSE(std::is_move_constructible_v<Operator>);

    utils::Args arguments = test_utils::setup_args();
    fem::FEM finiteElements = fem::setup_fem(arguments.mesh_file, arguments, 0);
    REQUIRE(finiteElements.okay());

    constexpr double densityValue = 1.37;
    constexpr double angularVelocity = 0.73;
    constexpr double targetAngularMomentum = 0.41;
    mfem::ParGridFunction densityField(finiteElements.densityFes.get());
    const mfem::Vector density = angular_momentum_test_utils::projectConstantDensity(
        finiteElements,
        densityValue,
        &densityField
    );
    mfem::Vector displacement(finiteElements.displacementFes->GetTrueVSize());
    displacement = 0.0;
    finiteElements.displacement->SetFromTrueDofs(displacement);

    auto dependencies = angular_momentum_test_utils::makeDependencies();
    operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        finiteElements,
        *finiteElements.domainMapperStateless
    );
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacement,
        dependencies
    );

    const models::CompiledFixedAngularMomentum originConstraint = models::compileConstraint(
        integral::FixedAngularMomentum({
            .Jtotal = dimensions::AngularMomentumValue{targetAngularMomentum},
            .axis = {0.0, 0.0, 4.0}
        })
    );
    Operator origin(
        finiteElements,
        *finiteElements.domainMapperStateless,
        gravityContext,
        originConstraint
    );
    const auto initial = origin.Prepare(angularVelocity, dependencies);
    CHECK(initial.rebuiltStaticPlan);
    CHECK(initial.refreshedGeometry);
    CHECK(initial.refreshedDensity);
    CHECK(initial.updatedAngularVelocity);
    CHECK(initial.assembledResidual);

    const double independentMoment = analysis::get_moment_of_inertia(finiteElements, densityField);
    CHECK(angular_momentum_test_utils::relativeError(origin.GetMomentOfInertia(), independentMoment) < 2.0e-13);
    CHECK(origin.GetCurrentAngularMomentum() ==
          Approx(angularVelocity * origin.GetMomentOfInertia()).epsilon(2.0e-15));
    CHECK(angular_momentum_test_utils::residual(origin) ==
          Approx(angularVelocity * origin.GetMomentOfInertia() - targetAngularMomentum).epsilon(2.0e-15));

    const auto report = origin.GetConstraintReport();
    CHECK(report.targetAngularMomentum == targetAngularMomentum);
    CHECK(report.achievedAngularMomentum == origin.GetCurrentAngularMomentum());
    CHECK(report.momentOfInertia == origin.GetMomentOfInertia());
    CHECK(report.angularVelocity == angularVelocity);
    const physics::RigidRotation rotation = origin.GetRotation();
    CHECK(rotation.angular_velocity()(0) == 0.0);
    CHECK(rotation.angular_velocity()(1) == 0.0);
    CHECK(rotation.angular_velocity()(2) == angularVelocity);

    constexpr double affineScale = 0.086;
    const mfem::Vector affineDisplacement =
        angular_momentum_test_utils::projectAffineDisplacement(finiteElements, affineScale);
    ++dependencies.displacement.revision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        affineDisplacement,
        dependencies
    );
    const auto affine = origin.Prepare(angularVelocity, dependencies);
    CHECK(affine.refreshedGeometry);
    CHECK_FALSE(affine.refreshedDensity);
    const double expectedAffineRatio = std::pow(1.0 + affineScale, 5);
    const double measuredAffineRatio = origin.GetMomentOfInertia() / independentMoment;
    INFO("Expected homothetic I ratio = " << expectedAffineRatio);
    INFO("Measured homothetic I ratio = " << measuredAffineRatio);
    CHECK(angular_momentum_test_utils::relativeError(measuredAffineRatio, expectedAffineRatio) < 7.0e-7);

    displacement = 0.0;
    ++dependencies.displacement.revision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacement,
        dependencies
    );
    origin.Prepare(angularVelocity, dependencies);
    constexpr std::array<double, 3> shiftedCenter{0.27, -0.19, 0.31};
    Operator shifted(
        finiteElements,
        *finiteElements.domainMapperStateless,
        gravityContext,
        models::compileConstraint(integral::FixedAngularMomentum({
            .Jtotal = dimensions::AngularMomentumValue{targetAngularMomentum},
            .axis = {0.0, 0.0, 1.0},
            .center = shiftedCenter
        }))
    );
    shifted.Prepare(angularVelocity, dependencies);

    const double mass = analysis::domain_integrate_grid_function(
        finiteElements,
        densityField,
        utils::DOMAINS::STELLAR,
        mapping::COORDINATE_SPACE::PHYSICAL
    );
    const mfem::Vector centerOfMass = analysis::get_com(finiteElements, densityField);
    const double expectedShiftedMoment = origin.GetMomentOfInertia() +
                                         mass * (shiftedCenter[0] * shiftedCenter[0] +
                                                 shiftedCenter[1] * shiftedCenter[1]) -
                                         2.0 * mass * (shiftedCenter[0] * centerOfMass(0) +
                                                       shiftedCenter[1] * centerOfMass(1));
    CHECK(angular_momentum_test_utils::relativeError(shifted.GetMomentOfInertia(), expectedShiftedMoment) < 3.0e-13);
}

TEST_CASE(
    "Prepared Angular Momentum Jacobian Matches Density Geometry And Angular Velocity Differences",
    "[fixed-angular-momentum][jacobian][accuracy]"
) {
    using namespace mean_field;

    utils::Args arguments = test_utils::setup_args();
    fem::FEM finiteElements = fem::setup_fem(arguments.mesh_file, arguments, 0);
    REQUIRE(finiteElements.okay());

    const mfem::Vector density = angular_momentum_test_utils::projectDensity(finiteElements, 0.31);
    const mfem::Vector densityDirection =
        angular_momentum_test_utils::projectDensityDirection(finiteElements, 0.67);
    const mfem::Vector displacement =
        angular_momentum_test_utils::projectDisplacementDirection(finiteElements, 0.43);
    const mfem::Vector displacementDirection =
        angular_momentum_test_utils::projectDisplacementDirection(finiteElements, -0.79);
    constexpr double angularVelocity = 0.63;
    constexpr double angularVelocityDirection = -0.37;

    auto dependencies = angular_momentum_test_utils::makeDependencies();
    operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        finiteElements,
        *finiteElements.domainMapperStateless
    );
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacement,
        dependencies
    );
    operators::PreparedAngularMomentumOperator operation(
        finiteElements,
        *finiteElements.domainMapperStateless,
        gravityContext,
        models::compileConstraint(
            integral::FixedAngularMomentum({.Jtotal = dimensions::AngularMomentumValue{0.81}})
        )
    );
    operation.Prepare(angularVelocity, dependencies);

    const mfem::Vector reducedDensityDirection = gravityContext.GetDensityMap().gather(densityDirection);
    const mfem::Vector reducedDisplacementDirection =
        gravityContext.GetDisplacementMap().gather(displacementDirection);
    mfem::Vector densityAction;
    mfem::Vector geometryAction;
    mfem::Vector angularVelocityAction;
    mfem::Vector completeAction;
    operation.ApplyDensityJacobianAction(reducedDensityDirection, densityAction);
    operation.ApplyDisplacementJacobianAction(reducedDisplacementDirection, geometryAction);
    operation.ApplyAngularVelocityJacobianAction(angularVelocityDirection, angularVelocityAction);
    operation.ApplyCompleteJacobianAction(
        reducedDensityDirection,
        reducedDisplacementDirection,
        angularVelocityDirection,
        completeAction
    );
    CHECK(angular_momentum_test_utils::relativeError(
              completeAction(0),
              densityAction(0) + geometryAction(0) + angularVelocityAction(0)
          ) < 3.0e-15);
    CHECK(angularVelocityAction(0) ==
          Catch::Approx(operation.GetMomentOfInertia() * angularVelocityDirection).epsilon(2.0e-15));

    constexpr double angularStep = 1.0e-6;
    ++dependencies.rotation.revision;
    operation.Prepare(angularVelocity + angularStep * angularVelocityDirection, dependencies);
    const double angularPlus = angular_momentum_test_utils::residual(operation);
    ++dependencies.rotation.revision;
    operation.Prepare(angularVelocity - angularStep * angularVelocityDirection, dependencies);
    const double angularMinus = angular_momentum_test_utils::residual(operation);
    const double angularDifference = (angularPlus - angularMinus) / (2.0 * angularStep);
    CHECK(angular_momentum_test_utils::relativeError(angularVelocityAction(0), angularDifference) < 2.0e-10);

    constexpr double densityStep = 1.0e-3;
    mfem::Vector densityPlus(density);
    densityPlus.Add(densityStep, densityDirection);
    ++dependencies.density.revision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        densityPlus,
        displacement,
        dependencies
    );
    operation.Prepare(angularVelocity, dependencies);
    const double densityPlusResidual = angular_momentum_test_utils::residual(operation);
    mfem::Vector densityMinus(density);
    densityMinus.Add(-densityStep, densityDirection);
    ++dependencies.density.revision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        densityMinus,
        displacement,
        dependencies
    );
    operation.Prepare(angularVelocity, dependencies);
    const double densityMinusResidual = angular_momentum_test_utils::residual(operation);
    const double densityDifference = (densityPlusResidual - densityMinusResidual) / (2.0 * densityStep);
    CHECK(angular_momentum_test_utils::relativeError(densityAction(0), densityDifference) < 4.0e-8);

    constexpr double geometryStep = 1.0e-6;
    mfem::Vector displacementPlus(displacement);
    displacementPlus.Add(geometryStep, displacementDirection);
    ++dependencies.density.revision;
    ++dependencies.displacement.revision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacementPlus,
        dependencies
    );
    operation.Prepare(angularVelocity, dependencies);
    const double geometryPlusResidual = angular_momentum_test_utils::residual(operation);
    mfem::Vector displacementMinus(displacement);
    displacementMinus.Add(-geometryStep, displacementDirection);
    ++dependencies.displacement.revision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacementMinus,
        dependencies
    );
    operation.Prepare(angularVelocity, dependencies);
    const double geometryMinusResidual = angular_momentum_test_utils::residual(operation);
    const double geometryDifference = (geometryPlusResidual - geometryMinusResidual) / (2.0 * geometryStep);
    INFO("Density angular-momentum derivative error = " <<
         angular_momentum_test_utils::relativeError(densityAction(0), densityDifference));
    INFO("Geometry angular-momentum derivative error = " <<
         angular_momentum_test_utils::relativeError(geometryAction(0), geometryDifference));
    CHECK(angular_momentum_test_utils::relativeError(geometryAction(0), geometryDifference) < 4.0e-7);
}

TEST_CASE(
    "Prepared Angular Momentum Refreshes Only Changed Runtime Data",
    "[fixed-angular-momentum][prepared][lifecycle]"
) {
    using namespace mean_field;

    utils::Args arguments = test_utils::setup_args();
    fem::FEM finiteElements = fem::setup_fem(arguments.mesh_file, arguments, 0);
    REQUIRE(finiteElements.okay());
    mfem::Vector density = angular_momentum_test_utils::projectDensity(finiteElements, 0.29);
    mfem::Vector displacement =
        angular_momentum_test_utils::projectDisplacementDirection(finiteElements, 0.41);
    auto dependencies = angular_momentum_test_utils::makeDependencies();
    std::uint64_t gravityPotentialRevision = 17;

    operators::context::gravity_field::GravityFieldLinearizationContext gravityContext(
        finiteElements,
        *finiteElements.domainMapperStateless
    );
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacement,
        dependencies,
        13,
        gravityPotentialRevision
    );
    operators::PreparedAngularMomentumOperator operation(
        finiteElements,
        *finiteElements.domainMapperStateless,
        gravityContext,
        models::compileConstraint(
            integral::FixedAngularMomentum({.Jtotal = dimensions::AngularMomentumValue{0.71}})
        )
    );
    operation.Prepare(0.52, dependencies);
    const auto preparationCount = operation.GetPreparationCount();
    const double moment = operation.GetMomentOfInertia();

    const auto repeated = operation.Prepare(0.52, dependencies);
    CHECK_FALSE(repeated.DidAnyWork());
    CHECK(operation.GetPreparationCount() == preparationCount);

    ++gravityPotentialRevision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacement,
        dependencies,
        13,
        gravityPotentialRevision
    );
    const auto unrelatedPotential = operation.Prepare(0.52, dependencies);
    CHECK_FALSE(unrelatedPotential.DidAnyWork());

    const double residualBeforeRotation = angular_momentum_test_utils::residual(operation);
    ++dependencies.rotation.revision;
    const auto rotationOnly = operation.Prepare(0.81, dependencies);
    CHECK(rotationOnly.updatedAngularVelocity);
    CHECK(rotationOnly.assembledResidual);
    CHECK_FALSE(rotationOnly.refreshedDensity);
    CHECK_FALSE(rotationOnly.refreshedGeometry);
    CHECK(operation.GetMomentOfInertia() == moment);
    CHECK(angular_momentum_test_utils::residual(operation) - residualBeforeRotation ==
          Catch::Approx((0.81 - 0.52) * moment).epsilon(3.0e-15));

    density = angular_momentum_test_utils::projectDensity(finiteElements, 0.83);
    ++dependencies.density.revision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacement,
        dependencies,
        13,
        gravityPotentialRevision
    );
    const auto densityOnly = operation.Prepare(0.81, dependencies);
    CHECK(densityOnly.refreshedDensity);
    CHECK_FALSE(densityOnly.refreshedGeometry);
    CHECK_FALSE(densityOnly.updatedAngularVelocity);

    displacement = angular_momentum_test_utils::projectDisplacementDirection(finiteElements, 0.87);
    ++dependencies.displacement.revision;
    angular_momentum_test_utils::prepareGravityContext(
        gravityContext,
        finiteElements,
        density,
        displacement,
        dependencies,
        13,
        gravityPotentialRevision
    );
    const auto geometryOnly = operation.Prepare(0.81, dependencies);
    CHECK(geometryOnly.refreshedGeometry);
    CHECK_FALSE(geometryOnly.refreshedDensity);
    CHECK_FALSE(geometryOnly.updatedAngularVelocity);
}