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

// PowerSeries.hpp
// 形式冪級数 (formal power series) テンプレートクラス
//
// f(x) = Σ_{n=0}^{N-1} aₙ xⁿ  (mod x^N で打ち切り)
//
// FLINT / NTL 参考、Newton 反復ベースの高速演算:
//   - 加算・減算: O(N)
//   - 乗算: O(N²) 直接畳み込み (小 N)
//   - 逆数: Newton 反復 O(M(N))  — g_{k+1} = g_k(2 - f·g_k)
//   - 対数: log f = ∫ f'/f dx  — O(M(N))
//   - 指数: Newton 反復 O(M(N))  — g_{k+1} = g_k(1 + f - log g_k)
//   - 平方根: Newton 反復 O(M(N)) — g_{k+1} = (g_k + f/g_k)/2
//   - 合成: f(g(x))  — Brent-Kung O(√N · M(N))
//   - 逆合成: compositional inverse — Newton 反復
//   - 微分・積分
//
// 設計:
//   - T = double, float, Int, Float, Rational, Complex<T> 等で使用可能
//   - Polynomial<T> との相互変換
//   - order (打ち切り次数) は各演算で保持・伝播

#pragma once

#include <vector>
#include <cmath>
#include <concepts>
#include <algorithm>
#include <initializer_list>
#include <cassert>
#include <iosfwd>
#include <sstream>
#include <string>

namespace calx {

// ================================================================
// コンセプト: 冪級数の係数型
// ================================================================

template<typename T>
concept PowerSeriesCoeff = requires(T a, T b, int n) {
    { a + b } -> std::convertible_to<T>;
    { a - b } -> std::convertible_to<T>;
    { a * b } -> std::convertible_to<T>;
    { -a } -> std::convertible_to<T>;
    { T(0) };
    { T(1) };
};

// 除算が必要な演算 (inverse, log, exp, sqrt) 用
template<typename T>
concept PowerSeriesDivisible = PowerSeriesCoeff<T> && requires(T a, T b) {
    { a / b } -> std::convertible_to<T>;
};

// ================================================================
// PowerSeries<T> クラス
// ================================================================

template<PowerSeriesCoeff T>
class PowerSeries {
private:
    std::vector<T> coeffs_;   // coeffs_[i] = x^i の係数
    size_t order_;            // 打ち切り次数 N: mod x^N で計算

    // N 項に切り詰め
    void truncate() {
        if (coeffs_.size() > order_)
            coeffs_.resize(order_);
    }

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

    // 零級数 (order=1)
    PowerSeries() : order_(1) {}

    // 零級数 (指定 order)
    explicit PowerSeries(size_t order) : order_(order) {}

    // 定数級数
    PowerSeries(const T& c, size_t order) : order_(order) {
        if (!(c == T(0))) coeffs_ = {c};
    }

    // 係数ベクトルから
    PowerSeries(std::vector<T> coeffs, size_t order)
        : coeffs_(std::move(coeffs)), order_(order) {
        truncate();
    }

    // 初期化子リストから
    // 例: {{1, 2, 3}, 5} → (1 + 2x + 3x²) mod x⁵
    PowerSeries(std::initializer_list<T> il, size_t order)
        : order_(order) {
        coeffs_.assign(il.begin(), il.end());
        truncate();
    }

    // ============================================================
    // ファクトリ
    // ============================================================

    // 単位級数 1
    [[nodiscard]] static PowerSeries one(size_t order) {
        return PowerSeries(T(1), order);
    }

    // x (生成元)
    [[nodiscard]] static PowerSeries x(size_t order) {
        assert(order >= 2);
        PowerSeries r(order);
        r.coeffs_.resize(2, T(0));
        r.coeffs_[1] = T(1);
        return r;
    }

    // x^k
    [[nodiscard]] static PowerSeries xPow(size_t k, size_t order) {
        if (k >= order) return PowerSeries(order);
        PowerSeries r(order);
        r.coeffs_.resize(k + 1, T(0));
        r.coeffs_[k] = T(1);
        return r;
    }

    // ============================================================
    // アクセサ
    // ============================================================

    [[nodiscard]] size_t order() const { return order_; }

    [[nodiscard]] bool isZero() const { return coeffs_.empty(); }

