import numpy as np
from time import perf_counter as pc
from time import sleep
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from collections import deque
import threading
from scipy.signal import butter, filtfilt, windows
from scipy.fft import fft, fftfreq
import warnings
warnings.filterwarnings('ignore')

# =====================
# CBHE-MD: Conjugate Black Hole Multi-Run Detector
# Fast measurement with ensemble smoothing
# =====================

class FastTimingCollector:
    """
    High-speed timing collector with sub-microsecond precision
    """
    def __init__(self, target_interval=np.pi/100):
        self.target_interval = target_interval
        self.d = 0.0
        self.kp = 0.02  # Proportional gain (tuned for speed)
        self.data = []
        self.running = False
        self.thread = None
        
    def single_measurement(self):
        """One high-precision timing measurement"""
        t = pc()
        sleep(self.target_interval - self.d)
        measured = 100 * (pc() - t)
        err = np.pi - measured
        self.d -= self.kp * err
        return self.d
    
    def collect_batch(self, n_samples, progress_callback=None):
        """Collect n samples as fast as possible"""
        self.data = np.zeros(n_samples)
        self.running = True
        
        for i in range(n_samples):
            if not self.running:
                break
            self.data[i] = self.single_measurement()
            
            if progress_callback and i % 1000 == 0:
                progress_callback(i / n_samples)
        
        return self.data.copy()
    
    def stop(self):
        self.running = False


class MultiRunSmoother:
    """
    Ensemble smoother - averages multiple runs to extract true signal
    """
    def __init__(self, max_runs=100):
        self.max_runs = max_runs
        self.all_runs = deque(maxlen=max_runs)
        self.ensemble_mean = None
        self.ensemble_std = None
        self.run_count = 0
        
        # Statistical tracking
        self.cumulative_sum = None
        self.cumulative_sq = None
        
    def add_run(self, data):
        """Add a new run to the ensemble"""
        # Normalize length
        if self.ensemble_mean is None:
            self.ensemble_mean = np.zeros_like(data)
            self.cumulative_sum = np.zeros_like(data)
            self.cumulative_sq = np.zeros_like(data)
        
        min_len = min(len(data), len(self.ensemble_mean))
        
        # Running statistics (Welford's algorithm for numerical stability)
        self.run_count += 1
        delta = data[:min_len] - self.ensemble_mean[:min_len]
        self.ensemble_mean[:min_len] += delta / self.run_count
        self.cumulative_sum[:min_len] += data[:min_len]
        self.cumulative_sq[:min_len] += data[:min_len]**2
        
        # Store the run
        self.all_runs.append(data[:min_len].copy())
        
    @property
    def smoothed_signal(self):
        """Return the smoothed ensemble mean"""
        return self.ensemble_mean.copy() if self.ensemble_mean is not None else None
    
    @property
    def uncertainty(self):
        """Return standard error of the mean"""
        if self.run_count < 2:
            return None
        # Variance using Bessel's correction
        variance = (self.cumulative_sq - self.cumulative_sum**2 / self.run_count) / (self.run_count - 1)
        return np.sqrt(np.maximum(variance, 0)) / np.sqrt(self.run_count)
    
    @property
    def snr_improvement(self):
        """Signal-to-noise ratio improvement from averaging"""
        if self.run_count < 2:
            return 1.0
        # Each run adds sqrt(n) to SNR
        return np.sqrt(self.run_count)


