module;

#include <algorithm>
#include <cmath>
#include <concepts>
#include <cstdint>
#include <functional>
#include <memory>
#include <optional>
#include <span>
#include <stdexcept>
#include <type_traits>
#include <utility>
#include <vector>

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

export module mean_field:solver.gravity_subsystem;

export import :solver.linear_backend;
export import :utils.blocks;

export namespace mean_field::solver {
    /* The maps are local true-dof maps. Gravity and retained coordinates need
     * not be contiguous or share the residual ordering. Only the discretized
     * sizes are runtime data; the production factory selects block identities
     * from the compiled equation form. */
    class GravityBlockPartition final {
    public:
        GravityBlockPartition(
            const int rootValueSize,
            const int rootResidualSize,
            std::vector<int> gravityValues,
            std::vector<int> gravityResiduals
        )
            : m_rootValueSize(rootValueSize),
              m_rootResidualSize(rootResidualSize),
              m_gravityValues(std::move(gravityValues)),
              m_gravityResiduals(std::move(gravityResiduals)),
              m_retainedValues(Complement(
                  rootValueSize,
                  m_gravityValues
              )),
              m_retainedResiduals(Complement(
                  rootResidualSize,
                  m_gravityResiduals
              )) {
            if (m_gravityValues.size() != m_gravityResiduals.size() ||
                m_retainedValues.size() != m_retainedResiduals.size()) {
                throw std::invalid_argument("Gravity elimination requires square gravity and retained blocks.");
            }
        }

        [[nodiscard]] int RootValueSize() const noexcept {
            return m_rootValueSize;
        }
        [[nodiscard]] int RootResidualSize() const noexcept {
            return m_rootResidualSize;
        }
        [[nodiscard]] int GravityValueSize() const noexcept {
            return static_cast<int>(m_gravityValues.size());
        }
        [[nodiscard]] int GravityResidualSize() const noexcept {
            return static_cast<int>(m_gravityResiduals.size());
        }
        [[nodiscard]] int RetainedValueSize() const noexcept {
            return static_cast<int>(m_retainedValues.size());
        }
        [[nodiscard]] int RetainedResidualSize() const noexcept {
            return static_cast<int>(m_retainedResiduals.size());
        }

        void GatherGravityValues(
            const mfem::Vector &full,
            mfem::Vector &compact
        ) const {
            Gather(m_gravityValues, m_rootValueSize, full, compact);
        }
        void GatherGravityResiduals(
            const mfem::Vector &full,
            mfem::Vector &compact
        ) const {
            Gather(m_gravityResiduals, m_rootResidualSize, full, compact);
        }
        void GatherRetainedValues(
            const mfem::Vector &full,
            mfem::Vector &compact
        ) const {
            Gather(m_retainedValues, m_rootValueSize, full, compact);
        }
        void GatherRetainedResiduals(
            const mfem::Vector &full,
            mfem::Vector &compact
        ) const {
            Gather(m_retainedResiduals, m_rootResidualSize, full, compact);
        }
        void InjectGravityValues(
            const mfem::Vector &compact,
            mfem::Vector &full
        ) const {
            Inject(m_gravityValues, m_rootValueSize, compact, full);
        }
        void InjectGravityResiduals(
            const mfem::Vector &compact,
            mfem::Vector &full
        ) const {
            Inject(m_gravityResiduals, m_rootResidualSize, compact, full);
        }
        void InjectRetainedValues(
            const mfem::Vector &compact,
            mfem::Vector &full
        ) const {
            Inject(m_retainedValues, m_rootValueSize, compact, full);
        }
        void InjectRetainedResiduals(
            const mfem::Vector &compact,
            mfem::Vector &full
        ) const {
            Inject(m_retainedResiduals, m_rootResidualSize, compact, full);
        }
        void AddGravityValues(
            const mfem::Vector &compact,
            mfem::Vector &full
        ) const {
            VerifySizes(m_gravityValues, m_rootValueSize, full, compact);
            for (std::size_t index = 0; index < m_gravityValues.size(); ++index) {
                full(m_gravityValues[index]) += compact(static_cast<int>(index));
            }
        }

