module;

#include <concepts>
#include <cmath>
#include <cstdint>
#include <limits>
#include <span>
#include <type_traits>
#include <utility>

export module mean_field:solver.newton_numerics;

export import :solver.enthalpy_admissibility;
export import :solver.linear_backend;

export namespace mean_field::solver::nonlinear {
    struct SampledEnthalpy final { };
    struct PolynomialEnthalpy final {
        EnthalpyAdmissibilityOptions options{};
    };

    struct MonolithicGravity final { };

    // These tolerances control an auxiliary equation solve, not the physical
    // model's acceptance threshold. The complete residual remains authoritative.
    struct GravitySolveOptions final {
        LinearSolveControl control{.relativeTolerance = 1.0e-10, .absoluteTolerance = 1.0e-14,
                                   .maximumIterations = 300};
        linear::FGMRESOptions backend{.restartLength = 80, .printLevel = -1};

        void Validate() const { control.Validate(); backend.Validate(); }
    };

    struct GravityCurvatureCorrection final {
        GravitySolveOptions gravity{};
    };

    struct EliminatedGravity final {
        GravitySolveOptions gravity{};
    };

    template <typename Candidate>
    concept EnthalpyNumerics = std::same_as<Candidate, SampledEnthalpy> ||
                              std::same_as<Candidate, PolynomialEnthalpy>;

    template <typename Candidate>
    concept GravityNumerics = std::same_as<Candidate, MonolithicGravity> ||
                              std::same_as<Candidate, GravityCurvatureCorrection> ||
                              std::same_as<Candidate, EliminatedGravity>;

    // Types select algorithms; tolerances and the mesh-dependent work sizes
    // remain runtime data. Adding a strategy cannot silently change a model's
    // equation set or impose an additional specification such as rho_c.
    template <EnthalpyNumerics Enthalpy = SampledEnthalpy, GravityNumerics Gravity = MonolithicGravity>
    struct NewtonNumerics final {
        using EnthalpyType = Enthalpy;
        using GravityType = Gravity;
        [[no_unique_address]] Enthalpy enthalpy{};
        [[no_unique_address]] Gravity gravity{};
    };

    template <typename Enthalpy, typename Gravity>
    NewtonNumerics(Enthalpy, Gravity) -> NewtonNumerics<Enthalpy, Gravity>;

    template <typename> struct IsNewtonNumerics : std::false_type { };
    template <typename Enthalpy, typename Gravity>
    struct IsNewtonNumerics<NewtonNumerics<Enthalpy, Gravity>> : std::true_type { };
    template <typename Candidate>
    concept NewtonNumericalConfiguration = IsNewtonNumerics<std::remove_cvref_t<Candidate>>::value;

    enum class LinearReportSpace : std::uint8_t { full_system, reduced_system, gravity_subsystem };

    struct NumericalStrategyDiagnostics final {
        int polynomialChecks{0};
        // Nonzero repair proposals include directions later rejected by a
        // descent check; this measures repair activity, not accepted steps.
        int directionRepairs{0};
        double lastRepairFraction{0.0};
        double lastMeritDirectionalDerivative{0.0};
        std::uint64_t gravitySolves{0};
        std::uint64_t gravityIterations{0};
        std::uint64_t gravityFailures{0};
        std::uint64_t gravityPreparedFailures{0};
        bool compactGravityOperator{false};
        bool compactGravityInverse{false};
        // Coordinates of the evaluation's lastLinearSolve report. A later
        // feasibility relift or trial solve has its own auxiliary diagnostics
        // and does not replace that report or relabel its stored norms.
        LinearReportSpace lastLinearReportSpace{LinearReportSpace::full_system};
        double gravitySolveSeconds{0.0};
        // Last completed auxiliary solve; a later failure before Solve leaves
        // this value and the cumulative counters unchanged.
        double lastGravityTrueResidual{0.0};
        // Current gravity attempt: NaN until that attempt measures the value.
        // Candidates rejected before invoking gravity leave its prior attempt
        // untouched because no auxiliary work was attempted.
        double lastGravityCorrectionNorm{0.0};
        double lastGravityPreparedDefect{0.0};
        double lastGravityAffineDisagreement{0.0};
        double lastGravityArithmeticAllowance{0.0};
        double lastGravityPreparedTargetThreshold{0.0};
        double initialGravityResidual{0.0};
        double completedInitialGravityResidual{0.0};
        double lastLiftedRelativeTrueResidual{0.0};
    };
}

export namespace mean_field::solver::detail {
    struct GravityPreparedResidualAssessment final {
        bool admissible{false};
        double arithmeticAllowance{std::numeric_limits<double>::quiet_NaN()};
        double targetThreshold{std::numeric_limits<double>::quiet_NaN()};
    };

    [[nodiscard]] inline GravityPreparedResidualAssessment AssessGravityPreparedResidual(
        const std::span<const double> contributionNorms,
        const double affineDisagreement,
        const double targetDefect,
        const double measuredInnerResidual
    ) {
        GravityPreparedResidualAssessment result;
        if (!std::isfinite(affineDisagreement) || affineDisagreement < 0.0 ||
            !std::isfinite(targetDefect) || targetDefect < 0.0 ||
            !std::isfinite(measuredInnerResidual) || measuredInnerResidual < 0.0)
            return result;
        long double scale = 0.0;
        for (const double norm : contributionNorms) {
            if (!std::isfinite(norm) || norm < 0.0) return result;
            scale += norm;
        }
        // This arithmetic estimate includes large terms that cancel in G.
        // It is not a rigorous error bound or a replacement for either the
        // independently checked inner residual or the full nonlinear merit.
        result.arithmeticAllowance = static_cast<double>(256.0L * std::numeric_limits<double>::epsilon() * scale);
        result.targetThreshold = measuredInnerResidual + result.arithmeticAllowance;
        result.admissible = std::isfinite(result.arithmeticAllowance) && std::isfinite(result.targetThreshold) &&
                            affineDisagreement <= result.arithmeticAllowance &&
                            targetDefect <= result.targetThreshold;
        return result;
    }
}