#pragma once

// Include after importing mean_field. This is an experiment-local inverse;
// it does not replace or prepare the authoritative normalized Jacobian.
#include <algorithm>
#include <array>
#include <cmath>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <memory>
#include <span>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

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

namespace experiment::plateau_diagnostics {
    struct CoupledShapeCoarseOptions final {
        double relativeRankTolerance{1.e-10};
        double minimumSubspaceOverlap{1.e-8};
    };

    enum class CoupledShapeCoarseFallback { none, no_independent_modes, insufficient_subspace_overlap };
    enum class CoupledShapeCoarseModeStatus { retained, dependent_source, dependent_filtered_image };

    struct CoupledShapeCoarseModeReport final {
        // QR admission is retained even if the complete space later falls back.
        CoupledShapeCoarseModeStatus status{CoupledShapeCoarseModeStatus::dependent_source};
        double normalizedInputNorm{0.0};
        double sourceImageNorm{0.0};
        double sourceIndependentFraction{0.0};
        double filteredImageNorm{0.0};
        double filteredIndependentFraction{0.0};
    };

    struct CoupledShapeCoarseStatistics final {
        int inputModes{0};
        int independentModes{0};
        int activeModes{0};
        int sourceDependentModes{0};
        int imageDependentModes{0};
        CoupledShapeCoarseFallback fallback{CoupledShapeCoarseFallback::no_independent_modes};
        double reciprocalOverlapCondition{std::numeric_limits<double>::quiet_NaN()};
        double minimumOverlapLowerBound{std::numeric_limits<double>::quiet_NaN()};
        std::uint64_t setups{0};
        std::uint64_t setupOperatorApplications{0};
        std::uint64_t setupInverseApplications{0};
        double lastSetupSeconds{0.0};
        double lastSetupOperatorSeconds{0.0};
        double lastSetupInverseSeconds{0.0};
        double totalSetupSeconds{0.0};
        // Application clocks are rank-local; aggregate them outside the timed
        // solve. The only application reductions validate and project vectors.
        std::uint64_t applications{0};
        std::uint64_t inverseApplications{0};
        std::uint64_t failedApplications{0};
        double applicationSecondsLocal{0.0};
        double inverseSecondsLocal{0.0};
        std::size_t persistentVectorBytes{0};
    };

    template <class Candidate>
    concept CoupledCoarseNormalizedOperator =
        std::derived_from<Candidate, mfem::Operator> &&
        requires(const Candidate &operation, const mfem::Vector &physical, mfem::Vector &normalized) {
            { operation.IsPrepared() } -> std::same_as<bool>;
            { operation.GetProblem().GetPreparationGeneration() } -> std::convertible_to<std::uint64_t>;
            { operation.NormalizeState(physical, normalized) } -> std::same_as<void>;
        };

    template <class Candidate>
    concept CoupledCoarseStationaryInverse =
        std::derived_from<Candidate, mfem::Solver> &&
        mean_field::solver::LinearPreconditionerApplicationContractAvailable<Candidate> &&
        (mean_field::solver::linearPreconditionerApplicationContract<Candidate> ==
         mean_field::preconditioning::ApplicationContract::stationary_linear) &&
        requires(const Candidate &inverse) {
            { inverse.IsCurrent() } -> std::same_as<bool>;
        };

    namespace coupled_shape_detail {
        template <class Inverse>
        using RootBlock = std::remove_cvref_t<decltype(std::declval<const Inverse &>().GetBlock())>;
        template <class Inverse>
        using MaterialBlock = typename RootBlock<Inverse>::StructureComponent::MaterialSurfaceComponent;