    private:
        [[nodiscard]] static std::vector<int> Complement(
            const int size,
            const std::vector<int> &selected
        ) {
            if (size < 0) {
                throw std::invalid_argument("A gravity partition requires non-negative root dimensions.");
            }
            std::vector<bool> used(static_cast<std::size_t>(size), false);
            for (const int index : selected) {
                if (index < 0 || index >= size || used[static_cast<std::size_t>(index)]) {
                    throw std::invalid_argument("Gravity partition indices must be unique and in range.");
                }
                used[static_cast<std::size_t>(index)] = true;
            }
            std::vector<int> complement;
            complement.reserve(static_cast<std::size_t>(size) - selected.size());
            for (int index = 0; index < size; ++index) {
                if (!used[static_cast<std::size_t>(index)]) {
                    complement.push_back(index);
                }
            }
            return complement;
        }

        static void VerifySizes(
            const std::vector<int> &indices,
            const int fullSize,
            const mfem::Vector &full,
            const mfem::Vector &compact
        ) {
            if (full.Size() != fullSize || compact.Size() != static_cast<int>(indices.size())) {
                throw std::invalid_argument("Gravity partition vectors must have their preallocated sizes.");
            }
            const std::less<const mfem::real_t *> before;
            if (compact.Size() != 0 && full.Size() != 0 && before(full.GetData(), compact.GetData() + compact.Size()) &&
                before(compact.GetData(), full.GetData() + full.Size())) {
                throw std::invalid_argument("Gravity gather and injection require distinct vector storage.");
            }
        }

        static void Gather(
            const std::vector<int> &indices,
            const int fullSize,
            const mfem::Vector &full,
            mfem::Vector &compact
        ) {
            VerifySizes(indices, fullSize, full, compact);
            for (std::size_t index = 0; index < indices.size(); ++index) {
                compact(static_cast<int>(index)) = full(indices[index]);
            }
        }

        static void Inject(
            const std::vector<int> &indices,
            const int fullSize,
            const mfem::Vector &compact,
            mfem::Vector &full
        ) {
            VerifySizes(indices, fullSize, full, compact);
            full = 0.0;
            for (std::size_t index = 0; index < indices.size(); ++index) {
                full(indices[index]) = compact(static_cast<int>(index));
            }
        }

        int m_rootValueSize;
        int m_rootResidualSize;
        std::vector<int> m_gravityValues;
        std::vector<int> m_gravityResiduals;
        std::vector<int> m_retainedValues;
        std::vector<int> m_retainedResiduals;
    };

    template <typename Form>
    concept GravityPartitionForm = requires {
        typename Form::value_blocks;
        typename Form::residual_blocks;
        requires utils::blocks::block_form_is_valid_v<Form>;
        requires utils::blocks::contains_type_v<utils::blocks::gravity::gradient::value, typename Form::value_blocks>;
        requires utils::blocks::contains_type_v<utils::blocks::gravity::poisson::value, typename Form::value_blocks>;
        requires utils::blocks::contains_type_v<
            utils::blocks::gravity::gradient::residual, typename Form::residual_blocks>;
        requires utils::blocks::contains_type_v<
            utils::blocks::gravity::poisson::residual, typename Form::residual_blocks>;
    };

    template <GravityPartitionForm Form>
    [[nodiscard]] GravityBlockPartition MakeGravityBlockPartition(const utils::blocks::form_layout<Form> &layout) {
        using namespace utils::blocks;
        std::vector<int> values;
        std::vector<int> residuals;
        const auto append = [](std::vector<int> &indices, const int offset, const int size) {
            for (int index = 0; index < size; ++index) {
                indices.push_back(offset + index);
            }
        };
        const auto value = [&]<typename Block>() {
            constexpr auto block = value_block<type_index_v<Block, typename Form::value_blocks>>{};
            append(values, layout.offset(block), layout.size(block));
        };
        const auto residual = [&]<typename Block>() {
            constexpr auto block = residual_block<type_index_v<Block, typename Form::residual_blocks>>{};
            append(residuals, layout.offset(block), layout.size(block));
        };
        value.template operator()<gravity::gradient::value>();
        value.template operator()<gravity::poisson::value>();
        residual.template operator()<gravity::gradient::residual>();
        residual.template operator()<gravity::poisson::residual>();
        return {
            layout.value_offsets().Last(), layout.residual_offsets().Last(), std::move(values), std::move(residuals)
        };
    }

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

    template <typename Preconditioner>
    concept GravitySubsystemPreconditioner =
        std::derived_from<Preconditioner, mfem::Solver> &&
        LinearPreconditionerApplicationContractAvailable<Preconditioner> && requires(const Preconditioner &inverse) {
            { inverse.IsCurrent() } -> std::same_as<bool>;
        };
} // namespace mean_field::solver

