#include <algorithm>
#include <cmath>
#include <concepts>
#include <limits>
#include <numbers>
#include <string>
#include <type_traits>
#include <utility>
#include <vector>

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

import mean_field;
import test_helpers;

namespace stellar_numerical_strategy_test {
    using namespace mean_field;
    namespace nonlinear = solver::nonlinear;

    struct UnregisteredEnthalpy { };
    struct UnregisteredGravity { };
    template <typename Enthalpy, typename Gravity>
    concept MakesNumerics = requires { typename nonlinear::NewtonNumerics<Enthalpy, Gravity>; };
    static_assert(!MakesNumerics<UnregisteredEnthalpy, nonlinear::MonolithicGravity>);
    static_assert(!MakesNumerics<nonlinear::SampledEnthalpy, UnregisteredGravity>);
    static_assert(!nonlinear::NewtonNumericalConfiguration<int>);
    static_assert(std::same_as<decltype(nonlinear::Newton{}), nonlinear::Newton<>>);
    static_assert(std::same_as<decltype(nonlinear::Newton(nonlinear::NewtonOptions{})), nonlinear::Newton<>>);
    static_assert(std::same_as<decltype(nonlinear::Newton(nonlinear::NewtonOptions{},
                                                         nonlinear::NormalizedResidualMetric{})), nonlinear::Newton<>>);
    static_assert(MakesNumerics<nonlinear::PolynomialEnthalpy, nonlinear::GravityCurvatureCorrection>);
    static_assert(MakesNumerics<nonlinear::PolynomialEnthalpy, nonlinear::EliminatedGravity>);

    [[nodiscard]] double matchingCentralDensity() {
        return std::numbers::pi_v<double> * utils::MASS / (4.0 * utils::RADIUS * utils::RADIUS * utils::RADIUS);
    }