class SpectralAnalyzer:
    """
    Advanced spectral analysis for BH signature detection
    """
    def __init__(self, sample_rate=100):
        self.sample_rate = sample_rate
        self.freqs = None
        self.psd_mean = None
        self.psd_all = []
        
        # Known BH signatures to look for
        self.signature_bands = {
            'Primordial BH (slow)': (1e-6, 1e-4),
            'Primordial BH (fast)': (1e-4, 1e-2),
            'Solar system effects': (1e-3, 1e-1),
            'Local gravitational': (1e-1, 1e+1),
            'LIGO band (GW)': (1e+1, 1e+3),
        }
        
    def compute_psd(self, data):
        """Compute Power Spectral Density"""
        n = len(data)
        
        # Apply window to reduce spectral leakage
        windowed = data * windows.hann(n)
        
        # FFT
        fft_vals = fft(windowed)
        freqs = fftfreq(n, 1/self.sample_rate)
        
        # Power spectrum (one-sided)
        positive_mask = freqs > 0
        freqs_pos = freqs[positive_mask]
        psd = 2.0/n * np.abs(fft_vals[positive_mask])**2
        
        return freqs_pos, psd
    
    def add_psd(self, freqs, psd):
        """Add PSD to ensemble"""
        self.psd_all.append(psd.copy())
        if self.psd_mean is None:
            self.psd_mean = np.zeros_like(psd)
        self.psd_mean = np.mean(self.psd_all, axis=0)
        self.freqs = freqs
    
    def detect_peaks(self, threshold_sigma=3):
        """Detect significant spectral peaks"""
        if self.psd_mean is None or len(self.psd_all) < 3:
            return [], []
        
        # Mean and std across runs
        psd_stack = np.array(self.psd_all)
        mean = np.mean(psd_stack, axis=0)
        std = np.std(psd_stack, axis=0)
        
        # Find peaks above threshold
        threshold = mean + threshold_sigma * std
        peaks = np.where(self.psd_mean > threshold)[0]
        
        return self.freqs[peaks], self.psd_mean[peaks]
    
    def analyze_signatures(self):
        """Analyze each signature band"""
        results = []
        
        for name, (f_min, f_max) in self.signature_bands.items():
            if self.freqs is None:
                continue
                
            mask = (self.freqs >= f_min) & (self.freqs < f_max)
            if np.any(mask):
                power = np.mean(self.psd_mean[mask])
                peak_freq = self.freqs[mask][np.argmax(self.psd_mean[mask])]
                results.append({
                    'name': name,
                    'mean_power': power,
                    'peak_freq': peak_freq,
                    'range': (f_min, f_max)
                })
        
        return results


