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

// GoldilocksNtt.hpp
// Goldilocks prime (p = 2^64 - 2^32 + 1) NTT
// CRT-free single-prime NTT. Modular reduction reduces to shifts and add/sub only.
// Inputs are split into 16-bit words (4x expansion) to avoid coefficient overflow.
//
// Constraints: NTT size <= 2^32 (the 2-power factor of p-1 is 2^32)
//              Maximum coefficient: N * (2^16 - 1)^2 < p (always holds for N < 2^32)
//
// ----------------------------------------------------------------
// Performance comparison (2026-03-16, AMD Ryzen Threadripper PRO 5995WX, Release x64)
//
//   Goldilocks NTT is 4-10x slower than Prime NTT (3-prime + CRT, AVX2).
//
//   n(limbs)  Goldilocks  PrimeNTT  NTT-size      ratio
//       500    399 us      97 us    4K / 1K       4.1x
//      1000    872 us     193 us    8K / 2K       4.5x
//      3000   3.93 ms     381 us   32K / 8K     10.3x
//     10000  18.54 ms    1.83 ms  128K / 32K    10.2x
//     50000  99.31 ms   21.16 ms  512K / 128K    4.7x
//
//   Causes:
//   - The 4x expansion from 16-bit splitting dominates. NTT size grows by 4x
//     and the O(N log N) work grows by roughly 4-5x.
//   - The CRT-free win (3 NTTs → 1 NTT) cannot offset the 4x expansion.
//   - Goldilocks modular operations (shift+add, 1x MULX) are lighter than
//     Montgomery (2x MULX), but Prime NTT is already 4-way AVX2-parallel
//     while Goldilocks is scalar only, narrowing the effective gap.
//
//   Room for improvement:
//   - 21-bit splitting (3x expansion, safe for N < 2^22) cuts NTT size by 25%
//   - AVX2 Goldilocks butterflies (4 VPMULUDQ vs Montgomery's 11 VPMULUDQ)
//   - It remains unclear whether this would beat the heavily optimized existing
//     Prime NTT (AVX2 + MT + cache blocking)
// ----------------------------------------------------------------

#pragma once

#include <cstdint>
#include <cstddef>
#include <cstring>
#include <vector>
#include <algorithm>

#include <math/core/mp/Int/compiler_intrinsics.hpp>

