module;

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

#include <mfem.hpp>

export module mean_field:preconditioning.backend_implementations;

export import :preconditioning.backend;

export namespace mean_field::preconditioning::backend {
    struct BackendStatistics final {
        std::uint64_t setups{0};
        std::uint64_t applications{0};
        std::uint64_t innerIterations{0};
        std::uint64_t lastInnerIterations{0};
    };

    namespace detail {
        inline void verifyApplicationDimensions(
            const mfem::Solver &solver,
            const mfem::Vector &rightHandSide,
            const mfem::Vector &action
        ) {
            if (rightHandSide.Size() != solver.Width() || action.Size() != solver.Height()) {
                throw std::invalid_argument(
                    "A prepared preconditioning backend requires compatible, preallocated input and output vectors."
                );
            }
        }

        inline void verifySquarePositiveSize(
            const int height,
            const int width
        ) {
            if (height <= 0 || height != width) {
                throw std::invalid_argument("A preconditioning backend requires a positive square operator.");
            }
        }

        inline void configure(
            mfem::HypreBoomerAMG &solver,
            const FixedCycles &mode
        ) {
            if (mode.cycles <= 0) {
                throw std::invalid_argument("Fixed-cycle AMG requires at least one cycle.");
            }
            solver.SetMaxIter(mode.cycles);
            solver.SetTol(0.0);
            solver.SetPrintLevel(0);
            solver.iterative_mode = false;
        }

        inline void configure(
            mfem::HypreBoomerAMG &solver,
            const SolveToTolerance &mode
        ) {
            if (!std::isfinite(mode.relativeTolerance) || mode.relativeTolerance <= 0.0 ||
                mode.relativeTolerance >= 1.0) {
                throw std::invalid_argument("Tolerance-driven AMG requires a finite relative tolerance in (0, 1).");
            }
            if (mode.maximumCycles <= 0) {
                throw std::invalid_argument("Tolerance-driven AMG requires at least one permitted cycle.");
            }
            solver.SetMaxIter(mode.maximumCycles);
            solver.SetTol(mode.relativeTolerance);
            solver.SetPrintLevel(0);
            solver.iterative_mode = false;
        }
    } // namespace detail

    class PreparedDiagonal final : public mfem::Solver {
    public:
        PreparedDiagonal(
            Diagonal configuration,
            const mfem::Vector &diagonal
        )
            : mfem::Solver(diagonal.Size()),
              m_configuration(std::move(configuration)) {
            Refresh(diagonal);
        }

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

        void SetOperator(const mfem::Operator &operation) override {
            if (operation.Height() != Height() || operation.Width() != Width()) {
                throw std::invalid_argument("The diagonal backend received an operator with incompatible dimensions.");
            }
        }

        void Mult(
            const mfem::Vector &rightHandSide,
            mfem::Vector &action
        ) const override {
            detail::verifyApplicationDimensions(*this, rightHandSide, action);
            for (int index = 0; index < Height(); ++index) {
                action(index) = m_inverseDiagonal(index) * rightHandSide(index);
            }
            ++m_statistics.applications;
        }

        void Refresh(const mfem::Vector &diagonal) {
            if (diagonal.Size() <= 0 || diagonal.Size() != Height()) {
                throw std::invalid_argument("The diagonal backend requires a positive diagonal of unchanged size.");
            }

            m_inverseDiagonal.SetSize(diagonal.Size());
            for (int index = 0; index < diagonal.Size(); ++index) {
                const double entry = diagonal(index);
                if (!std::isfinite(entry) || entry == 0.0) {
                    throw std::invalid_argument("The diagonal backend cannot invert a zero or non-finite entry.");
                }
                m_inverseDiagonal(index) = 1.0 / entry;
            }
            ++m_statistics.setups;
        }

        [[nodiscard]] const Diagonal &GetConfiguration() const noexcept {
            return m_configuration;
        }

        [[nodiscard]] const mfem::Vector &GetInverseDiagonal() const noexcept {
            return m_inverseDiagonal;
        }

        [[nodiscard]] const BackendStatistics &GetStatistics() const noexcept {
            return m_statistics;
        }

    private:
        Diagonal m_configuration;
        mfem::Vector m_inverseDiagonal;
        mutable BackendStatistics m_statistics;
    };

    class PreparedMatrixFreeChebyshev final : public mfem::Solver {
    public:
        PreparedMatrixFreeChebyshev(
            MatrixFreeChebyshev configuration,
            const mfem::Operator &operation,
            const MPI_Comm communicator
        )
            : mfem::Solver(operation.Height()),
              m_configuration(std::move(configuration)),
              m_communicator(communicator) {
            ValidateConfiguration();
            Refresh(operation);
        }

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

