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

// Padic.hpp
//
// p-進数 (p-adic Numbers)
//
// p-進展開、四則演算、指数関数、対数関数。
// 有限精度 (N 桁) の p-進整数 / p-進有理数を扱う。
//
// 主な API:
//   Padic(p, prec)               — ゼロ p-進数
//   Padic::fromInteger(n, p, N)  — 整数 → p-進数
//   Padic::fromRational(a, b, p, N) — 有理数 a/b → p-進数
//   operator+, -, *, /           — 四則演算
//   padicExp(x)                  — p-進指数関数
//   padicLog(x)                  — p-進対数関数
//   padicNorm(x)                 — p-進ノルム |x|_p
//   padicValuation(x)            — p-進付値 v_p(x)
//   teichmullerLift(a, p, N)     — Teichmüller 持ち上げ

#ifndef CALX_CORE_PADIC_HPP
#define CALX_CORE_PADIC_HPP

#include <vector>
#include <cstdint>
#include <stdexcept>
#include <cmath>
#include <algorithm>
#include <string>
#include <sstream>
#include <numeric>

namespace calx {

/**
 * @brief p-進数クラス (有限精度)
 *
 * 内部表現: digits[i] は p^(valuation + i) の係数 (0 ≤ digits[i] < p)
 * valuation は p-進付値 (最低次の p のべき)
 * 精度 N: digits の長さ (p^N まで)
 */
class Padic {
public:
    /// ゼロ p-進数
    Padic(int64_t p, int64_t prec)
        : p_(p), prec_(prec), valuation_(prec), digits_(prec, 0) {
        if (p < 2) throw std::invalid_argument("Padic: p must be >= 2");
        if (prec < 1) throw std::invalid_argument("Padic: precision must be >= 1");
    }

    /// 整数から構築
    static Padic fromInteger(int64_t n, int64_t p, int64_t prec) {
        Padic result(p, prec);
        if (n == 0) return result;

        // 負の場合: p-進数では -n = (p^N - n) mod p^N として表現
        bool neg = (n < 0);
        // p-進展開のため、uint で処理
        // 負数は補数表現: -n ≡ p^N - n (mod p^N)
        int64_t val = 0;

        if (neg) {
            // -|n| の p-進展開
            // まず |n| を展開してから補数を取る
            result = fromInteger(-n, p, prec);
            result = negate(result);
            return result;
        }

        // 正の整数の p-進展開
        int64_t v = n;
        // まず付値を求める
        while (v != 0 && v % p == 0) {
            v /= p;
            ++val;
        }

        result.valuation_ = val;
        std::fill(result.digits_.begin(), result.digits_.end(), 0);

        int64_t remaining = v;
        for (int64_t i = 0; i < prec && remaining != 0; ++i) {
            result.digits_[static_cast<std::size_t>(i)] =
                static_cast<int64_t>(((remaining % p) + p) % p);
            remaining = (remaining - result.digits_[static_cast<std::size_t>(i)]) / p;
        }

        return result;
    }

    /// 有理数 a/b から構築
    static Padic fromRational(int64_t a, int64_t b, int64_t p, int64_t prec) {
        if (b == 0) throw std::invalid_argument("Padic::fromRational: denominator is zero");
        if (a == 0) return Padic(p, prec);

        // a/b の p-進付値
        int64_t va = 0, vb = 0;
        int64_t aa = std::abs(a), bb = std::abs(b);
        while (aa % p == 0) { aa /= p; ++va; }
        while (bb % p == 0) { bb /= p; ++vb; }

        int64_t val = va - vb;

        // b の逆元を mod p^N で求めて a と掛ける
        // b' = bb (p を除いた部分) は p と互いに素
        Padic result(p, prec);
        result.valuation_ = val;

        // bb^{-1} mod p^prec を求める (拡張ユークリッド)
        // p^prec を計算 (オーバーフロー回避のため段階的に)
        // 代わりに digit ごとに Hensel lifting で計算

        // a' / b' の p-進展開 (a', b' は p を除いた部分)
        bool negative = ((a < 0) != (b < 0));

        // digit ごとに計算: |a'| / |b'| の p-進展開
        int64_t bbInvModP = modInverse(bb % p, p);
        if (bbInvModP < 0) bbInvModP += p;

        // digit by digit division (正の値で計算)
        int64_t carry = static_cast<int64_t>(aa);
        for (int64_t i = 0; i < prec; ++i) {
            int64_t cmod = ((carry % p) + p) % p;
            int64_t digit = (cmod * bbInvModP) % p;
            if (digit < 0) digit += p;
            result.digits_[static_cast<std::size_t>(i)] = digit;
            carry = (carry - digit * bb) / p;
        }

        if (negative) {
            result = negate(result);
        }

        return result;
    }