namespace sangi {
namespace goldilocks_ntt {

// ================================================================
// Goldilocks prime constants
// ================================================================

// p = 2^64 - 2^32 + 1
constexpr uint64_t GOLD_P = 0xFFFFFFFF00000001ULL;

// p-1 = 2^32 * (2^32 - 1) = 2^32 * 3 * 5 * 17 * 257 * 65537
constexpr int GOLD_MAX_S = 32;  // Maximum NTT length = 2^32

// Primitive root g = 7 (g^(p-1) = 1 mod p, g^((p-1)/q) != 1 for all prime q | p-1)
constexpr uint64_t GOLD_G = 7;

// ================================================================
// Goldilocks fast modular operations
// ================================================================

// (a + b) mod p — a, b < p
// a + b < 2p < 2^65. On overflow, add 2^64 mod p = 2^32 - 1.
// Proof: when carry=1, the wrapped sum < 2p - 2^64 = 2^64 - 2^33 + 2,
//        so sum + (2^32-1) < 2^64 - 2^32 + 1 = p → no second overflow.
inline uint64_t gold_add(uint64_t a, uint64_t b) {
    uint64_t sum = a + b;
    if (sum < a) {
        // carry=1: true value is sum + 2^64 ≡ sum + (2^32 - 1) (mod p)
        // sum + (2^32-1) < p, so no final normalization is required
        return sum + 0xFFFFFFFFULL;
    }
    return (sum >= GOLD_P) ? (sum - GOLD_P) : sum;
}

// (a - b) mod p — a, b < p
inline uint64_t gold_sub(uint64_t a, uint64_t b) {
    if (a >= b) return a - b;
    return a - b + GOLD_P;  // wrap: a - b + p
}

// Goldilocks reduction: (hi : lo) mod p
// p = 2^64 - 2^32 + 1, so 2^64 ≡ 2^32 - 1 (mod p)
// (hi * 2^64 + lo) mod p = (hi * (2^32 - 1) + lo) mod p
//
// Two-step reduction:
//   Step 1: compute hi*(2^32-1) = (hi<<32) - hi in 96 bits and add lo → {r1, r0}
//   Step 2: r1 < 2^32+1, so r1*(2^32-1) < 2^64. Compute r0 + r1*(2^32-1).
//           If it overflows, correct by +0xFFFFFFFF (no second overflow).
inline uint64_t gold_reduce(uint64_t lo, uint64_t hi) {
    // Step 1: {h1, h0} = hi << 32 (96bit value)
    uint64_t h0 = hi << 32;       // low 64 bits
    uint64_t h1 = hi >> 32;       // high 32 bits (< 2^32)

    // {t1, t0} = {h1, h0} - hi
    uint64_t borrow = (h0 < hi) ? 1ULL : 0ULL;
    uint64_t t0 = h0 - hi;
    uint64_t t1 = h1 - borrow;    // h1 >= borrow (proof: h1=0 → hi<2^32 → h0=hi<<32>=hi → borrow=0)

    // {r1, r0} = {t1, t0} + lo
    uint64_t r0 = t0 + lo;
    uint64_t carry = (r0 < t0) ? 1ULL : 0ULL;
    uint64_t r1 = t1 + carry;     // r1 <= 2^32 (t1 < 2^32, carry <= 1)

    // Step 2: result = r0 + r1 * (2^32 - 1) (mod p)
    // r1 * (2^32-1) <= 2^32 * (2^32-1) = 2^64 - 2^32, fits in uint64
    uint64_t adj = r1 * 0xFFFFFFFFULL;
    uint64_t result = r0 + adj;
    if (result < r0) {
        // carry: result += 2^32 - 1 (proof that there is no second overflow:
        // result_wrapped <= r0 + adj - 2^64 <= (2^64-1) + (2^64-2^32) - 2^64 = 2^64-2^32-1
        // result_wrapped + (2^32-1) <= 2^64 - 2 < 2^64)
        result += 0xFFFFFFFFULL;
    }

    // Final normalization
    if (result >= GOLD_P) result -= GOLD_P;
    return result;
}

// (a * b) mod p — Goldilocks fast reduction
// Compared with Montgomery's 2x MULX, this uses 1x MULX plus shifts/add/sub only
inline uint64_t gold_mul(uint64_t a, uint64_t b) {
#if defined(_MSC_VER) && defined(_M_X64)
    uint64_t hi;
    uint64_t lo = _umul128(a, b, &hi);
    return gold_reduce(lo, hi);
#elif defined(__GNUC__) || defined(__clang__)
    unsigned __int128 prod = (unsigned __int128)a * b;
    return gold_reduce((uint64_t)prod, (uint64_t)(prod >> 64));
#endif
}

// base^exp mod p
inline uint64_t gold_pow(uint64_t base, uint64_t exp) {
    uint64_t result = 1;
    base %= GOLD_P;
    while (exp > 0) {
        if (exp & 1)
            result = gold_mul(result, base);
        base = gold_mul(base, base);
        exp >>= 1;
    }
    return result;
}

// a^(-1) mod p (Fermat)
inline uint64_t gold_inv(uint64_t a) {
    return gold_pow(a, GOLD_P - 2);
}

// ================================================================
// NTT root table
// ================================================================

struct GoldRoots {
    size_t N = 0;
    std::vector<uint64_t> fwd_roots;  // omega^0, omega^1, ..., omega^(N-1)
    std::vector<uint64_t> inv_roots;  // omega_inv^0, ..., omega_inv^(N-1)
    uint64_t inv_n = 0;               // N^(-1) mod p

