#include <array>
#include <chrono>
#include <cmath>
#include <string_view>
#include <utility>

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

import mean_field;
import test_helpers;

namespace {
    class DenseLinearOperator final : public mfem::Operator {
    public:
        explicit DenseLinearOperator(mfem::DenseMatrix matrix)
            : mfem::Operator(
                  matrix.Height(),
                  matrix.Width()
              ),
              m_matrix(std::move(matrix)) {
        }

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override {
            m_matrix.Mult(input, output);
        }

    private:
        mfem::DenseMatrix m_matrix;
    };

    class DiagonalInversePreconditioner final : public mfem::Solver {
    public:
        explicit DiagonalInversePreconditioner(mfem::Vector diagonal)
            : mfem::Solver(diagonal.Size()),
              m_diagonal(std::move(diagonal)) {
        }

        void SetOperator(const mfem::Operator &operation) override {
            REQUIRE(operation.Height() == Height());
            REQUIRE(operation.Width() == Width());
        }

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override {
            REQUIRE(input.Size() == Width());
            output.SetSize(Height());
            for (int index = 0; index < Height(); ++index) {
                output(index) = input(index) / m_diagonal(index);
            }
        }

    private:
        mfem::Vector m_diagonal;
    };

    [[nodiscard]] mfem::DenseMatrix diagonal_matrix(
        const std::array<
            double,
            4> &diagonal
    ) {
        mfem::DenseMatrix matrix(4);
        matrix = 0.0;
        for (int index = 0; index < 4; ++index) {
            matrix(index, index) = diagonal[static_cast<std::size_t>(index)];
        }
        return matrix;
    }

    [[nodiscard]] mean_field::operators::RootBlockDescriptor residual_block(
        const std::string_view stableId,
        const int index,
        const int offset,
        const int size
    ) {
        using namespace mean_field::operators;
        return {
            .stableId       = stableId,
            .symbol         = stableId,
            .kind           = RootBlockKind::residual,
            .provenance     = RootBlockProvenance::physical_operator,
            .source         = "test",
            .rowInjection   = RootRowInjection::physical_equation,
            .columnPolicy   = RootColumnPolicy::no_column,
            .scalePolicy    = RootScalePolicy::unscaled,
            .canonicalIndex = index,
            .offset         = offset,
            .size           = size,
            .scale          = 1.0
        };
    }

    [[nodiscard]] bool contains_eigenvalue(
        const mean_field::solver::ArnoldiSpectralMeasurement &measurement,
        const double realPart,
        const double imaginaryPart,
        const double tolerance
    ) {
        for (const auto &value : measurement.ritzValues) {
            if (std::hypot(value.realPart - realPart, value.imaginaryPart - imaginaryPart) < tolerance) {
                return true;
            }
        }
        return false;
    }
} // namespace