    // アクセサ
    int64_t prime() const { return p_; }
    int64_t precision() const { return prec_; }
    int64_t valuation() const { return valuation_; }
    const std::vector<int64_t>& digits() const { return digits_; }

    /// p-進数がゼロかどうか
    bool isZero() const {
        return std::all_of(digits_.begin(), digits_.end(),
                          [](int64_t d) { return d == 0; });
    }

    /// 整数に変換 (可能な場合)
    int64_t toInteger() const {
        int64_t result = 0;
        int64_t pk = 1;
        // valuation から始まる
        for (int64_t i = 0; i < static_cast<int64_t>(digits_.size()); ++i) {
            if (valuation_ + i >= 0) {
                // p^(valuation+i) を計算
                int64_t power = 1;
                for (int64_t j = 0; j < valuation_ + i; ++j) power *= p_;
                result += digits_[static_cast<std::size_t>(i)] * power;
            }
        }
        return result;
    }

    /// 文字列表現 (... d_2 d_1 d_0 . d_{-1} d_{-2} ... の形)
    std::string toString() const {
        std::ostringstream oss;
        // digits を上位桁から表示
        bool started = false;
        for (int64_t i = static_cast<int64_t>(digits_.size()) - 1; i >= 0; --i) {
            if (digits_[static_cast<std::size_t>(i)] != 0) started = true;
            if (started) {
                if (oss.str().size() > 0 && i < static_cast<int64_t>(digits_.size()) - 1)
                    oss << " ";
                oss << digits_[static_cast<std::size_t>(i)];
            }
        }
        if (!started) oss << "0";

        if (valuation_ != 0) {
            oss << " * " << p_ << "^" << valuation_;
        }
        return oss.str();
    }

    // 四則演算
    friend Padic operator+(const Padic& a, const Padic& b) {
        checkSamePrime(a, b);
        int64_t p = a.p_;
        int64_t prec = std::min(a.prec_, b.prec_);
        int64_t minVal = std::min(a.valuation_, b.valuation_);

        Padic result(p, prec);
        result.valuation_ = minVal;

        int64_t carry = 0;
        for (int64_t i = 0; i < prec; ++i) {
            int64_t da = a.digitAt(minVal + i);
            int64_t db = b.digitAt(minVal + i);
            int64_t sum = da + db + carry;
            result.digits_[static_cast<std::size_t>(i)] = ((sum % p) + p) % p;
            carry = (sum - result.digits_[static_cast<std::size_t>(i)]) / p;
        }

        result.normalize();
        return result;
    }

    friend Padic operator-(const Padic& a, const Padic& b) {
        return a + negate(b);
    }

    friend Padic operator*(const Padic& a, const Padic& b) {
        checkSamePrime(a, b);
        int64_t p = a.p_;
        int64_t prec = std::min(a.prec_, b.prec_);

        Padic result(p, prec);
        result.valuation_ = a.valuation_ + b.valuation_;

        // 畳み込み
        for (int64_t i = 0; i < prec; ++i) {
            int64_t sum = 0;
            for (int64_t j = 0; j <= i; ++j) {
                int64_t k = i - j;
                if (j < static_cast<int64_t>(a.digits_.size()) &&
                    k < static_cast<int64_t>(b.digits_.size())) {
                    sum += a.digits_[static_cast<std::size_t>(j)] *
                           b.digits_[static_cast<std::size_t>(k)];
                }
            }
            result.digits_[static_cast<std::size_t>(i)] = sum;
        }

        // キャリー処理
        int64_t carry = 0;
        for (int64_t i = 0; i < prec; ++i) {
            int64_t v = result.digits_[static_cast<std::size_t>(i)] + carry;
            result.digits_[static_cast<std::size_t>(i)] = ((v % p) + p) % p;
            carry = (v - result.digits_[static_cast<std::size_t>(i)]) / p;
        }

        result.normalize();
        return result;
    }

