module;

#include <cstdint>
#include <span>
#include <string>
#include <string_view>
#include <vector>

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

export module mean_field:solver.preconditioning_diagnostics;

export import :operators.root_manifest;

export namespace mean_field::solver {
    struct OperatorApplicationStatistics final {
        std::uint64_t applications{0};
        double totalSeconds{0.0};
        double maximumSeconds{0.0};
    };

    struct PreconditionerLifecycleStatistics final {
        std::uint64_t setups{0};
        std::uint64_t refreshes{0};
        double setupSeconds{0.0};
        double refreshSeconds{0.0};
    };

    /*
     * A non-owning measurement wrapper. Statistics are local to an MPI rank;
     * cross-rank wall-clock reductions are performed when a solve report is
     * assembled. Krylov application is sequential, so counters intentionally
     * do not impose atomic overhead.
     */
    class InstrumentedOperator final : public mfem::Operator {
    public:
        explicit InstrumentedOperator(const mfem::Operator &operation);

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override;

        void ResetStatistics() const noexcept;
        [[nodiscard]] const OperatorApplicationStatistics &GetStatistics() const noexcept;
        [[nodiscard]] const mfem::Operator &GetOperation() const noexcept;

    private:
        const mfem::Operator *m_operation;
        mutable OperatorApplicationStatistics m_statistics;
    };

    class InstrumentedPreconditioner final : public mfem::Solver {
    public:
        explicit InstrumentedPreconditioner(mfem::Solver &preconditioner);

        void SetOperator(const mfem::Operator &operation) override;

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override;

        void ResetStatistics() const noexcept;
        [[nodiscard]] const OperatorApplicationStatistics &GetStatistics() const noexcept;
        [[nodiscard]] const PreconditionerLifecycleStatistics &GetLifecycleStatistics() const noexcept;
        [[nodiscard]] const mfem::Solver &GetPreconditioner() const noexcept;

    private:
        mfem::Solver *m_preconditioner;
        mutable OperatorApplicationStatistics m_statistics;
        PreconditionerLifecycleStatistics m_lifecycleStatistics;
    };

    class IdentityPreconditioner final : public mfem::Solver {
    public:
        explicit IdentityPreconditioner(int size);

        void SetOperator(const mfem::Operator &operation) override;

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

    /*
     * If the supplied solver applies M^{-1}, this operator represents the
     * fixed right-preconditioned product J M^{-1}. It is deliberately
     * independent of the Krylov implementation used in production.
     */
    class FixedRightPreconditionedOperator final : public mfem::Operator {
    public:
        FixedRightPreconditionedOperator(
            const mfem::Operator &jacobian,
            const mfem::Solver &inversePreconditioner
        );

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override;

        [[nodiscard]] const mfem::Operator &GetJacobian() const noexcept;
        [[nodiscard]] const mfem::Solver &GetInversePreconditioner() const noexcept;

    private:
        const mfem::Operator *m_jacobian;
        const mfem::Solver *m_inversePreconditioner;
        mutable mfem::Vector m_preconditionedDirection;
    };

    struct IterationResidualMeasurement final {
        int iteration;
        double reportedNorm;
        bool final;
    };

    class ResidualHistoryMonitor final : public mfem::IterativeSolverMonitor {
    public:
        void Reset() override;

        void MonitorResidual(
            int iteration,
            double norm,
            const mfem::Vector &residual,
            bool final
        ) override;

        [[nodiscard]] const std::vector<IterationResidualMeasurement> &GetHistory() const noexcept;

    private:
        std::vector<IterationResidualMeasurement> m_history;
    };

    struct ResidualBlockMeasurement final {
        std::string stableId;
        int size{0};
        double descriptorScale{1.0};
        double rightHandSideNorm{0.0};
        double trueResidualNorm{0.0};
        double blockRelativeResidual{0.0};
        double scaledRightHandSideNorm{0.0};
        double scaledTrueResidualNorm{0.0};
        double contributionToGlobalRelativeResidual{0.0};
        double fractionOfGlobalSquaredResidualNorm{0.0};
    };

    struct DirectResidualMeasurement final {
        double rightHandSideNorm{0.0};
        double trueResidualNorm{0.0};
        double relativeResidual{0.0};
        std::vector<ResidualBlockMeasurement> blocks;
    };

