"""
Explicit Formula Engine (from v3, enhanced for v_mix)
Uses the von Mangoldt explicit formula with 100,000 zeros.
"""

import numpy as np
from typing import List, Dict
from dataclasses import dataclass


@dataclass
class PrimePrediction:
    predicted_value: int
    confidence: float
    actual_is_prime: bool = False
    error: int = 0
    method: str = ""


class ExplicitFormulaEngine:
    def __init__(self, zeros: List[float]):
        self.zeros = np.array(zeros, dtype=np.float64)
        self._log = []

    def log(self, msg: str):
        self._log.append(msg)
        print(f"  [ExplicitFormula] {msg}")

    def compute_psi(self, x: float, n_zeros: int = None) -> float:
        if n_zeros is None:
            n_zeros = len(self.zeros)
        n_zeros = min(n_zeros, len(self.zeros))
        if x <= 1:
            return 0.0
        log_x = np.log(x)
        sqrt_x = np.sqrt(x)
        psi = x
        gamma_subset = self.zeros[:n_zeros]
        denom = 0.25 + gamma_subset ** 2
        cos_term = np.cos(gamma_subset * log_x) * 0.5
        sin_term = np.sin(gamma_subset * log_x) * gamma_subset
        psi -= np.sum(2 * sqrt_x * (cos_term + sin_term) / denom)
        psi -= np.log(2 * np.pi)
        if x > 1.01:
            psi -= 0.5 * np.log(1 - 1 / (x * x))
        return psi

    def compute_psi_vectorized(self, x_values: np.ndarray, n_zeros: int = None) -> np.ndarray:
        if n_zeros is None:
            n_zeros = min(10000, len(self.zeros))
        n_zeros = min(n_zeros, len(self.zeros))
        x_values = np.asarray(x_values, dtype=np.float64)
        log_x = np.log(np.maximum(x_values, 1.01))
        sqrt_x = np.sqrt(np.maximum(x_values, 0))
        psi = x_values.copy()
        gamma_arr = self.zeros[:n_zeros]
        for gamma in gamma_arr:
            denom = 0.25 + gamma * gamma
            cos_terms = np.cos(gamma * log_x) * 0.5
            sin_terms = np.sin(gamma * log_x) * gamma
            psi -= 2 * sqrt_x * (cos_terms + sin_terms) / denom
        psi -= np.log(2 * np.pi)
        mask = x_values > 1.01
        psi[mask] -= 0.5 * np.log(1 - 1 / (x_values[mask] ** 2))
        return psi

    def count_predicted_primes(self, x_max: int, n_zeros: int = None) -> int:
        if n_zeros is None:
            n_zeros = min(5000, len(self.zeros))
        x_range = np.arange(2, x_max + 1, dtype=np.float64)
        psi_prev = self.compute_psi_vectorized(x_range - 0.5, n_zeros=n_zeros)
        psi_next = self.compute_psi_vectorized(x_range + 0.5, n_zeros=n_zeros)
        jumps = psi_next - psi_prev
        expected = np.log(x_range)
        predicted = np.sum(jumps > expected * 0.5)
        return int(predicted)

    def minimum_zeros_for_primes(self, x_max: int, max_zeros: int = 100000) -> Dict:
        sieve = np.ones(x_max + 1, dtype=bool)
        sieve[:2] = False
        for i in range(2, int(x_max ** 0.5) + 1):
            if sieve[i]:
                sieve[i * i::i] = False
        actual_primes = np.sum(sieve)

        N_tests = [10, 50, 100, 200, 500, 1000, 2000, 5000,
                     10000, 20000, 50000, 100000]
        N_tests = [n for n in N_tests if n <= max_zeros]

        results = []
        for N in N_tests:
            predicted = self.count_predicted_primes(x_max, n_zeros=N)
            results.append({
                'N': N,
                'predicted_primes': predicted,
                'actual_primes': int(actual_primes),
                'accuracy': predicted / actual_primes if actual_primes > 0 else 0,
                'correct': predicted == actual_primes,
            })
            if predicted == actual_primes:
                break

        return {
            'x_max': x_max,
            'results': results,
            'minimum_N': next((r['N'] for r in results if r['correct']), None),
        }

    def prime_oscillation_contributions(self, x: float, n_zeros: int = 1000) -> Dict:
        n_zeros = min(n_zeros, len(self.zeros))
        log_x = np.log(x)
        sqrt_x = np.sqrt(x)
        contributions = []
        cumulative = 0.0
        for i in range(n_zeros):
            gamma = self.zeros[i]
            denom = 0.25 + gamma * gamma
            cos_term = np.cos(gamma * log_x) * 0.5
            sin_term = np.sin(gamma * log_x) * gamma
            actual = -2 * sqrt_x * (cos_term + sin_term) / denom
            cumulative += actual
            contributions.append({
                'index': i + 1,
                'gamma': gamma,
                'contribution': float(actual),
                'cumulative': float(cumulative),
            })
        return {
            'x': x,
            'total_contribution': float(cumulative),
            'contributions': contributions,
        }

    def get_log(self):
        return self._log