"""
Singularity Factorization via CCT-ODE Gradient Descent
=======================================================
Finds prime factors p1, p2 of c = p1 * p2 by flowing through
the continuous gap space until the singularity condition is met:
    
    d² + 4c = s²    (where d = p1 - p2)

The "void" between integers is filled by probability P(d),
flowing via gradient descent until collapse to discrete primes.
"""

import numpy as np
from dataclasses import dataclass, field
from typing import Optional, Tuple, List
import matplotlib.pyplot as plt


# =============================================================================
# CORE MATHEMATICS
# =============================================================================

def compute_singularity(d: float, c: int) -> Tuple[float, float, float]:
    """
    Compute the singularity proximity metric.
    
    Parameters:
        d: continuous gap (p1 - p2)
        c: composite number (p1 * p2)
    
    Returns:
        s: sqrt(d² + 4c) - the "square root variable"
        S: singularity error (should → 0 at collapse)
        w: weight (probability of being near singularity)
    """
    s = np.sqrt(d**2 + 4 * c)
    s_rounded = round(s)
    
    # Singularity error: deviation from perfect square
    S = d**2 + 4 * c - s_rounded**2
    
    # Weight: Gaussian proximity to singularity
    # σ² controls basin width - larger = wider search
    sigma_sq = max(1.0, c / 1e6)  # Adaptive sigma
    w = np.exp(-S**2 / (2 * sigma_sq))
    
    return s, S, w


def compute_gradient(d: float, c: int, eta: float = 0.1) -> float:
    """
    Compute gradient for gradient descent.
    
    dS/dd = 2d  (derivative of singularity error)
    
    The gradient points toward/away from singularity based on S sign.
    """
    s = np.sqrt(d**2 + 4 * c)
    s_rounded = round(s)
    S = d**2 + 4 * c - s_rounded**2
    
    # Gradient: dS/dd = 2d
    # But we want to minimize |S|, so:
    # If S > 0: need smaller d (d decreases)
    # If S < 0: need larger d (d increases)
    grad_S = 2 * d
    
    # Adaptive step size based on proximity to singularity
    step = eta * grad_S
    
    # Clamp to prevent overflow
    step = np.clip(step, -1e6, 1e6)
    
    return step


def is_prime(n: int) -> bool:
    """Primality test."""
    if n < 2:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    for i in range(3, int(np.sqrt(n)) + 1, 2):
        if n % i == 0:
            return False
    return True


def extract_factors(d: float, c: int) -> Optional[Tuple[int, int]]:
    """
    Extract prime factors from gap d.
    
    From:
        a * b = 4c
        d = (b - a) / 2
        s = sqrt(d² + 4c)
    
    We get:
        a = s - d
        b = s + d
        p2 = a / 2
        p1 = b / 2
    """
    s, _, _ = compute_singularity(d, c)
    s = round(s)
    d_round = round(d)
    
    a = s - d_round
    b = s + d_round
    
    # Check divisibility
    if a % 2 != 0 or b % 2 != 0:
        return None
    
    p2 = a // 2
    p1 = b // 2
    
    # Verify primality
    if is_prime(p1) and is_prime(p2):
        return (p1, p2)
    
    return None


# =============================================================================
# PROBABILITY DISTRIBUTION OVER GAP SPACE
# =============================================================================

@dataclass
class GapProbabilityField:
    """
    Probability field P(d) filling the void between integers.
    
    The gap space [0, sqrt(c)] is dense with probability mass
    that flows toward the singularity via gradient descent.
    """
    c: int
    num_samples: int = 1000
    samples: np.ndarray = field(init=False)
    weights: np.ndarray = field(init=False)
    
    def __post_init__(self):
        # Initialize samples as Gaussian centered near sqrt(c)/2
        # (prior: primes are roughly equal size)
        d_max = int(np.sqrt(self.c)) + 1
        mean = d_max / 4  # Expect small gap for random primes
        std = d_max / 2
        
        self.samples = np.abs(np.random.normal(mean, std, self.num_samples))
        self.weights = np.ones(self.num_samples)
    
    def compute_weights(self):
        """Update weights based on singularity proximity."""
        for i, d in enumerate(self.samples):
            _, _, w = compute_singularity(d, self.c)
            self.weights[i] = w
    
    def resample(self):
        """
        Resample from weighted distribution.
        This is the "probability flow" - mass moves toward singularity.
        """
        # Normalize weights
        probs = self.weights / np.sum(self.weights)
        
        # Resample
        indices = np.random.choice(
            self.num_samples, 
            size=self.num_samples, 
            p=probs,
            replace=True
        )
        self.samples = self.samples[indices]
        
        # Add small noise (exploration)
        self.samples += np.random.normal(0, 0.5, self.num_samples)
        self.samples = np.clip(self.samples, 0, int(np.sqrt(self.c)) + 1)
        
        # Reset weights
        self.weights = np.ones(self.num_samples)
    
    def mean_estimate(self) -> float:
        """Weighted mean estimate of the gap."""
        return np.sum(self.samples * self.weights) / np.sum(self.weights)