    template <bool CentralConstraint>
    [[nodiscard]] auto makeModel() {
        const auto equationOfState = eos::Polytrope({.n = 1.0,
            .K = 2.0 * utils::G * utils::RADIUS * utils::RADIUS / std::numbers::pi_v<double>});
        const auto pressureSurface = surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}});
        const auto mass = integral::FixedTotalMass({.Mtotal = dimensions::MassValue{utils::MASS}});
        const auto angularMomentum = integral::FixedAngularMomentum({.Jtotal = dimensions::AngularMomentumValue{0.05}});
        if constexpr (CentralConstraint)
            return model::StellarModel(equationOfState, pressureSurface, mass, angularMomentum,
                constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{matchingCentralDensity()}}));
        else
            return model::StellarModel(equationOfState, pressureSurface, mass, angularMomentum);
    }

    template <bool CentralConstraint>
    [[nodiscard]] auto makeContext() {
        const auto args = test_utils::setup_args();
        auto finiteElements = fem::setup_fem(args.mesh_file, args, 0);
        REQUIRE(finiteElements.okay());
        return solver::makeContext(
            makeModel<CentralConstraint>(), equilibrium::StellarDiscretization{std::move(finiteElements)},
            preconditioning::makePreconditioner(), solver::linear::FGMRES({.restartLength = 2, .printLevel = -1}),
            seed::LaneEmden({.centralDensity = dimensions::DensityValue{matchingCentralDensity()}}));
    }

    struct Snapshot {
        mfem::Vector physical;
        mfem::Vector normalized;
        mfem::Vector residual;
        std::vector<std::string> specifications;
        int retainedSize{};
        int gravitySize{};
    };

    template <typename Context>
    [[nodiscard]] Snapshot snapshot(Context &context) {
        Snapshot snapshot;
        solver::detail::StellarEquilibriumContextDiagnostics::WithState(context, [&](auto &state, const auto &) {
            using Form = typename std::remove_cvref_t<decltype(state.Problem())>::FormType;
            snapshot.physical = state.AcceptedPhysicalState();
            snapshot.normalized = state.acceptedNormalizedState;
            snapshot.residual = state.acceptedNormalizedResidual;
            const auto &manifest = state.Problem().GetManifest();
            for (const auto &descriptor : manifest.constraints())
                snapshot.specifications.emplace_back(descriptor.stableId);
            const auto partition = solver::MakeGravityBlockPartition<Form>(manifest.layout());
            snapshot.retainedSize = partition.RetainedValueSize();
            snapshot.gravitySize = partition.GravityValueSize();
            CHECK(partition.RetainedResidualSize() == partition.RetainedValueSize());
            CHECK(partition.RootValueSize() == state.Problem().StateSize());
            CHECK(partition.GravityResidualSize() == partition.GravityValueSize());
            CHECK(state.physicalInverse->IsCurrent());
            CHECK(state.normalizedInverse->IsCurrent());
        });
        return snapshot;
    }

    [[nodiscard]] double relativeDifference(const mfem::Vector &left, const mfem::Vector &right) {
        REQUIRE(left.Size() == right.Size());
        mfem::Vector difference(left);
        difference -= right;
        return difference.Norml2() / std::max(1.0, left.Norml2());
    }

    void checkPreserved(const Snapshot &before, const Snapshot &after) {
        CHECK(before.specifications == after.specifications);
        CHECK(before.retainedSize == after.retainedSize);
        CHECK(before.gravitySize == after.gravitySize);
        CHECK(relativeDifference(before.physical, after.physical) < 1.e-13);
        CHECK(relativeDifference(before.normalized, after.normalized) < 1.e-13);
        CHECK(relativeDifference(before.residual, after.residual) < 1.e-13);
    }

    template <typename Context>
    void checkGravityApplications(Context &context) {
        solver::detail::StellarEquilibriumContextDiagnostics::WithState(context, [&](auto &state, const auto &) {
            using Form = typename std::remove_cvref_t<decltype(state.Problem())>::FormType;
            const auto partition = solver::MakeGravityBlockPartition<Form>(state.Problem().GetManifest().layout());
            solver::PreparedGravitySubsystem compact(*state.normalizedOperator, *state.normalizedInverse,
                                                      partition, state.Problem().GetCommunicator());
            solver::PreparedGravitySubsystem reference(*state.normalizedOperator, *state.normalizedInverse,
                partition, state.Problem().GetCommunicator(), solver::linear::FGMRES{},
                solver::RootRestrictedGravityApplications{});
            REQUIRE(compact.UsesCompactOperator());
            REQUIRE(compact.UsesCompactPreconditioner());
            mfem::Vector direction(partition.GravityValueSize()), actual(direction.Size()), expected(direction.Size());
            // Dense directions expose ordering/scaling mistakes which a few
            // isolated basis vectors may miss in a distributed field block.
            for (int trial = 0; trial < 3; ++trial) {
                for (int i = 0; i < direction.Size(); ++i)
                    direction(i) = std::sin(.37*(i+1)*(trial+1)) + .4*std::cos(.11*(i+1)*(trial+2));
                compact.GetOperator().Mult(direction, actual);
                reference.GetOperator().Mult(direction, expected);
                CHECK(relativeDifference(expected, actual) < 2.e-12);
                compact.GetPreconditioner().Mult(direction, actual);
                reference.GetPreconditioner().Mult(direction, expected);
                CHECK(relativeDifference(expected, actual) < 2.e-12);
            }
        });
    }

    [[nodiscard]] nonlinear::NewtonOptions oneAttempt() {
        // A relative threshold of exactly ||b|| need not stop MFEM before
        // its first iteration. Use a separated absolute threshold so this
        // wiring test deliberately retains its zero initial correction.
        return {.relativeTolerance = 0.0, .absoluteTolerance = 0.0, .maximumIterations = 1,
                .linearSolve = {.relativeTolerance = 0.0, .absoluteTolerance = 1.e100, .maximumIterations = 1},
                .backtracking = {.maximumTrials = 1, .geometryPreflight = false}};
    }

    [[nodiscard]] nonlinear::GravitySolveOptions permissiveGravity() {
        // Accepting a zero correction exercises wiring and rollback cheaply;
        // the ordinary residual metric still rejects the unchanged trial.
        // These controls deliberately establish no physical accuracy claim.
        return {.control = {.relativeTolerance = 1.0, .absoluteTolerance = 1.e100, .maximumIterations = 1},
                .backend = {.restartLength = 2, .printLevel = -1}};
    }

    template <typename Context, typename Numerics>
    [[nodiscard]] auto evaluate(Context &context, Numerics numerics) {
        auto instance = solver::make(context, nonlinear::Newton(oneAttempt(),
            nonlinear::NormalizedResidualMetric{}, std::move(numerics)));
        return instance.evaluate();
    }
} // namespace stellar_numerical_strategy_test

