// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// RationalMatrix.hpp
//
// Rational matrix with a single cleared common denominator (FLINT fmpq_mat style).
//
// In Matrix<Rational> every element carries its own numerator and denominator.
// An N×N matrix multiply triggers N^3 rational reductions, so GCD computation
// dominates. RationalMatrix collects the denominator into a single Int and
// operates over an integer matrix, reducing the reduction count to one at the
// very end.
//
// Internal representation: M = num_ / den_  (num_ : Matrix<Int>, den_ : Int>0)
// Determinant: Bareiss method (fraction-free Gauss) applied to num_.
//          det(M) = det(num_) / den_^n.
//
// Design notes:
//   - No template specialization of Matrix<Rational> (the internal
//     representation is fundamentally different).
//   - Element access (operator()(i,j)) is read-only and returns a Rational.
//     Write access is not exposed (it would break the common denominator).
//   - Header-only implementation that reuses the existing Matrix<Int> / IntOps
//     functionality as-is.

#ifndef SANGI_LINALG_RATIONAL_MATRIX_HPP
#define SANGI_LINALG_RATIONAL_MATRIX_HPP

#include <math/core/matrix.hpp>
#include <math/core/mp/Int.hpp>
#include <math/core/mp/Rational.hpp>
#include <math/linalg/decomposition.hpp>
#include <math/linalg/dixon_solve.hpp>
#include <cstddef>
#include <stdexcept>
#include <utility>
#include <vector>

namespace sangi {

    /**
     * @brief Rational matrix with a single cleared common denominator.
     *
     * Internal representation: an integer matrix num_ and a positive common
     * denominator den_ represent the logical matrix as
     *   M(i, j) = num_(i, j) / den_
     * Unlike Matrix<Rational>, where each element is reduced individually,
     * operations reduce to integer-matrix arithmetic, drastically lowering the
     * reduction cost.
     */
    class RationalMatrix {
    public:
        using size_type = std::size_t;

    private:
        Matrix<Int> num_;
        Int den_;   // invariant: den_ > 0

    public:
        // ====================================================================
        // Construction
        // ====================================================================

        /// 0×0 matrix (den_ = 1)
        RationalMatrix() : num_(), den_(1) {}

        /// rows×cols zero matrix (den_ = 1)
        RationalMatrix(size_type rows, size_type cols)
            : num_(rows, cols), den_(1) {
            num_.zero();
        }

        /// Construct from a Matrix<Rational>. den_ is set to the LCM of all
        /// element denominators.
        explicit RationalMatrix(const Matrix<Rational>& src)
            : num_(src.rows(), src.cols()), den_(1) {
            const size_type R = src.rows();
            const size_type C = src.cols();

            // 1. Compute the LCM of all element denominators
            for (size_type i = 0; i < R; ++i) {
                for (size_type j = 0; j < C; ++j) {
                    const Int& d = src(i, j).denominator();
                    if (!d.isOne())
                        den_ = lcm(den_, d);
                }
            }

            // 2. Scale each element by den_/d to make it integral
            for (size_type i = 0; i < R; ++i) {
                for (size_type j = 0; j < C; ++j) {
                    const Rational& r = src(i, j);
                    if (r.isZero()) {
                        num_(i, j) = Int(0);
                    } else {
                        // num_(i,j) = numerator * (den_ / denominator)
                        Int scale = den_ / r.denominator();
                        num_(i, j) = r.numerator() * scale;
                    }
                }
            }
        }

        /// Construct directly from an integer matrix and a denominator. Throws
        /// if den is 0 or negative.
        /// reduce=true reduces at construction time (default false: fast
        /// construction).
        RationalMatrix(Matrix<Int> num, Int den, bool reduce = false)
            : num_(std::move(num)), den_(std::move(den)) {
            if (den_.isZero())
                throw std::invalid_argument("RationalMatrix: denominator must be nonzero");
            if (den_.isNegative()) {
                den_ = -den_;
                negateNumerator();
            }
            if (reduce) this->reduce();
        }

        RationalMatrix(const RationalMatrix&) = default;
        RationalMatrix(RationalMatrix&&) noexcept = default;
        RationalMatrix& operator=(const RationalMatrix&) = default;
        RationalMatrix& operator=(RationalMatrix&&) noexcept = default;

        // ====================================================================
        // Factories
        // ====================================================================

        /// rows×cols zero matrix
        static RationalMatrix zero(size_type rows, size_type cols) {
            return RationalMatrix(rows, cols);
        }

        /// n×n identity matrix
        static RationalMatrix identity(size_type n) {
            RationalMatrix m(n, n);
            for (size_type i = 0; i < n; ++i)
                m.num_(i, i) = Int(1);
            // den_ stays 1
            return m;
        }

        // ====================================================================
        // Accessors
        // ====================================================================

        size_type rows() const noexcept { return num_.rows(); }
        size_type cols() const noexcept { return num_.cols(); }
        bool isSquare() const noexcept { return num_.rows() == num_.cols(); }

