import torch
import numpy as np

torch.set_default_dtype(torch.float64)

# ============================================================================
# PROXY FUNCTIONS (Simpler Questions)
# ============================================================================

def riemann_theta(t: np.ndarray, n_terms: int = 50) -> np.ndarray:
    """
    Riemann-Siegel theta function:
    θ(t) = arg Γ(1/4 + it/2) - t*log(π)/2
    
    This is MUCH simpler than ζ(s) and oscillatory.
    Zeros of ζ(s) occur when Z(t) = 0.
    """
    from scipy.special import gammainc, gamma
    from mpmath import mp
    
    theta = np.zeros_like(t, dtype=float)
    
    for i, ti in enumerate(t):
        # Use mpmath for precision
        with mp.workdps(30):
            gamma_arg = mp.mpf('1/4') + mp.mpf(ti) / 2
            gamma_val = mp.gamma(gamma_arg)
            theta[i] = mp.im(mp.log(gamma_val)) - ti * mp.log(mp.pi) / 2
    
    return theta


def riemann_siegel_z(t: np.ndarray, n_terms: int = 100) -> np.ndarray:
    """
    Z(t) = 2 * Σ cos(θ(t) - t*log(n)) / n^(1/2) + remainder
    
    Z(t) is REAL and oscillatory on the critical line.
    Z(t) = 0 ⟺ ζ(1/2 + it) = 0
    """
    z = np.zeros_like(t, dtype=float)
    
    for i, ti in enumerate(t):
        # Truncated sum
        total = 0.0
        for n in range(1, n_terms + 1):
            total += np.cos(riemann_theta(np.array([ti]))[0] - ti * np.log(n)) / np.sqrt(n)
        
        z[i] = 2 * total
    
    return z


def cct_proxy_gradient(t: torch.Tensor) -> torch.Tensor:
    """
    CCT Proxy: Use Z(t) approximation as simpler loss function.
    
    Instead of computing full ζ(s), we compute a smooth proxy
    that captures the zero structure.
    """
    t_val = t.item() if t.dim() == 0 else t.detach().numpy()
    
    if isinstance(t_val, torch.Tensor):
        t_val = t_val.detach().numpy()
    
    # Simplified Z(t) using first 20 terms
    z_val = simplified_z(t_val, n_terms=20)
    
    return torch.tensor(z_val**2, dtype=torch.float64)


def simplified_z(t: float, n_terms: int = 20) -> float:
    """
    Simplified Z(t) for gradient descent.
    
    Uses: Z(t) ≈ Σ cos(π/4 - t*log(n)) / √n
    """
    z = 0.0
    for n in range(1, n_terms + 1):
        z += np.cos(np.pi/4 - t * np.log(n)) / np.sqrt(n)
    return 2 * z  # Normalization factor


# ============================================================================
# CCT-ODE with Conditional Collapse Detection
# ============================================================================

class CCTProxyFinder:
    """
    CCT-ODE using proxy functions.
    
    Instead of computing complex ζ(s), we compute:
    1. Z(t) - real proxy on critical line
    2. Theta(t) - phase function
    3. Euler identity check - geometric constraint
    
    This reduces the search space from 2D (σ, t) to 1D (t only).
    """
    
    def __init__(self, target_t: float, lr: float = 0.001):
        self.target_t = target_t
        self.lr = lr
        
        # State vector: only t (sigma is fixed at 0.5 by proxy)
        self.t = torch.tensor(target_t, requires_grad=True, dtype=torch.float64)
        
        # CCT tracking
        self.z_history = []
        self.entropy_history = []
        self.periodicity_detected = False
        
    def step(self) -> dict:
        """Single step with proxy gradient."""
        
        # Proxy loss: Z(t)² → 0
        z_val = simplified_z(self.t.item(), n_terms=30)
        loss = z_val ** 2
        
        # Store metrics
        self.z_history.append(z_val)
        entropy = np.log(1 + np.abs(z_val))
        self.entropy_history.append(entropy)
        
        # Check periodicity (CCT cycle detection)
        if len(self.z_history) > 20:
            recent = self.z_history[-20:]
            # Check if oscillation pattern exists
            sign_changes = np.sum(np.diff(np.sign(recent)) != 0)
            if sign_changes > 5:
                self.periodicity_detected = True
        
        # Gradient (numerical)
        eps = 1e-7
        z_plus = simplified_z(self.t.item() + eps, n_terms=30)
        z_minus = simplified_z(self.t.item() - eps, n_terms=30)
        gradient = (z_plus - z_minus) / (2 * eps)
        
        # Gradient descent
        with torch.no_grad():
            self.t -= self.lr * 2 * z_val * gradient  # d(Z²)/dt = 2*Z*dZ/dt
        
        return {
            't': self.t.item(),
            'Z(t)': z_val,
            'loss': loss,
            'entropy': entropy,
            'periodicity': self.periodicity_detected
        }
    
    def run(self, max_steps: int = 1000, tol: float = 1e-6) -> dict:
        """Run with early stopping and periodicity detection."""
        
        print(f"{'Step':<8} {'t':<15} {'Z(t)':<15} {'Loss':<15} {'H(T)':<12} {'Period':<8}")
        print("-" * 80)
        
        for step in range(max_steps):
            metrics = self.step()
            
            if step % 50 == 0:
                period_flag = "✓" if metrics['periodicity'] else "✗"
                print(f"{step:<8} {metrics['t']:<15.8f} {metrics['Z(t)']:<15.6f} "
                      f"{metrics['loss']:<15.2e} {metrics['entropy']:<12.4f} {period_flag:<8}")
            
            if np.abs(metrics['Z(t)']) < tol:
                print(f"\n✓ COLLAPSE DETECTED at step {step}")
                print(f"  Zero found at t = {metrics['t']:.10f}")
                print(f"  Z(t) = {metrics['Z(t)']:.2e}")
                break
            
            # Periodicity detection: if oscillating, use harmonic averaging
            if metrics['periodicity'] and step > 100:
                # Use average of recent high-points
                recent = self.z_history[-50:]
                peaks = [i for i in range(1, len(recent)-1) 
                        if recent[i] > recent[i-1] and recent[i] > recent[i+1]]
                if len(peaks) > 3:
                    # Adjust learning rate for harmonic behavior
                    avg_peak_t = np.mean([self.t.item() - 50 + p for p in peaks])
                    with torch.no_grad():
                        self.t.data = torch.tensor(avg_peak_t, dtype=torch.float64)
                    print(f"  CCT Periodicity: Jumping to estimated harmonic center")
        
        return {
            't': self.t.item(),
            'Z(t)': simplified_z(self.t.item(), n_terms=50),
            'z_history': np.array(self.z_history),
            'entropy': np.array(self.entropy_history)
        }