        template <mean_field::equilibrium::DiscretizedStellarEquilibriumProblem Problem>
        [[nodiscard]] mean_field::deformation::PreparedNodalRadialSurface
        MakeDefaultRadialSurface(const Problem &problem) {
            using namespace mean_field;
            const auto &physical    = problem.GetPhysicalOperator();
            const auto &domain      = physical.GetDomainDeformation();
            const auto &fem         = physical.GetHydrostaticOperator().GetFEM();
            const auto dependencies = domain.discretizationDependencies();
            const auto communicator = problem.GetCommunicator();
            const auto all          = [communicator](bool condition) {
                int local = condition, global = 0;
                if (MPI_Allreduce(&local, &global, 1, MPI_INT, MPI_MIN, communicator) != MPI_SUCCESS)
                    throw std::runtime_error("Could not validate the radial shape basis geometry.");
                return global != 0;
            };
            if (!all(
                    physical.IsPrepared() && fem.mesh && fem.surfaceDeformationFes && fem.displacementFes &&
                    dependencies.isCurrent() && dependencies.physicalMeshIdentity == fem.mesh.get() &&
                    dependencies.surfaceScalarSpaceIdentity == fem.surfaceDeformationFes.get() &&
                    dependencies.volumeDisplacementSpaceIdentity == fem.displacementFes.get()
                ))
                throw std::logic_error("Radial shape modes require the current physical deformation spaces.");
            const auto &scalar = *fem.surfaceDeformationFes;
            const auto &vector = *fem.displacementFes;
            mfem::Vector center(3);
            center = 0.0;
            const deformation::NodalRadialSurface prescription(std::move(center));
            if (!all(
                    fem.mesh->SpaceDimension() == 3 && scalar.GetVDim() == 1 && vector.GetVDim() == 3 &&
                    vector.GetOrdering() == mfem::Ordering::byNODES &&
                    vector.GetTrueVSize() == 3 * scalar.GetTrueVSize() && vector.FEColl() == scalar.FEColl() &&
                    domain.descriptor().surfaceDeformation == prescription.descriptor()
                ))
                throw std::invalid_argument(
                    "The shape benchmark requires the default radial surface and its scalar companion."
                );
            // StellarEquilibriumProblem currently compiles this zero-center
            // prescription, then erases its type. Reuse the same compiler and
            // exact FEM boundary map instead of reproducing its geometry math.
            auto radial = deformation::compileSurfaceDeformationPrescription(
                prescription,
                deformation::SurfaceDeformationCompilationContext{
                    *fem.surfaceDeformationFes, field::make_stellar_surface_scalar_dof_map(*fem.surfaceDeformationFes)
                }
            );
            if (!all(radial.parameterCount() == domain.parameterCount()))
                throw std::invalid_argument("The reconstructed radial map does not match the physical parameters.");
            // The runtime descriptor does not retain the center. Verify actual
            // surface motion against the reconstructed directions; metadata
            // alone must not admit a future shifted-center prescription.
            mfem::Vector unit(domain.parameterCount()), action(domain.volumeDisplacementSize());
            unit = 1.0;
            domain.applyJacobian(physical.GetSurfaceDeformationParameters(), unit, action);
            bool matches         = true;
            const int scalarSize = scalar.GetTrueVSize();
            for (int i = 0; i < radial.parameterCount(); ++i)
                for (int component = 0; component < 3; ++component) {
                    // MFEM byNODES stores component-sized blocks, as required
                    // by the existing radial extension's true-DOF mapping.
                    const int dof         = radial.surfaceDofMap().volume_true_dof(i) + component * scalarSize;
                    const double observed = action(dof), expected = radial.radialDirection(i, component);
                    matches = matches && std::isfinite(observed) && std::isfinite(expected) &&
                              std::abs(observed - expected) <= 512.0 * std::numeric_limits<double>::epsilon();
                }
            if (!all(matches && dependencies.isCurrent()))
                throw std::invalid_argument("The actual surface motion disagrees with the zero-center radial basis.");
            return radial;
        }

        template <class Problem, class Inverse>
        concept SupportsShapeModes =
            mean_field::equilibrium::DiscretizedStellarEquilibriumProblem<Problem> &&
            requires(const Problem &problem, const Inverse &inverse) {
                typename MaterialBlock<Inverse>::Factorization;
                requires std::same_as<
                    typename MaterialBlock<Inverse>::Factorization,
                    mean_field::preconditioning::SurfaceThenMaterialTriangular>;
                { inverse.GetProblem() } -> std::same_as<const Problem &>;
                { inverse.IsCurrent() } -> std::same_as<bool>;
                inverse.GetCoordinateMap().GetCorrectionRanges();
                inverse.GetGroupedPreconditioner().GetStructurePreconditioner().GetCrossCouplings();
                inverse.GetGroupedPreconditioner()
                    .GetStructurePreconditioner()
                    .GetMaterialSurfacePreconditioner()
                    .GetCoupledOperator();
                inverse.GetGroupedPreconditioner().GetStructurePreconditioner().GetGravityPreconditioner();
                { problem.GetPreparationGeneration() } -> std::convertible_to<std::uint64_t>;
                {
                    MakeDefaultRadialSurface(problem)
                } -> std::same_as<mean_field::deformation::PreparedNodalRadialSurface>;
            };
    } // namespace coupled_shape_detail