    // 係数アクセス (範囲外は T(0))
    [[nodiscard]] const T& operator[](size_t i) const {
        static const T zero = T(0);
        return (i < coeffs_.size()) ? coeffs_[i] : zero;
    }

    // 係数への参照アクセス
    T& operator[](size_t i) {
        if (i >= coeffs_.size())
            coeffs_.resize(std::min(i + 1, order_), T(0));
        return coeffs_[i];
    }

    [[nodiscard]] const std::vector<T>& coefficients() const { return coeffs_; }

    // 定数項
    [[nodiscard]] T constantTerm() const {
        return coeffs_.empty() ? T(0) : coeffs_[0];
    }

    // order を変更 (拡大または切り詰め)
    PowerSeries& setOrder(size_t newOrder) {
        order_ = newOrder;
        truncate();
        return *this;
    }

    // ============================================================
    // 比較演算子 (同 order の係数を比較)
    // ============================================================

    [[nodiscard]] bool operator==(const PowerSeries& rhs) const {
        size_t n = std::max(coeffs_.size(), rhs.coeffs_.size());
        for (size_t i = 0; i < n; ++i)
            if (!((*this)[i] == rhs[i])) return false;
        return true;
    }
    [[nodiscard]] bool operator!=(const PowerSeries& rhs) const {
        return !(*this == rhs);
    }

    // ============================================================
    // 単項演算子
    // ============================================================

    [[nodiscard]] PowerSeries operator+() const { return *this; }

    [[nodiscard]] PowerSeries operator-() const {
        PowerSeries result(order_);
        result.coeffs_.resize(coeffs_.size());
        for (size_t i = 0; i < coeffs_.size(); ++i)
            result.coeffs_[i] = -coeffs_[i];
        return result;
    }

    // ============================================================
    // 加算・減算: O(N)
    // ============================================================

    [[nodiscard]] PowerSeries operator+(const PowerSeries& rhs) const {
        size_t ord = std::min(order_, rhs.order_);
        size_t n = std::min(std::max(coeffs_.size(), rhs.coeffs_.size()), ord);
        PowerSeries result(ord);
        result.coeffs_.resize(n, T(0));
        for (size_t i = 0; i < n; ++i)
            result.coeffs_[i] = (*this)[i] + rhs[i];
        // 末尾のゼロを除去
        while (!result.coeffs_.empty() && result.coeffs_.back() == T(0))
            result.coeffs_.pop_back();
        return result;
    }

    [[nodiscard]] PowerSeries operator-(const PowerSeries& rhs) const {
        size_t ord = std::min(order_, rhs.order_);
        size_t n = std::min(std::max(coeffs_.size(), rhs.coeffs_.size()), ord);
        PowerSeries result(ord);
        result.coeffs_.resize(n, T(0));
        for (size_t i = 0; i < n; ++i)
            result.coeffs_[i] = (*this)[i] - rhs[i];
        while (!result.coeffs_.empty() && result.coeffs_.back() == T(0))
            result.coeffs_.pop_back();
        return result;
    }

    PowerSeries& operator+=(const PowerSeries& rhs) { *this = *this + rhs; return *this; }
    PowerSeries& operator-=(const PowerSeries& rhs) { *this = *this - rhs; return *this; }

    // スカラー演算
    [[nodiscard]] PowerSeries operator+(const T& c) const {
        PowerSeries result = *this;
        if (result.coeffs_.empty()) result.coeffs_.resize(1, T(0));
        result.coeffs_[0] = result.coeffs_[0] + c;
        return result;
    }
    [[nodiscard]] PowerSeries operator-(const T& c) const {
        return *this + (-c);
    }
    [[nodiscard]] PowerSeries operator*(const T& c) const {
        PowerSeries result(order_);
        result.coeffs_.resize(coeffs_.size());
        for (size_t i = 0; i < coeffs_.size(); ++i)
            result.coeffs_[i] = coeffs_[i] * c;
        return result;
    }
    PowerSeries& operator*=(const T& c) { *this = *this * c; return *this; }

    // ============================================================
    // 乗算: O(N²) 直接畳み込み mod x^N
    // ============================================================

