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

// IntSequence.hpp
// Integer sequences (Fibonacci numbers, Lucas numbers)
// Algorithm: square-only fast doubling, O(S(n) log n)
//   S(n) = squaring cost. Eliminates general multiplication entirely.

#pragma once

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

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

namespace sangi {

/**
 * @brief Fibonacci and Lucas numbers
 *
 * Square-only fast doubling: track the pair (F(k), F(k-1)) and
 * compute each step with two squarings plus addition/subtraction (no general multiplication).
 *
 * Formulas:
 *   F(2k-1) = F(k)^2 + F(k-1)^2
 *   F(2k+1) = 4*F(k)^2 - F(k-1)^2 + 2*(-1)^k
 *   F(2k)   = F(2k+1) - F(2k-1)
 *
 * GMP correspondence: mpz_fib_ui, mpz_fib2_ui, mpz_lucnum_ui, mpz_lucnum2_ui
 */
class IntSequence {
public:
    /**
     * @brief Fibonacci number F(n)
     *
     * F(0)=0, F(1)=1, F(n)=F(n-1)+F(n-2)
     *
     * @param n Non-negative integer
     * @return F(n)
     */
    static Int fibonacci(uint64_t n) {
        return fibonacci2(n).first;
    }

    /**
     * @brief Fibonacci-number pair {F(n), F(n-1)}
     *
     * @param n Non-negative integer (when 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) };

        // F(93) is the largest Fibonacci number that fits in uint64_t
        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,
        };

        // Small n: table lookup
        if (n <= 93) {
            return { Int(fib_table[n]), Int(fib_table[n > 0 ? n - 1 : 0]) };
        }

        // Use the direct mpn version (GMP-style 3 buffers + table start) for all sizes.
        // The old 50000 threshold fell back to the IntOps path, but Int construction overhead was large.
        // Large sizes require heap allocation, but it is still faster than IntOps.

        // Table start: halve n down to the table range to reduce iteration count
        // GMP method: nfirst = n >> steps, then double steps times from there
        uint64_t nfirst = n;
        int steps = 0;
        while (nfirst > 93) { nfirst /= 2; steps++; }

        // Size estimate for 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;

        // 3 buffers (fp, f1p, xp) + scratch
        // GMP method: fp=F[k], f1p=F[k-1], xp=scratch (no pointer swap needed)
        size_t alloc = 2 * max_limbs + 4;
        size_t scratch_sz = mpn::square_scratch_size(max_limbs);
        size_t total = 3 * alloc + scratch_sz;

        // Small sizes use a stack buffer (avoids heap allocation and zero initialization)
        // Large sizes use a thread_local arena to amortize reallocation costs (eliminates the zero-init cost on repeated calls)
        static constexpr size_t STACK_LIMIT = 4096;  // 32KB — fits in the stack up to n ~ 250000
        uint64_t stack_mem[STACK_LIMIT];
        thread_local std::vector<uint64_t> tls_arena;
        uint64_t* mem_ptr;
        if (total <= STACK_LIMIT) {
            mem_ptr = stack_mem;
        } else {
            if (tls_arena.size() < total) tls_arena.resize(total);
            mem_ptr = tls_arena.data();
        }

        uint64_t* fp  = mem_ptr;                // F[k]
        uint64_t* f1p = mem_ptr + alloc;         // F[k-1]
        uint64_t* xp  = mem_ptr + 2 * alloc;    // scratch (holds F[k]^2)
        uint64_t* scratch = mem_ptr + 3 * alloc;

        // Set the initial value from the table
        fp[0] = fib_table[nfirst];
        f1p[0] = fib_table[nfirst > 0 ? nfirst - 1 : 0];
        size_t size = 1;

        // GMP-style doubling loop:
        //   xp ← fp²          (F[k]²)
        //   fp <- f1p^2         (F[k-1]^2, overwriting fp - safe since the input is f1p)
        //   f1p ← xp + fp     (F[2k-1] = F[k]² + F[k-1]²)
        //   fp ← 4*xp - fp ± 2 (F[2k+1] = 4·F[k]² - F[k-1]² + 2·(-1)^k)
        //   bit=1: f1p ← fp - f1p  → (fp, f1p) = (F[2k+1], F[2k])
        //   bit=0: fp ← fp - f1p   → (fp, f1p) = (F[2k], F[2k-1])
        for (int i = steps - 1; i >= 0; --i) {
            bool k_odd = (n >> (i + 1)) & 1;  // corresponding bit of nfirst = LSB of the current k
            // Note: since k = n >> i (upper bits), the LSB of k is not the LSB of n >> (i+1),
            //     but the lowest bit above bit (i) of n. As in GMP, test with n & (1 << (i+1)).
            //     Note i+1 cannot reach 64 (n <= 50000).

            // xp = fp² (F[k]²)
            mpn::square(xp, fp, size, scratch);
            // fp = f1p^2 (F[k-1]^2) — safe because fp != f1p
            mpn::square(fp, f1p, size, scratch);
            size *= 2;
            size -= (xp[size - 1] == 0);

            // f1p = xp + fp (F[2k-1])
            f1p[size] = mpn::add(f1p, xp, size, fp, size);

            // F[2k+1] = 4·F[k]² - F[k-1]² + 2·(-1)^k
            // GMP bit-OR trick: bit 1 of a square is always 0, so |= 2 absorbs +/-2
            if (k_odd) {
                // (-1)^k = -1 -> -2: OR 2 into fp (= F[k-1]^2) to increase the minuend
                fp[0] |= 2;
            }
            uint64_t cy = mpn::lshift(xp, xp, size, 2);  // 4·F[k]²
            if (!k_odd) {
                // (-1)^k = +1 -> +2: OR 2 into xp (= 4 * F[k]^2)
                xp[0] |= 2;
            }
            // fp = xp - fp (= 4·F[k]² ± 2 - F[k-1]² ∓ 2 = F[2k+1])
            cy -= mpn::sub(fp, xp, size, fp, size);
            fp[size] = cy;
            size += (fp[size] != 0);

            // F[2k] = F[2k+1] - F[2k-1]: replace the one not needed
            if ((n >> i) & 1) {
                // bit=1: (fp, f1p) → (F[2k+1], F[2k])
                mpn::sub(f1p, fp, size, f1p, size);
            } else {
                // bit=0: (fp, f1p) → (F[2k], F[2k-1])
                mpn::sub(fp, fp, size, f1p, size);
                // Since F[2k] < F[2k+1], leading zeros may appear
                while (size > 0 && fp[size - 1] == 0) --size;
            }
        }

        // Build Int (final result only): fp = F(n), f1p = F(n-1)
        // Use span + PreNormalized to avoid temporary vector allocation
        while (size > 0 && fp[size - 1] == 0) --size;
        size_t f1size = size;
        while (f1size > 0 && f1p[f1size - 1] == 0) --f1size;

        Int result_fn, result_fn1;
        if (size > 0) {
            result_fn = Int::fromRawWordsPreNormalized(
                std::span<const uint64_t>(fp, size), 1);
        }
        if (f1size > 0) {
            result_fn1 = Int::fromRawWordsPreNormalized(
                std::span<const uint64_t>(f1p, f1size), 1);
        }
        return { std::move(result_fn), std::move(result_fn1) };
    }

private:
    // For large sizes: IntOps square-only fast doubling
    static std::pair<Int, Int> fibonacci2_intops(uint64_t n, int bits) {
        Int a(1), b(0);  // F(1), F(0)
        bool neg = true;
        Int a2, b2, f2km1, f2kp1, f2k;
        static const Int TWO(2);

        for (int i = bits - 2; i >= 0; --i) {
            // a, b start from Int(1)/Int(0) and are updated only via Unchecked ops -> always Normal
            IntOps::squareUnchecked(a, a2);
            IntOps::squareUnchecked(b, b2);

            // F(2k-1) = a² + b²
            IntOps::addUnchecked(a2, b2, f2km1);

            // F(2k+1) = 4*a² - b² + 2*(-1)^k
            IntOps::leftShiftUnchecked(a2, 2, f2kp1);
            IntOps::subUnchecked(f2kp1, b2, f2kp1);
            if (neg) {
                IntOps::subUnchecked(f2kp1, TWO, f2kp1);
            } else {
                IntOps::addUnchecked(f2kp1, TWO, f2kp1);
            }

            // F(2k) = F(2k+1) - F(2k-1)
            IntOps::subUnchecked(f2kp1, f2km1, f2k);

            if ((n >> i) & 1) {
                a = std::move(f2kp1);
                b = std::move(f2k);
                neg = true;
            } else {
                a = std::move(f2k);
                b = std::move(f2km1);
                neg = false;
            }
        }
        return { std::move(a), std::move(b) };
    }

public:
    /**
     * @brief Lucas number 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 Non-negative integer
     * @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-number pair {L(n), L(n-1)}
     *
     * @param n Non-negative integer (when 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 sangi