TEST_CASE("Typed numerical policies preserve the existing Newton configuration", "[solver][numerical_strategies]") {
    using namespace stellar_numerical_strategy_test;
    const nonlinear::Newton original;
    CHECK(original.options().maximumIterations == 50);
    const auto selected = nonlinear::Newton(nonlinear::NewtonOptions{}, nonlinear::NormalizedResidualMetric{},
        nonlinear::NewtonNumerics{nonlinear::SampledEnthalpy{}, nonlinear::GravityCurvatureCorrection{}});
    CHECK(selected.numerics().gravity.gravity.control.relativeTolerance == 1.e-10);
}

TEST_CASE("Production gravity strategies retain optional model borders and restore rejected trials",
          "[solver][numerical_strategies][gravity][rollback]") {
    using namespace stellar_numerical_strategy_test;
    const auto check = []<bool CentralConstraint>() {
        CAPTURE(CentralConstraint);
        auto context = makeContext<CentralConstraint>();
        const auto original = snapshot(context);
        checkGravityApplications(context);
        checkPreserved(original, snapshot(context));
        REQUIRE(original.gravitySize > 0);
        REQUIRE(original.retainedSize > 0);
        // The fixed mass and angular momentum are always retained. Central
        // density adds one specification only when the model requests it.
        CHECK(original.specifications.size() == (CentralConstraint ? 4 : 3));
        {
            const auto report = evaluate(context, nonlinear::NewtonNumerics{nonlinear::SampledEnthalpy{},
                nonlinear::GravityCurvatureCorrection{permissiveGravity()}});
            CHECK_FALSE(report.converged());
            CHECK(report.diagnostics().attemptedNonlinearIterations == 1);
            CHECK(report.diagnostics().acceptedNonlinearIterations == 0);
            REQUIRE(report.diagnostics().lastLinearSolve.has_value());
            CHECK(report.diagnostics().lastLinearSolve->iterations == 0);
            CHECK(report.diagnostics().numericalStrategies.gravitySolves > 0);
            CHECK(report.diagnostics().numericalStrategies.gravityFailures == 0);
            CHECK(std::isfinite(report.diagnostics().numericalStrategies.lastGravityCorrectionNorm));
            CHECK(std::isfinite(report.diagnostics().numericalStrategies.lastGravityPreparedDefect));
            CHECK(std::isfinite(report.diagnostics().numericalStrategies.lastGravityAffineDisagreement));
            CHECK(std::isfinite(report.diagnostics().numericalStrategies.lastGravityArithmeticAllowance));
            CHECK(std::isfinite(report.diagnostics().numericalStrategies.lastGravityPreparedTargetThreshold));
            CHECK(report.lastAcceptedCheckpointView().valid());
        }
        REQUIRE(context.isReady());
        CHECK_FALSE(context.hasActiveSolver());
        checkPreserved(original, snapshot(context));
        {
            const auto report = evaluate(context, nonlinear::NewtonNumerics{nonlinear::SampledEnthalpy{},
                nonlinear::EliminatedGravity{permissiveGravity()}});
            CHECK_FALSE(report.converged());
            CHECK(report.diagnostics().attemptedNonlinearIterations == 1);
            CHECK(report.diagnostics().acceptedNonlinearIterations == 0);
            REQUIRE(report.diagnostics().lastLinearSolve.has_value());
            CHECK(report.diagnostics().lastLinearSolve->iterations == 0);
            CHECK(report.diagnostics().numericalStrategies.gravitySolves > 0);
            CHECK(report.diagnostics().numericalStrategies.gravityFailures == 0);
            CHECK(std::isfinite(report.diagnostics().numericalStrategies.lastLiftedRelativeTrueResidual));
            CHECK(report.lastAcceptedCheckpointView().valid());
        }
        REQUIRE(context.isReady());
        checkPreserved(original, snapshot(context));
    };
    check.template operator()<false>();
    check.template operator()<true>();
}