        const Matrix<Int>& numerator() const noexcept { return num_; }
        const Int& denominator() const noexcept { return den_; }

        /// Element access (read-only). Constructs and returns a Rational.
        Rational operator()(size_type i, size_type j) const {
            return Rational(num_(i, j), den_);
        }

        // ====================================================================
        // Conversion
        // ====================================================================

        /// Convert back to a Matrix<Rational> (constructs a Rational per element)
        [[nodiscard]] Matrix<Rational> toRationalMatrix() const {
            Matrix<Rational> result(num_.rows(), num_.cols());
            for (size_type i = 0; i < num_.rows(); ++i)
                for (size_type j = 0; j < num_.cols(); ++j)
                    result(i, j) = Rational(num_(i, j), den_);
            return result;
        }

        // ====================================================================
        // Reduction
        // ====================================================================

        /**
         * @brief Reduce by the common factor of all elements and den_.
         *
         * Computes g = gcd(den_, gcd(all elements)) and divides num_ and den_
         * by g. Scans every element once, so it costs O(rows*cols * gcd_cost).
         * Call it explicitly when the size of an operation's result matters.
         */
        void reduce() {
            if (den_.isOne()) return;
            Int g = den_;
            for (size_type i = 0; i < num_.rows() && !g.isOne(); ++i) {
                for (size_type j = 0; j < num_.cols() && !g.isOne(); ++j) {
                    if (!num_(i, j).isZero())
                        g = gcd(g, num_(i, j));
                }
            }
            if (g.isOne()) return;
            for (size_type i = 0; i < num_.rows(); ++i)
                for (size_type j = 0; j < num_.cols(); ++j)
                    if (!num_(i, j).isZero())
                        num_(i, j) = num_(i, j) / g;
            den_ = den_ / g;
        }

        // ====================================================================
        // Arithmetic
        // ====================================================================

        /// Unary minus
        [[nodiscard]] RationalMatrix operator-() const {
            RationalMatrix r = *this;
            r.negateNumerator();
            return r;
        }

        /// Addition: (Na/da) + (Nb/db) = (Na*db + Nb*da) / (da*db)
        [[nodiscard]] RationalMatrix operator+(const RationalMatrix& rhs) const {
            checkSameShape(rhs, "operator+");
            // If the common denominators are equal, only an integer-matrix add
            if (den_ == rhs.den_) {
                return RationalMatrix(num_ + rhs.num_, den_);
            }
            // General case: g = gcd(da, db), da' = da/g, db' = db/g
            // result num = Na*db' + Nb*da', result den = da*db' = da'*db'*g
            Int g = gcd(den_, rhs.den_);
            Int da_g = den_ / g;
            Int db_g = rhs.den_ / g;
            Matrix<Int> n_lhs = scaleNum(num_, db_g);
            Matrix<Int> n_rhs = scaleNum(rhs.num_, da_g);
            return RationalMatrix(n_lhs + n_rhs, den_ * db_g);
        }

        /// Subtraction
        [[nodiscard]] RationalMatrix operator-(const RationalMatrix& rhs) const {
            checkSameShape(rhs, "operator-");
            if (den_ == rhs.den_) {
                return RationalMatrix(num_ - rhs.num_, den_);
            }
            Int g = gcd(den_, rhs.den_);
            Int da_g = den_ / g;
            Int db_g = rhs.den_ / g;
            Matrix<Int> n_lhs = scaleNum(num_, db_g);
            Matrix<Int> n_rhs = scaleNum(rhs.num_, da_g);
            return RationalMatrix(n_lhs - n_rhs, den_ * db_g);
        }

        /// Matrix multiplication: (Na/da) * (Nb/db) = (Na*Nb) / (da*db)
        [[nodiscard]] RationalMatrix operator*(const RationalMatrix& rhs) const {
            if (cols() != rhs.rows()) {
                throw DimensionError("RationalMatrix::operator*: dimension mismatch");
            }
            return RationalMatrix(num_ * rhs.num_, den_ * rhs.den_);
        }

        /// Scalar multiplication (Rational): (N/d) * (p/q) = (N*p) / (d*q)
        [[nodiscard]] RationalMatrix operator*(const Rational& s) const {
            if (s.isZero()) return RationalMatrix::zero(rows(), cols());
            return RationalMatrix(scaleNum(num_, s.numerator()), den_ * s.denominator());
        }

        /// Scalar multiplication (Int): (N/d) * k = (N*k) / d
        [[nodiscard]] RationalMatrix operator*(const Int& s) const {
            if (s.isZero()) return RationalMatrix::zero(rows(), cols());
            return RationalMatrix(scaleNum(num_, s), den_);
        }

        /// Scalar multiplication (int)
        [[nodiscard]] RationalMatrix operator*(int s) const {
            return *this * Int(s);
        }

