#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <iostream>
#include <map>
#include <numbers>
#include <ranges>
#include <span>
#include <string>
#include <utility>
#include <vector>

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

import mean_field;
import test_helpers;
import experiment;

namespace {
    using Clock = std::chrono::steady_clock;

    [[nodiscard]] const char *build_configuration() noexcept {
#ifdef NDEBUG
        return "release";
#else
        return "debug";
#endif
    }

    [[nodiscard]] double maximum_rank_seconds(
        const Clock::time_point start,
        const MPI_Comm communicator
    ) {
        const double localSeconds = std::chrono::duration<double>(Clock::now() - start).count();
        double maximumSeconds{0.0};
        MPI_Allreduce(&localSeconds, &maximumSeconds, 1, MPI_DOUBLE, MPI_MAX, communicator);
        return maximumSeconds;
    }

    void announce(
        const MPI_Comm communicator,
        const std::string &message
    ) {
        int rank{0};
        MPI_Comm_rank(communicator, &rank);
        if (rank == 0) {
            std::cout << message << std::endl;
        }
    }

    class ArnoldiProgressOperator final : public mfem::Operator {
    public:
        ArnoldiProgressOperator(
            const mfem::Operator &operation,
            const MPI_Comm communicator,
            const int expectedApplications,
            const int reportingInterval
        )
            : mfem::Operator(
                  operation.Height(),
                  operation.Width()
              ),
              m_operation(&operation),
              m_communicator(communicator),
              m_expectedApplications(expectedApplications),
              m_reportingInterval(reportingInterval) {
        }

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override {
            m_operation->Mult(input, output);
            ++m_completedApplications;
            if (m_completedApplications == 1 || m_completedApplications == m_expectedApplications ||
                m_completedApplications % m_reportingInterval == 0) {
                announce(
                    m_communicator, "Arnoldi progress: " + std::to_string(m_completedApplications) + "/" +
                                        std::to_string(m_expectedApplications) + " Jacobian applications"
                );
            }
        }

    private:
        const mfem::Operator *m_operation;
        MPI_Comm m_communicator;
        int m_expectedApplications;
        int m_reportingInterval;
        mutable int m_completedApplications{0};
    };

    [[nodiscard]] mean_field::operators::StellarEquilibriumDependencies make_dependencies() {
        return {
            .discretization     = {.identity = 8101, .revision = 1},
            .density            = {.identity = 8103, .revision = 1},
            .surfaceDeformation = {.identity = 8107, .revision = 1},
            .gravityGradient    = {.identity = 8111, .revision = 1},
            .gravityPotential   = {.identity = 8117, .revision = 1},
            .enthalpy           = {.identity = 8123, .revision = 1},
            .bernoulliConstant  = {.identity = 8129, .revision = 1},
            .rotation           = {.identity = 8131, .revision = 1},
            .targetMass         = {.identity = 8137, .revision = 1}
        };
    }

    [[nodiscard]] mean_field::physics::RigidRotation make_zero_rotation() {
        mfem::Vector angularVelocity(3);
        mfem::Vector center(3);
        angularVelocity = 0.0;
        center          = 0.0;
        return {angularVelocity, center};
    }

    [[nodiscard]] double global_norm(
        const mfem::Vector &vector,
        const MPI_Comm communicator
    ) {
        const double localSquaredNorm = vector * vector;
        double globalSquaredNorm{0.0};
        MPI_Allreduce(&localSquaredNorm, &globalSquaredNorm, 1, MPI_DOUBLE, MPI_SUM, communicator);
        return std::sqrt(std::max(globalSquaredNorm, 0.0));
    }

