#include <cmath>
#include <limits>
#include <catch2/catch_approx.hpp>

#include "serif/equations/angular_momentum.hpp"
#include "serif/equations/barotropic_density.hpp"
#include "serif/equations/primal_gravity.hpp"
#include "serif/equations/mass_normalization.hpp"
#include "serif/tests/equations/gravity_port.hpp"

namespace {
    using namespace serif;
    using namespace serif::tests::equations;
    using namespace serif::tests::equations::gravity;
    using Equation = serif::equations::AngularMomentumConstraint;
    using Prepared = VolumeEquationFixture<Equation, Schema>;
    using Vector = ::mfem::Vector;
    using Axis = serif::equations::RigidRotation::Vector;

    struct Values {
        Vector density, omega, displacement;
    };

    auto state(GravityFixture& f, const Values& values) {
        return backend::make_term_coefficients(backend::bind(f.density_binding, values.density), backend::bind(f.angular_velocity_binding, values.omega));
    }

    auto make_operator(GravityFixture& f, double target, Axis axis = {0.0, 0.0, 1.0}, Axis center = {}) {
        return Prepared{Equation{target, axis, center}, f.angular_velocity_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    }

    template <typename OperatorT> double residual(GravityFixture& f, OperatorT& op) {
        Vector result;
        op.residual(f.workspace, result);
        REQUIRE(result.Size() == 1);
        return result[0];
    }

    template <typename OperatorT> double residual(GravityFixture& f, OperatorT& op, const Values& values) {
        op.prepare(state(f, values), values.displacement, f.workspace);
        return residual(f, op);
    }

    template <typename OperatorT> double action(GravityFixture& f, OperatorT& op, const Values& direction) {
        Vector result;
        op.apply_jacobian(state(f, direction), direction.displacement, f.workspace, result);
        REQUIRE(result.Size() == 1);
        return result[0];
    }

    Values zeros(const GravityFixture& f) {
        return {constant(f.density_binding.size(), 0.0), constant(1, 0.0), constant(f.displacement_binding.size(), 0.0)};
    }

    // MeanField tests/operators/prepared_angular_momentum.cpp: preserve its projected profiles.
    Vector density(GravityFixture& f, double phase, bool direction = false) {
        return f.project(f.density_binding, [=](const Vector& x) {
            if (direction) { return 0.17 * std::sin(0.83 * x[0] + phase) - 0.12 * std::cos(0.79 * x[1] - phase) + 0.06 * x[2]; }
            return 0.94 + 0.08 * std::sin(0.71 * x[0] + phase) + 0.05 * std::cos(0.63 * x[1] - phase) + 0.03 * x[2] * x[2];
        });
    }

    Vector displacement(GravityFixture& f, double scale) {
        return gravity_port::project_vector(f.displacement_binding, [=](const Vector& x, Vector& value) {
            value.SetSize(3);
            value[0] = scale * (0.07 * x[0] + 0.018 * x[1] * x[2]);
            value[1] = scale * (-0.05 * x[1] + 0.013 * x[0] * x[2]);
            value[2] = scale * (0.04 * x[2] - 0.011 * x[0] * x[1]);
        });
    }

    double relative(double actual, double expected) {
        return std::abs(actual - expected) / std::max({std::abs(actual), std::abs(expected), 100.0 * std::numeric_limits<double>::epsilon()});
    }

    Values shifted(Values values, const Values& direction, double step) {
        values.density.Add(step, direction.density);
        values.omega.Add(step, direction.omega);
        values.displacement.Add(step, direction.displacement);
        return values;
    }

    struct Integrals {
        double moment = 0.0, mass = 0.0, action = 0.0;
        Axis first_moment{};
    };

    // Direct MFEM quadrature, independent of the equation and operator evaluators.
    // delta(I) = integral [delta(rho) s^2 + 2 rho r_perp . delta(x) + rho s^2 trace(A^-1 delta(A))] dV.
    Integrals reference(GravityFixture& f, const Values& values, Axis axis = {0.0, 0.0, 1.0}, Axis center = {}, const Values* direction = nullptr) {
        const double norm = std::hypot(axis[0], axis[1], axis[2]);
        for (double& component : axis) { component /= norm; }
        ::mfem::GridFunction rho(&f.density_space), u(&f.displacement_space), drho(&f.density_space), du(&f.displacement_space);
        rho = f.density_binding.scatter(values.density);
        u = f.displacement_binding.scatter(values.displacement);
        if (direction) {
            drho = f.density_binding.scatter(direction->density);
            du = f.displacement_binding.scatter(direction->displacement);
        }
        Integrals result;
        for (int element = 0; element < f.mesh.GetNE(); ++element) {
            if (!f.density_binding.supports(element)) { continue; }
            auto& transform = *f.mesh.GetElementTransformation(element);
            const auto& rule = ::mfem::IntRules.Get(f.mesh.GetElementGeometry(element), reference_order);
            for (int point = 0; point < rule.GetNPoints(); ++point) {
                const auto& ip = rule.IntPoint(point);
                transform.SetIntPoint(&ip);
                Vector x(3), displacement_value(3), dx(3);
                transform.Transform(ip, x);
                u.GetVectorValue(transform, ip, displacement_value);
                x += displacement_value;
                ::mfem::DenseMatrix A(3), dA(3), inverse(3);
                u.GetVectorGradient(transform, A);
                for (int i = 0; i < 3; ++i) { A(i, i) += 1.0; }
                double axial_distance = 0.0;
                for (int i = 0; i < 3; ++i) { axial_distance += axis[i] * (x[i] - center[i]); }
                Axis perpendicular{};
                double radius_squared = 0.0;
                for (int i = 0; i < 3; ++i) {
                    perpendicular[i] = x[i] - center[i] - axis[i] * axial_distance;
                    radius_squared += perpendicular[i] * perpendicular[i];
                }
                const double weight = ip.weight * transform.Weight() * A.Det();
                const double density_value = rho.GetValue(transform, ip);
                result.mass += weight * density_value;
                result.moment += weight * density_value * radius_squared;
                for (int i = 0; i < 3; ++i) { result.first_moment[i] += weight * density_value * x[i]; }
                if (direction) {
                    du.GetVectorValue(transform, ip, dx);
                    du.GetVectorGradient(transform, dA);
                    ::mfem::CalcInverse(A, inverse);
                    double trace = 0.0, radius_variation = 0.0;
                    for (int i = 0; i < 3; ++i) {
                        radius_variation += 2.0 * perpendicular[i] * dx[i];
                        for (int j = 0; j < 3; ++j) { trace += inverse(i, j) * dA(j, i); }
                    }
                    result.action += weight * (direction->omega[0] * density_value * radius_squared + values.omega[0] *
                        (drho.GetValue(transform, ip) * radius_squared + density_value * (radius_variation + radius_squared * trace)));
                }
            }
        }
        return result;
    }

    double moment(GravityFixture& f, Prepared& op) {
        auto direction = zeros(f);
        direction.omega = 1.0;
        return action(f, op, direction);
    }

    void check_columns(GravityFixture& f, Axis axis = {0.0, 0.0, 1.0}, Axis center = {}) {
        CAPTURE(f.mesh.GetNE(), f.mesh.GetElementGeometry(0), f.displacement_space.GetOrdering(), axis, center);
        const Values values{density(f, 0.31), constant(1, 0.63), displacement(f, 0.43)};
        const Values direction{density(f, 0.67, true), constant(1, -0.37), displacement(f, -0.79)};
        auto op = make_operator(f, 0.81, axis, center);
        (void)residual(f, op, values);
        const double complete = action(f, op, direction);
        CHECK(relative(complete, reference(f, values, axis, center, &direction).action) < 2.0e-11);
        auto density_direction = zeros(f), geometry_direction = zeros(f), omega_direction = zeros(f);
        density_direction.density = direction.density;
        geometry_direction.displacement = direction.displacement;
        omega_direction.omega = direction.omega;
        const double density_action = action(f, op, density_direction), geometry_action = action(f, op, geometry_direction), omega_action = action(f, op, omega_direction);
        CHECK(relative(complete, density_action + geometry_action + omega_action) < 3.0e-15);
        CHECK(omega_action == Catch::Approx(moment(f, op) * direction.omega[0]).epsilon(2.0e-15));
        backend::JacobianOperator adapter(op.cache());
        CHECK(adapter.Height() == 1);
        CHECK(adapter.Width() == direction.density.Size() + 1 + direction.displacement.Size());
        Vector result;
        adapter.Mult(packed(direction.density, direction.omega, direction.displacement), result);
        CHECK(result[0] == complete);

        // Preserve all three original difference steps and thresholds.
        const std::array<Values, 3> columns{density_direction, geometry_direction, omega_direction};
        const std::array<double, 3> actions{density_action, geometry_action, omega_action};
        const std::array<double, 3> steps{1.0e-3, 1.0e-6, 1.0e-6}, tolerances{4.0e-8, 4.0e-7, 2.0e-10};
        for (int column = 0; column < 3; ++column) {
            const double plus = residual(f, op, shifted(values, columns[column], steps[column]));
            const double minus = residual(f, op, shifted(values, columns[column], -steps[column]));
            const double difference = (plus - minus) / (2.0 * steps[column]);
            CAPTURE(column, actions[column], difference);
            CHECK(relative(actions[column], difference) < tolerances[column]);
        }
    }
}

TEST_CASE("Angular momentum retains moment scaling and the parallel axis theorem", "[equations][angular_momentum][integration]") {
    STATIC_CHECK_FALSE(std::is_copy_constructible_v<Prepared::Cache>);
    STATIC_CHECK(std::is_move_constructible_v<Prepared::Cache>);
    auto fixture = gravity_port::stellar_fixture();
    auto& f = *fixture;
    auto values = zeros(f);
    values.density = 1.37;
    values.omega = 0.73;
    auto op = make_operator(f, 0.41, {0.0, 0.0, 4.0});
    const double initial = residual(f, op, values), origin_moment = moment(f, op);
    const auto independent = reference(f, values);
    CHECK(relative(origin_moment, independent.moment) < 2.0e-13);
    CHECK(initial + 0.41 == Catch::Approx(0.73 * origin_moment).epsilon(2.0e-15));
    CHECK(initial == Catch::Approx(0.73 * origin_moment - 0.41).epsilon(2.0e-15));
    values.displacement = gravity_port::affine_displacement(f, 0.086);
    (void)residual(f, op, values);
    CHECK(relative(moment(f, op) / independent.moment, std::pow(1.0 + 0.086, 5)) < 7.0e-7);

    values.displacement = 0.0;
    const Axis center{0.27, -0.19, 0.31};
    auto shifted_op = make_operator(f, 0.41, {0.0, 0.0, 1.0}, center);
    (void)residual(f, shifted_op, values);
    const double expected = origin_moment + independent.mass * (center[0] * center[0] + center[1] * center[1]) -
        2.0 * (center[0] * independent.first_moment[0] + center[1] * independent.first_moment[1]);
    CHECK(relative(moment(f, shifted_op), expected) < 3.0e-13);
}

TEST_CASE("Angular momentum retains density geometry and angular velocity action accuracy", "[equations][angular_momentum][integration]") {
    for (const auto geometry : {::mfem::Element::TETRAHEDRON, ::mfem::Element::HEXAHEDRON}) {
        for (const auto ordering : {::mfem::Ordering::byNODES, ::mfem::Ordering::byVDIM}) {
            GravityFixture f(geometry, ordering);
            f.set_compactification();
            check_columns(f);
            check_columns(f, {0.17, -0.12, 0.43}, {0.031, -0.024, 0.018});
        }
    }
    auto fixture = gravity_port::stellar_fixture();
    check_columns(*fixture);
    check_columns(*fixture, {0.17, -0.12, 0.43}, {0.031, -0.024, 0.018});
}

TEST_CASE("Angular momentum snapshots and refreshed states preserve the old lifecycle numerics", "[equations][angular_momentum][integration]") {
    auto fixture = gravity_port::stellar_fixture();
    auto& f = *fixture;
    Values values{density(f, 0.29), constant(1, 0.52), displacement(f, 0.41)};
    const Values direction{density(f, 0.67, true), constant(1, -0.19), displacement(f, -0.53)};
    auto op = make_operator(f, 0.71);
    const double initial = residual(f, op, values), initial_moment = moment(f, op);
    auto check_action = [&] { CHECK(relative(action(f, op, direction), reference(f, values, {0.0, 0.0, 1.0}, {}, &direction).action) < 2.0e-11); };
    check_action();
    CHECK(residual(f, op, values) == initial);
    check_action();
    auto potential = f.potential();
    potential *= 1.7;
    op.prepare(backend::make_term_coefficients(backend::bind(f.density_binding, values.density), backend::bind(f.angular_velocity_binding, values.omega), backend::bind(f.potential_binding, potential)), values.displacement, f.workspace);
    CHECK(residual(f, op) == initial);
    check_action();

    values.omega = 0.81;
    CHECK(residual(f, op) == initial);
    CHECK(residual(f, op, values) - initial == Catch::Approx((0.81 - 0.52) * initial_moment).epsilon(3.0e-15));
    CHECK(moment(f, op) == initial_moment);
    check_action();
    const double before_density = residual(f, op);
    values.density = density(f, 0.83);
    CHECK(residual(f, op) == before_density);
    (void)residual(f, op, values);
    check_action();
    const double before_geometry = residual(f, op);
    values.displacement = displacement(f, 0.87);
    CHECK(residual(f, op) == before_geometry);
    (void)residual(f, op, values);
    check_action();
}

TEST_CASE("Angular momentum rejects inadmissible moments and angular speeds then recovers", "[equations][angular_momentum][unit]") {
    GravityFixture f;
    auto values = zeros(f);
    values.density = 0.91;
    values.omega = 0.52;
    auto op = make_operator(f, 0.71);
    const double initial = residual(f, op, values);
    for (const double omega : {0.52, 0.0, -0.52}) {
        values.omega = omega;
        values.density = -0.91;
        const auto count = op.preparation_count();
        const auto rejected = op.try_prepare(state(f, values), values.displacement, f.workspace);
        REQUIRE_FALSE(rejected.has_value());
        CHECK(rejected.error().reason == backend::PreparationRejectionReason::EQUATION);
        CHECK(op.preparation_count() == count);
        CHECK_FALSE(op.is_prepared());
        CHECK_THROWS_AS(op.prepare(state(f, values), values.displacement, f.workspace), std::domain_error);
        values.density = 0.91;
        values.omega = 0.52;
        CHECK(residual(f, op, values) == initial);
    }
    for (double bad : {std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::infinity()}) {
        values.omega = bad;
        const auto rejected = op.try_prepare(state(f, values), values.displacement, f.workspace);
        REQUIRE_FALSE(rejected.has_value());
        CHECK(rejected.error().reason == backend::PreparationRejectionReason::NON_FINITE);
        CHECK_FALSE(op.is_prepared());
        CHECK_THROWS_AS(op.prepare(state(f, values), values.displacement, f.workspace), std::domain_error);
        values.omega = 0.52;
        CHECK(residual(f, op, values) == initial);
    }
    // Local negative density is allowed when the total moment remains nonnegative.
    values.density = f.project(f.density_binding, [](const Vector& x) { return x[0] - 0.25; });
    REQUIRE(reference(f, values).moment > 0.0);
    CHECK(relative(residual(f, op, values), 0.52 * reference(f, values).moment - 0.71) < 2.0e-13);
}

TEST_CASE("Angular momentum target axis and center retain the old validity requirements", "[equations][angular_momentum][unit]") {
    for (double target : {-1.0, std::numeric_limits<double>::infinity(), std::numeric_limits<double>::quiet_NaN()}) {
        CHECK_THROWS_AS(Equation{target}, std::invalid_argument);
    }
    CHECK_NOTHROW(Equation{0.0});
    CHECK_THROWS_AS((Equation{1.0, {0.0, 0.0, 0.0}}), std::invalid_argument);
    CHECK_THROWS_AS((Equation{1.0, {0.0, 0.0, 1.0}, {0.0, std::numeric_limits<double>::quiet_NaN(), 0.0}}), std::invalid_argument);
}

TEST_CASE("Angular momentum rejects overflow in integrated moments and residuals", "[equations][angular_momentum][unit]") {
    // A long cylinder-like box keeps point values finite while their volume integral overflows.
    auto mesh = make_mesh(::mfem::Element::HEXAHEDRON);
    for (int vertex = 0; vertex < mesh.GetNV(); ++vertex) { mesh.GetVertex(vertex)[2] *= 100.0; }
    GravityFixture f(std::move(mesh));
    auto values = zeros(f);
    const double maximum = std::numeric_limits<double>::max();
    for (const double target : {0.0, maximum}) {
        auto op = make_operator(f, target);
        values.density = 1.0;
        values.omega = target == 0.0 ? 0.1 * maximum : -0.01 * maximum;
        const auto rejected = op.try_prepare(state(f, values), values.displacement, f.workspace);
        REQUIRE_FALSE(rejected.has_value());
        CHECK(rejected.error().reason == backend::PreparationRejectionReason::NON_FINITE);
        CHECK_FALSE(op.is_prepared());
    }
    auto op = make_operator(f, 0.0);
    values.density = 0.1 * maximum;
    values.omega = 0.0;
    const auto rejected = op.try_prepare(state(f, values), values.displacement, f.workspace);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().reason == backend::PreparationRejectionReason::EQUATION);
    CHECK_FALSE(op.is_prepared());
    values.density = 1.0;
    CHECK(residual(f, op, values) == 0.0);
}