    void build(size_t ntt_size) {
        N = ntt_size;
        fwd_roots.resize(N);
        inv_roots.resize(N);

        // omega = g^((p-1)/N) mod p
        uint64_t omega = gold_pow(GOLD_G, (GOLD_P - 1) / N);
        uint64_t omega_inv = gold_inv(omega);
        inv_n = gold_inv(N);

        fwd_roots[0] = 1;
        for (size_t i = 1; i < N; ++i)
            fwd_roots[i] = gold_mul(fwd_roots[i - 1], omega);

        inv_roots[0] = 1;
        for (size_t i = 1; i < N; ++i)
            inv_roots[i] = gold_mul(inv_roots[i - 1], omega_inv);
    }
};

// ================================================================
// NTT transforms (DIF / DIT)
// ================================================================

// Forward NTT (Decimation In Frequency)
inline void forward_ntt_gold(uint64_t* data, size_t N,
                              const uint64_t* roots) {
    for (size_t len = N; len >= 2; len >>= 1) {
        size_t half = len >> 1;
        size_t step = N / len;
        for (size_t i = 0; i < N; i += len) {
            for (size_t j = 0; j < half; ++j) {
                uint64_t u = data[i + j];
                uint64_t v = data[i + j + half];
                data[i + j] = gold_add(u, v);
                data[i + j + half] = gold_mul(gold_sub(u, v), roots[j * step]);
            }
        }
    }
}

// Inverse NTT (Decimation In Time)
inline void inverse_ntt_gold(uint64_t* data, size_t N,
                              const uint64_t* inv_roots, uint64_t inv_n) {
    for (size_t len = 2; len <= N; len <<= 1) {
        size_t half = len >> 1;
        size_t step = N / len;
        for (size_t i = 0; i < N; i += len) {
            for (size_t j = 0; j < half; ++j) {
                uint64_t u = data[i + j];
                uint64_t v = gold_mul(data[i + j + half], inv_roots[j * step]);
                data[i + j] = gold_add(u, v);
                data[i + j + half] = gold_sub(u, v);
            }
        }
    }
    // 1/N scaling
    for (size_t i = 0; i < N; ++i)
        data[i] = gold_mul(data[i], inv_n);
}

// Pointwise multiplication
inline void pointwise_mul_gold(uint64_t* a, const uint64_t* b, size_t N) {
    for (size_t i = 0; i < N; ++i)
        a[i] = gold_mul(a[i], b[i]);
}

// ================================================================
// 16-bit word splitting / reassembly
// ================================================================

// 64-bit words → split into 16-bit chunks, zero-padded
inline void pack_16bit(const uint64_t* src, size_t n,
                       uint64_t* dst, size_t N) {
    size_t out = 0;
    for (size_t i = 0; i < n; ++i) {
        dst[out++] = src[i] & 0xFFFF;
        dst[out++] = (src[i] >> 16) & 0xFFFF;
        dst[out++] = (src[i] >> 32) & 0xFFFF;
        dst[out++] = (src[i] >> 48) & 0xFFFF;
    }
    for (size_t i = out; i < N; ++i)
        dst[i] = 0;
}

// Reassemble the convolution result (base 2^16) into 64-bit words with carry propagation
inline void unpack_16bit(const uint64_t* conv, size_t conv_len,
                         uint64_t* rp, size_t rn) {
    std::memset(rp, 0, rn * sizeof(uint64_t));

    uint64_t carry = 0;
    for (size_t k = 0; k < conv_len; ++k) {
        // conv[k] + carry (maximum ~2^50; 128-bit arithmetic is unnecessary but used defensively)
#if defined(_MSC_VER) && defined(_M_X64)
        uint64_t sum_lo;
        unsigned char c = _addcarry_u64(0, conv[k], carry, &sum_lo);
        uint64_t sum_hi = c;
#elif defined(__GNUC__) || defined(__clang__)
        unsigned __int128 sum128 = (unsigned __int128)conv[k] + carry;
        uint64_t sum_lo = (uint64_t)sum128;
        uint64_t sum_hi = (uint64_t)(sum128 >> 64);
#endif

        // Emit the low 16 bits
        size_t word = k / 4;
        size_t shift = (k % 4) * 16;
        if (word < rn) {
            rp[word] |= (sum_lo & 0xFFFF) << shift;
        }

        carry = (sum_lo >> 16) | (sum_hi << 48);
    }

    // Emit the remaining carry
    size_t word = conv_len / 4;
    size_t shift = (conv_len % 4) * 16;
    while (carry > 0 && word < rn) {
        if (shift == 0 && word >= conv_len / 4) {
            // Words past conv_len should still be 0
        }
        rp[word] |= (carry & 0xFFFF) << shift;
        carry >>= 16;
        shift += 16;
        if (shift >= 64) { shift = 0; word++; }
    }
}

// ================================================================
// Utilities
// ================================================================

inline size_t next_pow2_gold(size_t n) {
    size_t p = 1;
    while (p < n) p <<= 1;
    return p;
}

// ================================================================
// Main multiplication function
// ================================================================

// rp[0..an+bn-1] = ap[0..an-1] * bp[0..bn-1]
inline void mul_goldilocks_ntt(uint64_t* rp,
                                const uint64_t* ap, size_t an,
                                const uint64_t* bp, size_t bn) {
    if (an < bn) { std::swap(ap, bp); std::swap(an, bn); }

    size_t rn = an + bn;
    size_t pack_a = 4 * an;
    size_t pack_b = 4 * bn;
    size_t conv_len = pack_a + pack_b - 1;

    size_t N = next_pow2_gold(conv_len + 1);

    if (N > (1ULL << GOLD_MAX_S)) return;

    // Root table (thread_local cache)
    thread_local GoldRoots roots;
    if (roots.N != N) {
        roots.build(N);
    }

    // Working buffer (thread_local)
    thread_local std::vector<uint64_t> work;
    size_t total = 2 * N;
    if (work.size() < total) work.resize(total);

    uint64_t* da = work.data();
    uint64_t* db = da + N;

    pack_16bit(ap, an, da, N);
    pack_16bit(bp, bn, db, N);

    forward_ntt_gold(da, N, roots.fwd_roots.data());
    forward_ntt_gold(db, N, roots.fwd_roots.data());
    pointwise_mul_gold(da, db, N);
    inverse_ntt_gold(da, N, roots.inv_roots.data(), roots.inv_n);

    unpack_16bit(da, conv_len, rp, rn);
}

// Square: rp[0..2*an-1] = ap[0..an-1]^2
inline void sqr_goldilocks_ntt(uint64_t* rp,
                                const uint64_t* ap, size_t an) {
    size_t rn = 2 * an;
    size_t pack_a = 4 * an;
    size_t conv_len = 2 * pack_a - 1;
    size_t N = next_pow2_gold(conv_len + 1);

    if (N > (1ULL << GOLD_MAX_S)) return;

    thread_local GoldRoots roots;
    if (roots.N != N) {
        roots.build(N);
    }

    thread_local std::vector<uint64_t> work;
    if (work.size() < N) work.resize(N);

    uint64_t* da = work.data();

    pack_16bit(ap, an, da, N);
    forward_ntt_gold(da, N, roots.fwd_roots.data());

    for (size_t i = 0; i < N; ++i)
        da[i] = gold_mul(da[i], da[i]);

    inverse_ntt_gold(da, N, roots.inv_roots.data(), roots.inv_n);

    unpack_16bit(da, conv_len, rp, rn);
}

} // namespace goldilocks_ntt
} // namespace sangi