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

// PadicNumber.hpp
// p進数テンプレートクラス (ALG-2)
//
// p進展開: 与えられた整数値を素数 p で展開し、p進数として保持する。
// p進付値 (valuation): p で割り切れる最大冪。
// 表示: "...d2d1d0" (最下位桁が右端)
//
// 用途: 数論、p進解析、Hensel の補題

#pragma once

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

namespace calx {

// ================================================================
// PadicNumber<T> クラス
// ================================================================
// p進数の有限精度表現。digits_[0] が最下位桁。
// valuation_ は p^v が元の値を割り切る最大の v。

template<typename T = int64_t>
class PadicNumber {
    int prime_;                  // 素数 p
    std::vector<int> digits_;   // p進展開の桁 (最下位桁が先頭)
    int valuation_;              // p進付値

public:
    // ================================================================
    // コンストラクタ
    // ================================================================

    /// デフォルトコンストラクタ (ゼロ)
    PadicNumber() : prime_(2), digits_{0}, valuation_(0) {}

    /// 素数 p と整数値から構築
    /// @param prime 素数 p (>= 2)
    /// @param value 展開する整数値
    /// @param precision p進展開の桁数
    PadicNumber(int prime, T value, int precision = 20)
        : prime_(prime), valuation_(0) {
        if (prime < 2) {
            throw std::invalid_argument("PadicNumber: prime must be >= 2");
        }

        if (value == T(0)) {
            digits_.assign(precision, 0);
            valuation_ = precision;  // 慣例: 0 の付値は無限大だが精度で制限
            return;
        }

        // 付値を計算: value を p で割れるだけ割る
        T v = value;
        bool negative = (v < T(0));
        if (negative) v = -v;

        while (v != T(0) && v % T(prime) == T(0)) {
            v /= T(prime);
            ++valuation_;
        }

        // p進展開
        // 負数の場合: p進数での補数表現
        // ...（p-1）（p-1）...（p-1）- |value| の p進展開
        if (negative) {
            // 負数の p進展開: p^precision - |value| の展開と等価
            // まず |value| / p^valuation_ の展開を求め、
            // p の補数を取る
            digits_.resize(precision, 0);
            T rem = v;
            int borrow = 0;
            for (int i = 0; i < precision; ++i) {
                int digit = static_cast<int>(rem % T(prime));
                rem /= T(prime);

                // 補数: (p - digit) with borrow
                int comp = prime - digit - borrow;
                if (i == 0 && digit != 0) {
                    // 最初の非零桁
                    comp = prime - digit;
                    borrow = 1;
                } else if (digit != 0 || borrow != 0) {
                    comp = prime - digit - borrow;
                    if (comp < 0) {
                        comp += prime;
                        borrow = 1;
                    } else {
                        borrow = (digit != 0 || borrow != 0) ? 1 : 0;
                    }
                } else {
                    comp = 0;
                }
                digits_[i] = ((comp % prime) + prime) % prime;
            }
            // 上位桁を (p-1) で埋める (負数の p進表現)
            for (int i = static_cast<int>(digits_.size()) - 1; i >= 0; --i) {
                if (rem == T(0) && borrow == 0) {
                    // 残りは (p-1) で埋める
                    for (int j = i + 1; j < precision; ++j) {
                        if (digits_[j] == 0) digits_[j] = prime - 1;
                    }
                    break;
                }
            }
        } else {
            // 正の数: 単純な p進展開
            digits_.reserve(precision);
            T rem = v;
            for (int i = 0; i < precision; ++i) {
                if (rem == T(0)) {
                    digits_.push_back(0);
                } else {
                    digits_.push_back(static_cast<int>(rem % T(prime)));
                    rem /= T(prime);
                }
            }
        }
    }

    // ================================================================
    // 算術演算
    // ================================================================

    /// p進加算
    PadicNumber operator+(const PadicNumber& rhs) const {
        checkCompatible(rhs);
        int prec = std::max(digits_.size(), rhs.digits_.size());
        PadicNumber result;
        result.prime_ = prime_;
        result.digits_.resize(prec, 0);
        result.valuation_ = 0;

        // 付値の差を考慮してアラインメント
        int shift_l = valuation_;
        int shift_r = rhs.valuation_;

        int carry = 0;
        for (int i = 0; i < prec; ++i) {
            int dl = getAlignedDigit(i, shift_l);
            int dr = rhs.getAlignedDigit(i, shift_r);
            int sum = dl + dr + carry;
            result.digits_[i] = sum % prime_;
            carry = sum / prime_;
        }

        // 付値を再計算
        result.computeValuation();
        return result;
    }