TEST_CASE(
    "Preconditioning Instrumentation Counts Work And Independently Measures The True Residual",
    tags::preconditioning_diagnostics_unit
) {
    using Catch::Approx;
    using namespace mean_field;

    constexpr std::array<double, 4> diagonalValues{2.0, 4.0, 8.0, 16.0};
    DenseLinearOperator jacobian(diagonal_matrix(diagonalValues));
    mfem::Vector diagonal(4);
    for (int index = 0; index < 4; ++index) {
        diagonal(index) = diagonalValues[static_cast<std::size_t>(index)];
    }
    DiagonalInversePreconditioner inversePreconditioner(std::move(diagonal));

    solver::InstrumentedOperator instrumentedJacobian(jacobian);
    solver::InstrumentedPreconditioner instrumentedPreconditioner(inversePreconditioner);
    solver::FixedRightPreconditionedOperator rightPreconditioned(instrumentedJacobian, instrumentedPreconditioner);

    mfem::Vector input({1.0, -2.0, 3.0, -4.0});
    mfem::Vector product(rightPreconditioned.Height());
    rightPreconditioned.Mult(input, product);
    REQUIRE(product.Size() == input.Size());
    for (int index = 0; index < input.Size(); ++index) {
        CHECK(product(index) == Approx(input(index)));
    }
    CHECK(instrumentedJacobian.GetStatistics().applications == 1);
    CHECK(instrumentedPreconditioner.GetStatistics().applications == 1);
    CHECK(instrumentedJacobian.GetStatistics().totalSeconds >= 0.0);
    CHECK(instrumentedPreconditioner.GetStatistics().totalSeconds >= 0.0);

    instrumentedJacobian.ResetStatistics();
    instrumentedPreconditioner.ResetStatistics();

    mfem::Vector exactSolution({0.25, -0.5, 0.75, -1.0});
    mfem::Vector rightHandSide(jacobian.Height());
    jacobian.Mult(exactSolution, rightHandSide);
    mfem::Vector computedSolution(4);
    computedSolution = 0.0;

    solver::ResidualHistoryMonitor monitor;
    mfem::FGMRESSolver krylov(MPI_COMM_WORLD);
    krylov.SetPreconditioner(instrumentedPreconditioner);
    krylov.SetOperator(instrumentedJacobian);
    krylov.SetMonitor(monitor);
    krylov.SetRelTol(1.0e-13);
    krylov.SetAbsTol(1.0e-15);
    krylov.SetMaxIter(20);
    krylov.SetKDim(10);
    krylov.SetPrintLevel(0);

    const auto start = std::chrono::steady_clock::now();
    krylov.Mult(rightHandSide, computedSolution);
    const double elapsed = std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count();

    const std::array residualBlocks{residual_block("first", 0, 0, 2), residual_block("second", 1, 2, 2)};
    const solver::LinearSolveMeasurement measurement = solver::measureLinearSolve(
        krylov, jacobian, rightHandSide, computedSolution, residualBlocks, instrumentedJacobian.GetStatistics(),
        instrumentedPreconditioner.GetStatistics(), instrumentedPreconditioner.GetLifecycleStatistics(), monitor,
        elapsed, MPI_COMM_WORLD
    );

    CHECK(measurement.solverConverged);
    CHECK(measurement.outerIterations > 0);
    CHECK(measurement.jacobian.applications > 0);
    CHECK(measurement.inversePreconditioner.applications > 0);
    CHECK(measurement.inversePreconditionerLifecycle.setups > 0);
    CHECK(measurement.solveSecondsMaximumRank >= 0.0);
    CHECK(measurement.solverReportedResidualReduction < 1.0e-12);
    CHECK(measurement.trueResidualDigitsReducedPerJacobianApplication > 0.0);
    CHECK(measurement.directResidual.relativeResidual < 1.0e-12);
    REQUIRE(measurement.directResidual.blocks.size() == 2);
    CHECK(measurement.directResidual.blocks[0].stableId == "first");
    CHECK(measurement.directResidual.blocks[1].stableId == "second");
    CHECK(measurement.directResidual.blocks[0].descriptorScale == 1.0);
    CHECK(measurement.directResidual.blocks[0].blockRelativeResidual < 1.0e-12);
    CHECK(measurement.directResidual.blocks[1].blockRelativeResidual < 1.0e-12);
    CHECK(measurement.directResidual.blocks[0].fractionOfGlobalSquaredResidualNorm >= 0.0);
    CHECK(measurement.directResidual.blocks[1].fractionOfGlobalSquaredResidualNorm >= 0.0);
    CHECK_FALSE(measurement.reportedResidualHistory.empty());
}

