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

// ntt.hpp — NTT (Number Theoretic Transform) public API
//
// The implementation is compiled into sangi_ntt.lib (src/math/fft/ntt.cpp).
// Callers must link against sangi_ntt.

#pragma once

#include <cstdint>
#include <cstddef>

namespace sangi {
namespace prime_ntt { struct NttPrime; }

class Ntt {
public:
    struct Prime {
        uint64_t p;
        uint64_t g;
        size_t maxNttSize;
    };

    // P1 =  30 * 2^47 + 1, g=19
    // P2 = 345 * 2^47 + 1, g=13
    // P3 = 420 * 2^47 + 1, g=17
    static constexpr Prime prime1() { return { 0x000F'0000'0000'0001ULL, 19, size_t(1) << 47 }; }
    static constexpr Prime prime2() { return { 0x00AC'8000'0000'0001ULL, 13, size_t(1) << 47 }; }
    static constexpr Prime prime3() { return { 0x00D2'0000'0000'0001ULL, 17, size_t(1) << 47 }; }

    static size_t nextSmoothSize(size_t n);

    static bool isSmooth(size_t n) {
        if (n == 0) return false;
        while (n % 2 == 0) n /= 2;
        return n == 1 || n == 3 || n == 5;
    }

    static bool isValidSize(size_t N, const Prime& prime) {
        return N > 0 && (prime.p - 1) % N == 0;
    }

    static void forward(uint64_t* data, size_t N, const Prime& prime);
    static void inverse(uint64_t* data, size_t N, const Prime& prime);
    static void pointwiseMul(uint64_t* r, const uint64_t* a, const uint64_t* b,
                              size_t N, const Prime& prime);
    static void polyMul(uint64_t* r, const uint64_t* a, size_t an,
                         const uint64_t* b, size_t bn, const Prime& prime);

private:
    static prime_ntt::NttPrime makePrime(const Prime& prime);
    static void forwardSmooth(uint64_t* data, size_t N, const Prime& prime);
    static void inverseSmooth(uint64_t* data, size_t N, const Prime& prime);
    static void bluestein(uint64_t* data, size_t N, const Prime& prime, bool is_inverse);
    static void forwardMont(uint64_t* data, size_t N, const prime_ntt::NttPrime& nttPrime);
    static void inverseMont(uint64_t* data, size_t N, const prime_ntt::NttPrime& nttPrime);
};

} // namespace sangi