#include <catch2/catch_test_macros.hpp>

#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <map>
#include <stdexcept>
#include <string>
#include <vector>

#include <mfem.hpp>
#include <mpi.h>

import experiment;
import experiment.stellar_null_space;
import mean_field;
import test_helpers;

namespace {
    struct DeterminantPolynomial final {
        double linear{0.0};
        double quadratic{0.0};
        double cubic{0.0};
    };

    struct CriticalAmplitude final {
        double magnitude{0.0};
        double determinant{1.0};
        bool searchLimitReached{false};
    };

    struct SymmetricFiniteDifferenceStep final {
        double step{0.0};
        double positiveMinimumDeterminant{0.0};
        double negativeMinimumDeterminant{0.0};
    };

    struct PolynomialRoots final {
        std::array<double, 3> values{
            std::numeric_limits<double>::quiet_NaN(), std::numeric_limits<double>::quiet_NaN(),
            std::numeric_limits<double>::quiet_NaN()
        };
        int count{0};
    };

    [[nodiscard]] double relative_difference(
        const mfem::Vector &computed,
        const mfem::Vector &reference,
        const MPI_Comm communicator
    ) {
        mfem::Vector difference(computed);
        difference -= reference;
        const double scale = std::max(
            {experiment::null_space::global_norm(computed, communicator),
             experiment::null_space::global_norm(reference, communicator), std::numeric_limits<double>::epsilon()}
        );
        return experiment::null_space::global_norm(difference, communicator) / scale;
    }

    void add_block_metrics(
        std::map<
            std::string,
            double> &metrics,
        const std::string &prefix,
        const std::array<
            double,
            6> &norms
    ) {
        for (std::size_t block = 0; block < norms.size(); ++block) {
            metrics.emplace(prefix + experiment::null_space::residualBlockNames[block] + "_norm", norms[block]);
        }
    }

    [[nodiscard]] std::vector<DeterminantPolynomial> collect_determinant_polynomials(
        const mean_field::fem::FEM &fem,
        const mfem::Vector &unitVolumeDirection
    ) {
        MFEM_VERIFY(fem.mesh->SpaceDimension() == 3, "The spherical-harmonic frequency probe requires 3D geometry.");

        mfem::ParGridFunction displacement(fem.displacementFes.get());
        displacement.SetFromTrueDofs(unitVolumeDirection);

        std::vector<DeterminantPolynomial> polynomials;
        polynomials.reserve(static_cast<std::size_t>(fem.mesh->GetNE()) * 64);
        for (int element = 0; element < fem.mesh->GetNE(); ++element) {
            mfem::ElementTransformation *transformation = fem.mesh->GetElementTransformation(element);
            const mfem::FiniteElement *finiteElement    = fem.displacementFes->GetFE(element);
            const int integrationOrder =
                std::max(finiteElement->GetOrder() + 2, 2 * fem.mesh->SpaceDimension() * finiteElement->GetOrder());
            const mfem::IntegrationRule &rule = mfem::IntRules.Get(transformation->GetGeometryType(), integrationOrder);

            for (int point = 0; point < rule.GetNPoints(); ++point) {
                transformation->SetIntPoint(&rule.IntPoint(point));
                mfem::DenseMatrix gradient;
                displacement.GetVectorGradient(*transformation, gradient);

                double trace        = 0.0;
                double traceSquared = 0.0;
                for (int row = 0; row < 3; ++row) {
                    trace += gradient(row, row);
                    for (int column = 0; column < 3; ++column) {
                        traceSquared += gradient(row, column) * gradient(column, row);
                    }
                }

                polynomials.push_back(
                    {.linear = trace, .quadratic = 0.5 * (trace * trace - traceSquared), .cubic = gradient.Det()}
                );
            }
        }
        return polynomials;
    }

