#pragma once

// Include after import mean_field. The experiment changes only the inverse
// supplied to Krylov; every column below comes from the authoritative Jacobian.
#include <algorithm>
#include <cmath>
#include <concepts>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <span>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

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

namespace experiment::plateau_diagnostics {
    struct SurfaceBorderSchurCoordinates final {
        // Equation and unknown ordering need not coincide in a generated form.
        std::vector<int> residualIndices;
        std::vector<int> correctionIndices;
    };

    template <class PhysicalInverse>
    [[nodiscard]] SurfaceBorderSchurCoordinates MakeSurfaceBorderSchurCoordinates(const PhysicalInverse &inverse) {
        using namespace mean_field;
        using namespace utils::blocks;
        using Root = std::remove_cvref_t<decltype(inverse.GetBlock())>;
        using Structure = typename Root::StructureComponent;
        constexpr auto shapeColumn = type_index_v<surface_deformation::parameters::value,
                                                  typename Root::CorrectionBlocks>;
        constexpr auto shapeRow = type_index_v<surface_deformation::shape_equilibrium::residual,
                                               typename Root::ResidualBlocks>;
        static_assert(Root::borderValueArity == Root::borderResidualArity);
        static_assert(shapeColumn < Structure::CorrectionBlocks::size);
        static_assert(shapeRow < Structure::ResidualBlocks::size);
        const auto &map = inverse.GetCoordinateMap();
        const auto &columns = map.GetCorrectionRanges();
        const auto &rows = map.GetResidualRanges();
        SurfaceBorderSchurCoordinates result;
        const auto append = [](const auto &range, std::vector<int> &indices) {
            for (int i = 0; i < range.size; ++i)
                indices.push_back(range.equilibriumOffset + i);
        };
        append(columns[shapeColumn], result.correctionIndices);
        append(rows[shapeRow], result.residualIndices);
        // Generated scalar blocks follow the complete structure descriptor.
        // This includes the optional central phase condition without naming a
        // particular model or interpreting a runtime stable-ID string.
        for (std::size_t i = Structure::CorrectionBlocks::size; i < columns.size(); ++i)
            append(columns[i], result.correctionIndices);
        for (std::size_t i = Structure::ResidualBlocks::size; i < rows.size(); ++i)
            append(rows[i], result.residualIndices);
        return result;
    }

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