    // These are geometry modes, not Newton-RHS samples. Complete their material
    // and gravity response with existing fixed inverses; neither completion is
    // an exact elimination. Scalar unknowns stay zero, while the later full-J
    // images retain every scalar equation, including an optional central border.
    // Construction performs nine explicit surface-to-material actions, material
    // inverse applications, material-to-gravity actions, and gravity applications.
    // The material inverse also performs its usual internal coupling actions.
    // One extra domain Jacobian action verifies the reconstructed radial basis.
    template <
        class Problem,
        class PhysicalInverse>
        requires coupled_shape_detail::SupportsShapeModes<
            Problem,
            PhysicalInverse>
    [[nodiscard]] std::vector<mfem::Vector> MakeCoupledShapeModes(
        const Problem &problem,
        const PhysicalInverse &inverse
    ) {
        using namespace mean_field;
        using namespace utils::blocks;
        using Material    = coupled_shape_detail::MaterialBlock<PhysicalInverse>;
        using Root        = coupled_shape_detail::RootBlock<PhysicalInverse>;
        using Corrections = typename Material::CorrectionBlocks;
        using Residuals   = typename Material::ResidualBlocks;
        using Gravity     = typename Root::StructureComponent::GravityComponent;
        // The borrowed coupling operator exposes this exact three-field packing.
        // Constraining the descriptor prevents a superficially compatible custom
        // material recipe from silently assigning a physical mode to wrong rows.
        static_assert(
            std::same_as<
                Corrections,
                type_list<density::mass::value, surface_deformation::parameters::value, enthalpy::specific::value>>
        );
        static_assert(std::same_as<
                      Residuals, type_list<
                                     density::mass::residual, surface_deformation::shape_equilibrium::residual,
                                     enthalpy::specific::residual>>);
        static_assert(
            std::same_as<
                typename Gravity::CorrectionBlocks, type_list<gravity::gradient::value, gravity::poisson::value>>
        );
        static_assert(
            std::same_as<
                typename Gravity::ResidualBlocks, type_list<gravity::gradient::residual, gravity::poisson::residual>>
        );
        constexpr auto densityRow    = type_index_v<density::mass::residual, Residuals>;
        constexpr auto enthalpyRow   = type_index_v<enthalpy::specific::residual, Residuals>;
        constexpr auto shapeColumn   = type_index_v<surface_deformation::parameters::value, Corrections>;
        constexpr auto materialFirst = type_index_v<density::mass::value, typename Root::CorrectionBlocks>;
        constexpr auto gravityFirst  = type_index_v<gravity::gradient::value, typename Root::CorrectionBlocks>;
        const auto communicator      = problem.GetCommunicator();
        if (communicator == MPI_COMM_NULL)
            throw std::invalid_argument("Physical shape modes require a valid communicator.");
        const auto all = [communicator](bool condition) {
            int local = condition, global = 0;
            if (MPI_Allreduce(&local, &global, 1, MPI_INT, MPI_MIN, communicator) != MPI_SUCCESS)
                throw std::runtime_error("Could not validate collective physical shape modes.");
            return global != 0;
        };
        if (!all(
                &inverse.GetProblem() == &problem && inverse.IsCurrent() &&
                problem.GetPhysicalOperator().IsPrepared() && !inverse.iterative_mode
            ))
            throw std::logic_error("Physical shape modes require the current inverse of this problem.");
        const auto generation      = problem.GetPreparationGeneration();
        const auto &structure      = inverse.GetGroupedPreconditioner().GetStructurePreconditioner();
        const auto &material       = structure.GetMaterialSurfacePreconditioner();
        const auto &gravityInverse = structure.GetGravityPreconditioner();
        const auto &couplings      = material.GetCoupledOperator();
        const auto &cross          = structure.GetCrossCouplings();
        const auto &offsets        = couplings.GetOffsets();
        const auto &coordinates    = inverse.GetCoordinateMap();
        const auto &ranges         = coordinates.GetCorrectionRanges();
        const auto radial          = coupled_shape_detail::MakeDefaultRadialSurface(problem);
        const auto materialOffset  = ranges[materialFirst].preconditionerOffset;
        const auto gravityOffset   = ranges[gravityFirst].preconditionerOffset;
        const auto shapeSize       = offsets[shapeColumn + 1] - offsets[shapeColumn];
        if (!all(
                radial.spatialDimension() == 3 && radial.parameterCount() == shapeSize &&
                material.Height() == material.Width() && gravityInverse.Height() == gravityInverse.Width() &&
                material.Height() == cross.MaterialSize() && gravityInverse.Height() == cross.GravitySize() &&
                materialOffset == 0 && gravityOffset == material.Height() &&
                gravityOffset + gravityInverse.Height() <= coordinates.PreconditionerCorrectionSize() &&
                !material.iterative_mode && !gravityInverse.iterative_mode
            ))
            throw std::invalid_argument("Physical shape modes require compatible radial three-dimensional blocks.");
        mfem::Vector materialRhs(material.Width()), materialAction(material.Height());
        mfem::Vector gravityRhs(gravityInverse.Width()), gravityAction(gravityInverse.Height());
        mfem::Vector grouped(coordinates.PreconditionerCorrectionSize());
        mfem::Vector shape(shapeSize);
        mfem::Vector densityRhs(materialRhs, offsets[densityRow], offsets[densityRow + 1] - offsets[densityRow]);
        mfem::Vector enthalpyRhs(materialRhs, offsets[enthalpyRow], offsets[enthalpyRow + 1] - offsets[enthalpyRow]);
        mfem::Vector shapeAction(materialAction, offsets[shapeColumn], shapeSize);
        mfem::Vector groupedMaterial(grouped, materialOffset, materialAction.Size());
        mfem::Vector groupedGravity(grouped, gravityOffset, gravityAction.Size());
        std::vector<mfem::Vector> result;
        result.reserve(9);
        for (int mode = 0; mode < 9; ++mode) {
            bool finite = true;
            for (int i = 0; i < shapeSize; ++i) {
                const double x = radial.radialDirection(i, 0), y = radial.radialDirection(i, 1),
                             z = radial.radialDirection(i, 2), radius = radial.referenceRadius(i);
                // All real harmonics through l=2, up to irrelevant constants.
                // Their span is rotation invariant; individual columns are not
                // asserted to be exact translation or homology null directions.
                const std::array<double, 9> angular{1.0, x, y, z, x * y, x * z, y * z, x * x - y * y, 3 * z * z - 1};
                shape(i) = radius * angular[mode];
                finite   = finite && std::isfinite(shape(i)) && radius > 0.0;
            }
            if (!all(finite))
                throw std::runtime_error("A physical radial shape mode is nonfinite or has invalid radius.");
            materialRhs = 0.0;
            couplings.ApplySurfaceToMaterial(shape, densityRhs, enthalpyRhs);
            densityRhs *= -1.0;
            enthalpyRhs *= -1.0;
            densityRhs.SyncAliasMemory(materialRhs);
            enthalpyRhs.SyncAliasMemory(materialRhs);
            // Zero shape RHS in this policy holds its returned shape at zero.
            // We then restore the prescribed shape rather than solving A_qq.
            material.Mult(materialRhs, materialAction);
            shapeAction = shape;
            shapeAction.SyncAliasMemory(materialAction);
            cross.ApplyMaterialToGravity(materialAction, gravityRhs);
            gravityRhs *= -1.0;
            gravityInverse.Mult(gravityRhs, gravityAction);
            grouped         = 0.0;
            groupedMaterial = materialAction;
            groupedGravity  = gravityAction;
            groupedMaterial.SyncAliasMemory(grouped);
            groupedGravity.SyncAliasMemory(grouped);
            result.emplace_back(problem.StateSize());
            coordinates.UnpackCorrection(grouped, result.back());
            finite = true;
            for (int i = 0; i < result.back().Size(); ++i)
                finite = finite && std::isfinite(result.back()(i));
            if (!all(finite && inverse.IsCurrent() && generation == problem.GetPreparationGeneration()))
                throw std::runtime_error("Physical shape completion failed or its preparation changed.");
        }
        return result;
    }