    [[nodiscard]] double global_minimum_determinant(
        const std::vector<DeterminantPolynomial> &polynomials,
        const double amplitude,
        const MPI_Comm communicator
    ) {
        double localMinimum = std::numeric_limits<double>::infinity();
        for (const DeterminantPolynomial &polynomial : polynomials) {
            const double determinant =
                1.0 +
                amplitude * (polynomial.linear + amplitude * (polynomial.quadratic + amplitude * polynomial.cubic));
            localMinimum = std::min(localMinimum, determinant);
        }

        double globalMinimum = std::numeric_limits<double>::infinity();
        MPI_Allreduce(&localMinimum, &globalMinimum, 1, MPI_DOUBLE, MPI_MIN, communicator);
        return globalMinimum;
    }

    [[nodiscard]] double evaluate(
        const DeterminantPolynomial &polynomial,
        const double amplitude
    ) {
        return 1.0 +
               amplitude * (polynomial.linear + amplitude * (polynomial.quadratic + amplitude * polynomial.cubic));
    }

    void append_root(
        PolynomialRoots &roots,
        const double root
    ) {
        if (roots.count < static_cast<int>(roots.values.size()) && std::isfinite(root)) {
            roots.values[static_cast<std::size_t>(roots.count++)] = root;
        }
    }

    [[nodiscard]] PolynomialRoots real_roots(const DeterminantPolynomial &polynomial) {
        PolynomialRoots roots;
        const double coefficientScale =
            std::max({1.0, std::abs(polynomial.linear), std::abs(polynomial.quadratic), std::abs(polynomial.cubic)});
        const double tolerance = 64.0 * std::numeric_limits<double>::epsilon() * coefficientScale;

        if (std::abs(polynomial.cubic) <= tolerance) {
            if (std::abs(polynomial.quadratic) <= tolerance) {
                if (std::abs(polynomial.linear) > tolerance) {
                    append_root(roots, -1.0 / polynomial.linear);
                }
                return roots;
            }

            const double discriminant = polynomial.linear * polynomial.linear - 4.0 * polynomial.quadratic;
            const double discriminantTolerance =
                64.0 * std::numeric_limits<double>::epsilon() * std::max(1.0, polynomial.linear * polynomial.linear);
            if (discriminant < -discriminantTolerance) {
                return roots;
            }

            const double squareRoot      = std::sqrt(std::max(0.0, discriminant));
            const double stableNumerator = -0.5 * (polynomial.linear + std::copysign(squareRoot, polynomial.linear));
            if (stableNumerator == 0.0) {
                append_root(roots, -polynomial.linear / (2.0 * polynomial.quadratic));
            } else {
                append_root(roots, stableNumerator / polynomial.quadratic);
                if (squareRoot > std::sqrt(discriminantTolerance)) {
                    append_root(roots, 1.0 / stableNumerator);
                }
            }
            return roots;
        }

        const double quadratic       = polynomial.quadratic / polynomial.cubic;
        const double linear          = polynomial.linear / polynomial.cubic;
        const double constant        = 1.0 / polynomial.cubic;
        const double depressedLinear = linear - quadratic * quadratic / 3.0;
        const double depressedConstant =
            2.0 * quadratic * quadratic * quadratic / 27.0 - quadratic * linear / 3.0 + constant;
        const double halfConstant = 0.5 * depressedConstant;
        const double thirdLinear  = depressedLinear / 3.0;
        const double discriminant = halfConstant * halfConstant + thirdLinear * thirdLinear * thirdLinear;
        const double discriminantTolerance =
            128.0 * std::numeric_limits<double>::epsilon() *
            std::max({1.0, std::abs(halfConstant * halfConstant), std::abs(thirdLinear * thirdLinear * thirdLinear)});
        const double shift = quadratic / 3.0;

        if (discriminant > discriminantTolerance) {
            const double squareRoot = std::sqrt(discriminant);
            append_root(roots, std::cbrt(-halfConstant + squareRoot) + std::cbrt(-halfConstant - squareRoot) - shift);
        } else if (std::abs(depressedLinear) <= tolerance || thirdLinear >= 0.0) {
            append_root(roots, std::cbrt(-depressedConstant) - shift);
        } else {
            const double radius         = 2.0 * std::sqrt(std::max(0.0, -thirdLinear));
            const double cosineArgument = std::clamp(
                -halfConstant / std::sqrt(std::max(0.0, -thirdLinear * thirdLinear * thirdLinear)), -1.0, 1.0
            );
            const double phase              = std::acos(cosineArgument) / 3.0;
            constexpr double twoPiOverThree = 2.0943951023931954923;
            for (int root = 0; root < 3; ++root) {
                append_root(roots, radius * std::cos(phase - twoPiOverThree * static_cast<double>(root)) - shift);
            }
        }

        for (int root = 0; root < roots.count; ++root) {
            double &value = roots.values[static_cast<std::size_t>(root)];
            for (int iteration = 0; iteration < 3; ++iteration) {
                const double derivative =
                    polynomial.linear + value * (2.0 * polynomial.quadratic + 3.0 * value * polynomial.cubic);
                if (std::abs(derivative) <= tolerance) {
                    break;
                }
                value -= evaluate(polynomial, value) / derivative;
            }
        }
        return roots;
    }