    friend Padic operator/(const Padic& a, const Padic& b) {
        checkSamePrime(a, b);
        if (b.isZero()) throw std::invalid_argument("Padic: division by zero");

        int64_t p = a.p_;
        int64_t prec = std::min(a.prec_, b.prec_);

        Padic result(p, prec);
        result.valuation_ = a.valuation_ - b.valuation_;

        // b の最初の非ゼロ桁の逆元
        int64_t b0Inv = modInverse(b.digits_[0], p);

        // 長除法
        std::vector<int64_t> rem(a.digits_.begin(), a.digits_.end());
        rem.resize(static_cast<std::size_t>(prec), 0);

        for (int64_t i = 0; i < prec; ++i) {
            int64_t qi = (((rem[static_cast<std::size_t>(i)] % p) + p) % p * b0Inv) % p;
            if (qi < 0) qi += p;
            result.digits_[static_cast<std::size_t>(i)] = qi;

            // rem -= qi * b.digits (shifted by i)
            int64_t carry = 0;
            for (int64_t j = 0; j < static_cast<int64_t>(b.digits_.size()) && i + j < prec; ++j) {
                rem[static_cast<std::size_t>(i + j)] -= qi * b.digits_[static_cast<std::size_t>(j)] + carry;
                // normalize
                int64_t v = rem[static_cast<std::size_t>(i + j)];
                int64_t normalized = ((v % p) + p) % p;
                carry = (v - normalized) / p;
                // carry は上位桁で処理するので rem にキャリーを伝播
                rem[static_cast<std::size_t>(i + j)] = normalized;
                // 実際は carry を次の桁に足す
            }
            // 残りのキャリーを伝播
            for (int64_t j = static_cast<int64_t>(b.digits_.size()); i + j < prec && carry != 0; ++j) {
                rem[static_cast<std::size_t>(i + j)] -= carry;
                int64_t v = rem[static_cast<std::size_t>(i + j)];
                int64_t normalized = ((v % p) + p) % p;
                carry = (v - normalized) / p;
                rem[static_cast<std::size_t>(i + j)] = normalized;
            }
        }

        return result;
    }

    friend bool operator==(const Padic& a, const Padic& b) {
        checkSamePrime(a, b);
        if (a.isZero() && b.isZero()) return true;
        if (a.valuation_ != b.valuation_) return false;
        int64_t n = std::min(static_cast<int64_t>(a.digits_.size()),
                             static_cast<int64_t>(b.digits_.size()));
        for (int64_t i = 0; i < n; ++i) {
            if (a.digits_[static_cast<std::size_t>(i)] != b.digits_[static_cast<std::size_t>(i)])
                return false;
        }
        return true;
    }

    friend bool operator!=(const Padic& a, const Padic& b) {
        return !(a == b);
    }

private:
    int64_t p_;       // 素数
    int64_t prec_;    // 精度 (桁数)
    int64_t valuation_; // p-進付値
    std::vector<int64_t> digits_; // 係数 (digits[i] = p^(valuation+i) の係数)

    /// i 番目の絶対位置の桁を取得
    int64_t digitAt(int64_t absPos) const {
        int64_t idx = absPos - valuation_;
        if (idx < 0 || idx >= static_cast<int64_t>(digits_.size())) return 0;
        return digits_[static_cast<std::size_t>(idx)];
    }

    /// 先頭のゼロを除去して valuation を調整
    void normalize() {
        while (!digits_.empty() && digits_.front() == 0) {
            digits_.erase(digits_.begin());
            ++valuation_;
        }
        if (digits_.empty()) {
            valuation_ = prec_;
        }
        // 末尾のゼロも除去 (精度範囲外)
        while (digits_.size() > 1 && digits_.back() == 0) {
            digits_.pop_back();
        }
        // 精度に合わせてサイズを制限
        if (static_cast<int64_t>(digits_.size()) > prec_) {
            digits_.resize(static_cast<std::size_t>(prec_));
        }
    }

    static void checkSamePrime(const Padic& a, const Padic& b) {
        if (a.p_ != b.p_)
            throw std::invalid_argument("Padic: different primes");
    }

    /// p-進数の符号反転 (補数)
    static Padic negate(const Padic& x) {
        if (x.isZero()) return x;
        int64_t p = x.p_;
        int64_t prec = x.prec_;

        Padic result(p, prec);
        result.valuation_ = x.valuation_;

        // p の補数: (p - 1 - d_i) に 1 を加算
        int64_t carry = 1;
        for (std::size_t i = 0; i < x.digits_.size(); ++i) {
            int64_t v = (p - 1 - x.digits_[i]) + carry;
            result.digits_[i] = v % p;
            carry = v / p;
        }
        // 残りの桁は p-1 (= ...1111 in p-adic)
        for (std::size_t i = x.digits_.size(); i < static_cast<std::size_t>(prec); ++i) {
            int64_t v = (p - 1) + carry;
            result.digits_[i] = v % p;
            carry = v / p;
        }

        return result;
    }

    /// mod p での逆元 (拡張ユークリッド)
    static int64_t modInverse(int64_t a, int64_t m) {
        a = ((a % m) + m) % m;
        if (a == 0) throw std::invalid_argument("Padic: no inverse (a=0 mod p)");
        int64_t g, x, y;
        extgcd(a, m, g, x, y);
        if (g != 1) throw std::invalid_argument("Padic: no inverse");
        return ((x % m) + m) % m;
    }