    /*
     * The caller supplies deterministic, physically completed shape directions
     * before exposing a Newton RHS. With A the full normalized Jacobian, form
     * V = orth(A T), U = P0 V, and A U = Q R. Then apply
     *
     *   P r = P0 [r + (V R^-1 - Q) Q^T r].
     *
     * This equals U R^-1 Q^T r + P0 (I-Q Q^T)r and is exact on range(A U)
     * in exact arithmetic. The P0 A filtering can change the supplied physical
     * correction space; this is not an exact completed-shape Schur inverse.
     *
     * For stationary P0, the extra right factor is nonsingular iff Q^T V is,
     * provided R is nonsingular. Orthonormal V and Q let the small overlap
     * matrix expose near-orthogonal subspaces. A failed overlap check selects
     * P0 explicitly; it never regularizes A or adds a physical constraint.
     * Invertibility/quality of P0 itself and full-system true-defect acceptance
     * remain the caller's responsibility. No solver improvement is assumed.
     */
    template <CoupledCoarseNormalizedOperator Operation, CoupledCoarseStationaryInverse Inverse>
    class PreparedCoupledShapeCoarse final : public mfem::Solver {
    public:
        static constexpr auto applicationContract = mean_field::preconditioning::ApplicationContract::stationary_linear;
        static constexpr int maximumModes         = 16;

        PreparedCoupledShapeCoarse(
            const Operation &operation,
            const Inverse &inverse,
            std::span<const mfem::Vector> physicalModes,
            MPI_Comm communicator,
            CoupledShapeCoarseOptions options = {}
        )
            : mfem::Solver(
                  operation.Width(),
                  operation.Height()
              ),
              m_operation(&operation),
              m_inverse(&inverse),
              m_communicator(communicator),
              m_options(options),
              m_residualWork(operation.Height()),
              m_actionWork(operation.Width()) {
            Refresh(physicalModes);
        }