        // ====================================================================
        // Transpose
        // ====================================================================

        [[nodiscard]] RationalMatrix transpose() const {
            return RationalMatrix(num_.transpose(), den_);
        }

        // ====================================================================
        // Determinant (Bareiss)
        // ====================================================================

        /**
         * @brief Determinant (Bareiss fraction-free Gauss applied to the
         * integer matrix).
         *
         * det(M) = det(num_) / den_^n, where n = rows().
         * Integer arithmetic only, and every division is exact (Bareiss
         * divisibility), so the reduction cost is dramatically lower than
         * Gaussian elimination over Matrix<Rational>.
         */
        [[nodiscard]] Rational determinant() const {
            if (!isSquare()) {
                throw DimensionError("RationalMatrix::determinant: matrix must be square");
            }
            const size_type n = rows();
            if (n == 0) return Rational(1);

            Int det_num = algorithms::bareiss_determinant(num_);
            // den_^n
            Int den_pow = pow(den_, static_cast<unsigned int>(n));
            return Rational(det_num, den_pow);
        }

        // ====================================================================
        // Linear system solving (Dixon p-adic lifting)
        // ====================================================================

        /**
         * @brief Solve A · X = b exactly via Dixon p-adic lifting.
         *
         * A = num_ / den_ (n×n), b = b.num_ / b.den_ (n×m). We build X by
         * solving, column by column with dixon::solve_int, the integer-matrix /
         * integer-right-hand-side system
         *   A_int · (b.den_ · X) = den_ · b.num_
         *
         * Complexity: on the order of O(n^3 log) (no intermediate GCD cost of
         * Gauss<Rational>). Faster than `solve_lq_continuous_laub` and similar
         * for large n or bit-rich coefficients.
         *
         * @param b right-hand side (n×m RationalMatrix)
         * @return X (n×m RationalMatrix)
         * @throws DimensionError shape mismatch / non-square
         * @throws std::runtime_error matrix is singular (det ≡ 0 for all
         *         candidate primes)
         */
        [[nodiscard]] RationalMatrix solve(const RationalMatrix& b) const {
            if (!isSquare())
                throw DimensionError("RationalMatrix::solve: A must be square");
            if (b.rows() != rows())
                throw DimensionError("RationalMatrix::solve: dimension mismatch");

            const size_type n = rows();
            const size_type m = b.cols();
            if (n == 0) return RationalMatrix(0, 0);

            // Build each element of the result X as a Rational, then collect
            // into a RationalMatrix at the end
            Matrix<Rational> X(n, m);
            const Int& d_a = den_;
            const Int& d_b = b.den_;
            // Since x_j = y_j / d_b, "solve with Dixon using d_a * b.num_ as the
            // right-hand side"; dividing the solution y by d_b yields x.
            const Rational inv_d_b(Int(1), d_b);

            for (size_type j = 0; j < m; ++j) {
                std::vector<Int> rhs(n);
                for (size_type i = 0; i < n; ++i) {
                    rhs[i] = d_a * b.num_(i, j);
                }
                std::vector<Rational> y = detail::dixon::solve_int(num_, rhs);
                for (size_type i = 0; i < n; ++i) {
                    X(i, j) = y[i] * inv_d_b;
                }
            }
            return RationalMatrix(X);
        }

    private:
        // ====================================================================
        // Internal helpers
        // ====================================================================

        void negateNumerator() {
            for (size_type i = 0; i < num_.rows(); ++i)
                for (size_type j = 0; j < num_.cols(); ++j)
                    if (!num_(i, j).isZero())
                        num_(i, j) = -num_(i, j);
        }

        void checkSameShape(const RationalMatrix& rhs, const char* op) const {
            if (rows() != rhs.rows() || cols() != rhs.cols()) {
                std::string msg = "RationalMatrix::";
                msg += op;
                msg += ": shape mismatch";
                throw DimensionError(msg);
            }
        }

        static Matrix<Int> scaleNum(const Matrix<Int>& m, const Int& k) {
            Matrix<Int> r(m.rows(), m.cols());
            for (size_type i = 0; i < m.rows(); ++i)
                for (size_type j = 0; j < m.cols(); ++j)
                    r(i, j) = m(i, j) * k;
            return r;
        }
    };

    // ========================================================================
    // Free functions / left-hand scalar
    // ========================================================================

    inline RationalMatrix operator*(const Rational& s, const RationalMatrix& m) {
        return m * s;
    }

    inline RationalMatrix operator*(const Int& s, const RationalMatrix& m) {
        return m * s;
    }

    inline RationalMatrix operator*(int s, const RationalMatrix& m) {
        return m * s;
    }

    /// Convert a Matrix<Rational> to a RationalMatrix (free-function version)
    inline RationalMatrix toRationalMatrix(const Matrix<Rational>& src) {
        return RationalMatrix(src);
    }

} // namespace sangi

#endif // SANGI_LINALG_RATIONAL_MATRIX_HPP