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

// IntSequence.hpp
// 整数列（Fibonacci 数、Lucas 数）
// アルゴリズム: fast doubling 法, O(M(n) log n)

#pragma once

#include "IntBase.hpp"
#include "IntOps.hpp"
#include "MpnOps.hpp"
#include "ScratchArena.hpp"

#include <cstdint>
#include <cstring>
#include <utility>

namespace calx {

/**
 * @brief Fibonacci 数・Lucas 数
 *
 * fast doubling 法により O(log n) 回の多倍長乗算で計算します。
 *
 * Fast doubling 公式 (F(k), F(k+1) から):
 *   F(2k)   = F(k) * [2*F(k+1) - F(k)]
 *   F(2k+1) = F(k)^2 + F(k+1)^2
 *
 * GMP 対応: mpz_fib_ui, mpz_fib2_ui, mpz_lucnum_ui, mpz_lucnum2_ui
 */
class IntSequence {
public:
    /**
     * @brief Fibonacci 数 F(n)
     *
     * F(0)=0, F(1)=1, F(n)=F(n-1)+F(n-2)
     *
     * @param n 非負整数
     * @return F(n)
     */
    static Int fibonacci(uint64_t n) {
        return fibonacci2(n).first;
    }

    /**
     * @brief Fibonacci 数のペア {F(n), F(n-1)}
     *
     * @param n 非負整数 (n=0 のとき {F(0), F(-1)} = {0, 1})
     * @return {F(n), F(n-1)}
     */
    static std::pair<Int, Int> fibonacci2(uint64_t n) {
        if (n == 0) return { Int(0), Int(1) };
        if (n == 1) return { Int(1), Int(0) };

        // 小さい n: テーブルルックアップ
        if (n <= 93) {
            // F(93) は uint64_t に収まる最大の Fibonacci 数
            static constexpr uint64_t fib_table[] = {
                0,1,1,2,3,5,8,13,21,34,55,89,144,233,377,610,987,
                1597,2584,4181,6765,10946,17711,28657,46368,75025,
                121393,196418,317811,514229,832040,1346269,2178309,
                3524578,5702887,9227465,14930352,24157817,39088169,
                63245986,102334155,165580141,267914296,433494437,
                701408733,1134903170,1836311903,2971215073ULL,
                4807526976ULL,7778742049ULL,12586269025ULL,
                20365011074ULL,32951280099ULL,53316291173ULL,
                86267571272ULL,139583862445ULL,225851433717ULL,
                365435296162ULL,591286729879ULL,956722026041ULL,
                1548008755920ULL,2504730781961ULL,4052739537881ULL,
                6557470319842ULL,10610209857723ULL,17167680177565ULL,
                27777890035288ULL,44945570212853ULL,72723460248141ULL,
                117669030460994ULL,190392490709135ULL,308061521170129ULL,
                498454011879264ULL,806515533049393ULL,1304969544928657ULL,
                2111485077978050ULL,3416454622906707ULL,5527939700884757ULL,
                8944394323791464ULL,14472334024676221ULL,23416728348467685ULL,
                37889062373143906ULL,61305790721611591ULL,99194853094755497ULL,
                160500643816367088ULL,259695496911122585ULL,420196140727489673ULL,
                679891637638612258ULL,1100087778366101931ULL,
                1779979416004714189ULL,2880067194370816120ULL,
                4660046610375530309ULL,7540113804746346429ULL,
                12200160415121876738ULL,
            };
            return { Int(fib_table[n]), Int(fib_table[n > 0 ? n - 1 : 0]) };
        }

        // fast doubling: top-down binary
        int bits = 0;
        for (uint64_t tmp = n; tmp > 0; tmp >>= 1) ++bits;

        // n <= 50000: mpn 直接版 (Int オブジェクト構築回避)
        // n > 50000: IntOps 3引数版 (大サイズでは multiply が支配的で差が小さい)
        if (n > 50000) {
            return fibonacci2_intops(n, bits);
        }

        // F(n) のサイズ推定: ~0.694*n bits = ceil(0.694*n/64) limbs
        size_t max_limbs = static_cast<size_t>(static_cast<double>(n) * 0.694 / 64.0) + 8;
        if (max_limbs < 8) max_limbs = 8;

        // バッファ: a=F(k), b=F(k+1) — 最終サイズ max_limbs + 余裕
        // t, prod, prod2 — 乗算/二乗の結果 (最大 2*max_limbs)
        size_t buf_sz = max_limbs + 2;
        size_t prod_sz = 2 * max_limbs + 4;
        size_t mul_scratch_sz = mpn::multiply_scratch_size(max_limbs, max_limbs);
        size_t sqr_scratch_sz = mpn::square_scratch_size(max_limbs);
        size_t scratch_sz = std::max(mul_scratch_sz, sqr_scratch_sz);
        // 全バッファを1つの vector で確保 (ScratchArena の容量制限を回避)
        size_t total = 2 * buf_sz + 3 * prod_sz + scratch_sz;
        std::vector<uint64_t> mem(total, 0);
        uint64_t* p = mem.data();
        uint64_t* a = p; p += buf_sz;
        uint64_t* b = p; p += buf_sz;
        uint64_t* t = p; p += prod_sz;
        uint64_t* prod = p; p += prod_sz;
        uint64_t* prod2 = p; p += prod_sz;
        uint64_t* scratch = p;

        std::memset(a, 0, max_limbs * sizeof(uint64_t));
        std::memset(b, 0, max_limbs * sizeof(uint64_t));
        // a = 0 (F(0)), b = 1 (F(1))
        b[0] = 1;
        size_t an = 0, bn = 1;  // 有効 limb 数

        for (int i = bits - 1; i >= 0; --i) {
            // doubling: F(2k), F(2k+1) from F(k), F(k+1)
            // c = F(2k)   = a * (2b - a)
            // d = F(2k+1) = a^2 + b^2

            if (an == 0) {
                // F(k) = 0: c = 0, d = b^2
                // bit=0: a=0, b=b^2; bit=1: a=b^2, b=b^2
                if (bn > 0) {
                    mpn::square(prod, b, bn, scratch);
                    size_t dn = 2 * bn;
                    while (dn > 0 && prod[dn - 1] == 0) --dn;
                    if ((n >> i) & 1) {
                        std::memcpy(a, prod, dn * sizeof(uint64_t));
                        an = dn;
                        // b = a (F(2k+1) = F(2k+1))
                        std::memcpy(b, a, an * sizeof(uint64_t));
                        bn = an;
                    } else {
                        // a stays 0
                        std::memcpy(b, prod, dn * sizeof(uint64_t));
                        bn = dn;
                    }
                }
                continue;
            }

            // t = 2b - a (t は正: 2*F(k+1) > F(k) for k >= 0)
            size_t tn;
            {
                // 2b (= b << 1)
                uint64_t carry = mpn::lshift(t, b, bn, 1);
                tn = bn;
                if (carry) { t[tn] = carry; tn++; }
                // t -= a
                if (an > 0) {
                    mpn::sub(t, t, tn, a, an);
                    while (tn > 0 && t[tn - 1] == 0) --tn;
                }
            }

            // prod = a * t → c (F(2k))
            size_t cn;
            if (an > 0 && tn > 0) {
                mpn::multiply(prod, a, an, t, tn, scratch);
                cn = an + tn;
                while (cn > 0 && prod[cn - 1] == 0) --cn;
            } else {
                cn = 0;
            }
            // c は prod に格納

            // d = a^2 + b^2 (F(2k+1))
            // a^2 を t に格納 (t は 2*max_limbs+2 サイズ)
            mpn::square(t, a, an, scratch);
            size_t a2n = 2 * an;
            while (a2n > 0 && t[a2n - 1] == 0) --a2n;

            // b^2 を prod2 に格納
            mpn::square(prod2, b, bn, scratch);
            size_t b2n = 2 * bn;
            while (b2n > 0 && prod2[b2n - 1] == 0) --b2n;

            // d = a^2 + b^2 → b に格納
            size_t dn;
            if (a2n >= b2n) {
                std::memcpy(b, t, a2n * sizeof(uint64_t));
                dn = a2n;
                uint64_t carry = (b2n > 0) ? mpn::add(b, b, dn, prod2, b2n) : 0;
                if (carry) { b[dn] = carry; dn++; }
            } else {
                std::memcpy(b, prod2, b2n * sizeof(uint64_t));
                dn = b2n;
                uint64_t carry = (a2n > 0) ? mpn::add(b, b, dn, t, a2n) : 0;
                if (carry) { b[dn] = carry; dn++; }
            }

            if ((n >> i) & 1) {
                // bit=1: a = d, b = c + d
                std::memcpy(a, b, dn * sizeof(uint64_t));
                an = dn;
                // b = c + d
                if (cn >= dn) {
                    std::memcpy(b, prod, cn * sizeof(uint64_t));
                    bn = cn;
                    uint64_t carry = (dn > 0) ? mpn::add(b, b, bn, a, dn) : 0;
                    if (carry) { b[bn] = carry; bn++; }
                } else {
                    // b already = d
                    bn = dn;
                    uint64_t carry = (cn > 0) ? mpn::add(b, b, bn, prod, cn) : 0;
                    if (carry) { b[bn] = carry; bn++; }
                }
            } else {
                // bit=0: a = c, b = d (d already in b)
                std::memcpy(a, prod, cn * sizeof(uint64_t));
                an = cn;
                bn = dn;
            }
        }

        // F(n-1) = F(n+1) - F(n) = b - a
        size_t fn1n;
        if (bn >= an) {
            std::memcpy(t, b, bn * sizeof(uint64_t));
            fn1n = bn;
            if (an > 0) {
                mpn::sub(t, t, fn1n, a, an);
                while (fn1n > 0 && t[fn1n - 1] == 0) --fn1n;
            }
        } else {
            fn1n = 0;  // 理論上到達しない
        }

        // Int 構築 (最終結果のみ)
        Int result_a, result_fn1;
        if (an > 0) {
            std::vector<uint64_t> aw(a, a + an);
            result_a = Int::fromRawWords(aw, 1);
        }
        if (fn1n > 0) {
            std::vector<uint64_t> tw(t, t + fn1n);
            result_fn1 = Int::fromRawWords(tw, 1);
        }
        return { std::move(result_a), std::move(result_fn1) };
    }

private:
    // 大サイズ用フォールバック: IntOps 3引数版 fast doubling
    static std::pair<Int, Int> fibonacci2_intops(uint64_t n, int bits) {
        Int a(0), b(1), c, d, t;
        for (int i = bits - 1; i >= 0; --i) {
            IntOps::addUnchecked(b, b, t);
            IntOps::subUnchecked(t, a, t);
            IntOps::mulUnchecked(a, t, c);
            IntOps::mulUnchecked(a, a, d);
            IntOps::addmul(d, b, b);
            if ((n >> i) & 1) {
                a = std::move(d);
                IntOps::addUnchecked(c, a, b);
            } else {
                a = std::move(c);
                b = std::move(d);
            }
        }
        IntOps::subUnchecked(b, a, t);
        return { std::move(a), std::move(t) };
    }

public:
    /**
     * @brief Lucas 数 L(n)
     *
     * L(0)=2, L(1)=1, L(n)=L(n-1)+L(n-2)
     * L(n) = F(n-1) + F(n+1) = 2*F(n+1) - F(n)
     *
     * @param n 非負整数
     * @return L(n)
     */
    static Int lucas(uint64_t n) {
        if (n == 0) return Int(2);
        auto [fn, fn_1] = fibonacci2(n);
        // L(n) = F(n) + 2*F(n-1)
        return fn + fn_1 + fn_1;
    }

    /**
     * @brief Lucas 数のペア {L(n), L(n-1)}
     *
     * @param n 非負整数 (n=0 のとき {L(0), L(-1)} = {2, -1})
     * @return {L(n), L(n-1)}
     */
    static std::pair<Int, Int> lucas2(uint64_t n) {
        if (n == 0) return { Int(2), Int(-1) };
        auto [fn, fn_1] = fibonacci2(n);
        // L(n) = F(n) + 2*F(n-1)
        // L(n-1) = 2*F(n) - F(n-1)
        Int ln = fn + fn_1 + fn_1;
        Int ln_1 = fn + fn - fn_1;
        return { std::move(ln), std::move(ln_1) };
    }
};

} // namespace calx
