module;

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

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

export module mean_field:solver.gravity_reduction;

export import :solver.gravity_subsystem;
export import :solver.linear_backend;

export namespace mean_field::solver {
    template <typename Candidate>
    concept PreparedGravityReductionSubsystem = requires(
        Candidate &subsystem,
        const Candidate &constantSubsystem,
        const mfem::Vector &rightHandSide,
        mfem::Vector &correction,
        const LinearSolveControl &control
    ) {
        { constantSubsystem.GetPartition() } -> std::same_as<const GravityBlockPartition &>;
        { constantSubsystem.IsCurrent() } -> std::same_as<bool>;
        { constantSubsystem.GetCommunicator() } -> std::same_as<MPI_Comm>;
        requires std::derived_from<
            std::remove_cvref_t<decltype(constantSubsystem.GetFullOperator())>, mfem::Operator>;
        { subsystem.Solve(rightHandSide, correction, control) } -> std::same_as<LinearSolveReport>;
    };

    /*
     * MFEM's void Mult boundary cannot return a failed inner solve. Keep that
     * failure distinguishable from a contract violation so an outer backend
     * can stop the numerical attempt without accepting an inaccurate Schur
     * action. The Try methods below retain the ordinary report-based path.
     */
    class GravityReductionFailure final : public std::runtime_error {
    public:
        explicit GravityReductionFailure(LinearSolveReport report)
            : std::runtime_error("A reduced-gravity action failed its inner true-residual tolerance."),
              m_report(std::move(report)) {
        }

        [[nodiscard]] const LinearSolveReport &GetReport() const noexcept {
            return m_report;
        }

    private:
        LinearSolveReport m_report;
    };

    /*
     * For J = [F_u F_z; G_u G_z], this applies F_u - F_z G_z^{-1} G_u.
     * Neither the full operator nor its preparation may change during an
     * application. All coordinates, including specification borders, come
     * from the same partition used by the gravity subsystem.
     *
     * A tolerance-stopped inner iteration approximates this linear operator;
     * flexible outer preconditioning does not make an inexact operator exact.
     * Its owner must verify a lifted correction with the full Jacobian and
     * choose an inner tolerance commensurate with the outer accuracy target.
     */
    template <PreparedGravityReductionSubsystem Subsystem>
    class GravityReducedOperator final : public mfem::Operator {
    public:
        GravityReducedOperator(
            Subsystem &subsystem,
            LinearSolveControl innerControl
        )
            : mfem::Operator(
                  subsystem.GetPartition().RetainedResidualSize(),
                  subsystem.GetPartition().RetainedValueSize()
              ),
              m_subsystem(std::addressof(subsystem)),
              m_innerControl(std::move(innerControl)),
              m_fullDirection(subsystem.GetPartition().RootValueSize()),
              m_fullGravityDirection(subsystem.GetPartition().RootValueSize()),
              m_fullAction(subsystem.GetPartition().RootResidualSize()),
              m_fullGravityAction(subsystem.GetPartition().RootResidualSize()),
              m_gravityRightHandSide(subsystem.GetPartition().GravityResidualSize()),
              m_gravityCorrection(subsystem.GetPartition().GravityValueSize()),
              m_retainedResult(subsystem.GetPartition().RetainedResidualSize()),
              m_retainedGravityAction(subsystem.GetPartition().RetainedResidualSize()) {
            m_innerControl.Validate();
            const auto &partition = GetPartition();
            const auto &operation = subsystem.GetFullOperator();
            if (operation.Width() != partition.RootValueSize() ||
                operation.Height() != partition.RootResidualSize()) {
                throw std::invalid_argument("The reduced-gravity partition does not match its full operator.");
            }
        }

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

        [[nodiscard]] const GravityBlockPartition &GetPartition() const noexcept {
            return m_subsystem->GetPartition();
        }

        [[nodiscard]] const Subsystem &GetSubsystem() const noexcept {
            return *m_subsystem;
        }

        [[nodiscard]] MPI_Comm GetCommunicator() const noexcept {
            return m_subsystem->GetCommunicator();
        }

        [[nodiscard]] bool IsCurrent() const noexcept {
            return m_subsystem->IsCurrent();
        }

        [[nodiscard]] const LinearSolveControl &GetInnerSolveControl() const noexcept {
            return m_innerControl;
        }

