clock-curve-math 1.1.3

#!/usr/bin/env python3
"""
Complete Complex FFT/NTT Implementation for Production Use

This module provides mathematically complete and robust implementations of:
1. Fast Fourier Transform (FFT) using Cooley-Tukey algorithm
2. Number Theoretic Transform (NTT) for modular arithmetic
3. Polynomial multiplication with proper numerical stability
4. Large integer multiplication using FFT/NTT

Features:
- Proper complex arithmetic handling
- Numerical stability and precision control
- Support for arbitrary polynomial degrees
- Modular arithmetic with NTT
- Comprehensive error handling and validation
"""

import math
import cmath
import numpy as np
from typing import List, Tuple, Optional, Union
from dataclasses import dataclass


@dataclass
class ComplexNumber:
    """Complex number representation for precise arithmetic."""
    real: float
    imag: float

    def __add__(self, other):
        if isinstance(other, ComplexNumber):
            return ComplexNumber(self.real + other.real, self.imag + other.imag)
        elif isinstance(other, (int, float)):
            return ComplexNumber(self.real + other, self.imag)
        return NotImplemented

    def __sub__(self, other):
        if isinstance(other, ComplexNumber):
            return ComplexNumber(self.real - other.real, self.imag - other.imag)
        elif isinstance(other, (int, float)):
            return ComplexNumber(self.real - other, self.imag)
        return NotImplemented

    def __mul__(self, other):
        if isinstance(other, ComplexNumber):
            return ComplexNumber(
                self.real * other.real - self.imag * other.imag,
                self.real * other.imag + self.imag * other.real
            )
        elif isinstance(other, (int, float)):
            return ComplexNumber(self.real * other, self.imag * other)
        return NotImplemented

    def __truediv__(self, other):
        if isinstance(other, (int, float)):
            return ComplexNumber(self.real / other, self.imag / other)
        return NotImplemented

    def conjugate(self):
        return ComplexNumber(self.real, -self.imag)

    def magnitude_squared(self):
        return self.real * self.real + self.imag * self.imag

    def magnitude(self):
        return math.sqrt(self.magnitude_squared())

    def phase(self):
        return math.atan2(self.imag, self.real)

    @classmethod
    def from_polar(cls, magnitude, phase):
        return cls(magnitude * math.cos(phase), magnitude * math.sin(phase))

    @classmethod
    def exp(cls, angle):
        return cls(math.cos(angle), math.sin(angle))

    def __str__(self):
        if self.imag >= 0:
            return f"{self.real:.6f} + {self.imag:.6f}i"
        else:
            return f"{self.real:.6f} - {-self.imag:.6f}i"

    def __repr__(self):
        return f"ComplexNumber({self.real}, {self.imag})"


