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

// dixon_solve.hpp
//
// Exact (rational) solution of a square integer linear system Ax = b via Dixon p-adic lifting.
//
// References:
//   - Dixon, "Exact solution of linear equations using p-adic expansions" (1982)
//   - Storjohann, "High-Order Lifting and Integrality Certification" (2003)
//
// Algorithm overview:
//   1. Pick a 30-bit prime p (one that does not divide det(A)) and compute
//      B = A^{-1} mod p via modular Gaussian elimination.
//   2. Starting from r_0 = b, iterate k times:
//        c_i = (B · r_i) mod p     (mod p vector, range [0, p))
//        r_{i+1} = (r_i - A · c_i) / p   (integer arithmetic; the mod p residual term cancels, so the division is exact)
//   3. y = Σ_{i=0}^{k-1} c_i p^i is an integer vector satisfying y ≡ x (mod p^k).
//   4. Determine k from the Hadamard bound so that p^k > 2 N D, then recover each
//      component as a rational via rational reconstruction reconstructRational(y_j, p^k, N, D).
//
// Complexity:
//   Each step is O(n^2) (modular matrix-vector + integer residual update), and with
//   k = O(n log n) iterations the total is O(n^3 log n).
//   Compared head-to-head with Gauss<Rational> at O(n^3), Dixon wins as n grows
//   (lower intermediate GCD cost + the bit-complexity is paid once in the final reconstruction).
//
// Scope:
//   - Single-prime version (Phase 1). If p divides det, retry with the next prime.
//   - Square matrices only. Rectangular systems require a different algorithm (e.g. least-squares).
//   - Single right-hand side. Multiple right-hand sides could be supported by reusing the same B = A^{-1} mod p.

#ifndef SANGI_LINALG_DIXON_SOLVE_HPP
#define SANGI_LINALG_DIXON_SOLVE_HPP

#include <math/core/matrix.hpp>
#include <math/core/mp/Int.hpp>
#include <math/core/mp/Rational.hpp>
#include <array>
#include <cstdint>
#include <stdexcept>
#include <vector>

namespace sangi::detail::dixon {

    // ====================================================================
    // Candidate primes (30-bit, so that p^2 fits in int63)
    // ====================================================================

    inline constexpr std::array<int64_t, 8> kCandidatePrimes = {
        1073741789LL,  // 2^30 - 35 (largest 30-bit prime)
        1073741783LL,
        1073741741LL,
        1073741723LL,
        1073741719LL,
        1073741717LL,
        1073741689LL,
        1073741671LL
    };

    // ====================================================================
    // Modular integer arithmetic (raw int64-based)
    // ====================================================================

    /// a + b (mod p), with a, b ∈ [0, p)
    inline int64_t add_mod(int64_t a, int64_t b, int64_t p) noexcept {
        int64_t s = a + b;
        return s >= p ? s - p : s;
    }

    /// a - b (mod p), with a, b ∈ [0, p)
    inline int64_t sub_mod(int64_t a, int64_t b, int64_t p) noexcept {
        int64_t d = a - b;
        return d < 0 ? d + p : d;
    }

    /// a * b (mod p), with a, b ∈ [0, p); assumes p < 2^31 (the product fits in int63)
    inline int64_t mul_mod(int64_t a, int64_t b, int64_t p) noexcept {
        return (a * b) % p;
    }

    /// Return a^{-1} mod p via the extended Euclidean algorithm (assumes a ∈ [1, p) and gcd(a, p) = 1)
    inline int64_t inv_mod(int64_t a, int64_t p) {
        int64_t old_r = a, r = p;
        int64_t old_s = 1, s = 0;
        while (r != 0) {
            int64_t q = old_r / r;
            int64_t tmp = old_r - q * r;
            old_r = r; r = tmp;
            tmp = old_s - q * s;
            old_s = s; s = tmp;
        }
        // old_r should be 1 here (a, p coprime)
        int64_t result = old_s % p;
        return result < 0 ? result + p : result;
    }

    // ====================================================================
    // Modular matrices (row-major) and basic operations
    // ====================================================================

    using ModRow = std::vector<int64_t>;
    using ModMat = std::vector<ModRow>;
    using ModVec = std::vector<int64_t>;

    /// Reduce a Matrix<Int> mod p (each entry in [0, p))
    inline ModMat reduce_int_matrix(const Matrix<Int>& A, int64_t p) {
        const std::size_t R = A.rows();
        const std::size_t C = A.cols();
        ModMat M(R, ModRow(C));
        Int p_int(p);
        for (std::size_t i = 0; i < R; ++i) {
            for (std::size_t j = 0; j < C; ++j) {
                Int rem = A(i, j) % p_int;
                if (rem.isNegative()) rem = rem + p_int;
                // rem ∈ [0, p) fits in 64 bits
                M[i][j] = static_cast<int64_t>(rem.toInt64());
            }
        }
        return M;
    }