        PreparedCoupledShapeCoarse(const PreparedCoupledShapeCoarse &)            = delete;
        PreparedCoupledShapeCoarse &operator=(const PreparedCoupledShapeCoarse &) = delete;
        PreparedCoupledShapeCoarse(PreparedCoupledShapeCoarse &&)                 = delete;
        PreparedCoupledShapeCoarse &operator=(PreparedCoupledShapeCoarse &&)      = delete;

        void SetOperator(const mfem::Operator &operation) override {
            if (&operation != m_operation)
                throw std::invalid_argument("A coupled coarse inverse cannot bind a different Jacobian.");
            RequireCurrent();
        }

        [[nodiscard]] bool IsCurrent() const {
            return m_ready && m_operation->IsPrepared() && m_inverse->IsCurrent() && !m_inverse->iterative_mode &&
                   m_generation == m_operation->GetProblem().GetPreparationGeneration();
        }

        [[nodiscard]] bool UsesCoarseCorrection() const noexcept {
            return !m_sources.empty();
        }
        [[nodiscard]] MPI_Comm GetCommunicator() const noexcept {
            return m_communicator;
        }
        [[nodiscard]] const CoupledShapeCoarseStatistics &GetStatistics() const noexcept {
            return m_statistics;
        }
        [[nodiscard]] std::span<const CoupledShapeCoarseModeReport> GetModeReports() const noexcept {
            return m_reports;
        }

        // These normalized residual-space columns permit independent checks:
        // for each v, compare P(A(P0 v)) with P0 v using the authoritative A.
        [[nodiscard]] std::span<const mfem::Vector> GetSourceBasis() const noexcept {
            return m_sources;
        }

        void Refresh(std::span<const mfem::Vector> physicalModes) {
            m_ready            = false;
            const double start = MPI_Wtime();
            ValidatePreparation(physicalModes);
            m_generation = m_operation->GetProblem().GetPreparationGeneration();
            m_sources.clear();
            m_images.clear();
            m_reports.clear();
            m_sources.reserve(physicalModes.size());
            m_images.reserve(physicalModes.size());
            m_reports.reserve(physicalModes.size());
            m_upper.fill(0.0);
            m_statistics.inputModes       = static_cast<int>(physicalModes.size());
            m_statistics.independentModes = m_statistics.activeModes = 0;
            m_statistics.sourceDependentModes = m_statistics.imageDependentModes = 0;
            m_statistics.fallback                                                = CoupledShapeCoarseFallback::none;
            m_statistics.reciprocalOverlapCondition = std::numeric_limits<double>::quiet_NaN();
            m_statistics.minimumOverlapLowerBound   = std::numeric_limits<double>::quiet_NaN();
            double operatorSeconds = 0.0, inverseSeconds = 0.0;
            mfem::Vector value(Height()), source(Width()), image(Width());

            for (const auto &physical : physicalModes) {
                CoupledShapeCoarseModeReport report;
                m_operation->NormalizeState(physical, value);
                RequireFinite(value);
                report.normalizedInputNorm = Norm(value);
                if (report.normalizedInputNorm == 0.0) {
                    ++m_statistics.sourceDependentModes;
                    m_reports.push_back(report);
                    continue;
                }
                value /= report.normalizedInputNorm;
                double actionStart = MPI_Wtime();
                m_operation->Mult(value, source);
                operatorSeconds += MPI_Wtime() - actionStart;
                ++m_statistics.setupOperatorApplications;
                RequireFinite(source);
                report.sourceImageNorm = Norm(source);
                std::array<double, maximumModes> ignored{};
                Reorthogonalize(source, m_sources, ignored);
                const double sourceNorm          = Norm(source);
                report.sourceIndependentFraction = Fraction(sourceNorm, report.sourceImageNorm);
                if (sourceNorm == 0.0 || report.sourceIndependentFraction <= m_options.relativeRankTolerance) {
                    ++m_statistics.sourceDependentModes;
                    m_reports.push_back(report);
                    continue;
                }
                source /= sourceNorm;
                actionStart = MPI_Wtime();
                m_inverse->Mult(source, value);
                inverseSeconds += MPI_Wtime() - actionStart;
                ++m_statistics.setupInverseApplications;
                RequireFinite(value);
                actionStart = MPI_Wtime();
                m_operation->Mult(value, image);
                operatorSeconds += MPI_Wtime() - actionStart;
                ++m_statistics.setupOperatorApplications;
                RequireFinite(image);
                report.filteredImageNorm = Norm(image);
                std::array<double, maximumModes> upperColumn{};
                Reorthogonalize(image, m_images, upperColumn);
                const double imageNorm             = Norm(image);
                report.filteredIndependentFraction = Fraction(imageNorm, report.filteredImageNorm);
                if (imageNorm == 0.0 || report.filteredIndependentFraction <= m_options.relativeRankTolerance) {
                    report.status = CoupledShapeCoarseModeStatus::dependent_filtered_image;
                    ++m_statistics.imageDependentModes;
                    m_reports.push_back(report);
                    continue;
                }
                const int column = static_cast<int>(m_sources.size());
                for (int row = 0; row < column; ++row)
                    m_upper[row * maximumModes + column] = upperColumn[row];
                m_upper[column * maximumModes + column] = imageNorm;
                image /= imageNorm;
                m_sources.emplace_back(source);
                m_images.emplace_back(image);
                report.status = CoupledShapeCoarseModeStatus::retained;
                m_reports.push_back(report);
            }
            m_statistics.independentModes = static_cast<int>(m_sources.size());
            if (m_sources.empty()) {
                m_statistics.fallback = CoupledShapeCoarseFallback::no_independent_modes;
            } else {
                const auto [condition, overlap]         = CheckOverlap();
                m_statistics.reciprocalOverlapCondition = condition;
                m_statistics.minimumOverlapLowerBound   = overlap;
                if (!std::isfinite(overlap) || overlap < m_options.minimumSubspaceOverlap) {
                    m_statistics.fallback = CoupledShapeCoarseFallback::insufficient_subspace_overlap;
                    m_sources.clear();
                    m_images.clear();
                }
            }
            m_statistics.activeModes = static_cast<int>(m_sources.size());
            // Payload bytes exclude allocator metadata and borrowed objects.
            m_statistics.persistentVectorBytes =
                sizeof(mfem::real_t) * (static_cast<std::size_t>(m_residualWork.Size()) + m_actionWork.Size() +
                                        2 * m_sources.size() * static_cast<std::size_t>(Width()));
            if (!All(
                    m_operation->IsPrepared() && m_inverse->IsCurrent() &&
                    m_generation == m_operation->GetProblem().GetPreparationGeneration()
                ))
                throw std::logic_error("The coupled coarse preparation changed during setup.");
            m_statistics.lastSetupOperatorSeconds = Maximum(operatorSeconds);
            m_statistics.lastSetupInverseSeconds  = Maximum(inverseSeconds);
            m_statistics.lastSetupSeconds         = Maximum(MPI_Wtime() - start);
            m_statistics.totalSetupSeconds += m_statistics.lastSetupSeconds;
            ++m_statistics.setups;
            m_ready = true;
        }