    [[nodiscard]] mfem::Vector make_block_balanced_direction(
        const int stateSize,
        const std::span<const mean_field::operators::RootBlockDescriptor> valueBlocks,
        const MPI_Comm communicator
    ) {
        mfem::Vector direction(stateSize);
        direction = 0.0;

        for (const mean_field::operators::RootBlockDescriptor &block : valueBlocks) {
            mfem::Vector values(direction.GetData() + block.offset, block.size);
            for (int index = 0; index < values.Size(); ++index) {
                const double ordinal = static_cast<double>(block.canonicalIndex + 1);
                values(index)        = std::sin(0.6180339887498948 * static_cast<double>(index + 1) + ordinal);
            }
            const double norm = global_norm(values, communicator);
            if (norm > 0.0) {
                values /= norm;
            }
        }
        return direction;
    }

    void require_finite(const double value) {
        REQUIRE(std::isfinite(value));
    }

    [[nodiscard]] std::map<
        std::string,
        std::string>
    common_parameters(
        const std::string &measurement,
        const int stateSize
    ) {
        return {
            {"build_configuration", build_configuration()},
            {"equation_of_state", "Polytrope(n=3)"},
            {"experiment_schema", "p0_extended_v2"},
            {"linearization_state", "projected_lane_emden"},
            {"measurement", measurement},
            {"mesh_file", test_utils::setup_args().mesh_file},
            {"preconditioner", "identity"},
            {"preconditioned_product", "J M^-1"},
            {"root_dimension", std::to_string(stateSize)}
        };
    }
} // namespace