        void SetOperator(const mfem::Operator &operation) override {
            Refresh(operation);
        }

        void Mult(
            const mfem::Vector &rightHandSide,
            mfem::Vector &action
        ) const override {
            detail::verifyApplicationDimensions(*this, rightHandSide, action);
            m_smoother->Mult(rightHandSide, action);
            ++m_statistics.applications;
            m_statistics.lastInnerIterations = static_cast<std::uint64_t>(m_configuration.order);
            m_statistics.innerIterations += static_cast<std::uint64_t>(m_configuration.order);
        }

        void Refresh(const mfem::Operator &operation) {
            detail::verifySquarePositiveSize(operation.Height(), operation.Width());
            if (operation.Height() != Height()) {
                throw std::invalid_argument("The matrix-free Chebyshev backend cannot change size during refresh.");
            }

            operation.AssembleDiagonal(m_diagonal);
            if (m_diagonal.Size() != Height()) {
                throw std::invalid_argument(
                    "The matrix-free Chebyshev backend received an incompatible assembled diagonal."
                );
            }
            for (int index = 0; index < m_diagonal.Size(); ++index) {
                if (!std::isfinite(m_diagonal(index)) || m_diagonal(index) <= 0.0) {
                    throw std::invalid_argument(
                        "The matrix-free Chebyshev backend requires a finite, strictly positive diagonal."
                    );
                }
            }

            m_operation = std::addressof(operation);
            m_essentialTrueDofs.SetSize(0);
            m_smoother = std::make_unique<mfem::OperatorChebyshevSmoother>(
                operation, m_diagonal, m_essentialTrueDofs, m_configuration.order, m_communicator,
                m_configuration.powerIterations, m_configuration.powerTolerance, m_configuration.powerSeed
            );
            m_smoother->iterative_mode = false;
            ++m_statistics.setups;
        }

        [[nodiscard]] const MatrixFreeChebyshev &GetConfiguration() const noexcept {
            return m_configuration;
        }

        [[nodiscard]] const mfem::Operator &GetOperator() const noexcept {
            return *m_operation;
        }

        [[nodiscard]] const mfem::Vector &GetDiagonal() const noexcept {
            return m_diagonal;
        }

        [[nodiscard]] const BackendStatistics &GetStatistics() const noexcept {
            return m_statistics;
        }

    private:
        void ValidateConfiguration() const {
            if (m_configuration.order <= 0 || m_configuration.order > 5) {
                throw std::invalid_argument("Matrix-free Chebyshev requires a polynomial order in [1, 5].");
            }
            if (m_configuration.powerIterations <= 0) {
                throw std::invalid_argument("Matrix-free Chebyshev requires at least one power iteration.");
            }
            if (!std::isfinite(m_configuration.powerTolerance) || m_configuration.powerTolerance <= 0.0 ||
                m_configuration.powerTolerance >= 1.0) {
                throw std::invalid_argument(
                    "Matrix-free Chebyshev requires a finite power-method tolerance strictly between zero and one."
                );
            }
            if (m_configuration.powerSeed <= 0) {
                throw std::invalid_argument("Matrix-free Chebyshev requires a strictly positive power-method seed.");
            }
        }

        MatrixFreeChebyshev m_configuration;
        MPI_Comm m_communicator;
        const mfem::Operator *m_operation{nullptr};
        mfem::Vector m_diagonal;
        mfem::Array<int> m_essentialTrueDofs;
        std::unique_ptr<mfem::OperatorChebyshevSmoother> m_smoother;
        mutable BackendStatistics m_statistics;
    };

    class PreparedDenseDirect final : public mfem::Solver {
    public:
        PreparedDenseDirect(
            DenseDirect configuration,
            const mfem::DenseMatrix &matrix
        )
            : mfem::Solver(matrix.Height()),
              m_configuration(std::move(configuration)) {
            Refresh(matrix);
        }

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

        void SetOperator(const mfem::Operator &operation) override {
            const auto *matrix = dynamic_cast<const mfem::DenseMatrix *>(&operation);
            if (matrix == nullptr) {
                throw std::invalid_argument("The dense-direct backend requires an mfem::DenseMatrix.");
            }
            Refresh(*matrix);
        }

        void Mult(
            const mfem::Vector &rightHandSide,
            mfem::Vector &action
        ) const override {
            detail::verifyApplicationDimensions(*this, rightHandSide, action);
            m_inverse->Mult(rightHandSide, action);
            ++m_statistics.applications;
        }