    template <class Candidate>
    concept SurfaceBorderSchurInverse = 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>;
        };

    // Reusing the complete scalar-bordered inverse here would impose zero
    // scalar-constraint RHS before elimination. Its field restriction can be
    // singular even when A_FF is invertible (a zero scalar diagonal suffices).
    // Borrow only the original structure inverse, preserving its typed packing
    // and the exact physical/normalized coordinate maps.
    template <SurfaceBorderSchurOperation Operation, SurfaceBorderSchurInverse PhysicalInverse>
    class NormalizedUnborderedStructureInverse final : public mfem::Solver {
        using Root = std::remove_cvref_t<decltype(std::declval<const PhysicalInverse &>().GetBlock())>;
        using Backend = typename Root::StructureComponent::BackendType;
    public:
        static constexpr auto applicationContract = mean_field::preconditioning::backend::applicationContract<Backend>;
        static_assert(applicationContract == mean_field::preconditioning::ApplicationContract::stationary_linear);

        NormalizedUnborderedStructureInverse(const Operation &operation, const PhysicalInverse &inverse)
            : mfem::Solver(operation.Width(), operation.Height()), m_operation(&operation), m_inverse(&inverse),
              m_generation(operation.GetProblem().GetPreparationGeneration()), m_physicalRhs(Width()),
              m_physicalAction(Height()), m_groupedRhs(Width()), m_groupedAction(Height()) {
            if (&operation.GetProblem() != &inverse.GetProblem() || inverse.Height() != Height() ||
                inverse.Width() != Width() || !IsCurrent())
                throw std::invalid_argument("The structure adapter requires the current inverse of the same problem.");
        }
        void SetOperator(const mfem::Operator &operation) override {
            if (&operation != m_operation || !IsCurrent())
                throw std::invalid_argument("The structure adapter cannot bind another or stale Jacobian.");
        }
        [[nodiscard]] bool IsCurrent() const {
            return m_operation->IsPrepared() && m_inverse->IsCurrent() && !m_inverse->iterative_mode &&
                   !m_inverse->GetGroupedPreconditioner().GetStructurePreconditioner().iterative_mode &&
                   m_generation == m_operation->GetProblem().GetPreparationGeneration();
        }
        [[nodiscard]] const Operation &GetNormalizedJacobian() const noexcept { return *m_operation; }
        void Mult(const mfem::Vector &rhs, mfem::Vector &action) const override {
            if (!IsCurrent() || rhs.Size() != Width() || action.Size() != Height())
                throw std::logic_error("The structure adapter requires current objects and compatible vectors.");
            const auto &structure = m_inverse->GetGroupedPreconditioner().GetStructurePreconditioner();
            const auto &map = m_inverse->GetCoordinateMap();
            m_operation->DenormalizeResidual(rhs, m_physicalRhs);
            map.PackResidual(m_physicalRhs, m_groupedRhs);
            m_groupedAction = 0.0;
            const mfem::Vector fieldRhs(m_groupedRhs, 0, structure.Width());
            mfem::Vector fieldAction(m_groupedAction, 0, structure.Height());
            structure.Mult(fieldRhs, fieldAction);
            fieldAction.SyncAliasMemory(m_groupedAction);
            map.UnpackCorrection(m_groupedAction, m_physicalAction);
            m_operation->NormalizeState(m_physicalAction, action);
        }
    private:
        const Operation *m_operation;
        const PhysicalInverse *m_inverse;
        std::uint64_t m_generation;
        mutable mfem::Vector m_physicalRhs, m_physicalAction, m_groupedRhs, m_groupedAction;
    };

    enum class SurfaceBorderSchurStatus { unprepared, ready, singular, ill_conditioned };

    struct SurfaceBorderSchurOptions final {
        double minimumReciprocalCondition{1.e-14};
    };

    struct SurfaceBorderSchurStatistics final {
        SurfaceBorderSchurStatus status{SurfaceBorderSchurStatus::unprepared};
        int coarseSize{0};
        std::uint64_t setups{0}, setupOperatorApplications{0}, setupInverseApplications{0};
        double reciprocalCondition{0.0};
        double minimumScaledPivot{0.0};
        double lastSetupSeconds{0.0}, setupOperatorSeconds{0.0}, setupInverseSeconds{0.0};
        std::uint64_t applications{0}, operatorApplications{0}, inverseApplications{0}, failedApplications{0};
        double applicationSeconds{0.0}, operatorSeconds{0.0}, inverseSeconds{0.0};
        std::size_t persistentVectorBytes{0};
    };

    /* With C the complete surface/scalar block and F every remaining field,
     * B = R_F P0 E_F approximates A_FF^{-1}. Store
     *
     *   Z = E_C - E_F B A_FC,       S = R_C A Z.
     *
     * Application gives y = E_F B r_F and x = y + Z S^{-1}(r_C-R_C A y).
     * The coarse equations are satisfied exactly (up to arithmetic), even
     * when B is approximate. The complete inverse is exact only when B is.
     * Restricting the existing inverse defines B; its internal border response
     * is not claimed to be an exact elimination of the frozen field block.
     * No symmetry, positive definiteness, or reduced physical equation is
     * assumed. A singular/unsafe S is reported and cannot be applied.
     */
    template <SurfaceBorderSchurOperation Operation, SurfaceBorderSchurInverse Inverse>
    class PreparedSurfaceBorderSchur final : public mfem::Solver {
    public:
        static constexpr auto applicationContract = mean_field::preconditioning::ApplicationContract::stationary_linear;
        static constexpr std::size_t maximumCoarseSize = 128;

        PreparedSurfaceBorderSchur(const Operation &operation, const Inverse &inverse,
                                   SurfaceBorderSchurCoordinates coordinates, MPI_Comm communicator,
                                   SurfaceBorderSchurOptions options = {})
            : mfem::Solver(operation.Width(), operation.Height()), m_operation(&operation), m_inverse(&inverse),
              m_coordinates(std::move(coordinates)), m_communicator(communicator), m_options(options),
              m_rhs(operation.Height()), m_fieldAction(operation.Width()), m_image(operation.Height()) {
            Validate();
            Refresh();
        }

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

        void SetOperator(const mfem::Operator &operation) override {
            if (&operation != m_operation)
                throw std::invalid_argument("The surface-border Schur inverse cannot bind a different Jacobian.");
            RequireCurrent();
        }

        [[nodiscard]] bool IsCurrent() const {
            return m_statistics.status == SurfaceBorderSchurStatus::ready && DependenciesCurrent();
        }
        [[nodiscard]] const SurfaceBorderSchurStatistics &GetStatistics() const noexcept { return m_statistics; }
        [[nodiscard]] const SurfaceBorderSchurCoordinates &GetCoordinates() const noexcept { return m_coordinates; }
        [[nodiscard]] const mfem::DenseMatrix &GetSchurMatrix() const noexcept { return m_schur; }
        [[nodiscard]] const mfem::DenseMatrix &GetUneliminatedCoarseMatrix() const noexcept { return m_coarse; }
        [[nodiscard]] std::span<const mfem::Vector> GetCompletedColumns() const noexcept { return m_columns; }

        void Refresh() {
            const double start = MPI_Wtime();
            Validate();
            m_statistics.status = SurfaceBorderSchurStatus::unprepared;
            m_generation = m_operation->GetProblem().GetPreparationGeneration();
            const int count = static_cast<int>(m_coordinates.correctionIndices.size());
            m_statistics.coarseSize = count;
            m_columns.clear();
            m_columns.reserve(count);
            m_schur.SetSize(count);
            m_coarse.SetSize(count);
            m_solveWork.SetSize(count);
            mfem::Vector basis(Height()), image(Width()), response(Height());
            for (int column = 0; column < count; ++column) {
                basis = 0.0;
                basis(m_coordinates.correctionIndices[column]) = 1.0;
                ApplyOperation(basis, image, true);
                for (int row = 0; row < count; ++row)
                    m_coarse(row, column) = image(m_coordinates.residualIndices[row]);
                ApplyFields(image, response, true);
                response *= -1.0;
                response += basis;
                m_columns.emplace_back(response);
                ApplyOperation(response, image, true);
                for (int row = 0; row < count; ++row)
                    m_schur(row, column) = image(m_coordinates.residualIndices[row]);
            }
            Factor();
            if (!DependenciesCurrent())
                throw std::logic_error("The surface-border Schur dependencies changed during preparation.");
            m_statistics.lastSetupSeconds = MPI_Wtime() - start;
            m_statistics.persistentVectorBytes = sizeof(mfem::real_t) *
                (static_cast<std::size_t>(count) * Height() + 3 * static_cast<std::size_t>(Height()) + count);
            ++m_statistics.setups;
        }

        void Mult(const mfem::Vector &rightHandSide, mfem::Vector &action) const override {
            const double start = MPI_Wtime();
            ++m_statistics.applications;
            try {
                RequireCurrent();
                if (rightHandSide.Size() != Width() || action.Size() != Height())
                    throw std::invalid_argument("The surface-border Schur inverse requires preallocated root vectors.");
                RequireFinite(rightHandSide);
                ApplyFields(rightHandSide, m_fieldAction, false);
                ApplyOperation(m_fieldAction, m_image, false);
                for (int row = 0; row < m_solveWork.Size(); ++row)
                    m_solveWork(row) = rightHandSide(m_coordinates.residualIndices[row]) -
                                       m_image(m_coordinates.residualIndices[row]);
                SolveFactored(m_solveWork);
                // No write to the caller's output occurs before all RHS uses,
                // so in-place application has the same algebra as separate buffers.
                action = m_fieldAction;
                for (int column = 0; column < m_solveWork.Size(); ++column)
                    action.Add(m_solveWork(column), m_columns[column]);
                RequireFinite(action);
                RequireCurrent();
                m_statistics.applicationSeconds += MPI_Wtime() - start;
            } catch (...) {
                ++m_statistics.failedApplications;
                m_statistics.applicationSeconds += MPI_Wtime() - start;
                throw;
            }
        }

    private:
        void Validate() const {
            int ranks = 0;
            if (m_communicator == MPI_COMM_NULL || MPI_Comm_size(m_communicator, &ranks) != MPI_SUCCESS || ranks != 1)
                throw std::invalid_argument("The surface-border Schur experiment currently requires one MPI rank.");
            if (!m_operation->IsPrepared() || !m_inverse->IsCurrent() || m_inverse->iterative_mode ||
                m_operation->Height() != m_operation->Width() || m_inverse->Height() != Height() ||
                m_inverse->Width() != Width())
                throw std::invalid_argument("The surface-border Schur experiment requires current compatible objects.");
            if constexpr (requires { m_inverse->GetNormalizedJacobian(); }) {
                if (&m_inverse->GetNormalizedJacobian() != m_operation)
                    throw std::invalid_argument("The inverse belongs to a different normalized Jacobian.");
            }
            if (!std::isfinite(m_options.minimumReciprocalCondition) ||
                m_options.minimumReciprocalCondition <= 0.0 || m_options.minimumReciprocalCondition >= 1.0)
                throw std::invalid_argument("The Schur reciprocal-condition threshold must lie strictly between zero and one.");
            if (m_coordinates.residualIndices.empty() ||
                m_coordinates.residualIndices.size() != m_coordinates.correctionIndices.size() ||
                m_coordinates.residualIndices.size() >= static_cast<std::size_t>(Height()))
                throw std::invalid_argument("The Schur partition requires nonempty fields and square coarse coordinates.");
            if (m_coordinates.residualIndices.size() > maximumCoarseSize)
                throw std::invalid_argument("The bounded surface-border experiment supports at most 128 coarse coordinates.");
            const auto verify = [this](const auto &indices) {
                std::vector<bool> used(Height(), false);
                for (int index : indices) {
                    if (index < 0 || index >= Height() || used[index])
                        throw std::invalid_argument("Schur coordinate indices must be distinct and in range.");
                    used[index] = true;
                }
            };
            verify(m_coordinates.residualIndices);
            verify(m_coordinates.correctionIndices);
        }

        [[nodiscard]] bool DependenciesCurrent() const {
            return m_operation->IsPrepared() && m_inverse->IsCurrent() && !m_inverse->iterative_mode &&
                   m_generation == m_operation->GetProblem().GetPreparationGeneration();
        }
        void RequireCurrent() const {
            if (!IsCurrent())
                throw std::logic_error("The surface-border Schur inverse is stale, singular, or ill-conditioned.");
        }
        static void RequireFinite(const mfem::Vector &vector) {
            for (int i = 0; i < vector.Size(); ++i)
                if (!std::isfinite(vector(i)))
                    throw std::runtime_error("The surface-border Schur experiment encountered a nonfinite vector.");
        }
        void ApplyOperation(const mfem::Vector &input, mfem::Vector &output, bool setup) const {
            const double start = MPI_Wtime();
            m_operation->Mult(input, output);
            const double seconds = MPI_Wtime() - start;
            if (setup) {
                ++m_statistics.setupOperatorApplications;
                m_statistics.setupOperatorSeconds += seconds;
            } else {
                ++m_statistics.operatorApplications;
                m_statistics.operatorSeconds += seconds;
            }
            RequireFinite(output);
        }
        void ApplyFields(const mfem::Vector &input, mfem::Vector &output, bool setup) const {
            m_rhs = input;
            for (int index : m_coordinates.residualIndices)
                m_rhs(index) = 0.0;
            const double start = MPI_Wtime();
            m_inverse->Mult(m_rhs, output);
            const double seconds = MPI_Wtime() - start;
            if (setup) {
                ++m_statistics.setupInverseApplications;
                m_statistics.setupInverseSeconds += seconds;
            } else {
                ++m_statistics.inverseApplications;
                m_statistics.inverseSeconds += seconds;
            }
            RequireFinite(output);
            for (int index : m_coordinates.correctionIndices)
                output(index) = 0.0;
        }

        void Factor() {
            const int count = m_schur.Height();
            m_lu = m_schur;
            m_pivots.resize(count);
            m_scale = 0.0;
            for (int i = 0; i < count * count; ++i) {
                if (!std::isfinite(m_lu.Data()[i]))
                    throw std::runtime_error("The coarse Schur matrix is nonfinite.");
                m_scale = std::max(m_scale, std::abs(m_lu.Data()[i]));
            }
            m_statistics.reciprocalCondition = 0.0;
            m_statistics.minimumScaledPivot = 0.0;
            m_statistics.status = SurfaceBorderSchurStatus::singular;
            if (m_scale == 0.0)
                return;
            for (int i = 0; i < count * count; ++i)
                m_lu.Data()[i] /= m_scale;
            double norm = 0.0;
            for (int row = 0; row < count; ++row) {
                double sum = 0.0;
                for (int column = 0; column < count; ++column)
                    sum += std::abs(m_lu(row, column));
                norm = std::max(norm, sum);
            }
            m_statistics.minimumScaledPivot = std::numeric_limits<double>::infinity();
            for (int column = 0; column < count; ++column) {
                int pivot = column;
                for (int row = column + 1; row < count; ++row)
                    if (std::abs(m_lu(row, column)) > std::abs(m_lu(pivot, column)))
                        pivot = row;
                m_pivots[column] = pivot;
                const double magnitude = std::abs(m_lu(pivot, column));
                m_statistics.minimumScaledPivot = std::min(m_statistics.minimumScaledPivot, magnitude);
                if (!std::isfinite(magnitude) || magnitude == 0.0)
                    return;
                if (pivot != column)
                    for (int j = 0; j < count; ++j)
                        std::swap(m_lu(column, j), m_lu(pivot, j));
                for (int row = column + 1; row < count; ++row) {
                    m_lu(row, column) /= m_lu(column, column);
                    for (int j = column + 1; j < count; ++j)
                        m_lu(row, j) -= m_lu(row, column) * m_lu(column, j);
                }
            }
            // A small explicit inverse supplies an infinity-norm condition
            // estimate without regularizing a physically important coarse mode.
            mfem::Vector column(count), sums(count);
            sums = 0.0;
            for (int j = 0; j < count; ++j) {
                column = 0.0;
                column(j) = m_scale;
                SolveFactored(column);
                for (int i = 0; i < count; ++i)
                    sums(i) += std::abs(column(i));
            }
            const double inverseNorm = sums.Max();
            m_statistics.reciprocalCondition = 1.0 / (norm * inverseNorm);
            m_statistics.status = std::isfinite(m_statistics.reciprocalCondition) &&
                                  m_statistics.reciprocalCondition >= m_options.minimumReciprocalCondition
                ? SurfaceBorderSchurStatus::ready : SurfaceBorderSchurStatus::ill_conditioned;
        }

        void SolveFactored(mfem::Vector &values) const {
            const int count = values.Size();
            values /= m_scale;
            for (int i = 0; i < count; ++i) {
                if (m_pivots[i] != i)
                    std::swap(values(i), values(m_pivots[i]));
            }
            for (int i = 0; i < count; ++i)
                for (int j = 0; j < i; ++j)
                    values(i) -= m_lu(i, j) * values(j);
            for (int i = count - 1; i >= 0; --i) {
                for (int j = i + 1; j < count; ++j)
                    values(i) -= m_lu(i, j) * values(j);
                values(i) /= m_lu(i, i);
            }
        }

        const Operation *m_operation;
        const Inverse *m_inverse;
        SurfaceBorderSchurCoordinates m_coordinates;
        MPI_Comm m_communicator;
        SurfaceBorderSchurOptions m_options;
        std::uint64_t m_generation{0};
        mfem::DenseMatrix m_schur, m_coarse, m_lu;
        std::vector<mfem::Vector> m_columns;
        std::vector<int> m_pivots;
        double m_scale{0.0};
        mutable mfem::Vector m_rhs, m_fieldAction, m_image, m_solveWork;
        mutable SurfaceBorderSchurStatistics m_statistics;
    };
} // namespace experiment::plateau_diagnostics