TEST_CASE(
    "Stellar Equilibrium P0 Identity Preconditioning Baseline",
    "[preconditioning][diagnostics][baseline][spectrum]"
) {
    using namespace mean_field;

    constexpr int arnoldiDimension          = 48;
    const MPI_Comm world                    = MPI_COMM_WORLD;
    const Clock::time_point experimentStart = Clock::now();
    announce(world, "P0 extended baseline: constructing the finite-element discretization");

    const Clock::time_point finiteElementSetupStart = Clock::now();
    utils::Args args                                = test_utils::setup_args();
    fem::FEM finiteElementModel                     = fem::setup_fem(args.mesh_file, args, 0);
    REQUIRE(finiteElementModel.okay());
    const MPI_Comm communicator            = finiteElementModel.mesh->GetComm();
    const double finiteElementSetupSeconds = maximum_rank_seconds(finiteElementSetupStart, communicator);
    announce(
        communicator, "P0 extended baseline: finite-element setup completed in " +
                          std::to_string(finiteElementSetupSeconds) + " seconds"
    );

    constexpr double stellarRadius                                  = utils::RADIUS;
    constexpr double targetMass                                     = utils::MASS;
    const Clock::time_point calibrationStart                        = Clock::now();
    const seed::DimensionlessLaneEmdenSolution dimensionlessProfile = seed::integrateLaneEmden(3.0, 10.0);
    REQUIRE(dimensionlessProfile.firstZeroCoordinate.has_value());
    const double surfaceCoordinate = *dimensionlessProfile.firstZeroCoordinate;
    const double surfaceDerivative =
        dimensionlessProfile.thetaDerivative(dimensionlessProfile.thetaDerivative.Size() - 1);
    const double dimensionlessMass = -surfaceCoordinate * surfaceCoordinate * surfaceDerivative;
    REQUIRE(dimensionlessMass > 0.0);

    const double massScale          = targetMass / (4.0 * std::numbers::pi_v<double> * dimensionlessMass);
    const double polytropicConstant = std::numbers::pi_v<double> * utils::G * std::pow(massScale, 2.0 / 3.0);
    const double radialScale        = stellarRadius / surfaceCoordinate;
    const double centralDensity =
        std::pow(polytropicConstant / (std::numbers::pi_v<double> * utils::G * radialScale * radialScale), 1.5);
    const double calibrationSeconds                  = maximum_rank_seconds(calibrationStart, communicator);

    const Clock::time_point problemConstructionStart = Clock::now();
    const auto stellarModel                          = model::StellarModel(
        eos::Polytrope({.n = 3.0, .K = polytropicConstant}),
        surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{targetMass}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{centralDensity}})
    );
    auto problem                            = equilibrium::discretize(stellarModel, std::move(finiteElementModel));
    const double problemConstructionSeconds = maximum_rank_seconds(problemConstructionStart, communicator);
    announce(communicator, "P0 extended baseline: projecting the Lane-Emden seed");

    const Clock::time_point seedProjectionStart = Clock::now();
    const auto projected = seed::makeProjectedEquilibriumState(problem, seed::LaneEmden({.radialSampleCount = 4096}));
    const double seedProjectionSeconds = maximum_rank_seconds(seedProjectionStart, communicator);
    announce(communicator, "P0 extended baseline: preparing the complete equilibrium operator");

    const Clock::time_point operatorPreparationStart = Clock::now();
    const auto preparation = problem.Prepare(projected.values, make_dependencies(), make_zero_rotation());
    REQUIRE(preparation.assembledResidual);
    const double operatorPreparationSeconds = maximum_rank_seconds(operatorPreparationStart, communicator);

    const mfem::Operator &rawJacobian       = problem.GetLinearizationOperator();
    mfem::Vector knownDirection =
        make_block_balanced_direction(problem.StateSize(), problem.GetManifest().valueBlocks(), communicator);
    mfem::Vector rightHandSide(problem.EquationSize());
    const Clock::time_point applicationStart = Clock::now();
    rawJacobian.Mult(knownDirection, rightHandSide);
    const double applicationSeconds = maximum_rank_seconds(applicationStart, communicator);
    REQUIRE(rightHandSide.Size() == problem.EquationSize());
    require_finite(global_norm(rightHandSide, communicator));
    announce(
        communicator, "P0 extended baseline: first prepared Jacobian application completed in " +
                          std::to_string(applicationSeconds) + " seconds"
    );

    if (std::getenv("MEANFIELD_SINGLE_JACOBIAN_BENCHMARK") != nullptr) {
        int rank{0};
        MPI_Comm_rank(communicator, &rank);
        if (rank == 0) {
            std::cout << "Single prepared Jacobian application: " << applicationSeconds << " seconds\n";
        }
        return;
    }

    solver::IdentityPreconditioner identity(problem.StateSize());
    solver::InstrumentedOperator instrumentedJacobian(rawJacobian);
    solver::InstrumentedPreconditioner instrumentedPreconditioner(identity);
    solver::ResidualHistoryMonitor monitor;
    mfem::FGMRESSolver krylov(communicator);
    krylov.SetPreconditioner(instrumentedPreconditioner);
    krylov.SetOperator(instrumentedJacobian);
    krylov.SetMonitor(monitor);
    krylov.SetRelTol(1.0e-8);
    krylov.SetAbsTol(1.0e-12);
    krylov.SetMaxIter(40);
    krylov.SetKDim(20);
    krylov.SetPrintLevel(1);

    mfem::Vector solution(problem.StateSize());
    solution = 0.0;
    const operators::PreparedStellarEquilibriumStatistics statisticsBeforeSolve =
        problem.GetPreparedOperator().GetPhysicalOperator().GetStatistics();
    announce(communicator, "P0 extended baseline: starting the 40-iteration identity-preconditioned FGMRES solve");
    const Clock::time_point solveStart = Clock::now();
    krylov.Mult(rightHandSide, solution);
    const double localSolveSeconds = std::chrono::duration<double>(Clock::now() - solveStart).count();
    const operators::PreparedStellarEquilibriumStatistics statisticsAfterSolve =
        problem.GetPreparedOperator().GetPhysicalOperator().GetStatistics();

    announce(communicator, "P0 extended baseline: independently reconstructing the true residual");
    const Clock::time_point directResidualStart           = Clock::now();
    const solver::LinearSolveMeasurement solveMeasurement = solver::measureLinearSolve(
        krylov, rawJacobian, rightHandSide, solution, problem.GetManifest().residualBlocks(),
        instrumentedJacobian.GetStatistics(), instrumentedPreconditioner.GetStatistics(),
        instrumentedPreconditioner.GetLifecycleStatistics(), monitor, localSolveSeconds, communicator
    );
    const double directResidualMeasurementSeconds = maximum_rank_seconds(directResidualStart, communicator);
    require_finite(solveMeasurement.directResidual.relativeResidual);
    require_finite(solveMeasurement.solveSecondsMaximumRank);

    std::map<std::string, double> solveMetrics{
        {"solver_converged", solveMeasurement.solverConverged ? 1.0 : 0.0},
        {"outer_iterations", static_cast<double>(solveMeasurement.outerIterations)},
        {"reported_initial_residual_norm", solveMeasurement.solverReportedInitialNorm},
        {"reported_final_residual_norm", solveMeasurement.solverReportedFinalNorm},
        {"reported_residual_reduction", solveMeasurement.solverReportedResidualReduction},
        {"true_residual_norm", solveMeasurement.directResidual.trueResidualNorm},
        {"true_relative_residual", solveMeasurement.directResidual.relativeResidual},
        {"rhs_norm", solveMeasurement.directResidual.rightHandSideNorm},
        {"true_residual_digits_per_jacobian_application",
         solveMeasurement.trueResidualDigitsReducedPerJacobianApplication},
        {"finite_element_setup_seconds", finiteElementSetupSeconds},
        {"lane_emden_calibration_seconds", calibrationSeconds},
        {"equilibrium_problem_construction_seconds", problemConstructionSeconds},
        {"seed_projection_seconds", seedProjectionSeconds},
        {"operator_preparation_seconds", operatorPreparationSeconds},
        {"initial_jacobian_application_seconds", applicationSeconds},
        {"direct_residual_measurement_seconds", directResidualMeasurementSeconds},
        {"solve_seconds_maximum_rank", solveMeasurement.solveSecondsMaximumRank},
        {"jacobian_applications", static_cast<double>(solveMeasurement.jacobian.applications)},
        {"jacobian_application_seconds", solveMeasurement.jacobian.totalSeconds},
        {"jacobian_maximum_application_seconds", solveMeasurement.jacobian.maximumSeconds},
        {"inverse_preconditioner_applications",
         static_cast<double>(solveMeasurement.inversePreconditioner.applications)},
        {"inverse_preconditioner_application_seconds", solveMeasurement.inversePreconditioner.totalSeconds},
        {"inverse_preconditioner_maximum_application_seconds", solveMeasurement.inversePreconditioner.maximumSeconds},
        {"inverse_preconditioner_setups", static_cast<double>(solveMeasurement.inversePreconditionerLifecycle.setups)},
        {"inverse_preconditioner_refreshes",
         static_cast<double>(solveMeasurement.inversePreconditionerLifecycle.refreshes)},
        {"inverse_preconditioner_setup_seconds", solveMeasurement.inversePreconditionerLifecycle.setupSeconds},
        {"inverse_preconditioner_refresh_seconds", solveMeasurement.inversePreconditionerLifecycle.refreshSeconds},
        {"prepared_residual_assemblies_during_solve",
         static_cast<double>(statisticsAfterSolve.residualAssemblies - statisticsBeforeSolve.residualAssemblies)},
        {"prepared_geometry_builds_during_solve",
         static_cast<double>(
             statisticsAfterSolve.generatedGeometryBuilds - statisticsBeforeSolve.generatedGeometryBuilds
         )},
        {"prepared_jacobian_applications_during_solve",
         static_cast<double>(statisticsAfterSolve.jacobianApplications - statisticsBeforeSolve.jacobianApplications)}
    };
    for (const solver::ResidualBlockMeasurement &block : solveMeasurement.directResidual.blocks) {
        const std::string prefix                                   = "residual_block." + block.stableId;
        solveMetrics[prefix + ".descriptor_scale"]                 = block.descriptorScale;
        solveMetrics[prefix + ".rhs_norm"]                         = block.rightHandSideNorm;
        solveMetrics[prefix + ".true_norm"]                        = block.trueResidualNorm;
        solveMetrics[prefix + ".block_relative_residual"]          = block.blockRelativeResidual;
        solveMetrics[prefix + ".scaled_rhs_norm"]                  = block.scaledRightHandSideNorm;
        solveMetrics[prefix + ".scaled_true_norm"]                 = block.scaledTrueResidualNorm;
        solveMetrics[prefix + ".fraction_global_squared_residual"] = block.fractionOfGlobalSquaredResidualNorm;
        solveMetrics[prefix + ".global_relative_contribution"]     = block.contributionToGlobalRelativeResidual;
    }
    experiment::record_experiment_result(
        "stellar_preconditioning_p0", "identity_linear_solve", common_parameters("linear_solve", problem.StateSize()),
        std::move(solveMetrics)
    );

    const double reportedInitialDenominator = std::max(solveMeasurement.solverReportedInitialNorm, 1.0e-300);
    for (std::size_t sample = 0; sample < solveMeasurement.reportedResidualHistory.size(); ++sample) {
        const solver::IterationResidualMeasurement &residual = solveMeasurement.reportedResidualHistory[sample];
        experiment::record_experiment_result(
            "stellar_preconditioning_p0", "identity_fgmres_history_" + std::to_string(sample),
            common_parameters("fgmres_residual_history", problem.StateSize()),
            {{"history_sample", static_cast<double>(sample)},
             {"iteration", static_cast<double>(residual.iteration)},
             {"reported_residual_norm", residual.reportedNorm},
             {"reported_relative_residual", residual.reportedNorm / reportedInitialDenominator},
             {"final_measurement", residual.final ? 1.0 : 0.0}}
        );
    }

    instrumentedJacobian.ResetStatistics();
    instrumentedPreconditioner.ResetStatistics();
    solver::FixedRightPreconditionedOperator rightPreconditionedProduct(
        instrumentedJacobian, instrumentedPreconditioner
    );
    ArnoldiProgressOperator progressOperator(rightPreconditionedProduct, communicator, arnoldiDimension, 4);
    announce(
        communicator,
        "P0 extended baseline: starting the " + std::to_string(arnoldiDimension) + "-vector Arnoldi measurement"
    );
    const solver::ArnoldiSpectralMeasurement spectrum = solver::measureArnoldiSpectrum(
        progressOperator, knownDirection, communicator,
        {.krylovDimension                  = arnoldiDimension,
         .breakdownRelativeTolerance       = 1.0e-13,
         .ritzConvergenceRelativeTolerance = 1.0e-7,
         .reorthogonalize                  = true}
    );
    require_finite(spectrum.projectedLargestSingularValue);
    require_finite(spectrum.centroidRealPart);
    require_finite(spectrum.rmsClusterRadius);

    experiment::record_experiment_result(
        "stellar_preconditioning_p0", "identity_arnoldi_summary",
        common_parameters("arnoldi_summary", problem.StateSize()),
        {{"requested_krylov_dimension", static_cast<double>(spectrum.requestedDimension)},
         {"achieved_krylov_dimension", static_cast<double>(spectrum.achievedDimension)},
         {"invariant_subspace_found", spectrum.invariantSubspaceFound ? 1.0 : 0.0},
         {"operator_applications", static_cast<double>(spectrum.operatorApplications)},
         {"arnoldi_operator_application_seconds", spectrum.operatorApplicationSecondsMaximumRank},
         {"arnoldi_operator_maximum_application_seconds", spectrum.operatorMaximumApplicationSecondsMaximumRank},
         {"arnoldi_measurement_seconds", spectrum.measurementSecondsMaximumRank},
         {"arnoldi_nonapplication_seconds", spectrum.nonApplicationSecondsMaximumRank},
         {"experiment_elapsed_through_arnoldi_seconds", maximum_rank_seconds(experimentStart, communicator)},
         {"converged_ritz_values", static_cast<double>(spectrum.convergedRitzValueCount)},
         {"negative_real_part_ritz_values", static_cast<double>(spectrum.negativeRealPartCount)},
         {"projected_largest_singular_value", spectrum.projectedLargestSingularValue},
         {"projected_smallest_singular_value", spectrum.projectedSmallestSingularValue},
         {"projected_condition_proxy", spectrum.projectedConditionProxy},
         {"ritz_centroid_real", spectrum.centroidRealPart},
         {"ritz_centroid_imaginary", spectrum.centroidImaginaryPart},
         {"ritz_rms_distance_from_one", spectrum.rmsDistanceFromOne},
         {"ritz_rms_cluster_radius", spectrum.rmsClusterRadius},
         {"ritz_minimum_magnitude", spectrum.minimumMagnitude},
         {"ritz_maximum_magnitude", spectrum.maximumMagnitude},
         {"ritz_minimum_real_part", spectrum.minimumRealPart},
         {"ritz_maximum_real_part", spectrum.maximumRealPart},
         {"ritz_maximum_absolute_imaginary_part", spectrum.maximumAbsoluteImaginaryPart},
         {"ritz_conjugate_pair_defect", spectrum.conjugatePairDefect},
         {"projected_departure_from_normality", spectrum.projectedDepartureFromNormality},
         {"projected_field_of_values_minimum_real_part", spectrum.projectedFieldOfValuesMinimumRealPart},
         {"projected_field_of_values_maximum_real_part", spectrum.projectedFieldOfValuesMaximumRealPart},
         {"measured_jacobian_applications", static_cast<double>(instrumentedJacobian.GetStatistics().applications)},
         {"measured_jacobian_application_seconds", instrumentedJacobian.GetStatistics().totalSeconds},
         {"measured_jacobian_maximum_application_seconds", instrumentedJacobian.GetStatistics().maximumSeconds},
         {"measured_inverse_preconditioner_applications",
          static_cast<double>(instrumentedPreconditioner.GetStatistics().applications)},
         {"measured_inverse_preconditioner_application_seconds",
          instrumentedPreconditioner.GetStatistics().totalSeconds}}
    );

    std::vector<solver::RitzValueMeasurement> orderedRitzValues = spectrum.ritzValues;
    std::ranges::sort(orderedRitzValues, [](const auto &left, const auto &right) {
        if (left.realPart != right.realPart) {
            return left.realPart < right.realPart;
        }
        return left.imaginaryPart < right.imaginaryPart;
    });
    for (std::size_t index = 0; index < orderedRitzValues.size(); ++index) {
        const solver::RitzValueMeasurement &ritz = orderedRitzValues[index];
        experiment::record_experiment_result(
            "stellar_preconditioning_p0", "identity_ritz_" + std::to_string(index),
            common_parameters("ritz_value", problem.StateSize()),
            {{"ritz_index", static_cast<double>(index)},
             {"ritz_real", ritz.realPart},
             {"ritz_imaginary", ritz.imaginaryPart},
             {"ritz_magnitude", ritz.magnitude},
             {"ritz_distance_from_one", ritz.distanceFromOne},
             {"ritz_residual_estimate", ritz.residualEstimate},
             {"ritz_relative_residual_estimate", ritz.relativeResidualEstimate},
             {"ritz_converged", ritz.converged ? 1.0 : 0.0}}
        );
    }

    int rank{0};
    MPI_Comm_rank(communicator, &rank);
    if (rank == 0) {
        std::cout << "P0 identity baseline: " << solveMeasurement.outerIterations << " FGMRES iterations, "
                  << spectrum.achievedDimension << " Arnoldi vectors, true relative residual "
                  << solveMeasurement.directResidual.relativeResidual << '\n';
    }
}