namespace mean_field::solver::detail {
    template <GravitySubsystemOperation Operation> class GravityRestrictedOperator final : public mfem::Operator {
    public:
        GravityRestrictedOperator(
            const Operation &operation,
            const GravityBlockPartition &partition
        )
            : mfem::Operator(
                  partition.GravityResidualSize(),
                  partition.GravityValueSize()
              ),
              m_operation(std::addressof(operation)),
              m_partition(std::addressof(partition)),
              m_fullDirection(partition.RootValueSize()),
              m_fullAction(partition.RootResidualSize()) {
            if (operation.Width() != partition.RootValueSize() || operation.Height() != partition.RootResidualSize()) {
                throw std::invalid_argument("The gravity partition does not match its root operator.");
            }
            Refresh();
        }

        void Refresh() {
            if (!m_operation->IsPrepared()) {
                throw std::logic_error("Gravity requires a prepared root linearization before refresh.");
            }
            m_generation = m_operation->GetProblem().GetPreparationGeneration();
        }

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

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override {
            if (!IsCurrent()) {
                throw std::logic_error("The gravity linearization must be refreshed after preparation.");
            }
            m_partition->InjectGravityValues(direction, m_fullDirection);
            m_operation->Mult(m_fullDirection, m_fullAction);
            m_partition->GatherGravityResiduals(m_fullAction, action);
        }

    private:
        const Operation *m_operation;
        const GravityBlockPartition *m_partition;
        std::uint64_t m_generation{0};
        mutable mfem::Vector m_fullDirection;
        mutable mfem::Vector m_fullAction;
    };

    template <GravitySubsystemOperation Operation, GravitySubsystemPreconditioner Preconditioner>
    class GravityRestrictedInverse final : public mfem::Solver {
    public:
        static constexpr auto applicationContract = linearPreconditionerApplicationContract<Preconditioner>;

        GravityRestrictedInverse(
            const Preconditioner &inverse,
            const GravityRestrictedOperator<Operation> &operation,
            const GravityBlockPartition &partition
        )
            : mfem::Solver(partition.GravityValueSize()),
              m_inverse(std::addressof(inverse)),
              m_operation(std::addressof(operation)),
              m_partition(std::addressof(partition)),
              m_fullResidual(partition.RootResidualSize()),
              m_fullCorrection(partition.RootValueSize()) {
            if (inverse.Width() != partition.RootResidualSize() || inverse.Height() != partition.RootValueSize()) {
                throw std::invalid_argument("The gravity partition does not match its root inverse.");
            }
        }

        [[nodiscard]] bool IsCurrent() const {
            return m_operation->IsCurrent() && m_inverse->IsCurrent();
        }

        void SetOperator(const mfem::Operator &operation) override {
            if (std::addressof(operation) != m_operation) {
                throw std::invalid_argument(
                    "A prepared gravity inverse cannot be "
                    "rebound to a different operator."
                );
            }
        }

        void Mult(
            const mfem::Vector &residual,
            mfem::Vector &correction
        ) const override {
            if (!IsCurrent()) {
                throw std::logic_error("The gravity inverse requires a current root inverse.");
            }
            m_partition->InjectGravityResiduals(residual, m_fullResidual);
            m_inverse->Mult(m_fullResidual, m_fullCorrection);
            m_partition->GatherGravityValues(m_fullCorrection, correction);
        }

    private:
        const Preconditioner *m_inverse;
        const GravityRestrictedOperator<Operation> *m_operation;
        const GravityBlockPartition *m_partition;
        mutable mfem::Vector m_fullResidual;
        mutable mfem::Vector m_fullCorrection;
    };
} // namespace mean_field::solver::detail

export namespace mean_field::solver {
    /* Restricting the supplied inverse is a preconditioner, not an assertion
     * that its principal block equals G_z^{-1}. The backend still verifies the
     * true gravity defect. Keeping the exact root operator as the source of
     * every action also retains its boundary rows and normalization. */
    template <
        GravitySubsystemOperation Operation,
        GravitySubsystemPreconditioner Preconditioner,
        LinearBackendConfiguration Backend = linear::FGMRES>
        requires LinearBackendRuntimeAvailableFor<
            Backend,
            detail::GravityRestrictedOperator<Operation>,
            detail::GravityRestrictedInverse<Operation, Preconditioner>>
    class PreparedGravitySubsystem final {
    public:
        using OperatorType       = Operation;
        using PreconditionerType = Preconditioner;
        using RestrictedOperator = detail::GravityRestrictedOperator<Operation>;
        using RestrictedInverse  = detail::GravityRestrictedInverse<Operation, Preconditioner>;
        using BackendType        = PreparedLinearBackendType<Backend, RestrictedOperator, RestrictedInverse>;

