"""
Visualization pipeline for v_mix — 15+ plot types.
"""

import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import os


class VisualizationPipeline:
    def __init__(self, output_dir: str = "./output"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        self.files = []

    def _save(self, name: str):
        path = os.path.join(self.output_dir, name)
        plt.savefig(path, dpi=150, bbox_inches='tight')
        plt.close()
        self.files.append(path)
        return path

    def plot_zero_distribution(self, zeros):
        fig, ax = plt.subplots(figsize=(10, 4))
        gammas = [z['imaginary_part'] for z in zeros]
        ax.plot(gammas, np.zeros(len(gammas)), '|', markersize=10, alpha=0.7)
        ax.set_xlabel('Imaginary part γ')
        ax.set_title('Zeta Zero Distribution on Critical Line')
        ax.set_yticks([])
        return self._save("vmix_zero_distribution.png")

    def plot_spacing_histogram(self, spacings):
        fig, ax = plt.subplots(figsize=(8, 5))
        ax.hist(spacings, bins=50, density=True, alpha=0.7, color='steelblue', edgecolor='black')
        # Wigner surmise overlay
        s = np.linspace(0, max(spacings), 200)
        wigner = (np.pi / 2) * s * np.exp(-np.pi * s**2 / 4)
        ax.plot(s, wigner, 'r--', linewidth=2, label='GUE Wigner surmise')
        ax.set_xlabel('Normalized spacing s')
        ax.set_ylabel('Probability density')
        ax.set_title('Nearest-Neighbor Spacing Distribution')
        ax.legend()
        return self._save("vmix_spacing_histogram.png")

    def plot_pair_correlation(self, pc_result):
        fig, ax = plt.subplots(figsize=(8, 5))
        alpha = np.array(pc_result['alpha'])
        emp = np.array(pc_result['empirical'])
        gue = np.array(pc_result['gue_prediction'])
        ax.plot(alpha, emp, label='Empirical', color='steelblue', linewidth=1.5)
        ax.plot(alpha, gue, '--', label='GUE', color='red', linewidth=1.5)
        ax.set_xlabel('α')
        ax.set_ylabel('F(α)')
        ax.set_title('Pair Correlation Function')
        ax.legend()
        ax.set_xlim(0, max(alpha))
        return self._save("vmix_pair_correlation.png")

    def plot_gue_convergence(self, result):
        fig, ax = plt.subplots(figsize=(8, 6))
        N = np.array(result['n_zeros'])
        dev = np.array(result['ks_distances'])
        ax.loglog(N, dev, 'o-', color='steelblue', linewidth=2, markersize=6)

        # Overlay fits
        if 'power_law_fit_N' in result:
            fit = result['power_law_fit_N']
            beta = fit['exponent_beta']
            c = np.exp(np.polyfit(np.log(N), np.log(dev), 1)[1])
            fit_line = c * N ** (-beta)
            ax.loglog(N, fit_line, '--', color='red', label=f"N^(-{beta:.3f})  R²={fit['fit_quality_r2']:.3f}")
        if 'power_law_fit_inv_sqrt' in result:
            fit = result['power_law_fit_inv_sqrt']
            ax.plot([], [], ' ', label=f"1/√N fit: R²={fit['fit_quality_r2']:.3f}")

        ax.set_xlabel('Number of zeros N')
        ax.set_ylabel('KS distance to Wigner surmise')
        ax.set_title('GUE Convergence Rate: KS Distance vs N (NOVEL)')
        ax.legend()
        ax.grid(True, which='both', linestyle='--', alpha=0.5)
        return self._save("vmix_gue_convergence.png")

    def plot_cramer_gaps(self, result):
        fig, ax = plt.subplots(figsize=(8, 5))
        thresholds = np.array(result['thresholds'])
        emp = np.array(result['empirical_survival'])
        cra = np.array(result['cramer_survival'])
        gra = np.array(result['granville_survival'])
        ax.semilogy(thresholds, emp, 'o', label='Empirical', markersize=4, alpha=0.7)
        ax.semilogy(thresholds, cra, '-', label='Cramér: e^{-λ}', color='red')
        ax.semilogy(thresholds, gra, '--', label='Granville: e^{-3.56λ}', color='green')
        ax.set_xlabel('λ = gap / (ln p)²')
        ax.set_ylabel('P(ratio > λ)')
        ax.set_title("Cramér Gap Tail Distribution (up to 5M)")
        ax.legend()
        ax.grid(True, which='both', linestyle='--', alpha=0.5)
        return self._save("vmix_cramer_gaps.png")

    def plot_lindeloef(self, result):
        fig, ax = plt.subplots(figsize=(8, 5))
        gamma = np.array(result['gamma_values'])
        ratios = np.array(result['ratios'])
        ax.plot(gamma, ratios, 'o', markersize=3, alpha=0.5, color='steelblue')
        ax.axhline(y=result['bourgain_bound'], color='red', linestyle='--',
                   label=f"Bourgain bound = {result['bourgain_bound']:.4f}")
        ax.set_xlabel('γ_n')
        ax.set_ylabel('log|ζ(1/2+iγ)| / log(γ)')
        ax.set_title('Lindelöf Hypothesis: Empirical Exponent θ')
        ax.legend()
        ax.set_xscale('log')
        return self._save("vmix_lindeloef.png")

    def plot_chebyshev_bias(self, result):
        fig, ax = plt.subplots(figsize=(8, 5))
        x = np.array(result['sample_x'])
        ratios = np.array(result['ratios_4'])
        ax.plot(x, ratios, '-', color='steelblue', linewidth=1)
        ax.axhline(y=1.0, color='red', linestyle='--', label='Equal counts')
        ax.set_xlabel('x')
        ax.set_ylabel('π(x;4,3) / π(x;4,1)')
        ax.set_title("Chebyshev Bias: Primes mod 4")
        ax.set_xscale('log')
        ax.legend()
        return self._save("vmix_chebyshev_bias.png")

    def plot_lehmer_phenomena(self, result):
        fig, ax = plt.subplots(figsize=(8, 5))
        bins = np.array(result['bin_centers'])
        hist = np.array(result['spacing_histogram'])
        wigner = np.array(result['wigner_prediction'])
        ax.plot(bins, hist, label='Empirical', color='steelblue', linewidth=1.5)
        ax.plot(bins, wigner, '--', label='GUE Wigner', color='red', linewidth=1.5)
        ax.set_xlabel('Normalized spacing s')
        ax.set_ylabel('Density')
        ax.set_title(f"Lehmer Phenomena: Spacing Distribution (min={result['min_normalized_spacing']:.5f})")
        ax.legend()
        return self._save("vmix_lehmer_phenomena.png")

    def plot_new_strategies_comparison(self, results_list):
        fig, ax = plt.subplots(figsize=(8, 5))
        names = [r.get('strategy', 'unknown') for r in results_list]
        values = [r.get('mae', 0) or r.get('mean_persistence_entropy', 0) for r in results_list]
        colors = ['steelblue', 'green', 'orange']
        ax.bar(range(len(names)), values, color=colors[:len(names)])
        ax.set_xticks(range(len(names)))
        ax.set_xticklabels([n.replace('_', '\n') for n in names], fontsize=8)
        ax.set_ylabel('Metric (lower = better for MAE)')
        ax.set_title('New Strategy Performance Comparison')
        return self._save("vmix_new_strategies_comparison.png")

    def plot_entropy_convergence(self, result):
        fig, ax = plt.subplots(figsize=(8, 5))
        sizes = np.array(result['window_sizes'])
        ent = np.array(result['entropies'])
        ax.plot(sizes, ent, 'o-', color='steelblue', linewidth=2, markersize=6)
        ax.set_xlabel('Window size (zeros)')
        ax.set_ylabel('Shannon entropy of spacings')
        ax.set_title('Entropy of Zero Spacing Distribution')
        ax.set_xscale('log')
        ax.grid(True, linestyle='--', alpha=0.5)
        return self._save("vmix_entropy_convergence.png")

    def plot_ktuple_comparison(self, result):
        fig, ax = plt.subplots(figsize=(8, 5))
        names = []
        errors = []
        for k, v in result['patterns'].items():
            names.append(k)
            errors.append(v['relative_error'])
        ax.barh(range(len(names)), errors, color='steelblue')
        ax.set_yticks(range(len(names)))
        ax.set_yticklabels(names)
        ax.set_xlabel('Relative error |observed - predicted| / predicted')
        ax.set_title('Hardy-Littlewood k-Tuple Accuracy')
        ax.set_xlim(0, max(errors) * 1.2)
        return self._save("vmix_ktuple_accuracy.png")

    def get_files(self):
        return self.files