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

// CRT.hpp — Chinese Remainder Theorem
//
// Provides the extended_gcd and chinese_remainder_theorem functions and the
// CRT class. The CRT class operates efficiently when repeatedly transforming
// over the same set of moduli by precomputing the inverses.

#ifndef SANGI_CRT_HPP
#define SANGI_CRT_HPP

#include <math/core/common.hpp>

#include <vector>
#include <cstdint>
#include <concepts>
#include <stdexcept>

namespace sangi {

// ============================================================================
// Extended Euclidean algorithm
// ============================================================================

/**
 * @brief Extended Euclidean algorithm
 *
 * Finds s, t, and gcd(a, b) satisfying a*s + b*t = gcd(a, b).
 *
 * @param a the first integer
 * @param b the second integer
 * @param s output: s in a*s + b*t = gcd(a, b)
 * @param t output: t in a*s + b*t = gcd(a, b)
 * @return the greatest common divisor of a and b
 */
inline int64_t extended_gcd(int64_t a, int64_t b, int64_t& s, int64_t& t) {
    s = 1, t = 0;
    int64_t u = 0, v = 1;

    while (b != 0) {
        int64_t q = a / b;
        int64_t b_old = b;
        b = a - q * b;
        a = b_old;

        int64_t u_old = u;
        u = s - q * u;
        s = u_old;

        int64_t v_old = v;
        v = t - q * v;
        t = v_old;
    }

    return a;  // greatest common divisor
}

// ============================================================================
// Chinese Remainder Theorem (function version)
// ============================================================================

/**
 * @brief Conversion across multiple moduli using the Chinese Remainder Theorem
 * @param remainders array of residue values for each modulus
 * @param moduli array of moduli
 * @return the smallest non-negative integer satisfying all congruences
 */
template <std::integral T = int64_t>
T chinese_remainder_theorem(const std::vector<T>& remainders, const std::vector<int>& moduli) {
    if (remainders.size() != moduli.size() || remainders.empty()) {
        throw std::invalid_argument("Remainders and moduli arrays must have the same non-zero size");
    }

    T result = 0;
    T M = 1;

    // Compute the product of all moduli
    for (int mod : moduli) {
        M *= mod;
    }

    // Compute for each residue
    for (size_t i = 0; i < moduli.size(); i++) {
        T m_i = moduli[i];
        T M_i = M / m_i;

        // Compute the inverse using the extended Euclidean algorithm
        int64_t s, t;
        int64_t g = extended_gcd(static_cast<int64_t>(M_i), static_cast<int64_t>(m_i), s, t);

        // If the moduli are not pairwise coprime
        if (g != 1) {
            throw MathError("Moduli must be pairwise coprime in Chinese remainder theorem");
        }

        // Accumulate into the result
        result = (result + remainders[i] * M_i * s) % M;
    }

    // Normalize to a non-negative value
    if (result < 0) result += M;

    return result;
}

// ============================================================================
// CRT class (precomputed version)
// ============================================================================

/**
 * @brief Chinese Remainder Theorem class (CRT)
 *
 * Efficient when CRT is applied repeatedly to the same set of moduli.
 * The constructor precomputes the inverses and performs forward/inverse
 * conversion in O(n).
 *
 *
 * @tparam T integer type (int64_t, etc.)
 */
template<typename T = int64_t>
class CRT {
    size_t m_count;
    T m_modulus;                // Product of all moduli
    std::vector<T> m_moduli;   // Each modulus
    std::vector<T> m_inv;      // Precomputed weights (M_i * M_i^{-1} mod m_i)

public:
    /**
     * @brief Constructor — precomputes inverses
     * @param moduli array of moduli (must be pairwise coprime)
     */
    explicit CRT(const std::vector<int>& moduli)
        : m_count(moduli.size())
        , m_moduli(moduli.begin(), moduli.end())
        , m_inv(moduli.size())
    {
        if (moduli.empty()) {
            throw std::invalid_argument("CRT requires at least one modulus");
        }

        // Product of all moduli
        m_modulus = T(1);
        for (auto m : moduli) {
            m_modulus *= T(m);
        }

        // Precompute weights: m_inv[i] = M_i * (M_i^{-1} mod m_i)
        // where M_i = m_modulus / m_i
        for (size_t i = 0; i < m_count; ++i) {
            T M_i = m_modulus / T(moduli[i]);

            // Compute the inverse of (M_i mod m_i) via the extended GCD
            int64_t M_i_mod = static_cast<int64_t>(M_i % T(moduli[i]));
            int64_t s, t;
            int64_t g = extended_gcd(M_i_mod, static_cast<int64_t>(moduli[i]), s, t);
            if (g != 1) {
                throw MathError("Moduli must be pairwise coprime for CRT");
            }
            m_inv[i] = M_i * T(s);
        }
    }

    /** @brief Return the product of all moduli */
    const T& modulus() const { return m_modulus; }

    /** @brief Return the number of moduli */
    size_t count() const { return m_count; }

    /**
     * @brief Convert an integer to its residues (forward conversion)
     * @param x integer to convert
     * @return residue values for each modulus
     */
    std::vector<int> toResidues(const T& x) const {
        std::vector<int> residues(m_count);
        for (size_t i = 0; i < m_count; ++i) {
            residues[i] = static_cast<int>(x % m_moduli[i]);
        }
        return residues;
    }

    /**
     * @brief Recover an integer from its residues (inverse conversion)
     * @param residues residue values for each modulus
     * @return the recovered integer (0 <= result < modulus())
     */
    T fromResidues(const std::vector<int>& residues) const {
        if (residues.size() != m_count) {
            throw std::invalid_argument("Residues size must match moduli count");
        }

        T result = T(0);
        for (size_t i = 0; i < m_count; ++i) {
            result = (result + T(residues[i]) * m_inv[i]) % m_modulus;
        }

        // Normalize to a non-negative value
        if (result < T(0)) result += m_modulus;

        return result;
    }
};

} // namespace sangi

#endif // SANGI_CRT_HPP