        PreparedGravitySubsystem(
            const Operation &operation,
            const Preconditioner &preconditioner,
            GravityBlockPartition partition,
            const MPI_Comm communicator,
            Backend backend = {}
        )
            : m_fullOperator(std::addressof(operation)),
              m_fullPreconditioner(std::addressof(preconditioner)),
              m_partition(std::move(partition)),
              m_operator(
                  operation,
                  m_partition
              ),
              m_inverse(
                  preconditioner,
                  m_operator,
                  m_partition
              ),
              m_backend(prepareLinearBackend(
                  std::move(backend),
                  m_operator,
                  m_inverse,
                  communicator
              )) {
            if (!IsCurrent()) {
                throw std::logic_error("Gravity subsystem preparation requires a current inverse.");
            }
        }

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

        void Refresh() {
            if (!m_fullPreconditioner->IsCurrent()) {
                throw std::logic_error("Refresh the root inverse before refreshing its gravity subsystem.");
            }
            m_operator.Refresh();
        }

        [[nodiscard]] bool IsCurrent() const {
            return m_operator.IsCurrent() && m_inverse.IsCurrent() && m_backend.IsReady();
        }
        [[nodiscard]] const GravityBlockPartition &GetPartition() const noexcept {
            return m_partition;
        }
        [[nodiscard]] const RestrictedOperator &GetOperator() const noexcept {
            return m_operator;
        }
        [[nodiscard]] const Operation &GetFullOperator() const noexcept {
            return *m_fullOperator;
        }
        [[nodiscard]] const Preconditioner &GetFullPreconditioner() const noexcept {
            return *m_fullPreconditioner;
        }
        [[nodiscard]] MPI_Comm GetCommunicator() const noexcept {
            return m_backend.GetCommunicator();
        }
        [[nodiscard]] std::uint64_t GetSolveCount() const noexcept {
            return m_solveCount;
        }
        [[nodiscard]] std::uint64_t GetIterationCount() const noexcept {
            return m_iterationCount;
        }
        [[nodiscard]] std::uint64_t GetFailureCount() const noexcept {
            return m_failureCount;
        }
        [[nodiscard]] double GetSolveSeconds() const noexcept {
            return m_solveSeconds;
        }
        [[nodiscard]] const std::optional<LinearSolveReport> &GetLastSolveReport() const noexcept {
            return m_lastSolve;
        }

        [[nodiscard]] LinearSolveReport Solve(
            const mfem::Vector &rightHandSide,
            mfem::Vector &correction,
            const LinearSolveControl &control
        ) {
            ++m_solveCount;
            try {
                const auto report = m_backend.Solve(rightHandSide, correction, control);
                m_iterationCount += static_cast<std::uint64_t>(report.iterations);
                m_solveSeconds += report.solveSeconds;
                if (!report.Converged()) {
                    ++m_failureCount;
                }
                m_lastSolve = report;
                return report;
            } catch (...) {
                ++m_failureCount;
                throw;
            }
        }

    private:
        const Operation *m_fullOperator;
        const Preconditioner *m_fullPreconditioner;
        GravityBlockPartition m_partition;
        RestrictedOperator m_operator;
        RestrictedInverse m_inverse;
        BackendType m_backend;
        // Times and iteration counts cover completed calls; failed contracts
        // are counted without attempting another collective during unwinding.
        std::uint64_t m_solveCount{0};
        std::uint64_t m_iterationCount{0};
        std::uint64_t m_failureCount{0};
        double m_solveSeconds{0.0};
        std::optional<LinearSolveReport> m_lastSolve;
    };

    template <typename Subsystem>
    concept GravitySubsystemRuntime = requires(
        Subsystem &subsystem,
        const Subsystem &constant,
        const mfem::Vector &rhs,
        mfem::Vector &correction,
        const LinearSolveControl &control
    ) {
        { constant.IsCurrent() } -> std::same_as<bool>;
        { constant.GetPartition() } -> std::same_as<const GravityBlockPartition &>;
        { constant.GetCommunicator() } -> std::same_as<MPI_Comm>;
        { subsystem.Solve(rhs, correction, control) } -> std::same_as<LinearSolveReport>;
        { constant.GetFullOperator().Mult(rhs, correction) } -> std::same_as<void>;
        { constant.GetFullOperator().GetProblem().GetPreparationGeneration() } -> std::convertible_to<std::uint64_t>;
    };