# =============================================================================
# SINGULARITY FACTORIZATION MAIN CLASS
# =============================================================================

@dataclass
class SingularityFactorization:
    """
    CCT-ODE Factorization via Singularity Gradient Descent.
    
    The algorithm:
    1. Initialize probability field P(d) over gap space
    2. Flow P(d) toward singularity via gradient descent
    3. When singularity condition met, collapse to discrete primes
    
    Attributes:
        c: Composite number to factor (p1 * p2)
        eta: Learning rate for gradient descent
        sigma: Gaussian width for weight computation
        max_iterations: Maximum gradient steps before failure
        verbose: Print progress
    """
    c: int
    eta: float = 0.1
    sigma: float = 1.0
    max_iterations: int = 10000
    verbose: bool = True
    track_history: bool = True
    
    # Internal state
    d: float = 0.0
    p1: int = 0
    p2: int = 0
    converged: bool = False
    history: List[dict] = field(default_factory=list)
    
    def __post_init__(self):
        # Initialize d near sqrt(c)/2 (prior for roughly equal primes)
        self.d = np.sqrt(self.c) / 2 + np.random.uniform(-10, 10)
        self.d = max(0, self.d)
        
        if self.verbose:
            print(f"Initializing factorization of c = {self.c}")
            print(f"Initial gap estimate: d ≈ {self.d:.2f}")
    
    def step(self) -> Tuple[bool, Optional[Tuple[int, int]]]:
        """
        Single gradient descent step.
        
        Returns:
            (converged, factors) - Tuple of convergence status and factors if found
        """
        # Compute singularity metrics
        s, S, w = compute_singularity(self.d, self.c)
        
        # Gradient descent update
        grad = compute_gradient(self.d, self.c, self.eta)
        self.d -= grad  # Descend toward singularity
        
        # Clamp to valid range
        self.d = np.clip(self.d, 0, int(np.sqrt(self.c)) + 1)
        
        # Track history
        if self.track_history:
            self.history.append({
                'iteration': len(self.history),
                'd': self.d,
                's': s,
                'S': S,
                'w': w,
                'gap_int': round(self.d)
            })
        
        # Check for convergence (singularity reached)
        if abs(S) < 1e-6:
            # Try to extract factors
            factors = extract_factors(self.d, self.c)
            if factors is not None:
                self.p1, self.p2 = factors
                self.converged = True
                return True, factors
        
        return False, None
    
    def factorize(self) -> Tuple[int, int]:
        """
        Run full factorization.
        
        Returns:
            (p1, p2) - The prime factors
            
        Raises:
            ValueError: If factorization fails (c is prime or insufficient iterations)
        """
        for i in range(self.max_iterations):
            converged, factors = self.step()
            
            if converged and factors is not None:
                if self.verbose:
                    print(f"\n✓ SINGULARITY COLLAPSED at iteration {i}")
                    print(f"  d* = {self.d:.6f} (gap between primes)")
                    print(f"  p1 = {self.p1}, p2 = {self.p2}")
                    print(f"  Verification: {self.p1} × {self.p2} = {self.p1 * self.p2}")
                return factors
            
            # Progress report
            if self.verbose and i % 1000 == 0:
                s, S, w = compute_singularity(self.d, self.c)
                print(f"  Iter {i:6d}: d = {self.d:.4f}, S = {S:.6f}, w = {w:.6f}")
        
        raise ValueError(f"Failed to factor {self.c} in {self.max_iterations} iterations")


# =============================================================================
# PROBABILISTIC (PARTICLE FILTER) VERSION
# =============================================================================