TEST_CASE("A failed initial gravity completion preserves the accepted stellar checkpoint",
          "[solver][numerical_strategies][gravity][failure]") {
    using namespace stellar_numerical_strategy_test;
    auto context = makeContext<false>();
    const auto original = snapshot(context);
    nonlinear::GravitySolveOptions strict{
        .control = {.relativeTolerance = 0.0, .absoluteTolerance = 0.0, .maximumIterations = 1},
        .backend = {.restartLength = 1, .printLevel = -1}};
    {
        const auto report = evaluate(context, nonlinear::NewtonNumerics{nonlinear::SampledEnthalpy{},
            nonlinear::EliminatedGravity{strict}});
        REQUIRE_FALSE(report.converged());
        CHECK(report.failure().reason == mean_field::solver::StellarEquilibriumFailureReason::linear_solve_failure);
        CHECK(report.diagnostics().numericalStrategies.initialGravityResidual > 0.0);
        CHECK(report.diagnostics().numericalStrategies.gravitySolves == 1);
        CHECK(report.diagnostics().numericalStrategies.gravityFailures == 1);
        const auto &auxiliary = report.diagnostics().numericalStrategies;
        // A completed but unsuccessful inner solve has a meaningful true
        // defect. No correction/prepared assessment occurred in this attempt.
        CHECK(std::isfinite(auxiliary.lastGravityTrueResidual));
        CHECK(std::isnan(auxiliary.lastGravityCorrectionNorm));
        CHECK(std::isnan(auxiliary.lastGravityPreparedDefect));
        CHECK(std::isnan(auxiliary.lastGravityAffineDisagreement));
        CHECK(std::isnan(auxiliary.lastGravityArithmeticAllowance));
        CHECK(std::isnan(auxiliary.lastGravityPreparedTargetThreshold));
        CHECK(report.diagnostics().acceptedNonlinearIterations == 0);
        CHECK(report.lastAcceptedCheckpointView().valid());
    }
    REQUIRE(context.isReady());
    CHECK_FALSE(context.hasActiveSolver());
    checkPreserved(original, snapshot(context));
}

TEST_CASE("Polynomial admissibility is invoked by each production gravity strategy",
          "[solver][numerical_strategies][enthalpy][rollback]") {
    using namespace stellar_numerical_strategy_test;
    auto context = makeContext<false>();
    const auto original = snapshot(context);
    // A bounded certification budget must fail closed if the seed is not
    // certifiable. If it is certifiable, the zero Newton direction cannot
    // satisfy the required strict merit descent. Either path preserves it.
    const nonlinear::PolynomialEnthalpy polynomial{{.maximumSubdivisionDepth = 4, .maximumSubcells = 64}};
    const auto check = [&](auto gravity) {
        const auto report = evaluate(context, nonlinear::NewtonNumerics{polynomial, gravity});
        CHECK_FALSE(report.converged());
        CHECK(report.diagnostics().numericalStrategies.polynomialChecks > 0);
        CHECK(report.diagnostics().acceptedNonlinearIterations == 0);
        CHECK(report.lastAcceptedCheckpointView().valid());
    };
    check(nonlinear::MonolithicGravity{});
    REQUIRE(context.isReady());
    checkPreserved(original, snapshot(context));
    check(nonlinear::GravityCurvatureCorrection{permissiveGravity()});
    REQUIRE(context.isReady());
    checkPreserved(original, snapshot(context));
    check(nonlinear::EliminatedGravity{permissiveGravity()});
    REQUIRE(context.isReady());
    checkPreserved(original, snapshot(context));
}