    template <GravitySubsystemRuntime Subsystem> class PreparedGravityCurvatureCorrection final {
    public:
        explicit PreparedGravityCurvatureCorrection(Subsystem &subsystem)
            : m_subsystem(std::addressof(subsystem)),
              m_baseResidual(subsystem.GetPartition().GravityResidualSize()),
              m_baseAction(subsystem.GetPartition().GravityResidualSize()),
              m_fullAction(subsystem.GetPartition().RootResidualSize()),
              m_rightHandSide(subsystem.GetPartition().GravityResidualSize()) {
        }

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

        // Capture must precede trial preparation: using a candidate Jacobian
        // here would introduce a first-order change into the Newton tangent.
        void CaptureBase(
            const mfem::Vector &baseResidual,
            const mfem::Vector &baseDirection
        ) {
            m_hasBase             = false;
            const auto &partition = m_subsystem->GetPartition();
            RequireCollectively(
                m_subsystem->IsCurrent() && baseResidual.Size() == partition.RootResidualSize() &&
                    baseDirection.Size() == partition.RootValueSize(),
                "Gravity curvature capture requires a current base and "
                "complete vectors."
            );
            partition.GatherGravityResiduals(baseResidual, m_baseResidual);
            m_subsystem->GetFullOperator().Mult(baseDirection, m_fullAction);
            partition.GatherGravityResiduals(m_fullAction, m_baseAction);
            m_baseGeneration = m_subsystem->GetFullOperator().GetProblem().GetPreparationGeneration();
            m_hasBase        = true;
        }

        void Invalidate() noexcept {
            m_hasBase = false;
        }
        [[nodiscard]] bool HasBase() const noexcept {
            return m_hasBase;
        }
        [[nodiscard]] std::uint64_t GetBaseGeneration() const noexcept {
            return m_baseGeneration;
        }

        // The caller owns candidate preparation and commits only after a new
        // complete residual evaluation. A failed inner solve never modifies
        // an accepted state; its output remains a diagnostic correction only.
        [[nodiscard]] LinearSolveReport SolveCorrection(
            const double alpha,
            const mfem::Vector &trialResidual,
            mfem::Vector &gravityCorrection,
            const LinearSolveControl &control
        ) {
            const auto &partition = m_subsystem->GetPartition();
            RequireCollectively(
                m_hasBase && m_subsystem->IsCurrent() && std::isfinite(alpha) && alpha >= 0.0 &&
                    trialResidual.Size() == partition.RootResidualSize() &&
                    gravityCorrection.Size() == partition.GravityValueSize(),
                "Gravity curvature correction requires a captured base, a current "
                "candidate, and valid vectors."
            );
            double minimumAlpha = 0.0;
            double maximumAlpha = 0.0;
            if (MPI_Allreduce(&alpha, &minimumAlpha, 1, MPI_DOUBLE, MPI_MIN, m_subsystem->GetCommunicator()) !=
                    MPI_SUCCESS ||
                MPI_Allreduce(&alpha, &maximumAlpha, 1, MPI_DOUBLE, MPI_MAX, m_subsystem->GetCommunicator()) !=
                    MPI_SUCCESS) {
                throw std::runtime_error(
                    "Gravity curvature correction could not verify "
                    "its distributed step length."
                );
            }
            if (minimumAlpha != maximumAlpha) {
                throw std::invalid_argument(
                    "Gravity curvature correction requires the "
                    "same step length on every rank."
                );
            }
            partition.GatherGravityResiduals(trialResidual, m_rightHandSide);
            for (int index = 0; index < m_rightHandSide.Size(); ++index) {
                m_rightHandSide(index) = m_baseResidual(index) + alpha * m_baseAction(index) - m_rightHandSide(index);
            }
            gravityCorrection = 0.0;
            return m_subsystem->Solve(m_rightHandSide, gravityCorrection, control);
        }

    private:
        void RequireCollectively(
            const bool valid,
            const char *message
        ) const {
            int local  = valid ? 1 : 0;
            int global = 0;
            if (MPI_Allreduce(&local, &global, 1, MPI_INT, MPI_MIN, m_subsystem->GetCommunicator()) != MPI_SUCCESS) {
                throw std::runtime_error(
                    "Gravity curvature correction could not verify "
                    "distributed validity."
                );
            }
            if (global == 0) {
                throw std::logic_error(message);
            }
        }

        Subsystem *m_subsystem;
        mfem::Vector m_baseResidual;
        mfem::Vector m_baseAction;
        mfem::Vector m_fullAction;
        mfem::Vector m_rightHandSide;
        std::uint64_t m_baseGeneration{0};
        bool m_hasBase{false};
    };
} // namespace mean_field::solver