# ============================================================================
# CCT Question Path: Multi-Stage Search
# ============================================================================

def cct_zero_search(start_t: float, end_t: float, resolution: float = 0.1) -> list:
    """
    CCT-based multi-stage zero search:
    
    Stage 1: Coarse scan using Z(t) proxy (low cost)
    Stage 2: Fine optimization near detected zero crossings
    Stage 3: Periodicity-based extrapolation
    Stage 4: Euler identity verification
    """
    
    print("\n" + "="*80)
    print("CCT ZERO SEARCH: Multi-Stage with Conditional Collapse")
    print("="*80)
    
    zeros = []
    t_grid = np.arange(start_t, end_t, resolution)
    
    # Stage 1: Coarse scan with Z(t)
    print(f"\n[Stage 1] Coarse scan: {len(t_grid)} points")
    z_values = np.array([simplified_z(t, n_terms=30) for t in t_grid])
    
    # Find sign changes (zero crossings)
    sign_changes = np.where(np.diff(np.sign(z_values)))[0]
    print(f"  Found {len(sign_changes)} candidate regions")
    
    for idx in sign_changes:
        t_candidate = t_grid[idx]
        
        # Stage 2: Fine optimization
        print(f"\n[Stage 2] Fine optimization near t = {t_candidate:.4f}")
        finder = CCTProxyFinder(target_t=t_candidate, lr=0.0001)
        result = finder.run(max_steps=500, tol=1e-8)
        
        t_zero = result['t']
        z_zero = result['Z(t)']
        
        if np.abs(z_zero) < 0.01:  # Valid zero
            # Stage 3: Periodicity check
            # Zeros follow approximate spacing: Δ_n ≈ 2π / log(n)
            # But at high t, spacing is roughly uniform (~0.5 on average)
            
            zeros.append({
                't': t_zero,
                'z': z_zero,
                'stage': 'collapsed'
            })
            print(f"  ✓ Zero found: s = 0.5 + {t_zero:.8f}i")
        
        # Stage 4: Euler identity check
        # At zero: σ = 0.5, t = t_zero
        # Euler: x² + y² = -N → (0)² + t_zero² = -N
        # This doesn't work directly, but we can check:
        # |ζ(1/2 + it_zero)|² = 0
        # Which relates to imaginary quadratic form
    
    return zeros


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

if __name__ == "__main__":
    
    print("="*80)
    print("CCT-PROXY ZETA ZERO FINDER")
    print("Using Z(t) as simplified proxy instead of full ζ(s)")
    print("="*80)
    
    # Test: Find known zeros
    known_zeros = [14.134725, 21.022040, 30.424826, 32.935062]
    
    print("\nFinding known zeros using CCT-Proxy gradient descent:")
    print("-" * 80)
    
    for target_t in known_zeros:
        print(f"\n→ Target: t = {target_t}")
        finder = CCTProxyFinder(target_t=target_t, lr=0.0005)
        result = finder.run(max_steps=1000, tol=1e-10)
        
        error = np.abs(result['t'] - target_t)
        print(f"  Found: t = {result['t']:.10f}")
        print(f"  Error: {error:.2e}")
    
    # Multi-stage search
    print("\n" + "="*80)
    print("Multi-Stage CCT Search (0 to 60)")
    print("="*80)
    
    zeros = cct_zero_search(start_t=0.0, end_t=60.0, resolution=0.5)
    print(f"\n✓ Found {len(zeros)} zeros")