    [[nodiscard]] CriticalAmplitude find_critical_amplitude(
        const std::vector<DeterminantPolynomial> &polynomials,
        const double sign,
        const MPI_Comm communicator
    ) {
        constexpr double maximumSearchMagnitude = 0.5;
        MFEM_VERIFY(sign == 1.0 || sign == -1.0, "The critical-amplitude direction must be positive or negative.");

        double localCriticalMagnitude = std::numeric_limits<double>::infinity();
        for (const DeterminantPolynomial &polynomial : polynomials) {
            const PolynomialRoots roots = real_roots(polynomial);
            for (int root = 0; root < roots.count; ++root) {
                const double signedMagnitude = sign * roots.values[static_cast<std::size_t>(root)];
                if (signedMagnitude > 0.0) {
                    localCriticalMagnitude = std::min(localCriticalMagnitude, signedMagnitude);
                }
            }
        }

        double globalCriticalMagnitude = std::numeric_limits<double>::infinity();
        MPI_Allreduce(&localCriticalMagnitude, &globalCriticalMagnitude, 1, MPI_DOUBLE, MPI_MIN, communicator);

        if (!std::isfinite(globalCriticalMagnitude) || globalCriticalMagnitude > maximumSearchMagnitude) {
            return {
                .magnitude   = maximumSearchMagnitude,
                .determinant = global_minimum_determinant(polynomials, sign * maximumSearchMagnitude, communicator),
                .searchLimitReached = true
            };
        }

        return {
            .magnitude          = globalCriticalMagnitude,
            .determinant        = global_minimum_determinant(polynomials, sign * globalCriticalMagnitude, communicator),
            .searchLimitReached = false
        };
    }

    [[nodiscard]] SymmetricFiniteDifferenceStep find_symmetric_finite_difference_step(
        const mean_field::deformation::PreparedDomainDeformationRuntime &deformation,
        const mfem::Vector &unitVolumeDirection
    ) {
        constexpr double requestedStep = 1.0e-4;
        constexpr double minimumStep   = 1.0e-10;

        mfem::Vector trialVolumeDirection(unitVolumeDirection.Size());
        for (double step = requestedStep; step >= minimumStep; step *= 0.25) {
            trialVolumeDirection = unitVolumeDirection;
            trialVolumeDirection *= step;
            const mean_field::deformation::DomainDeformationGeometryReport positive =
                deformation.inspectMappedGeometry(trialVolumeDirection);

            trialVolumeDirection *= -1.0;
            const mean_field::deformation::DomainDeformationGeometryReport negative =
                deformation.inspectMappedGeometry(trialVolumeDirection);

            if (positive.isOrientationPreserving() && negative.isOrientationPreserving()) {
                return {
                    .step                       = step,
                    .positiveMinimumDeterminant = positive.minimumJacobianDeterminant,
                    .negativeMinimumDeterminant = negative.minimumJacobianDeterminant
                };
            }
        }

        throw std::domain_error(
            "No symmetric orientation-preserving finite-difference step was found for the surface mode."
        );
    }
} // namespace