    [[nodiscard]] DirectResidualMeasurement measureDirectResidual(
        const mfem::Operator &jacobian,
        const mfem::Vector &rightHandSide,
        const mfem::Vector &solution,
        std::span<const operators::RootBlockDescriptor> residualBlocks,
        MPI_Comm communicator,
        double denominatorFloor = 1.0e-300
    );

    struct LinearSolveMeasurement final {
        bool solverConverged{false};
        int outerIterations{0};
        double solverReportedInitialNorm{0.0};
        double solverReportedFinalNorm{0.0};
        double solverReportedResidualReduction{0.0};
        double trueResidualDigitsReducedPerJacobianApplication{0.0};
        double solveSecondsMaximumRank{0.0};
        OperatorApplicationStatistics jacobian;
        OperatorApplicationStatistics inversePreconditioner;
        PreconditionerLifecycleStatistics inversePreconditionerLifecycle;
        DirectResidualMeasurement directResidual;
        std::vector<IterationResidualMeasurement> reportedResidualHistory;
    };

    [[nodiscard]] LinearSolveMeasurement measureLinearSolve(
        const mfem::IterativeSolver &iterativeSolver,
        const mfem::Operator &jacobian,
        const mfem::Vector &rightHandSide,
        const mfem::Vector &solution,
        std::span<const operators::RootBlockDescriptor> residualBlocks,
        const OperatorApplicationStatistics &jacobianStatistics,
        const OperatorApplicationStatistics &inversePreconditionerStatistics,
        const PreconditionerLifecycleStatistics &inversePreconditionerLifecycle,
        const ResidualHistoryMonitor &monitor,
        double localSolveSeconds,
        MPI_Comm communicator,
        double denominatorFloor = 1.0e-300
    );

    struct ArnoldiOptions final {
        int krylovDimension{40};
        double breakdownRelativeTolerance{1.0e-13};
        double ritzConvergenceRelativeTolerance{1.0e-8};
        bool reorthogonalize{true};
    };

    struct RitzValueMeasurement final {
        double realPart{0.0};
        double imaginaryPart{0.0};
        double magnitude{0.0};
        double distanceFromOne{0.0};
        double residualEstimate{0.0};
        double relativeResidualEstimate{0.0};
        bool converged{false};
    };

    enum class RitzValueOrdering { closest_to_zero, farthest_from_one, smallest_real_part, largest_magnitude };

    struct ArnoldiSpectralMeasurement final {
        int requestedDimension{0};
        int achievedDimension{0};
        bool invariantSubspaceFound{false};
        std::uint64_t operatorApplications{0};
        double operatorApplicationSecondsMaximumRank{0.0};
        double operatorMaximumApplicationSecondsMaximumRank{0.0};
        double measurementSecondsMaximumRank{0.0};
        double nonApplicationSecondsMaximumRank{0.0};
        int convergedRitzValueCount{0};
        int negativeRealPartCount{0};

        double projectedLargestSingularValue{0.0};
        double projectedSmallestSingularValue{0.0};
        double projectedConditionProxy{0.0};
        double centroidRealPart{0.0};
        double centroidImaginaryPart{0.0};
        double rmsDistanceFromOne{0.0};
        double rmsClusterRadius{0.0};
        double minimumMagnitude{0.0};
        double maximumMagnitude{0.0};
        double minimumRealPart{0.0};
        double maximumRealPart{0.0};
        double maximumAbsoluteImaginaryPart{0.0};
        double conjugatePairDefect{0.0};
        double projectedDepartureFromNormality{0.0};
        double projectedFieldOfValuesMinimumRealPart{0.0};
        double projectedFieldOfValuesMaximumRealPart{0.0};

        std::vector<RitzValueMeasurement> ritzValues;
    };

    [[nodiscard]] ArnoldiSpectralMeasurement measureArnoldiSpectrum(
        const mfem::Operator &operation,
        const mfem::Vector &initialDirection,
        MPI_Comm communicator,
        const ArnoldiOptions &options = {}
    );

    [[nodiscard]] std::vector<RitzValueMeasurement> selectRitzValues(
        const ArnoldiSpectralMeasurement &measurement,
        RitzValueOrdering ordering,
        int count
    );
} // namespace mean_field::solver