class RealTimeBHDetector:
    """
    Main detector class with real-time visualization
    """
    def __init__(self, samples_per_run=10000, num_runs=20, sample_rate=100):
        self.samples_per_run = samples_per_run
        self.num_runs = num_runs
        self.sample_rate = sample_rate
        
        # Components
        self.collector = FastTimingCollector()
        self.smoother = MultiRunSmoother(max_runs=num_runs)
        self.analyzer = SpectralAnalyzer(sample_rate=sample_rate)
        
        # Visualization data
        self.fig = None
        self.axes = None
        self.lines = None
        
    def run_single_collection(self):
        """Collect one run of data"""
        data = self.collector.collect_batch(self.samples_per_run)
        return data
    
    def process_run(self, data):
        """Process a single run: add to smoother and analyzer"""
        self.smoother.add_run(data)
        freqs, psd = self.analyzer.compute_psd(data)
        self.analyzer.add_psd(freqs, psd)
        
    def run_experiment(self, show_progress=True):
        """Run multiple experiments with ensemble smoothing"""
        print("="*70)
        print("CBHE-MD: Conjugate Black Hole Multi-Run Detector")
        print(f"Configuration: {self.samples_per_run} samples × {self.num_runs} runs @ {self.sample_rate}Hz")
        print("="*70)
        
        for run_idx in range(self.num_runs):
            # Collect data
            data = self.run_single_collection()
            
            # Process
            self.process_run(data)
            
            # Progress
            if show_progress:
                snr = self.smoother.snr_improvement
                print(f"Run {run_idx+1:3d}/{self.num_runs} | "
                      f"Mean: {np.mean(data):+.6f} | "
                      f"Std: {np.std(data):.6f} | "
                      f"SNR boost: {snr:.2f}×")
        
        print("\n" + "="*70)
        print("EXPERIMENT COMPLETE - ANALYZING ENSEMBLE")
        print("="*70)
        
        return self.generate_report()
    
    def generate_report(self):
        """Generate comprehensive analysis report"""
        # Get smoothed signal
        smoothed = self.smoother.smoothed_signal
        uncertainty = self.smoother.uncertainty
        snr = self.smoother.snr_improvement
        
        print(f"\n📊 ENSEMBLE STATISTICS:")
        print(f"   Total runs averaged: {self.smoother.run_count}")
        print(f"   Total samples: {self.smoother.run_count * len(smoothed):,}")
        print(f"   SNR improvement: {snr:.2f}×")
        print(f"   Mean timing deviation: {np.mean(smoothed):+.8f}")
        print(f"   Std of mean: {np.std(smoothed):.8f}")
        
        # Spectral analysis
        print(f"\n📈 SPECTRAL ANALYSIS:")
        signatures = self.analyzer.analyze_signatures()
        
        for sig in signatures:
            bar_len = min(50, int(np.log10(sig['mean_power'] + 1) * 10))
            bar = '█' * bar_len + '░' * (50 - bar_len)
            print(f"   {sig['name']:30} | power: {sig['mean_power']:.4e} | {bar}")
        
        # Peak detection
        peaks, powers = self.analyzer.detect_peaks(threshold_sigma=2)
        if len(peaks) > 0:
            print(f"\n🚨 DETECTED SIGNIFICANT PEAKS ({len(peaks)}):")
            top_peaks = sorted(zip(peaks, powers), key=lambda x: -x[1])[:10]
            for freq, power in top_peaks:
                print(f"   freq: {freq:.6e} Hz | power: {power:.4e}")
        else:
            print(f"\n○ No significant peaks detected (need more runs)")
        
        # Detection assessment
        print(f"\n🎯 BLACK HOLE DETECTION ASSESSMENT:")
        
        # Check each band
        for sig in signatures:
            if sig['mean_power'] > 1e-4:  # Threshold
                print(f"   ✓ {sig['name']}: POTENTIAL DETECTION")
            else:
                print(f"   ○ {sig['name']}: No significant signal")
        
        return {
            'smoothed': smoothed,
            'uncertainty': uncertainty,
            'snr': snr,
            'signatures': signatures,
            'peaks': (peaks, powers)
        }
    
    def plot_realtime(self, interval_ms=500):
        """Create real-time updating plot"""
        self.fig, self.axes = plt.subplots(2, 2, figsize=(14, 10))
        self.fig.suptitle('CBHE-MD: Real-Time Black Hole Detection', fontsize=14)
        
        # Subplot 1: Time series comparison
        ax1 = self.axes[0, 0]
        self.line_raw, = ax1.plot([], [], 'b-', alpha=0.3, label='Single run')
        self.line_smooth, = ax1.plot([], [], 'r-', lw=2, label='Ensemble mean')
        ax1.fill_between([], [], [], alpha=0.3, color='red', label='±σ')
        ax1.set_xlabel('Sample')
        ax1.set_ylabel('Timing Deviation (d)')
        ax1.set_title('Timing Deviation: Raw vs Smoothed')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # Subplot 2: PSD evolution
        ax2 = self.axes[0, 1]
        self.line_psd, = ax2.plot([], [], 'r-', lw=1)
        ax2.set_xlabel('Frequency (Hz)')
        ax2.set_ylabel('Power Spectral Density')
        ax2.set_title('Ensemble-Averaged PSD')
        ax2.set_xscale('log')
        ax2.set_yscale('log')
        ax2.grid(True, alpha=0.3)
        
        # Subplot 3: SNR growth
        ax3 = self.axes[1, 0]
        self.snr_history = []
        self.line_snr, = ax3.plot([], [], 'g-', lw=2)
        ax3.set_xlabel('Run Number')
        ax3.set_ylabel('SNR Improvement (×)')
        ax3.set_title('Ensemble SNR Growth')
        ax3.grid(True, alpha=0.3)
        
        # Subplot 4: Running statistics
        ax4 = self.axes[1, 1]
        self.mean_history = []
        self.std_history = []
        self.line_mean, = ax4.plot([], [], 'b-', label='Mean')
        ax4.set_xlabel('Run Number')
        ax4.set_ylabel('Value')
        ax4.set_title('Running Statistics')
        ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        # Animation function
        def update(frame):
            if self.smoother.run_count > 0:
                # Update time series
                smoothed = self.smoother.smoothed_signal
                if smoothed is not None:
                    x = np.arange(len(smoothed))
                    self.line_smooth.set_data(x, smoothed)
                    ax1.set_xlim(0, len(smoothed))
                    ax1.set_ylim(np.min(smoothed) - 0.1, np.max(smoothed) + 0.1)
                    
                    # Show most recent raw run
                    if len(self.smoother.all_runs) > 0:
                        raw = self.smoother.all_runs[-1]
                        self.line_raw.set_data(x, raw[:len(smoothed)])
                
                # Update PSD
                if self.analyzer.psd_mean is not None:
                    freqs = self.analyzer.freqs
                    self.line_psd.set_data(freqs, self.analyzer.psd_mean)
                    ax2.set_xlim(1e-4, self.sample_rate/2)
                    ax2.set_ylim(1e-10, np.max(self.analyzer.psd_mean)*2)
                
                # Update SNR history
                self.snr_history.append(self.smoother.snr_improvement)
                self.line_snr.set_data(range(len(self.snr_history)), self.snr_history)
                ax3.set_xlim(0, max(5, len(self.snr_history)))
                ax3.set_ylim(0, max(self.snr_history) * 1.2)
                
                # Update statistics
                self.mean_history.append(np.mean(smoothed) if smoothed is not None else 0)
                self.std_history.append(np.std(smoothed) if smoothed is not None else 0)
                ax4.clear()
                ax4.plot(self.mean_history, 'b-', label='Mean')
                ax4.plot(self.std_history, 'r-', label='Std')
                ax4.set_xlabel('Run Number')
                ax4.set_ylabel('Value')
                ax4.set_title(f'Running Statistics (n={len(self.mean_history)})')
                ax4.legend()
                ax4.grid(True, alpha=0.3)
            
            return self.line_smooth, self.line_psd, self.line_snr
        
        # Start animation
        self.ani = FuncAnimation(self.fig, update, interval=interval_ms, blit=True)
        plt.show()
    
    def plot_final_analysis(self):
        """Generate final comprehensive plots"""
        if self.smoother.run_count == 0:
            print("No data to plot. Run experiment first.")
            return
        
        fig = plt.figure(figsize=(16, 12))
        
        # 1. Ensemble time series
        ax1 = fig.add_subplot(3, 2, 1)
        smoothed = self.smoother.smoothed_signal
        uncertainty = self.smoother.uncertainty
        x = np.arange(len(smoothed))
        
        ax1.fill_between(x, smoothed - uncertainty, smoothed + uncertainty, 
                         alpha=0.3, color='red')
        ax1.plot(x, smoothed, 'r-', lw=1, label='Ensemble mean')
        
        # Overlay individual runs (faded)
        for i, run in enumerate(self.smoother.all_runs):
            ax1.plot(x, run[:len(smoothed)], 'b-', alpha=0.1, lw=0.5)
        
        ax1.set_xlabel('Sample')
        ax1.set_ylabel('Timing Deviation')
        ax1.set_title('Ensemble Average (Blue=individual runs, Red=smoothed)')
        ax1.grid(True, alpha=0.3)
        
        # 2. PSD comparison
        ax2 = fig.add_subplot(3, 2, 2)
        if self.analyzer.psd_mean is not None:
            freqs = self.analyzer.freqs
            
            # Show individual PSDs (faded)
            for psd in self.analyzer.psd_all[:5]:
                ax2.semilogy(freqs, psd, 'b-', alpha=0.1)
            
            # Ensemble mean
            ax2.semilogy(freqs, self.analyzer.psd_mean, 'r-', lw=2, label='Ensemble mean')
            
            # Mark signature bands
            colors = ['green', 'orange', 'purple', 'cyan', 'yellow']
            for (name, (f_min, f_max)), color in zip(self.analyzer.signature_bands.items(), colors):
                ax2.axvspan(f_min, f_max, alpha=0.1, color=color, label=name)
        
        ax2.set_xlabel('Frequency (Hz)')
        ax2.set_ylabel('PSD')
        ax2.set_title('Power Spectral Density (with BH bands)')
        ax2.set_xscale('log')
        ax2.legend(fontsize=7)
        ax2.grid(True, alpha=0.3)
        
        # 3. SNR convergence
        ax3 = fig.add_subplot(3, 2, 3)
        snr_runs = np.arange(1, self.smoother.run_count + 1)
        expected_snr = np.sqrt(snr_runs)
        ax3.plot(snr_runs, expected_snr, 'g--', label='Expected √N')
        ax3.axhline(self.smoother.snr_improvement, color='red', linestyle='--',
                    label=f'Actual: {self.smoother.snr_improvement:.2f}×')
        ax3.set_xlabel('Number of Runs')
        ax3.set_ylabel('SNR Improvement (×)')
        ax3.set_title('Signal-to-Noise Ratio Improvement')
        ax3.legend()
        ax3.grid(True, alpha=0.3)
        
        # 4. Variance map
        ax4 = fig.add_subplot(3, 2, 4)
        if len(self.smoother.all_runs) > 1:
            variance_map = np.var(self.smoother.all_runs, axis=0)
            im = ax4.imshow([variance_map], aspect='auto', cmap='hot')
            ax4.set_xlabel('Sample')
            ax4.set_title('Variance Map (bright = variable)')
            plt.colorbar(im, ax=ax4, orientation='horizontal')
        
        # 5. Frequency zoom for PBH band
        ax5 = fig.add_subplot(3, 2, 5)
        if self.analyzer.psd_mean is not None:
            mask = (self.analyzer.freqs > 1e-6) & (self.analyzer.freqs < 1e-2)
            if np.any(mask):
                ax5.semilogy(self.analyzer.freqs[mask], self.analyzer.psd_mean[mask], 'r-', lw=2)
                ax5.set_xlabel('Frequency (Hz)')
                ax5.set_ylabel('PSD')
                ax5.set_title('Primordial Black Hole Detection Band (10⁻⁶ - 10⁻² Hz)')
                ax5.grid(True, alpha=0.3)
                
                # Mark detection threshold
                threshold = np.mean(self.analyzer.psd_mean[mask]) + 2*np.std(self.analyzer.psd_mean[mask])
                ax5.axhline(threshold, color='green', linestyle='--', label='Detection threshold')
                ax5.legend()
        
        # 6. Summary stats
        ax6 = fig.add_subplot(3, 2, 6)
        ax6.axis('off')
        
        summary_text = f"""
        CBHE-MD DETECTION SUMMARY
        =========================
        
        Configuration:
        • Samples per run: {self.samples_per_run:,}
        • Number of runs: {self.smoother.run_count}
        • Sample rate: {self.sample_rate} Hz
        • Total integration time: {self.samples_per_run * self.smoother.run_count / self.sample_rate:.1f}s
        
        Ensemble Performance:
        • SNR improvement: {self.smoother.snr_improvement:.2f}×
        • Mean deviation: {np.mean(smoothed):+.8f}
        • Std deviation: {np.std(smoothed):.8f}
        
        Detection Status:
        """
        
        for sig in self.analyzer.analyze_signatures():
            status = "✓" if sig['mean_power'] > 1e-4 else "○"
            summary_text += f"\n        {status} {sig['name']}"
        
        ax6.text(0.1, 0.9, summary_text, transform=ax6.transAxes,
                fontsize=10, verticalalignment='top', fontfamily='monospace',
                bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        plt.tight_layout()
        plt.savefig('bh_detection_final_analysis.png', dpi=150, bbox_inches='tight')
        plt.show()
        
        print("\n✓ Final analysis saved to 'bh_detection_final_analysis.png'")


def quick_test(num_runs=5, samples_per_run=5000):
    """Quick demonstration run"""
    detector = RealTimeBHDetector(samples_per_run=samples_per_run, 
                                  num_runs=num_runs, 
                                  sample_rate=100)
    
    # Run experiment
    results = detector.run_experiment(show_progress=True)
    
    # Plot final analysis
    detector.plot_final_analysis()
    
    return detector, results


# =====================
# Main execution
# =====================
if __name__ == "__main__":
    import sys
    
    # Parse arguments
    n_runs = int(sys.argv[1]) if len(sys.argv) > 1 else 10
    n_samples = int(sys.argv[2]) if len(sys.argv) > 2 else 10000
    
    print(f"\n🌀 Running CBHE-MD with {n_runs} runs of {n_samples} samples each\n")
    
    detector, results = quick_test(num_runs=n_runs, samples_per_run=n_samples)
    
    print("\n" + "="*70)
    print("DETECTION COMPLETE")
    print("="*70)