TEST_CASE(
    "Reduced Surface Mode Reachability And Stellar Equilibrium Linearization",
    "[null_space][surface_modes][reachability][linearization]"
) {
    mean_field::utils::Args args = test_utils::setup_args();
    args.p.rtol                  = 1.0e-12;
    args.p.atol                  = std::min(args.p.atol, 1.0e-14);
    args.p.max_iters             = std::max(args.p.max_iters, 2000);

    experiment::null_space::N3Equilibrium fixture(std::move(args));
    const MPI_Comm communicator = fixture.fem().mesh->GetComm();
    int rank                    = 0;
    MPI_Comm_rank(communicator, &rank);

    const auto modes = experiment::null_space::make_surface_modes(fixture);
    constexpr std::array<double, 2> rotationFractions{0.0, 0.5};
    const int totalCases = static_cast<int>(rotationFractions.size() * modes.size());
    int completedCases   = 0;

    for (const double rotationFraction : rotationFractions) {
        const mean_field::physics::RigidRotation rotation = fixture.rotation(rotationFraction);
        fixture.prepare(fixture.state(), rotation);

        const mfem::Vector baseResidual = fixture.residual();
        REQUIRE(std::isfinite(experiment::null_space::global_norm(baseResidual, communicator)));

        for (const experiment::null_space::SurfaceMode &mode : modes) {
            experiment::null_space::report_progress(
                communicator, "probing " + mode.name + " at rotation fraction " + std::to_string(rotationFraction) +
                                  " (" + std::to_string(completedCases + 1) + "/" + std::to_string(totalCases) + ")"
            );

            fixture.prepare(fixture.state(), rotation);
            const mfem::Vector action          = fixture.jacobian_action(mode.direction);
            const mfem::Vector liftedDirection = fixture.lifted_surface_direction(mode.direction);

            const double inputNorm             = experiment::null_space::global_norm(mode.direction, communicator);
            const double actionNorm            = experiment::null_space::global_norm(action, communicator);
            const double liftNorm              = experiment::null_space::global_norm(liftedDirection, communicator);
            const SymmetricFiniteDifferenceStep coarseStep = find_symmetric_finite_difference_step(
                fixture.stellar_operator().GetDomainDeformation(), liftedDirection
            );
            const std::array<double, 2> finiteDifferenceSteps{coarseStep.step, 1.0e-2 * coarseStep.step};

            REQUIRE(inputNorm > 0.0);
            REQUIRE(liftNorm > 0.0);
            REQUIRE(std::isfinite(actionNorm));

            std::map<std::string, double> metrics{
                {"surface_parameter_input_norm", inputNorm},
                {"lifted_volume_displacement_norm", liftNorm},
                {"lift_amplification", liftNorm / inputNorm},
                {"root_jacobian_action_norm", actionNorm},
                {"root_action_per_surface_parameter_norm", actionNorm / inputNorm},
                {"root_action_per_lifted_volume_norm", actionNorm / liftNorm},
                {"base_residual_norm", experiment::null_space::global_norm(baseResidual, communicator)},
                {"surface_parameter_count",
                 static_cast<double>(fixture.stellar_operator().GetDomainDeformation().parameterCount())},
                {"volume_displacement_count",
                 static_cast<double>(fixture.stellar_operator().GetDomainDeformation().volumeDisplacementSize())},
                {"finite_difference_coarse_step", finiteDifferenceSteps[0]},
                {"finite_difference_fine_step", finiteDifferenceSteps[1]},
                {"coarse_step_positive_minimum_determinant", coarseStep.positiveMinimumDeterminant},
                {"coarse_step_negative_minimum_determinant", coarseStep.negativeMinimumDeterminant}
            };

            add_block_metrics(
                metrics, "root_",
                experiment::null_space::residual_block_norms(
                    action, fixture.stellar_operator().GetLayout(), communicator
                )
            );

            for (const double step : finiteDifferenceSteps) {
                mfem::Vector plusState(fixture.state());
                plusState.Add(step, mode.direction);
                fixture.prepare(plusState, rotation);
                const mfem::Vector plusResidual = fixture.residual();

                mfem::Vector minusState(fixture.state());
                minusState.Add(-step, mode.direction);
                fixture.prepare(minusState, rotation);
                const mfem::Vector minusResidual = fixture.residual();

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

                const std::string stepName = step == finiteDifferenceSteps.front() ? "coarse" : "fine";
                metrics.emplace(
                    "finite_difference_relative_error_" + stepName,
                    relative_difference(action, finiteDifference, communicator)
                );
            }

            fixture.prepare(fixture.state(), rotation);

            if (rank == 0) {
                experiment::record_experiment_result(
                    "reduced_surface_mode_reachability", mode.name,
                    {{"mode_kind", experiment::null_space::surface_mode_kind_name(mode.kind)},
                     {"axis", std::to_string(mode.axis)},
                     {"rotation_fraction_of_keplerian", std::to_string(rotationFraction)},
                     {"mesh_file", test_utils::setup_args().mesh_file},
                     {"local_state_dofs", std::to_string(fixture.stellar_operator().Width())}},
                    std::move(metrics)
                );
            }

            ++completedCases;
            experiment::null_space::report_progress(
                communicator, "completed " + std::to_string(completedCases) + "/" + std::to_string(totalCases) +
                                  " reduced surface-mode cases"
            );
        }
    }

    experiment::null_space::report_progress(communicator, "reduced surface-mode probe complete; writing CSV output");
}