    [[nodiscard]] PowerSeries operator*(const PowerSeries& rhs) const {
        if (coeffs_.empty() || rhs.coeffs_.empty())
            return PowerSeries(std::min(order_, rhs.order_));
        size_t ord = std::min(order_, rhs.order_);
        size_t n = std::min(coeffs_.size() + rhs.coeffs_.size() - 1, ord);
        PowerSeries result(ord);
        result.coeffs_.resize(n, T(0));
        for (size_t i = 0; i < coeffs_.size(); ++i) {
            if (coeffs_[i] == T(0)) continue;
            size_t jmax = std::min(rhs.coeffs_.size(), n - i);
            for (size_t j = 0; j < jmax; ++j)
                result.coeffs_[i + j] = result.coeffs_[i + j] + coeffs_[i] * rhs.coeffs_[j];
        }
        while (!result.coeffs_.empty() && result.coeffs_.back() == T(0))
            result.coeffs_.pop_back();
        return result;
    }

    PowerSeries& operator*=(const PowerSeries& rhs) { *this = *this * rhs; return *this; }

    // ============================================================
    // x による乗除 (シフト)
    // ============================================================

    // f(x) * x^k mod x^N
    [[nodiscard]] PowerSeries shift(int k) const {
        if (k == 0) return *this;
        PowerSeries result(order_);
        if (k > 0) {
            size_t uk = static_cast<size_t>(k);
            if (uk >= order_) return result;
            size_t n = std::min(coeffs_.size() + uk, order_);
            result.coeffs_.resize(n, T(0));
            for (size_t i = 0; i < coeffs_.size() && i + uk < order_; ++i)
                result.coeffs_[i + uk] = coeffs_[i];
        } else {
            size_t uk = static_cast<size_t>(-k);
            if (uk >= coeffs_.size()) return result;
            result.coeffs_.assign(coeffs_.begin() + uk, coeffs_.end());
        }
        while (!result.coeffs_.empty() && result.coeffs_.back() == T(0))
            result.coeffs_.pop_back();
        return result;
    }

    // ============================================================
    // 微分: f'(x) mod x^{N-1}
    // ============================================================

    [[nodiscard]] PowerSeries derivative() const {
        if (coeffs_.size() <= 1) return PowerSeries(order_ > 0 ? order_ - 1 : 0);
        PowerSeries result(order_ > 0 ? order_ - 1 : 0);
        result.coeffs_.resize(coeffs_.size() - 1);
        for (size_t i = 1; i < coeffs_.size(); ++i)
            result.coeffs_[i - 1] = coeffs_[i] * T(static_cast<int>(i));
        while (!result.coeffs_.empty() && result.coeffs_.back() == T(0))
            result.coeffs_.pop_back();
        return result;
    }

    // ============================================================
    // 積分: ∫ f(x) dx mod x^{N+1} (定数項 = 0)
    // ============================================================

    [[nodiscard]] PowerSeries integral() const
        requires PowerSeriesDivisible<T>
    {
        PowerSeries result(order_ + 1);
        result.coeffs_.resize(coeffs_.size() + 1, T(0));
        // coeffs_[0] → 積分の x^1 項: coeffs_[0] / 1
        for (size_t i = 0; i < coeffs_.size(); ++i)
            result.coeffs_[i + 1] = coeffs_[i] / T(static_cast<int>(i + 1));
        result.truncate();
        while (!result.coeffs_.empty() && result.coeffs_.back() == T(0))
            result.coeffs_.pop_back();
        return result;
    }

    // ============================================================
    // 逆数: 1/f(x) mod x^N  — Newton 反復
    // f(0) ≠ 0 が必要
    // g_{k+1} = g_k · (2 - f · g_k)
    // ============================================================

    [[nodiscard]] PowerSeries inverse() const
        requires PowerSeriesDivisible<T>
    {
        assert(!coeffs_.empty() && !(coeffs_[0] == T(0)) &&
               "inverse requires f(0) != 0");

        // 初期近似: g_0 = 1/f(0)
        T inv0 = T(1) / coeffs_[0];
        if (order_ <= 1) return PowerSeries(inv0, order_);

        // Newton 反復: 精度を倍増しながら
        size_t prec = 1;
        PowerSeries g(inv0, order_);

        while (prec < order_) {
            prec = std::min(prec * 2, order_);
            // f と g を prec 項で計算
            PowerSeries f_trunc = *this;
            f_trunc.setOrder(prec);
            g.setOrder(prec);
            // g = g * (2 - f * g)
            PowerSeries fg = f_trunc * g;
            PowerSeries two_minus_fg = PowerSeries(T(1) + T(1), prec) - fg;
            g = g * two_minus_fg;
            g.setOrder(prec);
        }
        g.setOrder(order_);
        return g;
    }

