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

// CRT.hpp — 中国剰余定理 (Chinese Remainder Theorem)
//
// extended_gcd, chinese_remainder_theorem 関数と CRT クラスを提供する。
// CRT クラスは同じモジュラス群に対して繰り返し変換を行う場合に、
// 逆元の事前計算により効率的に動作する。

#ifndef CALX_CRT_HPP
#define CALX_CRT_HPP

#include <math/core/common.hpp>

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

namespace calx {

// ============================================================================
// 拡張ユークリッドアルゴリズム
// ============================================================================

/**
 * @brief 拡張ユークリッドアルゴリズム
 *
 * a*s + b*t = gcd(a,b) を満たす s, t, gcd(a,b) を求める。
 *
 * @param a 第1の整数
 * @param b 第2の整数
 * @param s a*s + b*t = gcd(a,b)となるs
 * @param t a*s + b*t = gcd(a,b)となるt
 * @return a,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;  // 最大公約数
}

// ============================================================================
// 中国剰余定理 (関数版)
// ============================================================================

/**
 * @brief 中国剰余定理を使った複数のモジュラス間の変換
 * @param remainders 各モジュラスでの剰余値の配列
 * @param moduli モジュラスの配列
 * @return 全ての合同式を満たす最小の非負整数
 */
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;

    // 全モジュラスの積を計算
    for (int mod : moduli) {
        M *= mod;
    }

    // 各剰余に対して計算
    for (size_t i = 0; i < moduli.size(); i++) {
        T m_i = moduli[i];
        T M_i = M / m_i;

        // 拡張ユークリッドアルゴリズムで逆元を計算
        int64_t s, t;
        int64_t g = extended_gcd(static_cast<int64_t>(M_i), static_cast<int64_t>(m_i), s, t);

        // モジュラスが互いに素でない場合
        if (g != 1) {
            throw MathError("Moduli must be pairwise coprime in Chinese remainder theorem");
        }

        // 結果に足し合わせる
        result = (result + remainders[i] * M_i * s) % M;
    }

    // 正の値に正規化
    if (result < 0) result += M;

    return result;
}

// ============================================================================
// CRT クラス (事前計算版)
// ============================================================================

/**
 * @brief 中国剰余定理クラス (CRT)
 *
 * 同じモジュラス群に対して繰り返し CRT を適用する場合に効率的。
 * コンストラクタで逆元を事前計算し、変換/逆変換を O(n) で実行する。
 *
 * 移植元: lib++20 中国剰余定理.h (CxCRT)
 *
 * @tparam T 整数型 (int64_t 等)
 */
template<typename T = int64_t>
class CRT {
    size_t m_count;
    T m_modulus;                // 全モジュラスの積
    std::vector<T> m_moduli;   // 各モジュラス
    std::vector<T> m_inv;      // 事前計算された重み (M_i * M_i^{-1} mod m_i)

public:
    /**
     * @brief コンストラクタ — 逆元を事前計算
     * @param moduli モジュラスの配列 (互いに素でなければならない)
     */
    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");
        }

        // 全モジュラスの積
        m_modulus = T(1);
        for (auto m : moduli) {
            m_modulus *= T(m);
        }

        // 重みの事前計算: m_inv[i] = M_i * (M_i^{-1} mod m_i)
        // ここで M_i = m_modulus / m_i
        for (size_t i = 0; i < m_count; ++i) {
            T M_i = m_modulus / T(moduli[i]);

            // M_i mod m_i の逆元を拡張 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 全モジュラスの積を返す */
    const T& modulus() const { return m_modulus; }

    /** @brief モジュラスの個数を返す */
    size_t count() const { return m_count; }

    /**
     * @brief 整数を剰余系に変換 (順変換)
     * @param x 変換する整数
     * @return 各モジュラスでの剰余値
     */
    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 剰余系から整数に復元 (逆変換)
     * @param residues 各モジュラスでの剰余値
     * @return 復元された整数 (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;
        }

        // 正の値に正規化
        if (result < T(0)) result += m_modulus;

        return result;
    }
};

} // namespace calx

#endif // CALX_CRT_HPP