TEST_CASE("Angular momentum handles zero speed signed speed and default quadrature", "[equations][angular_momentum][integration]") {
    for (const auto geometry : {::mfem::Element::TETRAHEDRON, ::mfem::Element::HEXAHEDRON}) {
        GravityFixture f(geometry);
        Values values{density(f, 0.31), constant(1, 0.0), displacement(f, 0.43)};
        Prepared op{Equation{0.41}, f.angular_velocity_binding, f.displacement_binding, f.compactification, f.mapper};
        CHECK(residual(f, op, values) == -0.41);
        CHECK(relative(moment(f, op), reference(f, values).moment) < 2.0e-13);
        const Values direction{density(f, 0.67, true), constant(1, 0.0), displacement(f, -0.79)};
        CHECK(action(f, op, direction) == 0.0);
        values.omega = -0.63;
        CHECK(relative(residual(f, op, values), -0.63 * reference(f, values).moment - 0.41) < 2.0e-13);
        values.density = 0.0;
        CHECK(residual(f, op, values) == -0.41);
        CHECK(moment(f, op) == 0.0);
    }
}

TEST_CASE("Angular momentum preserves physical admissibility", "[equations][angular_momentum][unit]") {
    GravityFixture f;
    Values values{density(f, 0.31), constant(1, 0.63), displacement(f, 0.43)};
    VolumeEquationFixture op{Equation{0.41}, f.angular_velocity_binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    CHECK(relative(residual(f, op, values), 0.63 * reference(f, values).moment - 0.41) < 2.0e-13);
    const Values direction{density(f, 0.67, true), constant(1, -0.37), displacement(f, -0.79)};
    CHECK(relative(action(f, op, direction), reference(f, values, {0.0, 0.0, 1.0}, {}, &direction).action) < 2.0e-11);
    values.density = -0.91;
    const auto rejected = op.try_prepare(state(f, values), values.displacement, f.workspace);
    REQUIRE_FALSE(rejected.has_value());
    CHECK(rejected.error().reason == backend::PreparationRejectionReason::EQUATION);
    CHECK_FALSE(op.is_prepared());
}

TEST_CASE("One angular speed drives all supported equilibrium rows consistently", "[equations][angular_momentum][integration]") {
    // Port the physical columns of MeanField stellar_equilibrium_system.cpp:253.
    // Root manifests, surface projection, and Lane-Emden seed generation are separate infrastructure.
    auto fixture = gravity_port::stellar_fixture();
    auto& f = *fixture;
    Values values{density(f, 0.31), constant(1, 0.63), displacement(f, 0.43)};
    auto direction = zeros(f);
    direction.omega = -0.37;
    const auto h = f.enthalpy(), phi = f.potential(), c = constant(1, 0.39);
    const auto zero_h = constant(h.Size(), 0.0), zero_phi = constant(phi.Size(), 0.0), zero_c = constant(1, 0.0);
    auto bind_all = [&](const Values& v, const Vector& enthalpy, const Vector& potential, const Vector& bernoulli) {
        return backend::make_term_coefficients(backend::bind(f.density_binding, v.density), backend::bind(f.angular_velocity_binding, v.omega),
            backend::bind(f.enthalpy_binding, enthalpy), backend::bind(f.potential_binding, potential),
            backend::bind(f.constant_binding, bernoulli));
    };
    const auto base = bind_all(values, h, phi, c), delta = bind_all(direction, zero_h, zero_phi, zero_c);
    auto make_row = [&](auto equation, const auto& binding) {
        return VolumeEquationFixture{equation, binding, f.displacement_binding, f.compactification, f.mapper, reference_policy()};
    };
    const Axis axis{0.0, 0.0, 3.0};
    const eos::models::Polytrope eos{1.0, 0.25};
    auto momentum = make_row(Equation{0.1, axis}, f.angular_velocity_binding);
    auto density_balance = make_row(serif::equations::BarotropicDensityBalance{eos, dimensions::PressureValue{0.0}, axis}, f.density_binding);
    auto mass = make_row(serif::equations::MassNormalization{1.0}, f.constant_binding);
    auto closure = make_row(::serif::tests::equations::BarotropicClosure{eos}, f.density_binding);
    auto poisson = make_row(serif::equations::PrimalPoissonBalance{four_pi_G}, f.potential_binding);

    auto check_rotating_row = [&](auto& row, double tolerance, bool zero_at_rest) {
        for (double omega : {0.63, 0.0}) {
            values.omega = omega;
            row.prepare(base, values.displacement, f.workspace);
            Vector actual, plus, minus;
            row.apply_jacobian(delta, direction.displacement, f.workspace, actual);
            if (omega == 0.0 && zero_at_rest) { CHECK(actual.Norml2() == 0.0); }
            else { CHECK(actual.Norml2() > 0.0); }
            constexpr double step = 1.0e-5;
            values.omega = omega + step * direction.omega[0];
            row.prepare(base, values.displacement, f.workspace);
            row.residual(f.workspace, plus);
            values.omega = omega - step * direction.omega[0];
            row.prepare(base, values.displacement, f.workspace);
            row.residual(f.workspace, minus);
            plus -= minus;
            plus /= 2.0 * step;
            const double scale = std::max({1.0, actual.Norml2(), plus.Norml2()});
            plus -= actual;
            CHECK(plus.Norml2() / scale < tolerance);
        }
    };
    check_rotating_row(momentum, 3.0e-10, false);
    check_rotating_row(density_balance, 3.0e-7, true);
    values.omega = 0.63;
    auto check_independent_row = [&](auto& row) {
        row.prepare(base, values.displacement, f.workspace);
        Vector actual;
        row.apply_jacobian(delta, direction.displacement, f.workspace, actual);
        CHECK(actual.Norml2() == 0.0);
    };
    check_independent_row(mass);
    check_independent_row(closure);
    check_independent_row(poisson);
}