    // ============================================================
    // 除算: f / g mod x^N
    // ============================================================

    [[nodiscard]] PowerSeries operator/(const PowerSeries& rhs) const
        requires PowerSeriesDivisible<T>
    {
        return *this * rhs.inverse();
    }

    // スカラー除算
    [[nodiscard]] PowerSeries operator/(const T& c) const
        requires PowerSeriesDivisible<T>
    {
        PowerSeries result(order_);
        result.coeffs_.resize(coeffs_.size());
        for (size_t i = 0; i < coeffs_.size(); ++i)
            result.coeffs_[i] = coeffs_[i] / c;
        return result;
    }

    // ============================================================
    // 対数: log f(x) mod x^N
    // f(0) = 1 が必要
    // log f = ∫ f'/f dx
    // ============================================================

    [[nodiscard]] PowerSeries log() const
        requires PowerSeriesDivisible<T>
    {
        assert(!coeffs_.empty() && coeffs_[0] == T(1) &&
               "log requires f(0) = 1");
        if (order_ <= 1) return PowerSeries(order_);

        // log f = ∫ (f'/f) dx
        PowerSeries fprime = this->derivative();
        fprime.setOrder(order_);
        PowerSeries finv = this->inverse();
        PowerSeries quot = fprime * finv;
        quot.setOrder(order_);
        return quot.integral();
    }

    // ============================================================
    // 指数: exp f(x) mod x^N
    // f(0) = 0 が必要
    // Newton 反復: g_{k+1} = g_k · (1 + f - log g_k)
    // ============================================================

    [[nodiscard]] PowerSeries exp() const
        requires PowerSeriesDivisible<T>
    {
        assert((coeffs_.empty() || coeffs_[0] == T(0)) &&
               "exp requires f(0) = 0");
        if (order_ <= 1) return PowerSeries::one(order_);

        // 初期近似: g_0 = 1
        size_t prec = 1;
        PowerSeries g = PowerSeries::one(order_);

        while (prec < order_) {
            prec = std::min(prec * 2, order_);
            g.setOrder(prec);
            PowerSeries f_trunc = *this;
            f_trunc.setOrder(prec);
            // g = g * (1 + f - log(g))
            PowerSeries log_g = g.log();
            PowerSeries correction = PowerSeries::one(prec) + f_trunc - log_g;
            g = g * correction;
            g.setOrder(prec);
        }
        g.setOrder(order_);
        return g;
    }

    // ============================================================
    // 平方根: sqrt(f(x)) mod x^N
    // f(0) = 1 が必要 (一般化: f(0) が完全平方)
    // Newton 反復: g_{k+1} = (g_k + f/g_k) / 2
    // ============================================================

    [[nodiscard]] PowerSeries sqrt() const
        requires PowerSeriesDivisible<T>
    {
        assert(!coeffs_.empty() && coeffs_[0] == T(1) &&
               "sqrt requires f(0) = 1");
        if (order_ <= 1) return PowerSeries::one(order_);

        // 初期近似: g_0 = 1
        size_t prec = 1;
        PowerSeries g = PowerSeries::one(order_);
        T two = T(1) + T(1);

        while (prec < order_) {
            prec = std::min(prec * 2, order_);
            g.setOrder(prec);
            PowerSeries f_trunc = *this;
            f_trunc.setOrder(prec);
            // g = (g + f/g) / 2
            PowerSeries f_over_g = f_trunc / g;
            g = (g + f_over_g) / PowerSeries(two, prec);
            g.setOrder(prec);
        }
        g.setOrder(order_);
        return g;
    }

    // ============================================================
    // 冪乗: f(x)^n mod x^N
    // 整数冪: 二分累乗法
    // ============================================================