        [[nodiscard]] const LinearSolveReport &GetLastInnerSolve() const noexcept {
            return m_lastInnerSolve;
        }

        [[nodiscard]] LinearSolveReport TryMult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const {
            RequireCurrent();
            const auto &partition = GetPartition();
            partition.InjectRetainedValues(input, m_fullDirection);
            FullMult(m_fullDirection, m_fullAction);
            partition.GatherGravityResiduals(m_fullAction, m_gravityRightHandSide);
            if (!SolveGravity()) {
                return m_lastInnerSolve;
            }
            partition.InjectGravityValues(m_gravityCorrection, m_fullGravityDirection);
            FullMult(m_fullGravityDirection, m_fullGravityAction);
            partition.GatherRetainedResiduals(m_fullAction, m_retainedResult);
            partition.GatherRetainedResiduals(m_fullGravityAction, m_retainedGravityAction);
            m_retainedResult -= m_retainedGravityAction;
            output = m_retainedResult;
            return m_lastInnerSolve;
        }

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override {
            const auto report = TryMult(input, output);
            if (!report.Converged()) {
                throw GravityReductionFailure(report);
            }
        }

        /*
         * Condensing a nonzero gravity residual preserves equivalence to the
         * monolithic Newton equations. Dropping it merely because gravity was
         * approximately completed would turn inner error into an outer bias.
         */
        [[nodiscard]] LinearSolveReport TryCondenseRightHandSide(
            const mfem::Vector &fullRightHandSide,
            mfem::Vector &reducedRightHandSide
        ) const {
            RequireCurrent();
            const auto &partition = GetPartition();
            partition.GatherGravityResiduals(fullRightHandSide, m_gravityRightHandSide);
            if (!SolveGravity()) {
                return m_lastInnerSolve;
            }
            partition.InjectGravityValues(m_gravityCorrection, m_fullGravityDirection);
            FullMult(m_fullGravityDirection, m_fullGravityAction);
            partition.GatherRetainedResiduals(fullRightHandSide, m_retainedResult);
            partition.GatherRetainedResiduals(m_fullGravityAction, m_retainedGravityAction);
            m_retainedResult -= m_retainedGravityAction;
            reducedRightHandSide = m_retainedResult;
            return m_lastInnerSolve;
        }

        [[nodiscard]] LinearSolveReport TryLiftCorrection(
            const mfem::Vector &retainedCorrection,
            const mfem::Vector &fullRightHandSide,
            mfem::Vector &fullCorrection
        ) const {
            RequireCurrent();
            const auto &partition = GetPartition();
            partition.InjectRetainedValues(retainedCorrection, m_fullDirection);
            FullMult(m_fullDirection, m_fullAction);
            m_fullAction *= -1.0;
            if (fullRightHandSide.Size() != m_fullAction.Size()) {
                throw std::invalid_argument("A lifted gravity correction requires a full residual right-hand side.");
            }
            m_fullAction += fullRightHandSide;
            partition.GatherGravityResiduals(m_fullAction, m_gravityRightHandSide);
            return FinishLift(fullCorrection);
        }

        [[nodiscard]] LinearSolveReport TryLiftTangent(
            const mfem::Vector &retainedDirection,
            mfem::Vector &fullDirection
        ) const {
            RequireCurrent();
            const auto &partition = GetPartition();
            partition.InjectRetainedValues(retainedDirection, m_fullDirection);
            FullMult(m_fullDirection, m_fullAction);
            partition.GatherGravityResiduals(m_fullAction, m_gravityRightHandSide);
            m_gravityRightHandSide *= -1.0;
            return FinishLift(fullDirection);
        }

    private:
        void RequireCurrent() const {
            if (!IsCurrent()) {
                throw std::logic_error("A reduced-gravity action requires the current prepared gravity subsystem.");
            }
        }

        void FullMult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const {
            m_subsystem->GetFullOperator().Mult(input, output);
            RequireCurrent();
        }

