#pragma once

#include <cmath>
#include <concepts>
#include <stdexcept>
#include <type_traits>

#include "plateau_preconditioner_benchmark.hpp"

// Include after import mean_field. These owning recipes belong to the
// experiment driver; the public physics-facing solver factory stays unchanged.
namespace experiment::plateau_diagnostics {
    template <class Policy, bool RefinedGravity>
    concept NewtonBenchmarkPolicy =
        (std::same_as<Policy, mean_field::preconditioning::IndependentStellarSubsystems> && RefinedGravity) ||
        (std::same_as<Policy, mean_field::preconditioning::ApproximateStellarBlockLDU> && !RefinedGravity);

    template <class Policy, bool RefinedGravity>
        requires NewtonBenchmarkPolicy<Policy, RefinedGravity>
    class NewtonBenchmarkPreconditioner final
        : public mean_field::preconditioning::StellarPreconditionerPrescriptionTag {
    public:
        explicit NewtonBenchmarkPreconditioner(mean_field::preconditioning::GravityConsistentSchurLDU gravity)
            : m_gravity(gravity) {
            if (gravity.correctionSweeps < 0 || !std::isfinite(gravity.relaxation) || gravity.relaxation <= 0.0)
                throw std::invalid_argument("Newton benchmark gravity controls must be finite and valid.");
        }

        [[nodiscard]] mean_field::preconditioning::GravityConsistentSchurLDU GravityPolicy() const noexcept {
            return m_gravity;
        }

    private:
        // A recipe can outlive command-line parsing. Keep numeric controls by
        // value; prepared objects acquire their problem binding only via ADL.
        mean_field::preconditioning::GravityConsistentSchurLDU m_gravity;
    };

    template <
        class Policy,
        bool RefinedGravity,
        mean_field::preconditioning::DefaultStellarPreconditionerAvailableFor Problem>
        requires NewtonBenchmarkPolicy<
            Policy,
            RefinedGravity>
    [[nodiscard]] auto prepareStellarPreconditioner(
        NewtonBenchmarkPreconditioner<
            Policy,
            RefinedGravity> prescription,
        const Problem &problem
    ) {
        PreconditionerBenchmarkOptions options;
        options.schurCorrectionSweeps = prescription.GravityPolicy().correctionSweeps;
        options.schurRelaxation       = prescription.GravityPolicy().relaxation;
        // Returning the benchmark's concrete prepared type lets the existing
        // context own refreshes and rebuilt scalar-border completion unchanged.
        return preconditioner_benchmark_detail::Prepare<Policy, false, RefinedGravity>(problem, options);
    }

    using GravitySchurNewtonPreconditioner =
        NewtonBenchmarkPreconditioner<mean_field::preconditioning::IndependentStellarSubsystems, true>;
    using StellarLDUNewtonPreconditioner =
        NewtonBenchmarkPreconditioner<mean_field::preconditioning::ApproximateStellarBlockLDU, false>;

    static_assert(mean_field::preconditioning::StellarPreconditionerPrescription<GravitySchurNewtonPreconditioner>);
    static_assert(mean_field::preconditioning::StellarPreconditionerPrescription<StellarLDUNewtonPreconditioner>);
    static_assert(std::is_trivially_copyable_v<GravitySchurNewtonPreconditioner>);
    static_assert(std::is_trivially_copyable_v<StellarLDUNewtonPreconditioner>);
} // namespace experiment::plateau_diagnostics