TEST_CASE(
    "Arnoldi Diagnostics Recover Real And Complex Conjugate Eigenvalue Clusters",
    tags::preconditioning_spectral_unit
) {
    using Catch::Approx;
    using namespace mean_field;

    mfem::DenseMatrix matrix(4);
    matrix       = 0.0;
    matrix(0, 0) = 2.0;
    matrix(1, 1) = 3.0;
    matrix(2, 3) = -1.0;
    matrix(3, 2) = 1.0;
    DenseLinearOperator operation(std::move(matrix));

    const mfem::Vector initialDirection({1.0, 2.0, 3.0, 4.0});
    const solver::ArnoldiSpectralMeasurement measurement = solver::measureArnoldiSpectrum(
        operation, initialDirection, MPI_COMM_WORLD,
        {.krylovDimension                  = 4,
         .breakdownRelativeTolerance       = 1.0e-12,
         .ritzConvergenceRelativeTolerance = 1.0e-9,
         .reorthogonalize                  = true}
    );

    REQUIRE(measurement.achievedDimension == 4);
    REQUIRE(measurement.ritzValues.size() == 4);
    CHECK(measurement.operatorApplications == 4);
    CHECK(measurement.operatorApplicationSecondsMaximumRank >= 0.0);
    CHECK(measurement.operatorMaximumApplicationSecondsMaximumRank >= 0.0);
    CHECK(measurement.measurementSecondsMaximumRank >= measurement.operatorApplicationSecondsMaximumRank);
    CHECK(measurement.nonApplicationSecondsMaximumRank >= 0.0);
    CHECK(measurement.invariantSubspaceFound);
    CHECK(contains_eigenvalue(measurement, 2.0, 0.0, 1.0e-10));
    CHECK(contains_eigenvalue(measurement, 3.0, 0.0, 1.0e-10));
    CHECK(contains_eigenvalue(measurement, 0.0, 1.0, 1.0e-10));
    CHECK(contains_eigenvalue(measurement, 0.0, -1.0, 1.0e-10));
    CHECK(measurement.conjugatePairDefect < 1.0e-10);
    CHECK(measurement.projectedLargestSingularValue == Approx(3.0).margin(1.0e-10));
    CHECK(measurement.projectedSmallestSingularValue == Approx(1.0).margin(1.0e-10));
    CHECK(measurement.projectedConditionProxy == Approx(3.0).margin(1.0e-10));
    CHECK(measurement.maximumAbsoluteImaginaryPart == Approx(1.0).margin(1.0e-10));
}

TEST_CASE(
    "Arnoldi Diagnostics Distinguish Exact Preconditioning From Nonnormal Clustering",
    tags::preconditioning_spectral_unit
) {
    using Catch::Approx;
    using namespace mean_field;

    DenseLinearOperator jacobian(diagonal_matrix({2.0, 4.0, 8.0, 16.0}));
    mfem::Vector diagonal({2.0, 4.0, 8.0, 16.0});
    DiagonalInversePreconditioner inversePreconditioner(std::move(diagonal));
    solver::FixedRightPreconditionedOperator exactProduct(jacobian, inversePreconditioner);
    const mfem::Vector initialDirection({1.0, -1.0, 2.0, -2.0});

    const solver::ArnoldiSpectralMeasurement exact = solver::measureArnoldiSpectrum(
        exactProduct, initialDirection, MPI_COMM_WORLD, {.krylovDimension = 4, .breakdownRelativeTolerance = 1.0e-12}
    );
    REQUIRE(exact.achievedDimension == 1);
    REQUIRE(exact.ritzValues.size() == 1);
    CHECK(exact.ritzValues[0].realPart == Approx(1.0).margin(1.0e-12));
    CHECK(exact.ritzValues[0].imaginaryPart == Approx(0.0).margin(1.0e-12));
    CHECK(exact.projectedConditionProxy == Approx(1.0).margin(1.0e-12));
    CHECK(exact.rmsDistanceFromOne < 1.0e-12);

    mfem::DenseMatrix jordan(4);
    jordan = 0.0;
    for (int index = 0; index < 4; ++index) {
        jordan(index, index) = 1.0;
    }
    jordan(0, 1) = 4.0;
    jordan(1, 2) = 4.0;
    jordan(2, 3) = 4.0;
    DenseLinearOperator nonnormal(std::move(jordan));
    const solver::ArnoldiSpectralMeasurement nonnormalMeasurement = solver::measureArnoldiSpectrum(
        nonnormal, mfem::Vector({1.0, 2.0, 3.0, 5.0}), MPI_COMM_WORLD,
        {.krylovDimension = 4, .breakdownRelativeTolerance = 1.0e-12}
    );
    CHECK(nonnormalMeasurement.projectedDepartureFromNormality > 0.1);
    CHECK(nonnormalMeasurement.projectedConditionProxy > 1.0);

    const std::vector<solver::RitzValueMeasurement> closest =
        solver::selectRitzValues(nonnormalMeasurement, solver::RitzValueOrdering::closest_to_zero, 2);
    CHECK(closest.size() <= 2);
}