        void Refresh(const mfem::DenseMatrix &matrix) {
            detail::verifySquarePositiveSize(matrix.Height(), matrix.Width());
            if (matrix.Height() != Height()) {
                throw std::invalid_argument("The dense-direct backend cannot change size during refresh.");
            }

            m_matrix  = matrix;
            m_inverse = std::make_unique<mfem::DenseMatrixInverse>(m_matrix);
            ++m_statistics.setups;
        }

        [[nodiscard]] const DenseDirect &GetConfiguration() const noexcept {
            return m_configuration;
        }

        [[nodiscard]] const mfem::DenseMatrix &GetDenseSurrogate() const noexcept {
            return m_matrix;
        }

        [[nodiscard]] const BackendStatistics &GetStatistics() const noexcept {
            return m_statistics;
        }

    private:
        DenseDirect m_configuration;
        mfem::DenseMatrix m_matrix;
        std::unique_ptr<mfem::DenseMatrixInverse> m_inverse;
        mutable BackendStatistics m_statistics;
    };

    template <ApplicationMode Mode> class PreparedHypreBoomerAMG final : public mfem::Solver {
    public:
        using Configuration = HypreBoomerAMG<Mode>;

        PreparedHypreBoomerAMG(
            Configuration configuration,
            const mfem::HypreParMatrix &matrix
        )
            : mfem::Solver(matrix.Height()),
              m_configuration(std::move(configuration)) {
            Refresh(matrix);
        }

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

        void SetOperator(const mfem::Operator &operation) override {
            const auto *matrix = dynamic_cast<const mfem::HypreParMatrix *>(&operation);
            if (matrix == nullptr) {
                throw std::invalid_argument("The BoomerAMG backend requires an mfem::HypreParMatrix surrogate.");
            }
            Refresh(*matrix);
        }

        void Mult(
            const mfem::Vector &rightHandSide,
            mfem::Vector &action
        ) const override {
            detail::verifyApplicationDimensions(*this, rightHandSide, action);
            m_solver->Mult(rightHandSide, action);

            int iterations = 0;
            m_solver->GetNumIterations(iterations);
            ++m_statistics.applications;
            m_statistics.lastInnerIterations = static_cast<std::uint64_t>(iterations);
            m_statistics.innerIterations += static_cast<std::uint64_t>(iterations);
        }

        void Refresh(const mfem::HypreParMatrix &matrix) {
            detail::verifySquarePositiveSize(matrix.Height(), matrix.Width());
            if (matrix.Height() != Height()) {
                throw std::invalid_argument("The BoomerAMG backend cannot change size during refresh.");
            }

            m_sparseSurrogate = std::addressof(matrix);
            m_solver          = std::make_unique<mfem::HypreBoomerAMG>(matrix);
            detail::configure(*m_solver, m_configuration.application);

            mfem::Vector setupRightHandSide(Width());
            mfem::Vector setupAction(Height());
            setupRightHandSide = 0.0;
            setupAction        = 0.0;
            m_solver->Setup(setupRightHandSide, setupAction);
            ++m_statistics.setups;
        }

        [[nodiscard]] const Configuration &GetConfiguration() const noexcept {
            return m_configuration;
        }

        [[nodiscard]] const mfem::HypreParMatrix &GetSparseSurrogate() const noexcept {
            return *m_sparseSurrogate;
        }

        [[nodiscard]] const BackendStatistics &GetStatistics() const noexcept {
            return m_statistics;
        }

    private:
        Configuration m_configuration;
        const mfem::HypreParMatrix *m_sparseSurrogate{nullptr};
        std::unique_ptr<mfem::HypreBoomerAMG> m_solver;
        mutable BackendStatistics m_statistics;
    };

    [[nodiscard]] inline PreparedDiagonal prepare(
        Diagonal configuration,
        const mfem::Vector &diagonal
    ) {
        return PreparedDiagonal{std::move(configuration), diagonal};
    }

    [[nodiscard]] inline PreparedMatrixFreeChebyshev prepare(
        MatrixFreeChebyshev configuration,
        const mfem::Operator &operation,
        const MPI_Comm communicator
    ) {
        return PreparedMatrixFreeChebyshev{std::move(configuration), operation, communicator};
    }

    [[nodiscard]] inline PreparedDenseDirect prepare(
        DenseDirect configuration,
        const mfem::DenseMatrix &matrix
    ) {
        return PreparedDenseDirect{std::move(configuration), matrix};
    }

    template <ApplicationMode Mode>
    [[nodiscard]] PreparedHypreBoomerAMG<Mode> prepare(
        HypreBoomerAMG<Mode> configuration,
        const mfem::HypreParMatrix &matrix
    ) {
        return PreparedHypreBoomerAMG<Mode>{std::move(configuration), matrix};
    }
} // namespace mean_field::preconditioning::backend