    /// p進減算
    PadicNumber operator-(const PadicNumber& rhs) const {
        checkCompatible(rhs);
        int prec = static_cast<int>(std::max(digits_.size(), rhs.digits_.size()));
        PadicNumber result;
        result.prime_ = prime_;
        result.digits_.resize(prec, 0);
        result.valuation_ = 0;

        int shift_l = valuation_;
        int shift_r = rhs.valuation_;

        int borrow = 0;
        for (int i = 0; i < prec; ++i) {
            int dl = getAlignedDigit(i, shift_l);
            int dr = rhs.getAlignedDigit(i, shift_r);
            int diff = dl - dr - borrow;
            if (diff < 0) {
                diff += prime_;
                borrow = 1;
            } else {
                borrow = 0;
            }
            result.digits_[i] = diff;
        }

        result.computeValuation();
        return result;
    }

    /// p進乗算
    PadicNumber operator*(const PadicNumber& rhs) const {
        checkCompatible(rhs);
        int prec = static_cast<int>(std::max(digits_.size(), rhs.digits_.size()));
        PadicNumber result;
        result.prime_ = prime_;
        result.digits_.resize(prec, 0);
        result.valuation_ = valuation_ + rhs.valuation_;

        // 筆算方式の乗算 (精度 prec 桁まで)
        std::vector<int> product(prec, 0);
        for (int i = 0; i < prec; ++i) {
            if (digits_[i] == 0) continue;
            int carry = 0;
            for (int j = 0; j < prec - i; ++j) {
                int idx = i + j;
                int val = product[idx] + digits_[i] * rhs.digit(j) + carry;
                product[idx] = val % prime_;
                carry = val / prime_;
            }
        }

        result.digits_ = std::move(product);
        result.computeValuation();
        return result;
    }

    // ================================================================
    // クエリ
    // ================================================================

    /// p進付値を返す
    int valuation() const { return valuation_; }

    /// 素数を返す
    int prime() const { return prime_; }

    /// 精度 (桁数) を返す
    int precision() const { return static_cast<int>(digits_.size()); }

    /// 整数値に変換 (有限桁の場合のみ正確)
    T toInteger() const {
        T result = T(0);
        T power = T(1);

        // valuation_ 分の p の冪を掛ける
        for (int i = 0; i < valuation_; ++i) {
            power *= T(prime_);
        }

        for (int i = 0; i < static_cast<int>(digits_.size()); ++i) {
            result += T(digits_[i]) * power;
            power *= T(prime_);
        }
        return result;
    }

    /// 文字列表現 "...d2d1d0" (最下位桁が右端)
    std::string toString() const {
        std::ostringstream oss;
        oss << "...";
        // 上位桁から順に出力
        bool leadingZero = true;
        for (int i = static_cast<int>(digits_.size()) - 1; i >= 0; --i) {
            if (leadingZero && digits_[i] == 0 && i > 0) continue;
            leadingZero = false;
            if (prime_ > 10) {
                if (i < static_cast<int>(digits_.size()) - 1 && !leadingZero) {
                    oss << '|';
                }
                oss << digits_[i];
            } else {
                oss << digits_[i];
            }
        }
        if (valuation_ > 0) {
            oss << " * " << prime_ << "^" << valuation_;
        }
        return oss.str();
    }

    /// i 番目の桁 (範囲外は 0)
    int digit(int i) const {
        if (i < 0 || i >= static_cast<int>(digits_.size())) return 0;
        return digits_[i];
    }

    /// 等価判定
    bool operator==(const PadicNumber& rhs) const {
        if (prime_ != rhs.prime_) return false;
        if (valuation_ != rhs.valuation_) return false;
        int prec = std::max(digits_.size(), rhs.digits_.size());
        for (int i = 0; i < prec; ++i) {
            if (digit(i) != rhs.digit(i)) return false;
        }
        return true;
    }

    bool operator!=(const PadicNumber& rhs) const {
        return !(*this == rhs);
    }

private:
    /// 互換性チェック
    void checkCompatible(const PadicNumber& rhs) const {
        if (prime_ != rhs.prime_) {
            throw std::invalid_argument("PadicNumber: mismatched primes");
        }
    }

    /// 付値でアラインした桁を取得
    int getAlignedDigit(int i, int shift) const {
        int idx = i - shift;
        if (idx < 0 || idx >= static_cast<int>(digits_.size())) return 0;
        return digits_[idx];
    }

    /// 付値を再計算
    void computeValuation() {
        valuation_ = 0;
        for (int i = 0; i < static_cast<int>(digits_.size()); ++i) {
            if (digits_[i] != 0) break;
            ++valuation_;
        }
    }
};

} // namespace calx