    [[nodiscard]] PowerSeries pow(int n) const
        requires PowerSeriesDivisible<T>
    {
        if (n == 0) return PowerSeries::one(order_);
        if (n < 0) return this->inverse().pow(-n);
        if (n == 1) return *this;

        // 二分累乗法
        PowerSeries result = PowerSeries::one(order_);
        PowerSeries base = *this;
        unsigned int exp = static_cast<unsigned int>(n);
        while (exp > 0) {
            if (exp & 1) result = result * base;
            base = base * base;
            exp >>= 1;
        }
        return result;
    }

    // 一般冪: f(x)^α = exp(α · log f(x))  — f(0) = 1 が必要
    [[nodiscard]] PowerSeries pow(const T& alpha) const
        requires PowerSeriesDivisible<T>
    {
        assert(!coeffs_.empty() && coeffs_[0] == T(1) &&
               "general pow requires f(0) = 1");
        return (this->log() * alpha).exp();
    }

    // ============================================================
    // 合成: f(g(x)) mod x^N  — Brent-Kung アルゴリズム
    // g(0) = 0 が必要
    // ============================================================

    [[nodiscard]] PowerSeries compose(const PowerSeries& g) const {
        assert((g.coeffs_.empty() || g.coeffs_[0] == T(0)) &&
               "compose requires g(0) = 0");

        size_t ord = std::min(order_, g.order_);
        size_t n = coeffs_.size();
        if (n == 0) return PowerSeries(ord);

        // Brent-Kung: m = √n ブロックに分割
        // f(g) = Σ_{i=0}^{m-1} h_i(g) · (g^b)^i
        // h_i(g) = Σ_{j=0}^{b-1} a_{i·b+j} · g^j
        size_t b = 1;
        while (b * b < n) ++b;
        size_t m = (n + b - 1) / b;

        // g の累乗テーブル: g^0, g^1, ..., g^{b-1}
        std::vector<PowerSeries> gpow(b, PowerSeries(ord));
        gpow[0] = PowerSeries::one(ord);
        if (b > 1) {
            PowerSeries gg = g;
            gg.setOrder(ord);
            gpow[1] = gg;
            for (size_t j = 2; j < b; ++j)
                gpow[j] = gpow[j - 1] * gg;
        }

        // g^b = g^{b-1} * g
        PowerSeries gg = g;
        gg.setOrder(ord);
        PowerSeries gb = gpow.back() * gg;
        gb.setOrder(ord);

        // Horner 法でブロック単位に評価 (逆順)
        PowerSeries result(ord);
        for (size_t i = m; i-- > 0;) {
            // h_i = Σ_{j=0}^{b-1} a_{i*b+j} · g^j
            PowerSeries hi(ord);
            for (size_t j = 0; j < b; ++j) {
                size_t idx = i * b + j;
                if (idx < n && !((*this)[idx] == T(0)))
                    hi += gpow[j] * (*this)[idx];
            }
            if (i + 1 < m)
                result = result * gb + hi;
            else
                result = hi;
        }
        return result;
    }

    // ============================================================
    // 逆合成 (compositional inverse): f(g(x)) = x
    // f(0) = 0, f'(0) ≠ 0 が必要
    // Newton 反復: g_{k+1} = g_k + (x - f(g_k)) / f'(g_k)
    // ============================================================

    [[nodiscard]] PowerSeries composeInverse() const
        requires PowerSeriesDivisible<T>
    {
        assert(coeffs_.size() >= 2 &&
               coeffs_[0] == T(0) && !(coeffs_[1] == T(0)) &&
               "composeInverse requires f(0) = 0, f'(0) != 0");

        size_t prec = 1;
        // 初期近似: g_0 = x / f'(0)
        PowerSeries g = PowerSeries::x(order_) / PowerSeries(coeffs_[1], order_);
        PowerSeries xps = PowerSeries::x(order_);

        while (prec < order_) {
            prec = std::min(prec * 2, order_);
            g.setOrder(prec);
            PowerSeries f_trunc = *this;
            f_trunc.setOrder(prec);
            xps.setOrder(prec);

            // f(g(x)) を計算
            PowerSeries fg = f_trunc.compose(g);
            // f'(g(x)) を計算
            PowerSeries fprime = f_trunc.derivative();
            fprime.setOrder(prec);
            PowerSeries fprime_g = fprime.compose(g);

            // g = g + (x - f(g)) / f'(g)
            PowerSeries residual = xps - fg;
            g = g + residual / fprime_g;
            g.setOrder(prec);
        }
        g.setOrder(order_);
        return g;
    }