    static void extgcd(int64_t a, int64_t b, int64_t& g, int64_t& x, int64_t& y) {
        if (a == 0) { g = b; x = 0; y = 1; return; }
        int64_t g1, x1, y1;
        extgcd(b % a, a, g1, x1, y1);
        g = g1;
        x = y1 - (b / a) * x1;
        y = x1;
    }
};

// =====================================================================
// p-進関数
// =====================================================================

/**
 * @brief p-進付値 v_p(x)
 */
[[nodiscard]]
inline int64_t padicValuation(const Padic& x) {
    return x.valuation();
}

/**
 * @brief p-進ノルム |x|_p = p^{-v_p(x)}
 */
[[nodiscard]]
inline double padicNorm(const Padic& x) {
    if (x.isZero()) return 0.0;
    return std::pow(static_cast<double>(x.prime()),
                    -static_cast<double>(x.valuation()));
}

/**
 * @brief p-進指数関数 exp(x) = Σ x^n / n!
 *
 * 収束条件: v_p(x) ≥ 1 (p ≥ 3) または v_p(x) ≥ 2 (p = 2)
 * つまり x ∈ pZ_p (p≥3) または x ∈ 4Z_2 (p=2)
 */
[[nodiscard]]
inline Padic padicExp(const Padic& x) {
    int64_t p = x.prime();
    int64_t prec = x.precision();

    // 収束条件チェック
    int64_t minVal = (p == 2) ? 2 : 1;
    if (!x.isZero() && x.valuation() < minVal) {
        throw std::domain_error("padicExp: x not in convergence domain");
    }

    if (x.isZero()) return Padic::fromInteger(1, p, prec);

    // exp(x) = Σ x^n / n!
    Padic result = Padic::fromInteger(1, p, prec);
    Padic term = Padic::fromInteger(1, p, prec);  // x^n / n!

    for (int64_t n = 1; n < prec * 2; ++n) {
        term = term * x;
        // n! で割る: term / n
        Padic nPadic = Padic::fromInteger(n, p, prec);
        term = term / nPadic;
        result = result + term;
    }

    return result;
}

/**
 * @brief p-進対数関数 log(1 + x) = Σ (-1)^{n+1} x^n / n
 *
 * 収束条件: v_p(x) ≥ 1
 */
[[nodiscard]]
inline Padic padicLog(const Padic& x) {
    int64_t p = x.prime();
    int64_t prec = x.precision();

    // x = 1 + t, t の付値 ≥ 1 が必要
    // 入力は 1+t ではなく t として受け取る仕様にする
    if (!x.isZero() && x.valuation() < 1) {
        throw std::domain_error("padicLog: argument not in convergence domain (v_p(x) >= 1 required)");
    }

    if (x.isZero()) return Padic(p, prec);

    // log(1+x) = x - x²/2 + x³/3 - ...
    Padic result(p, prec);
    Padic power = x;  // x^n

    for (int64_t n = 1; n < prec * 2; ++n) {
        Padic nPadic = Padic::fromInteger(n, p, prec);
        Padic term = power / nPadic;
        if (n % 2 == 1) {
            result = result + term;
        } else {
            result = result - term;
        }
        power = power * x;
    }

    return result;
}

/**
 * @brief Teichmüller 持ち上げ
 *
 * F_p の元 a (0 ≤ a < p) を Z_p に持ち上げる。
 * ω(a)^p = ω(a), ω(a) ≡ a (mod p)
 *
 * Newton 法: x_{n+1} = x_n - (x_n^p - x_n) / (p·x_n^{p-1} - 1)
 * 実際には: x_{n+1} = x_n^p (mod p^{2^n}) が二次収束
 */
[[nodiscard]]
inline Padic teichmullerLift(int64_t a, int64_t p, int64_t prec) {
    if (a < 0 || a >= p)
        throw std::invalid_argument("teichmullerLift: a must be in [0, p)");
    if (a == 0) return Padic(p, prec);

    // x^p mod p^N を繰り返す
    Padic x = Padic::fromInteger(a, p, prec);
    for (int64_t iter = 0; iter < prec + 2; ++iter) {
        // x = x^p
        Padic xp = x;
        for (int64_t i = 1; i < p; ++i) {
            xp = xp * x;
        }
        x = xp;
    }

    return x;
}

/**
 * @brief 整数の p-進付値
 */
[[nodiscard]]
inline int64_t padicValuation(int64_t n, int64_t p) {
    if (n == 0) return -1;  // 慣習: v_p(0) = ∞ → -1 で表現
    n = std::abs(n);
    int64_t v = 0;
    while (n % p == 0) { n /= p; ++v; }
    return v;
}

/**
 * @brief 整数の p-進ノルム
 */
[[nodiscard]]
inline double padicNorm(int64_t n, int64_t p) {
    if (n == 0) return 0.0;
    int64_t v = padicValuation(n, p);
    return std::pow(static_cast<double>(p), -static_cast<double>(v));
}

} // namespace calx

#endif // CALX_CORE_PADIC_HPP