        void Mult(
            const mfem::Vector &rightHandSide,
            mfem::Vector &action
        ) const override {
            const double start = MPI_Wtime();
            ++m_statistics.applications;
            try {
                RequireCurrent();
                if (!All(rightHandSide.Size() == Width() && action.Size() == Height()))
                    throw std::invalid_argument("A coupled coarse application requires preallocated root vectors.");
                RequireFinite(rightHandSide);
                m_residualWork  = rightHandSide;
                const int count = static_cast<int>(m_sources.size());
                if (count != 0) {
                    std::array<double, maximumModes> projection{}, coefficients{};
                    Project(rightHandSide, m_images, projection);
                    coefficients = projection;
                    for (int row = count - 1; row >= 0; --row) {
                        for (int column = row + 1; column < count; ++column)
                            coefficients[row] -= m_upper[row * maximumModes + column] * coefficients[column];
                        coefficients[row] /= m_upper[row * maximumModes + row];
                    }
                    for (int column = 0; column < count; ++column) {
                        m_residualWork.Add(coefficients[column], m_sources[column]);
                        m_residualWork.Add(-projection[column], m_images[column]);
                    }
                    RequireFinite(m_residualWork);
                }
                const double inverseStart = MPI_Wtime();
                ++m_statistics.inverseApplications;
                m_inverse->Mult(m_residualWork, m_actionWork);
                m_statistics.inverseSecondsLocal += MPI_Wtime() - inverseStart;
                RequireFinite(m_actionWork);
                RequireCurrent();
                action = m_actionWork;
                m_statistics.applicationSecondsLocal += MPI_Wtime() - start;
            } catch (...) {
                ++m_statistics.failedApplications;
                m_statistics.applicationSecondsLocal += MPI_Wtime() - start;
                throw;
            }
        }

    private:
        [[nodiscard]] static double Fraction(
            double independent,
            double original
        ) {
            return original == 0.0 ? 0.0 : independent / original;
        }

        [[nodiscard]] bool All(bool value) const {
            int local = value, global = 0;
            if (MPI_Allreduce(&local, &global, 1, MPI_INT, MPI_MIN, m_communicator) != MPI_SUCCESS)
                throw std::runtime_error("Could not validate a collective coupled coarse operation.");
            return global != 0;
        }

