#include #include #include #include #include #include #include namespace sangi {

// Multiplication-algorithm thresholds (in words)
// Tuned according to benchmark results (2026-02):
//   - Basecase → Karatsuba: 44% faster at 128 words
//   - Karatsuba → Toom-Cook: 15% faster at 256 words
//   - FFT (double-based): slower than Toom-Cook at every size (1.5x-2.8x)
//     → Will be re-evaluated once an NTT implementation lands. Currently unused in operator*.
namespace MultiplicationThresholds {
    constexpr size_t KARATSUBA_THRESHOLD = 128;  // 128 words (8K-bit) or larger uses Karatsuba
    constexpr size_t TOOMCOOK_THRESHOLD = 256;   // 256 words (16K-bit) or larger uses Toom-Cook-3
    constexpr size_t TOOMCOOK4_THRESHOLD = 640;  // 640 words (40K-bit) or larger uses Toom-Cook-4
    constexpr size_t FFT_THRESHOLD = 256;        // 256 words — currently unused (FFT is slower than Toom-Cook)
}

class IntMultiplication {
public:
    // Karatsuba multiplication (mpn-based)
    // Complexity: O(n^1.585) where n = max(X.words, Y.words)
    // Operates on raw limb pointers internally; Int is materialized only at entry and exit
    static Int karatsubaMultiply(const Int& X, const Int& Y) {
        int result_sign = X.getSign() * Y.getSign();
        if (result_sign == 0) return Int(0);
        if (X.isSpecialState() || Y.isSpecialState()) {
            return basecaseMultiply(X, Y);
        }

        const uint64_t* a = X.m_words.data();
        size_t an = X.m_words.size();
        const uint64_t* b = Y.m_words.data();
        size_t bn = Y.m_words.size();

        // For small inputs use basecase (mpn-based)
        if (an < MultiplicationThresholds::KARATSUBA_THRESHOLD ||
            bn < MultiplicationThresholds::KARATSUBA_THRESHOLD) {
            return basecaseMultiply(X, Y);
        }

        // Allocate the result buffer and scratch from the arena
        ScratchScope scope;
        size_t rn = an + bn;
        uint64_t* r = getThreadArena().alloc_limbs(rn);
        size_t scratch_size = mpn::mul_karatsuba_scratch_size(std::max(an, bn));
        uint64_t* scratch = getThreadArena().alloc_limbs(scratch_size);

        // Run mpn Karatsuba
        mpn::mul_karatsuba(r, a, an, b, bn, scratch);

        // Convert the result to Int (only one copy happens here)
        size_t actual_rn = mpn::normalized_size(r, rn);
        if (actual_rn == 0) return Int(0);

        return Int::fromRawWords(std::span (r, actual_rn), result_sign);
    }

    // Toom-Cook-3 multiplication (mpn-based)
    // Complexity: O(n^1.465) where n = max(X.words, Y.words)
    // Evaluation points {0, 1, -1, 2, ∞} (GMP convention)
    // Operates on raw limb pointers internally; Int is materialized only at entry and exit
    static Int toomCook3Multiply(const Int& X, const Int& Y) {
        int result_sign = X.getSign() * Y.getSign();
        if (result_sign == 0) return Int(0);
        if (X.isSpecialState() || Y.isSpecialState()) {
            return basecaseMultiply(X, Y);
        }

        const uint64_t* a = X.m_words.data();
        size_t an = X.m_words.size();
        const uint64_t* b = Y.m_words.data();
        size_t bn = Y.m_words.size();

        if (an < MultiplicationThresholds::TOOMCOOK_THRESHOLD ||
            bn < MultiplicationThresholds::TOOMCOOK_THRESHOLD) {
            return karatsubaMultiply(X, Y);
        }

        // Allocate the result buffer and scratch from the arena
        ScratchScope scope;
        size_t rn = an + bn;
        uint64_t* r = getThreadArena().alloc_limbs(rn);
        size_t scratch_size = mpn::mul_toomcook3_scratch_size(std::max(an, bn));
        uint64_t* scratch = getThreadArena().alloc_limbs(scratch_size);

        // Run mpn Toom-Cook-3
        mpn::mul_toomcook3(r, a, an, b, bn, scratch);

        // Convert the result to Int
        size_t actual_rn = mpn::normalized_size(r, rn);
        if (actual_rn == 0) return Int(0);

        return Int::fromRawWords(std::span (r, actual_rn), result_sign);
    }