class FFTComplex:
    """Complete FFT implementation with complex arithmetic."""

    @staticmethod
    def is_power_of_two(n: int) -> bool:
        """Check if n is a power of 2."""
        return n > 0 and (n & (n - 1)) == 0

    @staticmethod
    def next_power_of_two(n: int) -> int:
        """Find the smallest power of 2 >= n."""
        if n <= 1:
            return 1
        return 1 << (n - 1).bit_length()

    @staticmethod
    def bit_reverse_permutation(arr: List[ComplexNumber]) -> List[ComplexNumber]:
        """Perform bit-reversal permutation on array."""
        n = len(arr)
        result = arr.copy()

        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]

        return result

    def cooley_tukey_fft(self, coeffs: List[Union[int, float, ComplexNumber]],
                        inverse: bool = False) -> List[ComplexNumber]:
        """
        Complete Cooley-Tukey FFT implementation.

        Args:
            coeffs: Input coefficients (real or complex)
            inverse: If True, compute inverse FFT

        Returns:
            FFT result as ComplexNumber list
        """
        # Convert to ComplexNumber
        complex_coeffs = []
        for c in coeffs:
            if isinstance(c, ComplexNumber):
                complex_coeffs.append(c)
            elif isinstance(c, complex):
                complex_coeffs.append(ComplexNumber(c.real, c.imag))
            else:
                complex_coeffs.append(ComplexNumber(float(c), 0.0))

        n = len(complex_coeffs)

        # Pad to power of 2
        if not self.is_power_of_two(n):
            next_n = self.next_power_of_two(n)
            complex_coeffs.extend([ComplexNumber(0, 0) for _ in range(next_n - n)])
            n = next_n

        # Bit-reversal permutation
        result = self.bit_reverse_permutation(complex_coeffs)

        # Iterative Cooley-Tukey algorithm
        size = 2
        while size <= n:
            half_size = size // 2
            angle_step = (2.0 if inverse else -2.0) * math.pi / size

            for i in range(0, n, size):
                # Compute twiddle factors for this block
                for j in range(half_size):
                    angle = angle_step * j
                    w = ComplexNumber.exp(angle)

                    u = result[i + j]
                    v = result[i + j + half_size] * w

                    result[i + j] = u + v
                    result[i + j + half_size] = u - v

            size *= 2

        # Scale for inverse transform
        if inverse:
            scale = 1.0 / n
            result = [x * scale for x in result]

        return result

    def inverse_fft(self, coeffs: List[Union[int, float, ComplexNumber]]) -> List[ComplexNumber]:
        """Compute inverse FFT."""
        return self.cooley_tukey_fft(coeffs, inverse=True)

    def polynomial_multiply_fft(self, a: List[Union[int, float]],
                               b: List[Union[int, float]]) -> List[int]:
        """
        Multiply polynomials using FFT with proper precision handling.

        Args:
            a, b: Polynomial coefficients

        Returns:
            Product polynomial coefficients
        """
        # Determine FFT size (next power of 2 >= len(a) + len(b) - 1)
        n = self.next_power_of_two(len(a) + len(b) - 1)

        # Pad coefficients with zeros
        a_padded = a + [0] * (n - len(a))
        b_padded = b + [0] * (n - len(b))

        # Forward FFT
        a_fft = self.cooley_tukey_fft(a_padded)
        b_fft = self.cooley_tukey_fft(b_padded)

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

        # Inverse FFT
        c_complex = self.inverse_fft(c_fft)

        # Extract real parts with proper rounding
        result = []
        max_error = 0.0

        for z in c_complex:
            real_val = z.real
            rounded = round(real_val)
            error = abs(real_val - rounded)
            max_error = max(max_error, error)

            # Check for significant precision errors
            if error > 1e-10:
                print(f"Warning: High precision error at coefficient: {real_val} (error: {error})")

            result.append(rounded)

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

        print(f"FFT multiplication completed with max precision error: {max_error}")
        return result


class NTTModular:
    """Number Theoretic Transform for modular arithmetic."""

    def __init__(self, modulus: int):
        """
        Initialize NTT with given modulus.

        Args:
            modulus: Prime modulus for NTT
        """
        self.modulus = modulus
        self.params = self._find_ntt_parameters()

    def _find_ntt_parameters(self) -> dict:
        """Find NTT parameters for the modulus."""
        modulus = self.modulus

        # Find primitive root
        for g in range(2, modulus):
            if pow(g, modulus - 1, modulus) == 1:
                # Check if primitive root
                is_primitive = True
                order = modulus - 1
                for i in range(2, int(math.sqrt(order)) + 1):
                    if order % i == 0:
                        if pow(g, i, modulus) == 1 or pow(g, order // 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 NTT length
        max_length = 1
        while (modulus - 1) % (2 * max_length) == 0:
            max_length *= 2

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

    def ntt(self, coeffs: List[int], inverse: bool = False) -> List[int]:
        """
        Number Theoretic Transform.

        Args:
            coeffs: Input coefficients
            inverse: If True, compute inverse NTT

        Returns:
            NTT result
        """
        n = len(coeffs)
        if not FFTComplex.is_power_of_two(n):
            raise ValueError("NTT length must be power of 2")

        if n > self.params['max_length']:
            raise ValueError(f"NTT length {n} exceeds maximum {self.params['max_length']}")

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

        # Compute root of unity
        if inverse:
            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)

        # Bit-reversal permutation
        result = coeffs.copy()
        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]

        # Iterative NTT
        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 for inverse
        if inverse:
            n_inv = pow(n, modulus - 2, modulus)
            result = [(x * n_inv) % modulus for x in result]

        return result

    def polynomial_multiply_ntt(self, a: List[int], b: List[int]) -> List[int]:
        """
        Multiply polynomials using NTT.

        Args:
            a, b: Polynomial coefficients

        Returns:
            Product polynomial coefficients
        """
        n = FFTComplex.next_power_of_two(len(a) + len(b) - 1)

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

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

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

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

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

        return c


