clock-curve-math 1.1.3

#!/usr/bin/env python3
"""
Complete FFT/NTT Implementation for Large Integer Multiplication

This module provides complete implementations of:
1. Fast Fourier Transform (FFT) using Cooley-Tukey algorithm
2. Number Theoretic Transform (NTT) for modular arithmetic
3. Polynomial multiplication for large integer multiplication

These algorithms achieve O(n log n) complexity for multiplication of n-bit numbers.
"""

import math
import cmath


class FFTNTTComplete:
    """Complete FFT/NTT implementation for large integer multiplication."""

    def __init__(self, ntt_modulus=None):
        """
        Initialize FFT/NTT multiplication.

        Args:
            ntt_modulus: Prime modulus for NTT. If None, uses FFT instead.
        """
        self.ntt_modulus = ntt_modulus
        if ntt_modulus:
            self.ntt_params = self._find_ntt_parameters(ntt_modulus)

    def _find_ntt_parameters(self, modulus):
        """
        Find NTT parameters for a given modulus.

        Returns dict with:
        - modulus: The prime modulus
        - primitive_root: A primitive root modulo modulus
        - max_length: Maximum transform length supported
        """
        # Find a primitive root
        for g in range(2, modulus):
            if pow(g, modulus-1, modulus) == 1:
                # Check if it's a primitive root
                is_primitive = True
                for i in range(2, modulus-1):
                    if pow(g, i, modulus) == 1:
                        is_primitive = False
                        break
                if is_primitive:
                    primitive_root = g
                    break
        else:
            raise ValueError(f"No primitive root found for modulus {modulus}")

        # Find maximum power of 2 that divides modulus-1
        max_length = 1
        while (modulus - 1) % (2 * max_length) == 0:
            max_length *= 2

        return {
            'modulus': modulus,
            'primitive_root': primitive_root,
            'max_length': max_length
        }

    def cooley_tukey_fft(self, coeffs, inverse=False):
        """
        Cooley-Tukey Fast Fourier Transform using iterative algorithm.

        Args:
            coeffs: List of complex coefficients
            inverse: If True, compute inverse FFT

        Returns:
            FFT result
        """
        n = len(coeffs)
        if n <= 1:
            return coeffs

        # Pad to next power of 2 if necessary
        if n & (n - 1) != 0:
            next_power = 1 << (n - 1).bit_length()
            coeffs = coeffs + [0] * (next_power - n)
            n = next_power

        # Bit-reversal permutation
        result = [0] * n
        for i in range(n):
            j = 0
            for k in range(int(math.log2(n))):
                if i & (1 << k):
                    j |= 1 << (int(math.log2(n)) - 1 - k)
            result[j] = coeffs[i]

        # Cooley-Tukey iterative algorithm
        size = 2
        while size <= n:
            half_size = size // 2
            angle = (2 if inverse else -2) * math.pi * 1j / size

            for i in range(0, n, size):
                w = 1 + 0j
                for j in range(half_size):
                    u = result[i + j]
                    v = result[i + j + half_size] * w
                    result[i + j] = u + v
                    result[i + j + half_size] = u - v
                    w *= cmath.exp(angle)

            size *= 2

        if inverse:
            result = [x / n for x in result]

        return result

    def number_theoretic_transform(self, coeffs, inverse=False):
        """
        Number Theoretic Transform (NTT) for modular arithmetic.

        Args:
            coeffs: Input coefficients (integers)
            inverse: If True, compute inverse NTT

        Returns:
            NTT of input coefficients
        """
        if not self.ntt_modulus:
            raise ValueError("NTT modulus not set")

        n = len(coeffs)
        if n <= 1:
            return coeffs

        # Check if n is power of 2
        if n & (n - 1) != 0:
            raise ValueError("NTT length must be power of 2")

        params = self.ntt_params
        modulus = params['modulus']
        primitive_root = params['primitive_root']

        # Compute root of unity
        if inverse:
            # For inverse NTT, use w^(-1) where w is the n-th root of unity
            root = pow(primitive_root, (modulus - 1) // n, modulus)
            root = pow(root, modulus - 2, modulus)  # Modular inverse
        else:
            root = pow(primitive_root, (modulus - 1) // n, modulus)

        # Iterative Cooley-Tukey NTT
        result = coeffs[:]  # Copy input

        # Bit-reversal permutation
        for i in range(n):
            j = 0
            for k in range(int(math.log2(n))):
                if i & (1 << k):
                    j |= 1 << (int(math.log2(n)) - 1 - k)
            if j > i:
                result[i], result[j] = result[j], result[i]

        # Cooley-Tukey butterfly operations
        size = 2
        while size <= n:
            half_size = size // 2
            wlen = pow(root, n // size, modulus)

            for i in range(0, n, size):
                w = 1
                for j in range(half_size):
                    u = result[i + j]
                    v = (result[i + j + half_size] * w) % modulus
                    result[i + j] = (u + v) % modulus
                    result[i + j + half_size] = (u - v) % modulus
                    w = (w * wlen) % modulus

            size *= 2

        # Scale by n^(-1) for inverse transform
        if inverse:
            n_inv = pow(n, modulus - 2, modulus)  # Modular inverse of n
            result = [(x * n_inv) % modulus for x in result]

        return result

    def polynomial_multiply_fft(self, a_coeffs, b_coeffs):
        """
        Multiply two polynomials using FFT.

        Args:
            a_coeffs, b_coeffs: Coefficient lists

        Returns:
            Product polynomial coefficients
        """
        n = 1
        while n < len(a_coeffs) + len(b_coeffs) - 1:
            n *= 2

        # Pad coefficients
        a_padded = a_coeffs + [0] * (n - len(a_coeffs))
        b_padded = b_coeffs + [0] * (n - len(b_coeffs))

        # Convert to Python complex
        a_complex = [complex(float(x), 0.0) for x in a_padded]
        b_complex = [complex(float(x), 0.0) for x in b_padded]

        # FFT
        a_fft = self.cooley_tukey_fft(a_complex)
        b_fft = self.cooley_tukey_fft(b_complex)

        # Pointwise multiply
        c_fft = [x * y for x, y in zip(a_fft, b_fft)]

        # Inverse FFT
        c_complex = self.cooley_tukey_fft(c_fft, inverse=True)

        # Extract real parts and round
        result = []
        for x in c_complex:
            real_val = x.real
            rounded = round(real_val)
            result.append(rounded)

        # Remove trailing zeros
        while result and result[-1] == 0:
            result.pop()

        return result

    def polynomial_multiply_ntt(self, a_coeffs, b_coeffs):
        """
        Multiply two polynomials using NTT.

        Args:
            a_coeffs, b_coeffs: Coefficient lists

        Returns:
            Product polynomial coefficients
        """
        if not self.ntt_modulus:
            raise ValueError("NTT modulus not set")

        n = 1
        while n < len(a_coeffs) + len(b_coeffs) - 1:
            n *= 2

        # Pad coefficients
        a_padded = a_coeffs + [0] * (n - len(a_coeffs))
        b_padded = b_coeffs + [0] * (n - len(b_coeffs))

        # NTT
        a_ntt = self.number_theoretic_transform(a_padded)
        b_ntt = self.number_theoretic_transform(b_padded)

        # Pointwise multiplication modulo NTT modulus
        modulus = self.ntt_modulus
        c_ntt = [(x * y) % modulus for x, y in zip(a_ntt, b_ntt)]

        # Inverse NTT
        c = self.number_theoretic_transform(c_ntt, inverse=True)

        # Remove trailing zeros
        while c and c[-1] == 0:
            c.pop()

        return c

    def multiply_big_ints(self, a, b, method='fft', base=10**9):
        """
        Multiply two large integers using specified method.

        Args:
            a, b: Integers to multiply
            method: 'fft' or 'ntt'
            base: Base for coefficient conversion

        Returns:
            Product a * b
        """
        # Convert to base-10^9 digits for better precision
        def to_digits(n):
            if n == 0:
                return [0]
            digits = []
            while n > 0:
                digits.append(n % base)
                n //= base
            return digits

        a_coeffs = to_digits(a)
        b_coeffs = to_digits(b)

        if method == 'fft':
            product_coeffs = self.polynomial_multiply_fft(a_coeffs, b_coeffs)
        elif method == 'ntt':
            product_coeffs = self.polynomial_multiply_ntt(a_coeffs, b_coeffs)
        else:
            raise ValueError(f"Unknown method: {method}")

        # Convert back to integer
        result = 0
        for i, coeff in enumerate(product_coeffs):
            result += coeff * (base ** i)

        return result


def test_correctness():
    """Test correctness of multiplication algorithms."""
    print("Testing FFT/NTT multiplication correctness...")

    # Test FFT
    fft_mult = FFTNTTComplete()

    print("\nTesting FFT:")
    test_cases = [
        (0, 0),
        (1, 1),
        (123, 456),
        (1234, 5678),
        (10**10, 10**10),
    ]

    for a, b in test_cases:
        expected = a * b
        result = fft_mult.multiply_big_ints(a, b, method='fft')
        correct = (result == expected)
        print(f"  {a} * {b}: {'✓' if correct else '✗'}")
        if not correct:
            print(f"    Expected: {expected}")
            print(f"    Got:      {result}")

    # Test NTT with different moduli
    print("\nTesting NTT:")
    ntt_moduli = [13, 17, 97, 257]  # NTT-friendly primes

    for modulus in ntt_moduli:
        try:
            ntt_mult = FFTNTTComplete(ntt_modulus=modulus)
            print(f"\n  Modulus {modulus}:")
            for a, b in [(12, 34), (56, 78)]:  # Smaller test cases for NTT
                expected = a * b
                result = ntt_mult.multiply_big_ints(a, b, method='ntt')
                correct = (result == expected)
                print(f"    {a} * {b}: {'✓' if correct else '✗'}")
        except Exception as e:
            print(f"    Modulus {modulus}: Error - {e}")


def benchmark():
    """Benchmark multiplication algorithms."""
    print("\nBenchmarking multiplication algorithms...")
    print("Digits | FFT Time (s) | NTT Time (s)")
    print("-" * 35)

    import time

    fft_mult = FFTNTTComplete()
    ntt_mult = FFTNTTComplete(ntt_modulus=998244353)  # NTT-friendly prime

    sizes = [10, 100, 1000]

    for digits in sizes:
        # Generate test numbers
        a = 10**digits - 1  # All 9's
        b = 10**(digits//2) - 1  # Smaller number

        # Benchmark FFT
        start = time.time()
        fft_result = fft_mult.multiply_big_ints(a, b, method='fft')
        fft_time = time.time() - start

        # Benchmark NTT
        start = time.time()
        ntt_result = ntt_mult.multiply_big_ints(a, b, method='ntt')
        ntt_time = time.time() - start

        # Verify correctness
        expected = a * b
        fft_correct = (fft_result == expected)
        ntt_correct = (ntt_result == expected)

        print(">6"
              ">12.3f"
              ">12.3f"
              f"{' ✓' if fft_correct and ntt_correct else ' ✗'}")


if __name__ == "__main__":
    test_correctness()
    benchmark()