        [[nodiscard]] double Maximum(double value) const {
            double result = 0.0;
            if (MPI_Allreduce(&value, &result, 1, MPI_DOUBLE, MPI_MAX, m_communicator) != MPI_SUCCESS)
                throw std::runtime_error("Could not measure coupled coarse setup time.");
            return result;
        }

        void RequireCurrent() const {
            if (!All(IsCurrent()))
                throw std::logic_error("The coupled coarse inverse is stale; refresh its dependencies and basis.");
        }

        void RequireFinite(const mfem::Vector &vector) const {
            bool finite = true;
            for (int i = 0; i < vector.Size(); ++i)
                finite = finite && std::isfinite(vector(i));
            if (!All(finite))
                throw std::runtime_error("A coupled coarse vector is nonfinite.");
        }

        [[nodiscard]] double Norm(const mfem::Vector &vector) const {
            const double local = vector * vector;
            double global      = 0.0;
            if (MPI_Allreduce(&local, &global, 1, MPI_DOUBLE, MPI_SUM, m_communicator) != MPI_SUCCESS)
                throw std::runtime_error("Could not compute a coupled coarse norm.");
            if (!std::isfinite(global) || global < 0.0)
                throw std::runtime_error("A coupled coarse norm is nonfinite.");
            return std::sqrt(global);
        }

        void Project(
            const mfem::Vector &vector,
            const std::vector<mfem::Vector> &basis,
            std::array<
                double,
                maximumModes> &coefficients
        ) const {
            std::array<double, maximumModes> local{};
            for (std::size_t i = 0; i < basis.size(); ++i)
                local[i] = basis[i] * vector;
            if (MPI_Allreduce(
                    local.data(), coefficients.data(), static_cast<int>(basis.size()), MPI_DOUBLE, MPI_SUM,
                    m_communicator
                ) != MPI_SUCCESS)
                throw std::runtime_error("Could not project a coupled coarse vector.");
            for (std::size_t i = 0; i < basis.size(); ++i)
                if (!std::isfinite(coefficients[i]))
                    throw std::runtime_error("A coupled coarse projection is nonfinite.");
        }

        void Reorthogonalize(
            mfem::Vector &vector,
            const std::vector<mfem::Vector> &basis,
            std::array<
                double,
                maximumModes> &coefficients
        ) const {
            for (int pass = 0; pass < 2; ++pass) {
                std::array<double, maximumModes> projection{};
                Project(vector, basis, projection);
                for (std::size_t i = 0; i < basis.size(); ++i) {
                    coefficients[i] += projection[i];
                    vector.Add(-projection[i], basis[i]);
                }
            }
        }

        void ValidatePreparation(std::span<const mfem::Vector> modes) const {
            if (m_communicator == MPI_COMM_NULL)
                throw std::invalid_argument("A coupled coarse inverse requires a valid communicator.");
            const int count  = modes.size() <= maximumModes ? static_cast<int>(modes.size()) : maximumModes + 1;
            int minimumCount = 0, maximumCount = 0;
            if (MPI_Allreduce(&count, &minimumCount, 1, MPI_INT, MPI_MIN, m_communicator) != MPI_SUCCESS ||
                MPI_Allreduce(&count, &maximumCount, 1, MPI_INT, MPI_MAX, m_communicator) != MPI_SUCCESS)
                throw std::runtime_error("Could not compare coupled coarse mode counts.");
            const std::array<double, 2> options{m_options.relativeRankTolerance, m_options.minimumSubspaceOverlap};
            std::array<double, 2> minimumOptions{}, maximumOptions{};
            if (MPI_Allreduce(options.data(), minimumOptions.data(), 2, MPI_DOUBLE, MPI_MIN, m_communicator) !=
                    MPI_SUCCESS ||
                MPI_Allreduce(options.data(), maximumOptions.data(), 2, MPI_DOUBLE, MPI_MAX, m_communicator) !=
                    MPI_SUCCESS)
                throw std::runtime_error("Could not compare coupled coarse preparation options.");
            bool valid = count > 0 && count <= maximumModes && minimumCount == maximumCount &&
                         minimumOptions == maximumOptions && std::isfinite(options[0]) && options[0] > 0.0 &&
                         options[0] < 1.0 && std::isfinite(options[1]) && options[1] > 0.0 && options[1] < 1.0 &&
                         m_operation->IsPrepared() && m_inverse->IsCurrent() && !m_inverse->iterative_mode &&
                         m_operation->Height() == m_operation->Width() && m_inverse->Height() == Height() &&
                         m_inverse->Width() == Width();
            if constexpr (requires { m_inverse->GetNormalizedJacobian(); })
                valid = valid && &m_inverse->GetNormalizedJacobian() == m_operation;
            if (!All(valid))
                throw std::invalid_argument(
                    "Coupled coarse setup needs current compatible objects and common valid options."
                );
            const long long localSize = Width();
            long long globalSize      = 0;
            if (MPI_Allreduce(&localSize, &globalSize, 1, MPI_LONG_LONG, MPI_SUM, m_communicator) != MPI_SUCCESS ||
                globalSize <= 0)
                throw std::invalid_argument("A coupled coarse system must be globally nonempty.");
            for (const auto &mode : modes) {
                if (!All(mode.Size() == Height()))
                    throw std::invalid_argument("Coupled coarse modes must use complete physical root coordinates.");
                RequireFinite(mode);
            }
        }

