import numpy as np
import pandas as pd
import warnings

# Suppress matplotlib warnings before importing pyplot
warnings.filterwarnings('ignore', category=UserWarning)
warnings.filterwarnings('ignore')

import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
from scipy import signal
from scipy.stats import entropy

class CCT_LiveStream_Classifier:
    """
    Conditional Collapse Theory Classifier for Real-World Data Streams
    Predicts Red Star (Stable) vs Black Hole (Crash) Events using Trained ML Seed θ*
    """
    
    def __init__(self, trained_theta: float = None, 
                 collapse_threshold: float = 0.3,
                 critical_escape_prob: float = 0.5,
                 window_size: int = 50,
                 prediction_horizon: int = 10):
        """
        Args:
            trained_theta: The ML Seed θ* learned from 200×200 training
            collapse_threshold: H_c - entropy level for successful collapse
            critical_escape_prob: P_critical - minimum escape probability for Red Star
            window_size: Lookback window for entropy calculation
            prediction_horizon: Steps ahead to predict
        """
        # Load trained seed (default: optimized value from training)
        self.theta = trained_theta if trained_theta is not None else 0.7342
        
        # CCT Parameters
        self.H_COLLAPSE = collapse_threshold
        self.P_CRITICAL = critical_escape_prob
        self.WINDOW = window_size
        self.HORIZON = prediction_horizon
        
        # 16-Element Semantic State (for explainability)
        self.element_names = [
            "E01_Trend_Stability",    "E02_Volatility_Bound",
            "E03_Entropy_Gradient",   "E04_Correlation_Structure",
            "E05_Feedback_Loop",      "E06_Liquidity_Flow",
            "E07_Shockwave_Damping",  "E08_Phase_Transition",
            "E09_Pattern_Recognition","E10_Anomaly_Detection",
            "E11_Escape_Potential",   "E12_Singularity_Risk",
            "E13_Energy_Budget",      "E14_Prompt_Responsiveness",
            "E15_Mutual_Destruction", "E16_Final_Collapse"
        ]
        
        # State tracking
        self.entropy_history = []
        self.prediction_log = []
        self.element_activations = np.zeros(16)
        
    def compute_semantic_entropy(self, data_window: np.ndarray) -> float:
        """
        Computes Semantic Entropy H(T) for a data window.
        Uses multi-scale analysis to capture both local and global uncertainty.
        """
        if len(data_window) < 10:
            return 1.0  # High entropy for insufficient data
            
        # Method 1: Distributional Entropy (Shannon)
        hist, _ = np.histogram(data_window, bins='auto', density=True)
        hist = hist[hist > 0]  # Remove zeros for log
        H_shannon = -np.sum(hist * np.log(hist + 1e-10))
        
        # Method 2: Spectral Entropy (Fourier-based complexity)
        fft_vals = np.abs(np.fft.rfft(data_window - np.mean(data_window)))
        fft_probs = fft_vals / (np.sum(fft_vals) + 1e-10)
        H_spectral = -np.sum(fft_probs * np.log(fft_probs + 1e-10))
        
        # Method 3: Predictability Entropy (AR model residuals)
        if len(data_window) > 20:
            from scipy.stats import linregress
            x = np.arange(len(data_window))
            slope, intercept, r_value, p_value, std_err = linregress(x, data_window)
            residuals = data_window - (slope * x + intercept)
            H_residual = np.std(residuals)  # Proxy for unpredictability
        else:
            H_residual = 1.0
            
        # Combine: Weighted average emphasizing spectral complexity
        H_total = 0.3 * H_shannon + 0.5 * H_spectral + 0.2 * H_residual
        return np.clip(H_total, 0.0, 2.0)  # Normalize range
    
    def compute_escape_probability(self, data_window: np.ndarray, theta: float) -> float:
        """
        Computes Escape Probability P_esc(θ) using the trained ML Seed.
        Based on Red Star Theory fractal topology.
        """
        # Extract key features from data window
        features = self._extract_features(data_window)
        
        # Red Star ODE Dynamics (from Appendix A)
        # dy/dt = y² + μ + θ·D(t) where D(t) is mutual destruction term
        y0 = features['initial_state']
        mu = features['control_parameter']
        
        # Simulate short trajectory to estimate escape likelihood
        y = y0
        escaped = False
        max_steps = 20
        
        for t in range(max_steps):
            # Mutual destruction term (learned via θ)
            D_term = theta * np.exp(-np.abs(y)**2) * features['prompt_energy']
            
            # ODE step
            dy = (y**2 + mu + D_term) * 0.1
            y = y + dy
            
            # Escape condition: bounded trajectory or divergence with control
            if np.abs(y) < 10.0 and t > 5:
                escaped = True
                break
            if np.abs(y) > 100.0:  # Uncontrolled divergence = Black Hole
                break
        
        # Base escape probability from trajectory
        P_base = 1.0 if escaped else 0.2
        
        # Modulate by feature alignment with trained seed
        alignment = self._compute_feature_alignment(features, theta)
        P_esc = P_base * (0.5 + 0.5 * alignment)
        
        return np.clip(P_esc, 0.0, 1.0)
    
    def _extract_features(self, data: np.ndarray) -> dict:
        """Extract CCT-relevant features from data window"""
        features = {}
        
        # Initial state (normalized)
        features['initial_state'] = (data[0] - np.mean(data)) / (np.std(data) + 1e-10)

        # Control parameter (trend strength)
        if len(data) > 2:
            from scipy.stats import linregress
            x = np.arange(len(data))
            slope, intercept, r_value, p_value, std_err = linregress(x, data)
            features['control_parameter'] = np.clip(slope * len(data) / (np.std(data) + 1e-10), -2, 2)
        else:
            features['control_parameter'] = 0.0
        
        # Volatility (prompt energy proxy)
        features['prompt_energy'] = np.std(np.diff(data)) / (np.abs(np.mean(data)) + 1e-10)
        
        # Correlation structure (feedback loop indicator)
        if len(data) > 10:
            autocorr = np.correlate(data - np.mean(data), data - np.mean(data), mode='full')
            autocorr = autocorr[len(autocorr)//2:]
            features['feedback_strength'] = np.max(autocorr[1:10]) / (autocorr[0] + 1e-10)
        else:
            features['feedback_strength'] = 0.5
            
        # Anomaly score (deviation from expected pattern)
        features['anomaly_score'] = np.abs(data[-1] - np.mean(data[:-5])) / (np.std(data) + 1e-10)
        
        return features
    
    def _compute_feature_alignment(self, features: dict, theta: float) -> float:
        """
        Computes how well current features align with the trained seed θ*.
        Higher alignment = higher confidence in prediction.
        """
        # Simple alignment: features that θ* was trained to recognize
        alignment = 0.0
        
        # θ* learned to stabilize high-feedback, moderate-volatility regimes
        if 0.3 < features['feedback_strength'] < 0.8:
            alignment += 0.3
        if 0.1 < features['prompt_energy'] < 0.5:
            alignment += 0.3
        if features['anomaly_score'] < 2.0:  # Not extreme outlier
            alignment += 0.2
        if np.abs(features['control_parameter']) < 1.5:  # Moderate trend
            alignment += 0.2
            
        return np.clip(alignment, 0.0, 1.0)
    
    def compute_red_star_potential(self, data_window: np.ndarray) -> dict:
        """
        Computes full Red Star Potential Ψ_Red per the Red Star Equation.
        """
        # Component 1: Escape Probability
        P_esc = self.compute_escape_probability(data_window, self.theta)
        
        # Component 2: Entropy Collapse Indicator
        H_final = self.compute_semantic_entropy(data_window[-self.WINDOW//2:])
        entropy_indicator = 1.0 if H_final < self.H_COLLAPSE else 0.0
        
        # Component 3: Energy Cost Term
        E_destroy = self._estimate_destruction_energy(data_window)
        E_available = np.std(data_window) * len(data_window) * 0.1  # Proxy
        energy_term = np.exp(-E_destroy / (E_available + 1e-10))
        
        # Red Star Equation
        Psi_Red = P_esc * entropy_indicator * energy_term
        
        return {
            'Psi_Red': np.clip(Psi_Red, 0.0, 1.0),
            'P_esc': P_esc,
            'H_final': H_final,
            'entropy_indicator': entropy_indicator,
            'energy_term': energy_term,
            'E_destroy': E_destroy
        }
    
    def _estimate_destruction_energy(self, data: np.ndarray) -> float:
        """Estimates energy required for mutual destruction (scales as 1/δ³)"""
        # δ = distance to singularity (estimated from trajectory curvature)
        if len(data) < 10:
            return 10.0  # High energy needed for uncertain data
            
        # Estimate curvature (second derivative)
        curvature = np.abs(np.gradient(np.gradient(data)))
        delta = 1.0 / (np.mean(curvature[-10:]) + 0.1)  # Inverse curvature
        
        # Energy scales as 1/δ³
        E_destroy = 1.0 / (delta**3 + 0.01)
        return np.clip(E_destroy, 0.1, 100.0)
    
    def predict_event(self, data_stream: np.ndarray, current_idx: int) -> dict:
        """
        Predicts Red Star vs Black Hole event at current index.
        """
        # Extract lookback window
        start_idx = max(0, current_idx - self.WINDOW)
        window = data_stream[start_idx:current_idx + 1]
        
        if len(window) < self.WINDOW // 2:
            return {'prediction': 'INSUFFICIENT_DATA', 'confidence': 0.0}
        
        # Compute Red Star Potential
        psi_result = self.compute_red_star_potential(window)
        
        # Make prediction
        if psi_result['Psi_Red'] > self.P_CRITICAL and psi_result['entropy_indicator'] == 1.0:
            prediction = 'RED_STAR'
            confidence = psi_result['Psi_Red']
        else:
            prediction = 'BLACK_HOLE'
            confidence = 1.0 - psi_result['Psi_Red']
        
        # Update 16-element activations (for explainability)
        self._update_element_activations(window, psi_result, prediction)
        
        # Log prediction
        self.prediction_log.append({
            'index': current_idx,
            'prediction': prediction,
            'confidence': confidence,
            'Psi_Red': psi_result['Psi_Red'],
            'H_final': psi_result['H_final']
        })
        
        return {
            'prediction': prediction,
            'confidence': confidence,
            'Psi_Red': psi_result['Psi_Red'],
            'components': psi_result,
            'elements': dict(zip(self.element_names, self.element_activations))
        }
    
    def _update_element_activations(self, window: np.ndarray, 
                                   psi_result: dict, prediction: str):
        """Updates 16-element semantic state based on current analysis"""
        features = self._extract_features(window)
        
        # Update elements based on feature values and prediction
        self.element_activations[0] = 1.0 - np.abs(features['control_parameter']) / 2.0  # E01_Trend_Stability
        self.element_activations[1] = 1.0 - np.clip(np.std(window) / 2.0, 0, 1)  # E02_Volatility_Bound
        self.element_activations[2] = 1.0 - psi_result['H_final'] / 2.0  # E03_Entropy_Gradient
        self.element_activations[3] = features['feedback_strength']  # E04_Correlation_Structure
        self.element_activations[4] = features['feedback_strength'] * 0.8  # E05_Feedback_Loop
        self.element_activations[5] = 1.0 - features['prompt_energy']  # E06_Liquidity_Flow
        self.element_activations[6] = 1.0 - features['anomaly_score'] / 3.0  # E07_Shockwave_Damping
        self.element_activations[7] = psi_result['energy_term']  # E08_Phase_Transition
        self.element_activations[8] = psi_result['P_esc']  # E09_Pattern_Recognition
        self.element_activations[9] = features['anomaly_score'] / 3.0  # E10_Anomaly_Detection
        self.element_activations[10] = psi_result['P_esc']  # E11_Escape_Potential
        self.element_activations[11] = 1.0 - psi_result['P_esc']  # E12_Singularity_Risk
        self.element_activations[12] = psi_result['energy_term']  # E13_Energy_Budget
        self.element_activations[13] = features['prompt_energy']  # E14_Prompt_Responsiveness
        self.element_activations[14] = 1.0 if prediction == 'RED_STAR' else 0.0  # E15_Mutual_Destruction
        self.element_activations[15] = psi_result['Psi_Red']  # E16_Final_Collapse
        
        # Normalize
        self.element_activations = np.clip(self.element_activations, 0.0, 1.0)
    
    def simulate_data_stream(self, domain: str, length: int = 500, 
                            inject_events: bool = True) -> np.ndarray:
        """
        Generates realistic simulated data streams for testing.
        Domains: 'market', 'weather', 'traffic'
        """
        np.random.seed(42)
        t = np.arange(length)
        
        if domain == 'market':
            # Geometric Brownian Motion with occasional crashes
            data = np.cumsum(np.random.normal(0, 0.02, length))
            data = 100 * np.exp(data * 0.1)  # Price-like
            
            if inject_events:
                # Inject Red Star: orderly correction
                data[150:180] *= np.linspace(1.0, 0.85, 30)  # Gradual decline
                data[180:200] *= np.linspace(0.85, 0.95, 20)  # Recovery
                
                # Inject Black Hole: flash crash
                data[350:360] *= np.linspace(1.0, 0.6, 10)  # Sharp drop
                data[360:370] *= np.linspace(0.6, 0.7, 10)   # Partial recovery
                
        elif domain == 'weather':
            # Sinusoidal base + noise + extreme events
            data = 20 + 10 * np.sin(2 * np.pi * t / 100) + np.random.normal(0, 2, length)
            
            if inject_events:
                # Red Star: stable front passage
                data[200:230] += np.linspace(0, 5, 30)  # Gradual warming
                data[230:250] += np.linspace(5, 0, 20)   # Return to normal
                
                # Black Hole: chaotic bifurcation
                data[400:420] += np.random.normal(0, 8, 20)  # Extreme volatility
                
        elif domain == 'traffic':
            # Flow rate with congestion dynamics
            base_flow = 80 + 20 * np.sin(2 * np.pi * t / 150)
            data = base_flow + np.random.normal(0, 5, length)
            
            if inject_events:
                # Red Star: shockwave dissipation
                data[100:130] -= np.linspace(0, 30, 30)  # Buildup
                data[130:160] += np.linspace(30, 0, 30)   # Smooth recovery
                
                # Black Hole: gridlock cascade
                data[300:315] -= np.linspace(0, 60, 15)   # Rapid collapse
                data[315:340] = np.random.uniform(10, 30, 25)  # Chaotic low flow
                
        else:
            data = np.cumsum(np.random.normal(0, 1, length))
            
        return data
    
    def run_realtime_simulation(self, domain: str = 'market', 
                              duration: int = 500,
                              plot_results: bool = True):
        """
        Runs end-to-end simulation of CCT classifier on live-like data stream.
        """
        print(f"🛸 CCT Live Stream Classifier: {domain.upper()} Domain")
        print(f"Trained Seed θ*: {self.theta:.4f}")
        print(f"Prediction Horizon: {self.HORIZON} steps")
        print("-" * 70)
        
        # Generate data stream
        data = self.simulate_data_stream(domain, duration, inject_events=True)
        
        # Run predictions
        predictions = []
        red_star_indices = []
        black_hole_indices = []
        
        for idx in range(self.WINDOW, len(data)):
            result = self.predict_event(data, idx)
            predictions.append(result['prediction'])
            
            if result['prediction'] == 'RED_STAR':
                red_star_indices.append(idx)
            elif result['prediction'] == 'BLACK_HOLE':
                black_hole_indices.append(idx)
                
            # Progress indicator
            if idx % 100 == 0:
                print(f"  Processed {idx}/{len(data)} | "
                      f"Red Stars: {len(red_star_indices)} | "
                      f"Black Holes: {len(black_hole_indices)}")
        
        # Compute metrics
        total_predictions = len(predictions)
        red_star_pct = len(red_star_indices) / total_predictions * 100
        black_hole_pct = len(black_hole_indices) / total_predictions * 100
        
        print("-" * 70)
        print(f"📊 SIMULATION RESULTS:")
        print(f"  Total Predictions: {total_predictions}")
        print(f"  Red Star Events: {len(red_star_indices)} ({red_star_pct:.1f}%)")
        print(f"  Black Hole Events: {len(black_hole_indices)} ({black_hole_pct:.1f}%)")
        print(f"  Average Confidence: {np.mean([p['confidence'] for p in self.prediction_log]):.3f}")
        
        # Plot results
        if plot_results:
            self._plot_simulation_results(data, predictions, red_star_indices, black_hole_indices)
            
        return {
            'data': data,
            'predictions': predictions,
            'red_star_indices': red_star_indices,
            'black_hole_indices': black_hole_indices,
            'metrics': {
                'red_star_pct': red_star_pct,
                'black_hole_pct': black_hole_pct,
                'avg_confidence': np.mean([p['confidence'] for p in self.prediction_log])
            }
        }
    
    def _plot_simulation_results(self, data, predictions, red_star_idx, black_hole_idx):
        """Visualizes prediction results"""
        fig, axs = plt.subplots(2, 2, figsize=(16, 10))
        
        # Plot 1: Data Stream with Event Markers
        ax = axs[0, 0]
        ax.plot(data, linewidth=1, label='Data Stream', color='steelblue')
        if red_star_idx:
            ax.scatter(red_star_idx, [data[i] for i in red_star_idx],
                      c='green', s=50, label='Red Star (Stable)', zorder=5, marker='*')
        if black_hole_idx:
            ax.scatter(black_hole_idx, [data[i] for i in black_hole_idx],
                      c='red', s=50, label='Black Hole (Crash)', zorder=5, marker='o')
        ax.set_title("Data Stream with CCT Event Classification")
        ax.set_xlabel("Time Step")
        ax.set_ylabel("Value")
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        # Plot 2: Prediction Timeline
        ax = axs[0, 1]
        pred_numeric = [1 if p == 'RED_STAR' else -1 if p == 'BLACK_HOLE' else 0 
                       for p in predictions]
        ax.plot(range(self.WINDOW, self.WINDOW + len(pred_numeric)), 
               pred_numeric, linewidth=2, color='purple')
        ax.axhline(y=0, color='gray', linestyle='--', alpha=0.5)
        ax.set_title("Prediction Timeline (1=Red Star, -1=Black Hole)")
        ax.set_xlabel("Time Step")
        ax.set_ylabel("Prediction")
        ax.grid(True, alpha=0.3)
        
        # Plot 3: Red Star Potential & Entropy
        ax = axs[1, 0]
        psi_values = [p['Psi_Red'] for p in self.prediction_log]
        h_values = [p['H_final'] for p in self.prediction_log]
        time_steps = range(self.WINDOW, self.WINDOW + len(psi_values))
        
        ax.plot(time_steps, psi_values, label='Ψ_Red (Red Star Potential)', 
               color='red', linewidth=2)
        ax.plot(time_steps, h_values, label='H(T) (Semantic Entropy)', 
               color='orange', linewidth=2, alpha=0.7)
        ax.axhline(y=self.P_CRITICAL, color='red', linestyle=':', 
                  label=f'P_critical ({self.P_CRITICAL})')
        ax.axhline(y=self.H_COLLAPSE, color='green', linestyle=':', 
                  label=f'H_collapse ({self.H_COLLAPSE})')
        ax.set_title("Red Star Components")
        ax.set_xlabel("Time Step")
        ax.set_ylabel("Value")
        ax.legend(fontsize=8)
        ax.grid(True, alpha=0.3)
        
        # Plot 4: 16-Element Activation Snapshot (latest)
        ax = axs[1, 1]
        elements = np.arange(16)
        ax.bar(elements, self.element_activations, color='steelblue', alpha=0.7)
        ax.set_xticks(elements)
        ax.set_xticklabels([f"E{i+1:02d}" for i in elements], rotation=45, ha='right', fontsize=7)
        ax.set_title("16-Element Semantic State (Latest)")
        ax.set_ylabel("Activation")
        ax.set_ylim(0, 1.1)
        ax.grid(True, alpha=0.3, axis='y')
        
        plt.tight_layout()
        plt.savefig('cct_simulation_results.png', dpi=150, bbox_inches='tight')
        print(f"\n📊 Plot saved to: cct_simulation_results.png")
        plt.close()


if __name__ == "__main__":
    # Initialize classifier with default trained seed
    classifier = CCT_LiveStream_Classifier()
    
    # Run simulation on market data
    results = classifier.run_realtime_simulation(domain='market', duration=500, plot_results=True)