@dataclass
class ProbabilisticSingularityFactorization:
    """
    Particle-filter version: maintains probability field P(d) over many samples.
    
    More robust than single-point gradient descent, but slower.
    """
    c: int
    num_particles: int = 100
    max_iterations: int = 5000
    exploration_rate: float = 0.1
    verbose: bool = True
    
    particles: np.ndarray = field(init=False)
    weights: np.ndarray = field(init=False)
    
    def __post_init__(self):
        d_max = int(np.sqrt(self.c)) + 1
        self.particles = np.random.uniform(0, d_max, self.num_particles)
        self.weights = np.ones(self.num_particles)
        
        if self.verbose:
            print(f"Initializing {self.num_particles} particles for c = {self.c}")
    
    def step(self) -> Tuple[bool, Optional[Tuple[int, int]]]:
        """Update all particles toward singularity."""
        new_particles = []
        new_weights = []
        
        for particle in self.particles:
            s, S, w = compute_singularity(particle, self.c)
            
            # Gradient descent for this particle
            grad = 2 * particle  # dS/dd
            new_d = particle - 0.05 * grad * (1 - w)  # Stronger push when far from singularity
            
            # Add exploration noise
            new_d += np.random.normal(0, self.exploration_rate)
            new_d = np.clip(new_d, 0, int(np.sqrt(self.c)) + 1)
            
            new_particles.append(new_d)
            new_weights.append(w * self.weights[self.particles.tolist().index(particle)])
        
        self.particles = np.array(new_particles)
        self.weights = np.array(new_weights) / np.sum(new_weights)
        
        # Resample periodically
        if len(self.history if hasattr(self, 'history') else []) % 100 == 0:
            indices = np.random.choice(self.num_particles, size=self.num_particles, p=self.weights, replace=True)
            self.particles = self.particles[indices]
            self.weights = np.ones(self.num_particles)
        
        # Check best particle for convergence
        best_idx = np.argmax(self.weights)
        best_d = self.particles[best_idx]
        factors = extract_factors(best_d, self.c)
        
        return factors is not None, factors
    
    def factorize(self) -> Tuple[int, int]:
        """Run probabilistic factorization."""
        self.history = []
        
        for i in range(self.max_iterations):
            converged, factors = self.step()
            
            if self.verbose and i % 500 == 0:
                best_d = self.particles[np.argmax(self.weights)]
                s, S, w = compute_singularity(best_d, self.c)
                print(f"Iter {i}: mean_d = {np.mean(self.particles):.2f}, best_d = {best_d:.2f}, S = {S:.4f}")
            
            if converged and factors is not None:
                if self.verbose:
                    print(f"✓ COLLAPSED at iteration {i}: {factors}")
                return factors
        
        raise ValueError(f"Failed to factor {self.c}")


# =============================================================================
# VISUALIZATION
# =============================================================================

def plot_singularity_landscape(c: int, d_true: float = None):
    """
    Plot the singularity landscape for visualization.
    
    Shows: S(d) = d² + 4c - round(sqrt(d² + 4c))²
    The singularities are where S(d) = 0.
    """
    d_values = np.linspace(0, np.sqrt(c) + 10, 1000)
    S_values = []
    w_values = []
    
    for d in d_values:
        s, _, w = compute_singularity(d, c)
        S = d**2 + 4*c - round(s)**2
        S_values.append(S)
        w_values.append(w)
    
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 8))
    
    # Singularity error
    ax1.plot(d_values, S_values, 'b-', linewidth=1)
    ax1.axhline(y=0, color='r', linestyle='--', label='Singularity (S=0)')
    ax1.xlabel('Gap d')
    ax1.ylabel('Singularity Error S = d² + 4c - s²')
    ax1.title(f'Singularity Landscape for c = {c}')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    if d_true:
        ax1.axvline(x=d_true, color='g', linestyle=':', linewidth=2, label=f'True gap d* = {d_true}')
        ax1.legend()
    
    # Weight (probability of singularity)
    ax2.plot(d_values, w_values, 'r-', linewidth=1)
    ax2.xlabel('Gap d')
    ax2.ylabel('Weight w = exp(-S²/2σ²)')
    ax2.title('Singularity Probability Field (fills the void)')
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(f'singularity_landscape_c{c}.png', dpi=150)
    plt.show()