        [[nodiscard]] std::pair<
            double,
            double>
        CheckOverlap() const {
            const int count = static_cast<int>(m_sources.size());
            std::array<double, maximumModes * maximumModes> local{}, matrix{}, inverse{};
            for (int row = 0; row < count; ++row)
                for (int column = 0; column < count; ++column)
                    local[row * count + column] = m_images[row] * m_sources[column];
            if (MPI_Allreduce(local.data(), matrix.data(), count * count, MPI_DOUBLE, MPI_SUM, m_communicator) !=
                MPI_SUCCESS)
                throw std::runtime_error("Could not assemble the coupled coarse overlap matrix.");
            for (int index = 0; index < count * count; ++index)
                if (!std::isfinite(matrix[index]))
                    return {0.0, 0.0};
            double norm = 0.0;
            for (int row = 0; row < count; ++row) {
                inverse[row * count + row] = 1.0;
                double sum                 = 0.0;
                for (int column = 0; column < count; ++column)
                    sum += std::abs(matrix[row * count + column]);
                norm = std::max(norm, sum);
            }
            if (!std::isfinite(norm) || norm == 0.0)
                return {0.0, 0.0};
            for (int column = 0; column < count; ++column) {
                int pivot = column;
                for (int row = column + 1; row < count; ++row)
                    if (std::abs(matrix[row * count + column]) > std::abs(matrix[pivot * count + column]))
                        pivot = row;
                if (!std::isfinite(matrix[pivot * count + column]) ||
                    std::abs(matrix[pivot * count + column]) <= std::numeric_limits<double>::epsilon() * count * norm)
                    return {0.0, 0.0};
                for (int j = 0; j < count; ++j) {
                    std::swap(matrix[column * count + j], matrix[pivot * count + j]);
                    std::swap(inverse[column * count + j], inverse[pivot * count + j]);
                }
                const double diagonal = matrix[column * count + column];
                for (int j = 0; j < count; ++j) {
                    matrix[column * count + j] /= diagonal;
                    inverse[column * count + j] /= diagonal;
                }
                for (int row = 0; row < count; ++row)
                    if (row != column) {
                        const double factor = matrix[row * count + column];
                        for (int j = 0; j < count; ++j) {
                            matrix[row * count + j] -= factor * matrix[column * count + j];
                            inverse[row * count + j] -= factor * inverse[column * count + j];
                        }
                    }
            }
            double inverseOne = 0.0, inverseInfinity = 0.0;
            for (int index = 0; index < count * count; ++index)
                if (!std::isfinite(inverse[index]))
                    return {0.0, 0.0};
            for (int i = 0; i < count; ++i) {
                double rowSum = 0.0, columnSum = 0.0;
                for (int j = 0; j < count; ++j) {
                    rowSum += std::abs(inverse[i * count + j]);
                    columnSum += std::abs(inverse[j * count + i]);
                }
                inverseInfinity = std::max(inverseInfinity, rowSum);
                inverseOne      = std::max(inverseOne, columnSum);
            }
            if (!std::isfinite(inverseOne) || !std::isfinite(inverseInfinity))
                return {0.0, 0.0};
            // Condition number alone misses epsilon*I. Since ||S^-1||_2 is
            // bounded by sqrt(||S^-1||_1 ||S^-1||_inf), this also checks the
            // absolute overlap of the two orthonormal subspaces. The computed
            // inverse supplies a floating-point estimate, not an interval proof.
            return {1.0 / (norm * inverseInfinity), (1.0 / std::sqrt(inverseOne)) / std::sqrt(inverseInfinity)};
        }

        const Operation *m_operation;
        const Inverse *m_inverse;
        MPI_Comm m_communicator;
        CoupledShapeCoarseOptions m_options;
        std::uint64_t m_generation{0};
        bool m_ready{false};
        std::vector<mfem::Vector> m_sources, m_images;
        std::vector<CoupledShapeCoarseModeReport> m_reports;
        std::array<double, maximumModes * maximumModes> m_upper{};
        mutable mfem::Vector m_residualWork, m_actionWork;
        mutable CoupledShapeCoarseStatistics m_statistics;
    };
} // namespace experiment::plateau_diagnostics