    /// Reduce an Int vector mod p
    inline ModVec reduce_int_vec(const std::vector<Int>& v, int64_t p) {
        ModVec out(v.size());
        Int p_int(p);
        for (std::size_t i = 0; i < v.size(); ++i) {
            Int rem = v[i] % p_int;
            if (rem.isNegative()) rem = rem + p_int;
            out[i] = static_cast<int64_t>(rem.toInt64());
        }
        return out;
    }

    /// y = M · v mod p
    inline ModVec mod_matvec(const ModMat& M, const ModVec& v, int64_t p) {
        const std::size_t R = M.size();
        const std::size_t C = v.size();
        ModVec out(R, 0);
        for (std::size_t i = 0; i < R; ++i) {
            int64_t acc = 0;
            for (std::size_t j = 0; j < C; ++j) {
                acc = add_mod(acc, mul_mod(M[i][j], v[j], p), p);
            }
            out[i] = acc;
        }
        return out;
    }

    /// Compute A^{-1} mod p via Gauss-Jordan elimination.
    /// Returns true on success, false if a zero pivot is hit (= det ≡ 0 mod p).
    inline bool modular_inverse(const ModMat& A, int64_t p, ModMat& outInv) {
        const std::size_t n = A.size();
        if (n == 0 || A[0].size() != n) return false;
        // augmented [A | I]
        std::vector<std::vector<int64_t>> aug(n, std::vector<int64_t>(2 * n, 0));
        for (std::size_t i = 0; i < n; ++i) {
            for (std::size_t j = 0; j < n; ++j) aug[i][j] = A[i][j];
            aug[i][n + i] = 1;
        }
        for (std::size_t col = 0; col < n; ++col) {
            // partial pivoting
            std::size_t pivot_row = n;  // sentinel
            for (std::size_t r = col; r < n; ++r) {
                if (aug[r][col] != 0) { pivot_row = r; break; }
            }
            if (pivot_row == n) return false;
            if (pivot_row != col) std::swap(aug[col], aug[pivot_row]);

            int64_t pivot = aug[col][col];
            int64_t pivot_inv = inv_mod(pivot, p);
            // normalize pivot row
            for (std::size_t j = 0; j < 2 * n; ++j) {
                aug[col][j] = mul_mod(aug[col][j], pivot_inv, p);
            }
            // eliminate other rows
            for (std::size_t r = 0; r < n; ++r) {
                if (r == col) continue;
                int64_t factor = aug[r][col];
                if (factor == 0) continue;
                for (std::size_t j = 0; j < 2 * n; ++j) {
                    aug[r][j] = sub_mod(aug[r][j], mul_mod(factor, aug[col][j], p), p);
                }
            }
        }
        // extract
        outInv.assign(n, ModRow(n));
        for (std::size_t i = 0; i < n; ++i) {
            for (std::size_t j = 0; j < n; ++j) {
                outInv[i][j] = aug[i][n + j];
            }
        }
        return true;
    }

    // ====================================================================
    // Hadamard bound
    // ====================================================================

    /// Compute Π_i (||row_i(A)||_2^2 + b_i^2) as an Int.
    /// The square root H_Ab of this value bounds |det(A_j)| (the Cramer matrix with column j replaced by b).
    /// The bound on |det(A)| is the square root H_A of Π_i ||row_i(A)||_2^2, and H_A ≤ H_Ab.
    /// Hence |x_j numerator|, |x_j denominator| ≤ H_Ab.
    /// Returns H_Ab^2 (Int).
    inline Int hadamard_bound_squared(const Matrix<Int>& A, const std::vector<Int>& b) {
        const std::size_t R = A.rows();
        const std::size_t C = A.cols();
        Int prod(1);
        for (std::size_t i = 0; i < R; ++i) {
            Int row_norm_sq;
            for (std::size_t j = 0; j < C; ++j) {
                const Int& a = A(i, j);
                row_norm_sq = row_norm_sq + a * a;
            }
            Int with_b = row_norm_sq + b[i] * b[i];
            // if the factor is 0, use 1 as the base instead
            if (with_b.isZero()) with_b = Int(1);
            prod = prod * with_b;
        }
        return prod;
    }