    // Toom-Cook-4 multiplication (mpn-based)
    // Complexity: O(n^1.404) where n = max(X.words, Y.words)
    // Evaluation points {0, 1, -1, 2, -2, 3, ∞}
    // Operates on raw limb pointers internally; Int is materialized only at entry and exit
    static Int toomCook4Multiply(const Int& X, const Int& Y) {
        int result_sign = X.getSign() * Y.getSign();
        if (result_sign == 0) return Int(0);
        if (X.isSpecialState() || Y.isSpecialState()) {
            return basecaseMultiply(X, Y);
        }

        const uint64_t* a = X.m_words.data();
        size_t an = X.m_words.size();
        const uint64_t* b = Y.m_words.data();
        size_t bn = Y.m_words.size();

        if (an < MultiplicationThresholds::TOOMCOOK4_THRESHOLD ||
            bn < MultiplicationThresholds::TOOMCOOK4_THRESHOLD) {
            return toomCook3Multiply(X, Y);
        }

        // Allocate the result buffer and scratch from the arena
        ScratchScope scope;
        size_t rn = an + bn;
        uint64_t* r = getThreadArena().alloc_limbs(rn);
        size_t scratch_size = mpn::mul_toomcook4_scratch_size(std::max(an, bn));
        uint64_t* scratch = getThreadArena().alloc_limbs(scratch_size);

        // Run mpn Toom-Cook-4
        mpn::mul_toomcook4(r, a, an, b, bn, scratch);

        // Convert the result to Int
        size_t actual_rn = mpn::normalized_size(r, rn);
        if (actual_rn == 0) return Int(0);

        return Int::fromRawWords(std::span (r, actual_rn), result_sign);
    }

    // General-purpose multiplication (handles unbalanced inputs)
    // Selects the algorithm automatically based on size, and uses a chunked
    // strategy when the size ratio is large.
    static Int multiply(const Int& X, const Int& Y) {
        int result_sign = X.getSign() * Y.getSign();
        if (result_sign == 0) return Int(0);
        if (X.isSpecialState() || Y.isSpecialState()) {
            return basecaseMultiply(X, Y);
        }

        const uint64_t* a = X.m_words.data();
        size_t an = X.m_words.size();
        const uint64_t* b = Y.m_words.data();
        size_t bn = Y.m_words.size();

        ScratchScope scope;
        size_t rn = an + bn;
        uint64_t* r = getThreadArena().alloc_limbs(rn);
        size_t scratch_size = mpn::multiply_scratch_size(an, bn);
        uint64_t* scratch = (scratch_size > 0)
            ? getThreadArena().alloc_limbs(scratch_size) : nullptr;

        mpn::multiply(r, a, an, b, bn, scratch);

        size_t actual_rn = mpn::normalized_size(r, rn);
        if (actual_rn == 0) return Int(0);

        return Int::fromRawWords(std::span (r, actual_rn), result_sign);
    }

    // Compute only the top rn words of the product (short multiplication)
    // Uses mulhigh_basecase only when the size is at or below basecase.
    // Otherwise falls back to the full product followed by extracting the top words.
    static Int multiplyHigh(const Int& X, const Int& Y, size_t rn) {
        int result_sign = X.getSign() * Y.getSign();
        if (result_sign == 0) return Int(0);

        const uint64_t* a = X.m_words.data();
        size_t an = X.m_words.size();
        const uint64_t* b = Y.m_words.data();
        size_t bn = Y.m_words.size();

        size_t total = an + bn;
        if (rn >= total) {
            return multiply(X, Y);  // The full product suffices
        }

        ScratchScope scope;
        uint64_t* rp = getThreadArena().alloc_limbs(rn);

        if (std::min(an, bn) < mpn::KARATSUBA_THRESHOLD) {
            // Basecase: skip the low words via mulhigh
            mpn::mulhigh_basecase(rp, a, an, b, bn, rn);
        } else {
            // Karatsuba or larger: compute the full product and extract the top portion
            size_t scratch_size = mpn::multiply_scratch_size(an, bn);
            uint64_t* full = getThreadArena().alloc_limbs(total);
            uint64_t* scratch = (scratch_size > 0)
                ? getThreadArena().alloc_limbs(scratch_size) : nullptr;
            mpn::multiply(full, a, an, b, bn, scratch);
            std::memcpy(rp, full + (total - rn), rn * sizeof(uint64_t));
        }

        size_t actual_rn = mpn::normalized_size(rp, rn);
        if (actual_rn == 0) return Int(0);
        return Int::fromRawWords(std::span (rp, actual_rn), result_sign);
    }