class LargeIntegerMultiplier:
    """Complete large integer multiplication using FFT/NTT."""

    def __init__(self, use_ntt: bool = False, ntt_modulus: Optional[int] = None):
        """
        Initialize multiplier.

        Args:
            use_ntt: Whether to use NTT instead of FFT
            ntt_modulus: Modulus for NTT (required if use_ntt=True)
        """
        self.use_ntt = use_ntt
        self.fft = FFTComplex()

        if use_ntt:
            if ntt_modulus is None:
                ntt_modulus = 998244353  # Good default NTT prime
            self.ntt = NTTModular(ntt_modulus)
        else:
            self.ntt = None

    def _big_int_to_coeffs(self, n: int, base: int = 10**9) -> List[int]:
        """Convert big integer to coefficient array."""
        if n == 0:
            return [0]

        coeffs = []
        while n > 0:
            coeffs.append(n % base)
            n //= base
        return coeffs

    def _coeffs_to_big_int(self, coeffs: List[int], base: int = 10**9) -> int:
        """Convert coefficient array back to big integer."""
        result = 0
        for i, coeff in enumerate(coeffs):
            result += coeff * (base ** i)
        return result

    def multiply(self, a: int, b: int) -> int:
        """
        Multiply two large integers.

        Args:
            a, b: Integers to multiply

        Returns:
            Product a * b
        """
        if self.use_ntt and self.ntt:
            return self._multiply_ntt(a, b)
        else:
            return self._multiply_fft(a, b)

    def _multiply_fft(self, a: int, b: int) -> int:
        """Multiply using FFT."""
        base = 10**9  # Use 10^9 for good precision

        a_coeffs = self._big_int_to_coeffs(a, base)
        b_coeffs = self._big_int_to_coeffs(b, base)

        product_coeffs = self.fft.polynomial_multiply_fft(a_coeffs, b_coeffs)

        return self._coeffs_to_big_int(product_coeffs, base)

    def _multiply_ntt(self, a: int, b: int) -> int:
        """Multiply using NTT."""
        # For NTT, we need to choose base carefully
        modulus = self.ntt.modulus
        base = modulus // 100  # Conservative base

        a_coeffs = self._big_int_to_coeffs(a, base)
        b_coeffs = self._big_int_to_coeffs(b, base)

        product_coeffs = self.ntt.polynomial_multiply_ntt(a_coeffs, b_coeffs)

        return self._coeffs_to_big_int(product_coeffs, base)


def comprehensive_test():
    """Comprehensive test of the FFT/NTT implementation."""
    print("=== Comprehensive FFT/NTT Test ===\n")

    # Test ComplexNumber
    print("Testing ComplexNumber arithmetic...")
    z1 = ComplexNumber(1, 2)
    z2 = ComplexNumber(3, 4)
    print(f"z1 = {z1}")
    print(f"z2 = {z2}")
    print(f"z1 + z2 = {z1 + z2}")
    print(f"z1 * z2 = {z1 * z2}")
    print(f"z1 magnitude = {z1.magnitude()}")
    print()

    # Test FFT
    print("Testing FFT polynomial multiplication...")
    fft = FFTComplex()

    test_cases = [
        ([1, 2, 3], [4, 5]),  # Expected: [4, 13, 22, 15]
        ([1], [1]),            # Expected: [1]
        ([1, 1], [1, -1]),     # Expected: [1, 0, -1]
    ]

    for a, b in test_cases:
        result = fft.polynomial_multiply_fft(a, b)
        print(f"  {a} * {b} = {result}")

    print()

    # Test large integer multiplication
    print("Testing large integer multiplication...")

    # FFT multiplier
    fft_mult = LargeIntegerMultiplier(use_ntt=False)

    big_tests = [
        (123, 456),
        (1234, 5678),
        (10**6, 10**6),
        (2**50, 2**50),
    ]

    print("FFT Results:")
    for a, b in big_tests:
        expected = a * b
        result = fft_mult.multiply(a, b)
        correct = (result == expected)
        print(f"  {a} * {b} = {result} {'✓' if correct else '✗'}")

    print()

    # NTT multiplier
    print("NTT Results:")
    ntt_mult = LargeIntegerMultiplier(use_ntt=True, ntt_modulus=998244353)

    for a, b in [(12, 34), (56, 78), (123, 456)]:
        expected = a * b
        result = ntt_mult.multiply(a, b)
        correct = (result == expected)
        print(f"  {a} * {b} = {result} {'✓' if correct else '✗'}")

    print("\n=== Test Complete ===")


if __name__ == "__main__":
    comprehensive_test()