    /// Return the smallest k satisfying p^k > 2 H^2 (where H_squared = H^2)
    inline std::size_t compute_lift_count(const Int& H_squared, int64_t p) {
        // required bit count = bitLength(2 H^2) = bitLength(H_squared) + 1
        // ↑ strictly, since H_squared is itself H^2, the bit count of 2 H^2 is bitLength(H_squared)+1
        // divide by the bit length of p and take the ceiling
        const std::size_t need_bits = H_squared.bitLength() + 2;  // +2: factor 2 + safety
        std::size_t log2_p = 0;
        int64_t pp = p;
        while (pp > 1) { pp >>= 1; ++log2_p; }
        return (need_bits + log2_p - 1) / log2_p;
    }

    // ====================================================================
    // Main: Dixon solver for an integer matrix and integer right-hand side
    // ====================================================================

    /// Solve for x ∈ ℚ^n given A: an n×n integer square matrix, b: an n-element integer vector.
    /// Throws std::runtime_error if A is singular (det ≡ 0 mod p for every candidate prime).
    inline std::vector<Rational> solve_int(const Matrix<Int>& A,
                                            const std::vector<Int>& b) {
        const std::size_t n = A.rows();
        if (A.cols() != n)
            throw std::invalid_argument("dixon::solve_int: matrix must be square");
        if (b.size() != n)
            throw std::invalid_argument("dixon::solve_int: rhs size mismatch");
        if (n == 0) return {};

        // 1. Prime selection + B = A^{-1} mod p
        int64_t p = 0;
        ModMat B_mod;
        for (int64_t cand : kCandidatePrimes) {
            ModMat A_mod = reduce_int_matrix(A, cand);
            ModMat Binv;
            if (modular_inverse(A_mod, cand, Binv)) {
                p = cand;
                B_mod = std::move(Binv);
                break;
            }
        }
        if (p == 0)
            throw std::runtime_error("dixon::solve_int: matrix appears singular (all candidate primes divide det)");

        // 2. Hadamard bound
        Int H_sq = hadamard_bound_squared(A, b);
        const std::size_t k = compute_lift_count(H_sq, p);

        // 3. p and p^k as Int
        Int p_int(p);

        // 4. Dixon iteration: r initialized to b, y = 0
        std::vector<Int> r = b;
        std::vector<Int> y(n);  // accumulator
        Int p_pow_i(1);  // p^i

        for (std::size_t i = 0; i < k; ++i) {
            // c_i = B · (r mod p) mod p
            ModVec r_mod = reduce_int_vec(r, p);
            ModVec c_mod = mod_matvec(B_mod, r_mod, p);

            // y += c_i · p^i  (integer addition)
            for (std::size_t j = 0; j < n; ++j) {
                Int cj(static_cast<int64_t>(c_mod[j]));
                y[j] = y[j] + cj * p_pow_i;
            }

            // r_{i+1} = (r - A · c_i) / p   (exact division)
            // A · c_i (integer matrix × int64 vector → Int vector)
            std::vector<Int> Ac(n);
            for (std::size_t row = 0; row < n; ++row) {
                Int acc;
                for (std::size_t col = 0; col < n; ++col) {
                    Int cj(static_cast<int64_t>(c_mod[col]));
                    acc = acc + A(row, col) * cj;
                }
                Ac[row] = acc;
            }
            for (std::size_t row = 0; row < n; ++row) {
                Int diff = r[row] - Ac[row];
                // diff is exactly divisible by p (Dixon's invariant)
                r[row] = diff / p_int;
            }

            // p_pow_i ← p_pow_i · p
            p_pow_i = p_pow_i * p_int;
        }

        // 5. Recover each y_j as a Rational
        // p^k = p_pow_i (at loop exit)
        const Int& p_pow_k = p_pow_i;
        // N = D = sqrt(H_sq) ≤ ceil bound. Rational reconstruction seeks |p| ≤ N, q ≤ D.
        // Since H_sq is itself H_Ab^2, even N = D = H_sq is safe (just overly generous).
        // More precisely N = D = ceil(sqrt(H_sq)), but absent an Int sqrt we use H_sq itself
        // (acceptable: it only widens the range over which reconstruction succeeds, without affecting correctness).
        Int N_bound = Int(1) << static_cast<int>((H_sq.bitLength() + 1) / 2 + 1);
        Int D_bound = N_bound;
        // p^k > 2 N D is guaranteed by compute_lift_count

        std::vector<Rational> x(n);
        for (std::size_t j = 0; j < n; ++j) {
            x[j] = reconstructRational(y[j], p_pow_k, N_bound, D_bound);
        }
        return x;
    }

}  // namespace sangi::detail::dixon

#endif  // SANGI_LINALG_DIXON_SOLVE_HPP