TEST_CASE(
    "Spherical Harmonic Surface Frequencies Preserve Orientation Up To Measured Critical Amplitudes",
    "[surface_modes][frequency_limit][geometry][spherical_harmonic]"
) {
    mean_field::utils::Args args = test_utils::setup_args();
    mean_field::fem::FEM fem     = mean_field::fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(fem.okay());

    experiment::null_space::Model model = experiment::null_space::make_model();
    auto deformation                    = model.compileDomainDeformation(fem);
    const auto &surface                 = deformation.surfaceDeformationPrescription();
    const MPI_Comm communicator         = fem.mesh->GetComm();
    int rank                            = 0;
    MPI_Comm_rank(communicator, &rank);

    constexpr std::array<int, 13> angularDegrees{0, 1, 2, 3, 4, 5, 6, 8, 10, 12, 14, 16, 20};
    mfem::Vector zeroParameters(surface.parameterCount());
    zeroParameters = 0.0;

    for (std::size_t degreeIndex = 0; degreeIndex < angularDegrees.size(); ++degreeIndex) {
        const int angularDegree = angularDegrees[degreeIndex];
        experiment::null_space::report_progress(
            communicator, "measuring zonal spherical-harmonic degree " + std::to_string(angularDegree) + " (" +
                              std::to_string(degreeIndex + 1) + "/" + std::to_string(angularDegrees.size()) + ")"
        );

        mfem::Vector parameters(surface.parameterCount());
        double localMaximumAngularMagnitude = 0.0;
        for (int parameter = 0; parameter < parameters.Size(); ++parameter) {
            const double angularValue =
                experiment::null_space::zonal_legendre(angularDegree, surface.radialDirection(parameter, 2));
            parameters(parameter)        = surface.referenceRadius(parameter) * angularValue;
            localMaximumAngularMagnitude = std::max(localMaximumAngularMagnitude, std::abs(angularValue));
        }

        double globalMaximumAngularMagnitude = 0.0;
        MPI_Allreduce(
            &localMaximumAngularMagnitude, &globalMaximumAngularMagnitude, 1, MPI_DOUBLE, MPI_MAX, communicator
        );
        REQUIRE(globalMaximumAngularMagnitude > 0.0);
        parameters /= globalMaximumAngularMagnitude;

        mfem::Vector unitVolumeDirection(deformation.volumeDisplacementSize());
        deformation.applyJacobian(zeroParameters, parameters, unitVolumeDirection);
        const std::vector<DeterminantPolynomial> determinantPolynomials =
            collect_determinant_polynomials(fem, unitVolumeDirection);

        long long localSampleCount  = static_cast<long long>(determinantPolynomials.size());
        long long globalSampleCount = 0;
        MPI_Allreduce(&localSampleCount, &globalSampleCount, 1, MPI_LONG_LONG, MPI_SUM, communicator);
        REQUIRE(globalSampleCount > 0);

        const CriticalAmplitude positiveCritical = find_critical_amplitude(determinantPolynomials, 1.0, communicator);
        const CriticalAmplitude negativeCritical = find_critical_amplitude(determinantPolynomials, -1.0, communicator);

        const double determinantPositive1e4 = global_minimum_determinant(determinantPolynomials, 1.0e-4, communicator);
        const double determinantNegative1e4 = global_minimum_determinant(determinantPolynomials, -1.0e-4, communicator);
        const double determinantPositive1e3 = global_minimum_determinant(determinantPolynomials, 1.0e-3, communicator);
        const double determinantNegative1e3 = global_minimum_determinant(determinantPolynomials, -1.0e-3, communicator);
        const double determinantPositive1e2 = global_minimum_determinant(determinantPolynomials, 1.0e-2, communicator);
        const double determinantNegative1e2 = global_minimum_determinant(determinantPolynomials, -1.0e-2, communicator);

        if (angularDegree == 12) {
            mfem::Vector directInspectionDirection(unitVolumeDirection);
            directInspectionDirection *= 1.0e-3;
            const mean_field::deformation::DomainDeformationGeometryReport directInspection =
                deformation.inspectMappedGeometry(directInspectionDirection);
            const double comparisonScale = std::max(
                {1.0, std::abs(directInspection.minimumJacobianDeterminant), std::abs(determinantPositive1e3)}
            );
            CHECK(
                std::abs(directInspection.minimumJacobianDeterminant - determinantPositive1e3) <=
                1.0e-11 * comparisonScale
            );
        }

        REQUIRE(std::isfinite(positiveCritical.magnitude));
        REQUIRE(std::isfinite(negativeCritical.magnitude));
        REQUIRE(positiveCritical.magnitude > 0.0);
        REQUIRE(negativeCritical.magnitude > 0.0);

        if (rank == 0) {
            experiment::record_experiment_result(
                "spherical_harmonic_surface_frequency_limit", "zonal_l" + std::to_string(angularDegree),
                {{"angular_degree", std::to_string(angularDegree)},
                 {"azimuthal_order", "0"},
                 {"positive_limit_censored", positiveCritical.searchLimitReached ? "true" : "false"},
                 {"negative_limit_censored", negativeCritical.searchLimitReached ? "true" : "false"},
                 {"mesh_file", test_utils::setup_args().mesh_file}},
                {{"positive_critical_fractional_amplitude", positiveCritical.magnitude},
                 {"negative_critical_fractional_amplitude", negativeCritical.magnitude},
                 {"positive_critical_determinant", positiveCritical.determinant},
                 {"negative_critical_determinant", negativeCritical.determinant},
                 {"minimum_determinant_positive_1e-4", determinantPositive1e4},
                 {"minimum_determinant_negative_1e-4", determinantNegative1e4},
                 {"minimum_determinant_positive_1e-3", determinantPositive1e3},
                 {"minimum_determinant_negative_1e-3", determinantNegative1e3},
                 {"minimum_determinant_positive_1e-2", determinantPositive1e2},
                 {"minimum_determinant_negative_1e-2", determinantNegative1e2},
                 {"surface_parameter_norm", experiment::null_space::global_norm(parameters, communicator)},
                 {"lifted_volume_displacement_norm",
                  experiment::null_space::global_norm(unitVolumeDirection, communicator)},
                 {"global_geometry_sample_count", static_cast<double>(globalSampleCount)}}
            );
        }
    }

    experiment::null_space::report_progress(
        communicator, "spherical-harmonic frequency-limit probe complete; writing CSV output"
    );
}