"""
Unified Zeta Engine (v1 + v2 combined)
- Zero computation (v1)
- Spectral analysis: pair correlation, GUE, Δ₃, form factor (v2)
- Mertens function (v2)
- Chebyshev ψ(x) (v2)
- Operator candidate search (v2)
"""

import numpy as np
from mpmath import mp, mpf, mpc, zetazero, zeta, gamma, pi, log, exp, sqrt, fabs
from mpmath import re as mp_re, im as mp_im, siegeltheta
from typing import List, Dict, Tuple
from dataclasses import dataclass, field
from datetime import datetime
import json


@dataclass
class ZetaZero:
    index: int
    imaginary_part: float
    real_part: float = 0.5
    verified: bool = False
    precision: int = 50
    def to_complex(self):
        return complex(self.real_part, self.imaginary_part)


class UnifiedZetaEngine:
    def __init__(self, precision: int = 50, zeros_file: str = None):
        mp.dps = precision
        self.precision = precision
        self._zeros_cache: Dict[int, ZetaZero] = {}
        self._log: List[str] = []
        self.zeros_file = zeros_file
        self.loaded_zeros = []
        if zeros_file:
            self._load_zeros_file(zeros_file)

    def log(self, msg: str):
        self._log.append(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}")

    def _load_zeros_file(self, filepath: str):
        with open(filepath) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#'):
                    try:
                        self.loaded_zeros.append(float(line))
                    except ValueError:
                        pass
        self.log(f"Loaded {len(self.loaded_zeros)} zeros from {filepath}")

    def compute_zeros(self, n: int = 100) -> List[ZetaZero]:
        zeros = []
        for k in range(1, n + 1):
            if k in self._zeros_cache:
                zeros.append(self._zeros_cache[k])
                continue
            z = zetazero(k)
            zero = ZetaZero(
                index=k,
                imaginary_part=float(mp_im(z)),
                real_part=float(mp_re(z)),
                verified=abs(float(mp_re(z)) - 0.5) < 1e-10,
                precision=self.precision
            )
            zeros.append(zero)
            self._zeros_cache[k] = zero
        return zeros

    def get_loaded_zeros(self, n: int = None) -> List[ZetaZero]:
        if n is None:
            n = len(self.loaded_zeros)
        zeros = []
        for i, gamma in enumerate(self.loaded_zeros[:n]):
            zeros.append(ZetaZero(index=i + 1, imaginary_part=gamma, real_part=0.5, verified=True))
        return zeros

    def pair_correlation_function(self, zeros: List[ZetaZero], alpha_max: float = 5.0,
                                   n_bins: int = 200, use_normalization: bool = True) -> Dict:
        gamma = np.array([z.imaginary_part for z in zeros])
        n = len(gamma)
        spacings = np.diff(gamma)
        mean_s = np.mean(spacings)
        diffs = []
        max_pairs = 500000
        count = 0
        step = max(1, n // 2000)
        for i in range(0, n, step):
            for j in range(i + 1, min(i + 500, n)):
                d = (gamma[j] - gamma[i]) / mean_s if use_normalization else (gamma[j] - gamma[i])
                if d < alpha_max:
                    diffs.append(d)
                count += 1
                if count >= max_pairs:
                    break
            if count >= max_pairs:
                break

        diffs = np.array(diffs)
        hist, edges = np.histogram(diffs, bins=n_bins, range=(0, alpha_max))
        bin_width = edges[1] - edges[0]
        density = hist / (len(diffs) * bin_width)
        alpha_centers = (edges[:-1] + edges[1:]) / 2
        gue = 1.0 - (np.sin(np.pi * alpha_centers) / (np.pi * alpha_centers + 1e-15)) ** 2

        return {
            'alpha': alpha_centers.tolist(),
            'empirical': density.tolist(),
            'gue_prediction': gue.tolist(),
            'n_pairs': len(diffs),
            'n_zeros': n,
        }

    def compute_gue_deviation(self, zeros: List[ZetaZero]) -> float:
        pc = self.pair_correlation_function(zeros, alpha_max=3.0, n_bins=100)
        emp = np.array(pc['empirical'])
        gue = np.array(pc['gue_prediction'])
        dx = pc['alpha'][1] - pc['alpha'][0] if len(pc['alpha']) > 1 else 0.1
        l2 = float(np.sqrt(np.sum((emp - gue) ** 2) * dx))
        return l2

    def mertens_function(self, x_max: int) -> Dict:
        mu = [0] * (x_max + 1)
        mu[1] = 1
        is_prime = [True] * (x_max + 1)
        primes = []
        for i in range(2, x_max + 1):
            if is_prime[i]:
                primes.append(i)
                mu[i] = -1
            for p in primes:
                if i * p > x_max:
                    break
                is_prime[i * p] = False
                if i % p == 0:
                    mu[i * p] = 0
                    break
                else:
                    mu[i * p] = -mu[i]

        M = np.cumsum(mu)
        x_arr = np.arange(1, x_max + 1)
        ratios = np.abs(M[1:]) / np.sqrt(x_arr)
        return {
            'x': x_arr.tolist(),
            'M': M[1:].tolist(),
            'ratios': ratios.tolist(),
            'max_ratio': float(np.max(ratios)),
            'mean_ratio': float(np.mean(ratios)),
        }

    def chebyshev_psi(self, x_max: int) -> Dict:
        is_prime = np.ones(x_max + 1, dtype=bool)
        is_prime[:2] = False
        for i in range(2, int(x_max ** 0.5) + 1):
            if is_prime[i]:
                is_prime[i * i::i] = False
        primes = np.where(is_prime)[0]
        lambda_n = np.zeros(x_max + 1)
        for p in primes:
            log_p = np.log(p)
            pk = p
            while pk <= x_max:
                lambda_n[pk] = log_p
                pk *= p
        psi = np.cumsum(lambda_n)
        x_arr = np.arange(1, x_max + 1)
        errors = np.abs(psi[1:] - x_arr)
        bounds = x_arr ** 0.5 * (np.log(x_arr + 1)) ** 2
        return {
            'x': x_arr.tolist(),
            'psi': psi[1:].tolist(),
            'errors': errors.tolist(),
            'bounds': bounds.tolist(),
            'within_bound': bool(np.all(errors < bounds)),
        }

    def get_log(self):
        return self._log