def plot_convergence_history(history: List[dict], c: int):
    """Plot the convergence trajectory of the factorization."""
    iterations = [h['iteration'] for h in history]
    d_values = [h['d'] for h in history]
    S_values = [h['S'] for h in history]
    w_values = [h['w'] for h in history]
    
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(14, 10))
    
    # Gap trajectory
    ax1.plot(iterations, d_values, 'b-', linewidth=0.8)
    ax1.set_xlabel('Iteration')
    ax1.set_ylabel('Gap d')
    ax1.set_title('Gap d Convergence (flowing to singularity)')
    ax1.grid(True, alpha=0.3)
    
    # Singularity error
    ax2.plot(iterations, S_values, 'r-', linewidth=0.8)
    ax2.set_xlabel('Iteration')
    ax2.set_ylabel('Singularity Error S')
    ax2.set_title('Singularity Error → 0 (collapse)')
    ax2.grid(True, alpha=0.3)
    
    # Weight (probability)
    ax3.plot(iterations, w_values, 'g-', linewidth=0.8)
    ax3.set_xlabel('Iteration')
    ax3.set_ylabel('Weight w')
    ax3.set_title('Weight w → 1 (probability concentrated)')
    ax3.grid(True, alpha=0.3)
    
    # Phase space
    ax4.scatter(d_values, S_values, c=iterations, cmap='viridis', s=1, alpha=0.5)
    ax4.set_xlabel('Gap d')
    ax4.set_ylabel('Singularity Error S')
    ax4.set_title('Phase Space Trajectory')
    ax4.grid(True, alpha=0.3)
    
    plt.suptitle(f'CCT-ODE Factorization of c = {c}', fontsize=14)
    plt.tight_layout()
    plt.savefig(f'convergence_c{c}.png', dpi=150)
    plt.show()


# =============================================================================
# DEMONSTRATIONS
# =============================================================================

def demo_single_factorization():
    """Demonstrate factorization of a single number."""
    print("=" * 60)
    print("SINGULARITY FACTORIZATION DEMO")
    print("=" * 60)
    
    # Test cases
    test_cases = [
        77,      # 7 × 11
        143,     # 11 × 13
        1001,    # 7 × 11 × 13
        9991,    # 97 × 103
        2701,    # 37 × 73
        99991,   # 317 × 313 (close primes)
    ]
    
    for c in test_cases:
        print(f"\n{'─' * 40}")
        print(f"Factoring c = {c}")
        print(f"{'─' * 40}")
        
        try:
            sf = SingularityFactorization(c, verbose=False, max_iterations=5000)
            p1, p2 = sf.factorize()
            print(f"✓ Result: {p1} × {p2} = {p1 * p2}")
            
            if sf.track_history:
                print(f"  Iterations: {len(sf.history)}")
        except ValueError as e:
            print(f"✗ {e}")


def demo_with_visualization():
    """Factor with full visualization."""
    c = 1001  # 7 × 11 × 13 (will find first pair)
    
    print(f"\nVisualizing factorization of c = {c}")
    print(f"Expected factors: 7 × 11 or 7 × 13 or 11 × 13\n")
    
    # Show the singularity landscape
    plot_singularity_landscape(c)
    
    # Run factorization with history tracking
    sf = SingularityFactorization(c, verbose=True, max_iterations=5000, track_history=True)
    p1, p2 = sf.factorize()
    
    # Plot convergence
    plot_convergence_history(sf.history, c)


def demo_probabilistic():
    """Demonstrate probabilistic (particle filter) version."""
    print("\n" + "=" * 60)
    print("PROBABILISTIC (PARTICLE FILTER) FACTORIZATION")
    print("=" * 60)
    
    c = 99991  # 317 × 313
    
    print(f"\nFactoring c = {c} (larger number)")
    print(f"Expected: 313 × 317 = {313 * 317}\n")
    
    pf = ProbabilisticSingularityFactorization(c, num_particles=200, verbose=True)
    
    try:
        p1, p2 = pf.factorize()
        print(f"\n✓ Result: {p1} × {p2} = {p1 * p2}")
    except ValueError as e:
        print(f"✗ {e}")


# =============================================================================
# MAIN
# =============================================================================

if __name__ == "__main__":
    print("""
    ╔══════════════════════════════════════════════════════════════╗
    ║           SINGULARITY FACTORIZATION: CCT-ODE                 ║
    ║                                                              ║
    ║  The void between integers is filled by probability P(d)     ║
    ║  which flows via gradient descent to the singularity         ║
    ║  where d² + 4c = s² (perfect square)                        ║
    ║                                                              ║
    ║  At collapse: d* → integer → factors (p1, p2)               ║
    ╚══════════════════════════════════════════════════════════════╝
    """)
    
    # Run demonstrations
    demo_single_factorization()
    
    # Uncomment for visualization:
    # demo_with_visualization()
    # demo_probabilistic()