    // FFT multiplication
    // Complexity: O(n log n) where n = FFT size
    // Precision limit: N < 2^17 points (about 630k digits) due to the 53-bit double mantissa
    static Int fftMultiply(const Int& X, const Int& Y) {
        // Save the sign and compute on the absolute values
        int result_sign = X.getSign() * Y.getSign();

        // Take the absolute values
        Int absX = (X.getSign() < 0) ? -X : X;
        Int absY = (Y.getSign() < 0) ? -Y : Y;

        size_t Nx = absX.words().size();
        size_t Ny = absY.words().size();

        // Small numbers use Toom-Cook
        if (Nx < MultiplicationThresholds::FFT_THRESHOLD ||
            Ny < MultiplicationThresholds::FFT_THRESHOLD) {
            Int result = toomCook3Multiply(absX, absY);
            if (result_sign < 0) result = -result;
            return result;
        }

        // Decompose the Int's 64-bit words into 16-bit words
        // (Use a smaller word size to preserve FFT precision)
        constexpr size_t WORD_BITS = 16;
        constexpr size_t WORDS_PER_U64 = 64 / WORD_BITS;  // = 4

        size_t x_u16_words = Nx * WORDS_PER_U64;
        size_t y_u16_words = Ny * WORDS_PER_U64;

        // Determine the FFT size (next power of two)
        size_t result_u16_words = x_u16_words + y_u16_words;
        int nfft = 1;
        while (static_cast (nfft) < result_u16_words) {
            nfft *= 2;
        }

        // Precision-limit check (N <= 2^16)
        if (nfft > (1 << 16)) {
            // Fall back to Toom-Cook when the FFT precision limit is exceeded
            Int result = toomCook3Multiply(absX, absY);
            if (result_sign < 0) result = -result;
            return result;
        }

        // Initialize the FFT
        FFTEngine fft(nfft, FFTDivMode::Inverse);

        // Convert X to a 16-bit word array
        std::vector x_data(nfft, 0.0);
        auto x_words = absX.words();
        for (size_t i = 0; i < Nx; ++i) {
            uint64_t word = x_words[i];
            for (size_t j = 0; j < WORDS_PER_U64; ++j) {
                x_data[i * WORDS_PER_U64 + j] = static_cast (word & 0xFFFF);
                word >>= WORD_BITS;
            }
        }

        // Convert Y to a 16-bit word array
        std::vector y_data(nfft, 0.0);
        auto y_words = absY.words();
        for (size_t i = 0; i < Ny; ++i) {
            uint64_t word = y_words[i];
            for (size_t j = 0; j < WORDS_PER_U64; ++j) {
                y_data[i * WORDS_PER_U64 + j] = static_cast (word & 0xFFFF);
                word >>= WORD_BITS;
            }
        }

        // Real-FFT forward transform
        auto X_freq = fft.real_transform(x_data);
        auto Y_freq = fft.real_transform(y_data);

        // Pointwise multiplication in the frequency domain
        std::vector > Z_freq(nfft / 2 + 1);
        for (int i = 0; i <= nfft / 2; ++i) {
            Z_freq[i] = X_freq[i] * Y_freq[i];
        }

        // Real-FFT inverse transform
        auto z_data = fft.real_inverse(Z_freq);

        // Convert from 16-bit words to 64-bit words (with carry handling)
        std::vector result_words;
        result_words.reserve((nfft / WORDS_PER_U64) + 1);

        uint64_t carry = 0;
        uint64_t current_word = 0;
        size_t shift = 0;

        for (int i = 0; i < nfft; ++i) {
            // Round to the nearest integer
            uint64_t val = static_cast (z_data[i] + 0.5) + carry;

            // Process as a 16-bit word
            uint64_t low16 = val & 0xFFFF;
            carry = val >> WORD_BITS;

            // Pack into the 64-bit word
            current_word |= (low16 << shift);
            shift += WORD_BITS;

            if (shift == 64) {
                result_words.push_back(current_word);
                current_word = 0;
                shift = 0;
            }
        }

        // Handle the final carry
        if (shift > 0) {
            current_word |= (carry << shift);
            result_words.push_back(current_word);
            // Handle the case where the high bits of carry overflow by shift
            if (shift < 64) {
                uint64_t remaining = carry >> (64 - shift);
                if (remaining > 0) {
                    result_words.push_back(remaining);
                }
            }
        } else if (carry > 0) {
            result_words.push_back(carry);
        }

        // Strip leading zeros
        while (!result_words.empty() && result_words.back() == 0) {
            result_words.pop_back();
        }

        if (result_words.empty()) {
            return Int(0);
        }

        Int result = Int::fromRawWords(result_words, result_sign);
        return result;
    }

    // Basic multiplication (O(n^2) long multiplication, mpn-based)
    // Used for small inputs or those below the Karatsuba threshold
    static Int basecaseMultiply(const Int& X, const Int& Y) {
        int result_sign = X.getSign() * Y.getSign();
        if (result_sign == 0) return Int(0);

        const uint64_t* a = X.m_words.data();
        size_t an = X.m_words.size();
        const uint64_t* b = Y.m_words.data();
        size_t bn = Y.m_words.size();

        if (an == 0 || bn == 0) return Int(0);

        std::vector product(an + bn);
        mpn::mul_basecase(product.data(), a, an, b, bn);

        size_t actual = mpn::normalized_size(product.data(), product.size());
        if (actual == 0) return Int(0);
        product.resize(actual);

        return Int::fromRawWords(product, result_sign);
    }

};

} // namespace sangi