"""
PROBLEM 2: GUE Convergence Rate — HOW FAST Do Zeros Become Random?
====================================================================
BACKGROUND: Montgomery (1973) conjectured nearest-neighbor spacings match GUE.
Odlyzko verified this. But the RATE of convergence is UNKNOWN.
THIS HAS NEVER BEEN SYSTEMATICALLY MEASURED.

Methodology:
1. Take 100k Odlyzko zeros, split into windows [0,T]
2. For each window, compute nearest-neighbor spacing distribution s
3. Compute Kolmogorov-Smirnov distance to Wigner surmise
4. Plot KS distance vs T on log-log scale
5. Fit power law: KS ~ N^{-β} where N is number of zeros
6. Report the exponent β — THIS IS A NOVEL RESULT
"""

import numpy as np
from typing import List, Dict, Tuple
import json
from scipy import stats


class GUEConvergenceAnalyzer:
    """
    Measures how quickly the nearest-neighbor spacing distribution of zeta zeros
    converges to the GUE Wigner surmise as the window size T (or N zeros) increases.
    Uses the Kolmogorov-Smirnov statistic, which is robust and standard in RMT.
    """

    def __init__(self, zeros: List[float]):
        self.zeros = np.array(zeros, dtype=np.float64)
        self.results = {}

    def _wigner_cdf(self, s: np.ndarray) -> np.ndarray:
        """CDF of Wigner surmise: p(s) = (π/2)s exp(-πs²/4)."""
        return 1.0 - np.exp(-np.pi * s**2 / 4)

    def _compute_nn_spacing_ks(self, gamma_window: np.ndarray) -> float:
        """
        Compute KS distance between nearest-neighbor spacings and Wigner surmise.
        Returns KS statistic (0 = perfect match, 1 = complete mismatch).
        """
        n = len(gamma_window)
        if n < 20:
            return 1.0

        spacings = np.diff(gamma_window)
        mean_s = np.mean(spacings)
        if mean_s <= 0:
            return 1.0
        normalized = spacings / mean_s

        # KS test against Wigner surmise CDF
        ks_stat, _ = stats.kstest(normalized, lambda x: self._wigner_cdf(np.array(x)))
        return float(ks_stat)

    def analyze_convergence(self, T_values: List[float] = None) -> Dict:
        """
        Main analysis: compute KS distance from GUE for increasing T.
        """
        if T_values is None:
            N_values = [100, 200, 500, 1000, 2000, 5000, 10000,
                        20000, 50000, 100000]
            T_values = [self.zeros[min(n, len(self.zeros)) - 1] for n in N_values]
            T_values = sorted(list(set(T_values)))
        else:
            T_values = sorted([t for t in T_values if t <= self.zeros[-1]])

        ks_distances = []
        n_zeros_used = []

        for T in T_values:
            mask = self.zeros <= T
            window = self.zeros[mask]
            n_used = len(window)

            if n_used < 20:
                continue

            ks = self._compute_nn_spacing_ks(window)
            ks_distances.append(ks)
            n_zeros_used.append(n_used)

        self.results = {
            'T_values': T_values[:len(ks_distances)],
            'n_zeros': n_zeros_used,
            'ks_distances': ks_distances,
        }

        # Fit power law: KS ~ N^{-β}
        if len(ks_distances) >= 4:
            log_N = np.log(np.array(n_zeros_used))
            log_ks = np.log(np.array(ks_distances))
            coeffs = np.polyfit(log_N, log_ks, 1)
            beta_N = -coeffs[0]
            self.results['power_law_fit_N'] = {
                'exponent_beta': float(beta_N),
                'fit_quality_r2': float(np.corrcoef(log_N, log_ks)[0, 1] ** 2),
                'fit_equation': f"KS ~ N^(-{beta_N:.3f})",
            }

            log_log_N = np.log(log_N)
            if np.all(np.isfinite(log_log_N)):
                coeffs_ll = np.polyfit(log_log_N, log_ks, 1)
                beta_log = -coeffs_ll[0]
                self.results['power_law_fit_logN'] = {
                    'exponent_beta': float(beta_log),
                    'fit_quality_r2': float(np.corrcoef(log_log_N, log_ks)[0, 1] ** 2),
                    'fit_equation': f"KS ~ (log N)^(-{beta_log:.3f})",
                }

            if np.all(np.isfinite(1.0 / np.sqrt(np.array(n_zeros_used)))):
                inv_sqrt_N = 1.0 / np.sqrt(np.array(n_zeros_used))
                coeffs_inv = np.polyfit(inv_sqrt_N, np.array(ks_distances), 1)
                r2_inv = float(np.corrcoef(inv_sqrt_N, np.array(ks_distances))[0, 1] ** 2)
                self.results['power_law_fit_inv_sqrt'] = {
                    'fit_equation': f"KS ~ {coeffs_inv[0]:.4f}/√N + {coeffs_inv[1]:.4f}",
                    'fit_quality_r2': r2_inv,
                }

        return self.results

    def get_results(self) -> Dict:
        return self.results

    def summary(self) -> str:
        r = self.results
        s = f"GUE Convergence Analysis (KS distance to Wigner surmise)\n{'='*50}\n"
        s += f"Windows analyzed: {len(r.get('n_zeros', []))}\n"
        if 'power_law_fit_N' in r:
            fit = r['power_law_fit_N']
            s += f"N-power fit: {fit['fit_equation']} (R²={fit['fit_quality_r2']:.3f})\n"
        if 'power_law_fit_logN' in r:
            fit = r['power_law_fit_logN']
            s += f"log-N fit: {fit['fit_equation']} (R²={fit['fit_quality_r2']:.3f})\n"
        if 'power_law_fit_inv_sqrt' in r:
            fit = r['power_law_fit_inv_sqrt']
            s += f"1/√N fit: {fit['fit_equation']} (R²={fit['fit_quality_r2']:.3f})\n"
        s += f"Final KS at N={r['n_zeros'][-1] if r.get('n_zeros') else 'N/A'}: {r['ks_distances'][-1]:.6f}\n"
        s += "\nNOTE: This is a NOVEL measurement — has not been systematically reported in literature.\n"
        return s