        [[nodiscard]] bool SolveGravity() const {
            // A previous Schur action is not a suitable initial guess for an
            // unrelated Krylov direction and can make the action history dependent.
            m_gravityCorrection = 0.0;
            m_lastInnerSolve = m_subsystem->Solve(m_gravityRightHandSide, m_gravityCorrection, m_innerControl);
            RequireCurrent();
            if (m_lastInnerSolve.Converged()) {
                if (!std::isfinite(m_lastInnerSolve.rightHandSideNorm) ||
                    !std::isfinite(m_lastInnerSolve.trueResidualNorm)) {
                    m_lastInnerSolve.status = LinearSolveStatus::non_finite;
                } else if (m_lastInnerSolve.rightHandSideNorm < 0.0 || m_lastInnerSolve.trueResidualNorm < 0.0 ||
                           m_lastInnerSolve.trueResidualNorm >
                               m_innerControl.ConvergenceThreshold(m_lastInnerSolve.rightHandSideNorm)) {
                    m_lastInnerSolve.status = LinearSolveStatus::backend_failure;
                }
            }
            return m_lastInnerSolve.Converged();
        }

        [[nodiscard]] LinearSolveReport FinishLift(mfem::Vector &fullCorrection) const {
            if (!SolveGravity()) {
                return m_lastInnerSolve;
            }
            GetPartition().AddGravityValues(m_gravityCorrection, m_fullDirection);
            fullCorrection = m_fullDirection;
            return m_lastInnerSolve;
        }

        Subsystem *m_subsystem;
        LinearSolveControl m_innerControl;
        mutable mfem::Vector m_fullDirection;
        mutable mfem::Vector m_fullGravityDirection;
        mutable mfem::Vector m_fullAction;
        mutable mfem::Vector m_fullGravityAction;
        mutable mfem::Vector m_gravityRightHandSide;
        mutable mfem::Vector m_gravityCorrection;
        mutable mfem::Vector m_retainedResult;
        mutable mfem::Vector m_retainedGravityAction;
        mutable LinearSolveReport m_lastInnerSolve;
    };

    /*
     * The retained block of an exact full inverse equals the inverse Schur
     * complement. Restricting a supplied approximate inverse preserves that
     * relationship without assembling a second approximate physics operator.
     */
    template <PreparedGravityReductionSubsystem Subsystem, typename FullInverse>
        requires std::derived_from<FullInverse, mfem::Solver> &&
                 LinearPreconditionerApplicationContractAvailable<FullInverse>
    class GravityReducedPreconditioner final : public mfem::Solver {
    public:
        static constexpr auto applicationContract = linearPreconditionerApplicationContract<FullInverse>;

        GravityReducedPreconditioner(
            const GravityReducedOperator<Subsystem> &operation,
            FullInverse &fullInverse
        )
            : mfem::Solver(operation.Width(), operation.Height(), false),
              m_operation(std::addressof(operation)),
              m_fullInverse(std::addressof(fullInverse)),
              m_fullRightHandSide(operation.GetPartition().RootResidualSize()),
              m_fullCorrection(operation.GetPartition().RootValueSize()) {
            if (fullInverse.Width() != m_fullRightHandSide.Size() ||
                fullInverse.Height() != m_fullCorrection.Size()) {
                throw std::invalid_argument("A reduced-gravity preconditioner requires a compatible full inverse.");
            }
        }

        void SetOperator(const mfem::Operator &operation) override {
            if (std::addressof(operation) != m_operation) {
                throw std::invalid_argument("The reduced-gravity preconditioner cannot bind a different operator.");
            }
        }

        void Mult(
            const mfem::Vector &input,
            mfem::Vector &output
        ) const override {
            if (!m_operation->IsCurrent()) {
                throw std::logic_error("The reduced-gravity preconditioner requires the current prepared subsystem.");
            }
            const auto &partition = m_operation->GetPartition();
            partition.InjectRetainedResiduals(input, m_fullRightHandSide);
            m_fullInverse->Mult(m_fullRightHandSide, m_fullCorrection);
            if (!m_operation->IsCurrent()) {
                throw std::logic_error("The full inverse changed the reduced-gravity preparation.");
            }
            partition.GatherRetainedValues(m_fullCorrection, output);
        }

        [[nodiscard]] const FullInverse &GetFullInverse() const noexcept {
            return *m_fullInverse;
        }

    private:
        const GravityReducedOperator<Subsystem> *m_operation;
        FullInverse *m_fullInverse;
        mutable mfem::Vector m_fullRightHandSide;
        mutable mfem::Vector m_fullCorrection;
    };
} // namespace mean_field::solver