    // ============================================================
    // 評価: f(x0) — Horner 法
    // ============================================================

    [[nodiscard]] T evaluate(const T& x0) const {
        if (coeffs_.empty()) return T(0);
        T result = coeffs_.back();
        for (size_t i = coeffs_.size() - 1; i > 0; --i)
            result = result * x0 + coeffs_[i - 1];
        return result;
    }

    // ============================================================
    // Revert (係数を反転): rev(f) = x^{n-1} · f(1/x)
    // ============================================================

    [[nodiscard]] PowerSeries reverse(size_t len) const {
        PowerSeries result(order_);
        result.coeffs_.resize(len, T(0));
        for (size_t i = 0; i < len && i < coeffs_.size(); ++i)
            result.coeffs_[len - 1 - i] = coeffs_[i];
        while (!result.coeffs_.empty() && result.coeffs_.back() == T(0))
            result.coeffs_.pop_back();
        return result;
    }

    // ============================================================
    // Polynomial<T> との相互変換
    // ============================================================

    // Polynomial から (order を指定)
    template<typename Poly>
    [[nodiscard]] static PowerSeries fromPolynomial(const Poly& p, size_t order) {
        PowerSeries result(order);
        size_t n = std::min(static_cast<size_t>(p.degree() + 1), order);
        result.coeffs_.resize(n);
        for (size_t i = 0; i < n; ++i)
            result.coeffs_[i] = p[i];
        return result;
    }

    // ============================================================
    // 文字列表現
    // ============================================================

    [[nodiscard]] std::string toString(const std::string& var = "x") const {
        if (coeffs_.empty()) return "O(" + var + "^" + std::to_string(order_) + ")";

        std::ostringstream oss;
        bool first = true;
        for (size_t i = 0; i < coeffs_.size(); ++i) {
            if (coeffs_[i] == T(0)) continue;
            if (!first) oss << " + ";
            first = false;
            if (i == 0) {
                oss << coeffs_[i];
            } else if (i == 1) {
                if (coeffs_[i] == T(1))
                    oss << var;
                else
                    oss << coeffs_[i] << "*" << var;
            } else {
                if (coeffs_[i] == T(1))
                    oss << var << "^" << i;
                else
                    oss << coeffs_[i] << "*" << var << "^" << i;
            }
        }
        if (first) oss << "0";
        oss << " + O(" << var << "^" << order_ << ")";
        return oss.str();
    }

    // ストリーム出力
    friend std::ostream& operator<<(std::ostream& os, const PowerSeries& ps) {
        return os << ps.toString();
    }
};

// ================================================================
// フリー関数: 数学関数
// ================================================================

// スカラー × 級数
template<PowerSeriesCoeff T>
[[nodiscard]] PowerSeries<T> operator*(const T& c, const PowerSeries<T>& f) {
    return f * c;
}

// 逆数
template<PowerSeriesDivisible T>
[[nodiscard]] PowerSeries<T> inverse(const PowerSeries<T>& f) {
    return f.inverse();
}

// 対数
template<PowerSeriesDivisible T>
[[nodiscard]] PowerSeries<T> log(const PowerSeries<T>& f) {
    return f.log();
}

// 指数
template<PowerSeriesDivisible T>
[[nodiscard]] PowerSeries<T> exp(const PowerSeries<T>& f) {
    return f.exp();
}

// 平方根
template<PowerSeriesDivisible T>
[[nodiscard]] PowerSeries<T> sqrt(const PowerSeries<T>& f) {
    return f.sqrt();
}

// 合成
template<PowerSeriesCoeff T>
[[nodiscard]] PowerSeries<T> compose(const PowerSeries<T>& f, const PowerSeries<T>& g) {
    return f.compose(g);
}

// 微分
template<PowerSeriesCoeff T>
[[nodiscard]] PowerSeries<T> derivative(const PowerSeries<T>& f) {
    return f.derivative();
}

// 積分
template<PowerSeriesDivisible T>
[[nodiscard]] PowerSeries<T> integral(const PowerSeries<T>& f) {
    return f.integral();
}

} // namespace calx
