import numpy as np
import pandas as pd
import yfinance as yf
import warnings

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

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy import signal
from scipy.stats import entropy
from datetime import datetime, timedelta

class CCT_MarketCrash_Predictor:
    """
    Conditional Collapse Theory Classifier for Real Market Data
    Uses yfinance to fetch real market data and predict crash events
    using the trained ML Seed θ* from Fractal Escape Time Theory
    """

    def __init__(self, trained_theta: float = 0.7342,
                 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
        """
        self.theta = trained_theta
        self.H_COLLAPSE = collapse_threshold
        self.P_CRITICAL = critical_escape_prob
        self.WINDOW = window_size
        self.HORIZON = prediction_horizon

        # 16-Element Semantic State
        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"
        ]

        self.entropy_history = []
        self.prediction_log = []
        self.element_activations = np.zeros(16)

    def fetch_market_data(self, ticker: str = "^GSPC", 
                         period: str = "2y",
                         interval: str = "1d") -> pd.DataFrame:
        """
        Fetches real market data using yfinance
        Default: S&P 500 (^GSPC) for 2 years, daily data
        """
        print(f"📊 Fetching {ticker} data ({period}, {interval})...")
        
        stock = yf.Ticker(ticker)
        df = stock.history(period=period, interval=interval)
        
        if df.empty:
            raise ValueError(f"No data fetched for {ticker}")
        
        # Use closing prices
        prices = df['Close'].values
        dates = df.index
        
        print(f"✓ Retrieved {len(prices)} data points from {dates[0].strftime('%Y-%m-%d')} to {dates[-1].strftime('%Y-%m-%d')}")
        
        return df, prices, dates

    def compute_semantic_entropy(self, data_window: np.ndarray) -> float:
        """
        Computes Semantic Entropy H(T) for market data window
        Uses multi-scale analysis for uncertainty measurement
        """
        # Filter out NaN values
        data_window = data_window[~np.isnan(data_window)]
        
        if len(data_window) < 10:
            return 1.0

        # Method 1: Distributional Entropy (Shannon)
        hist, _ = np.histogram(data_window, bins='auto', density=True)
        hist = hist[hist > 0]
        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)
        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)

    def compute_escape_probability(self, data_window: np.ndarray, theta: float) -> float:
        """
        Computes Escape Probability P_esc(θ) using trained ML Seed
        Based on Red Star Theory fractal topology
        """
        features = self._extract_features(data_window)

        # Red Star ODE Dynamics
        y0 = features['initial_state']
        mu = features['control_parameter']

        # Simulate short trajectory
        y = y0
        escaped = False
        max_steps = 20

        for t in range(max_steps):
            D_term = theta * np.exp(-np.abs(y)**2) * features['prompt_energy']
            dy = (y**2 + mu + D_term) * 0.1
            y = y + dy

            if np.abs(y) < 10.0 and t > 5:
                escaped = True
                break
            if np.abs(y) > 100.0:
                break

        P_base = 1.0 if escaped else 0.2
        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 market data"""
        features = {}

        features['initial_state'] = (data[0] - np.mean(data)) / (np.std(data) + 1e-10)

        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

        features['prompt_energy'] = np.std(np.diff(data)) / (np.abs(np.mean(data)) + 1e-10)

        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

        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 alignment with trained seed θ*"""
        alignment = 0.0

        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:
            alignment += 0.2
        if np.abs(features['control_parameter']) < 1.5:
            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"""
        P_esc = self.compute_escape_probability(data_window, self.theta)

        H_final = self.compute_semantic_entropy(data_window[-self.WINDOW//2:])
        # Soft entropy indicator (sigmoid function instead of binary)
        entropy_indicator = 1.0 / (1.0 + np.exp(10 * (H_final - self.H_COLLAPSE)))

        E_destroy = self._estimate_destruction_energy(data_window)
        E_available = np.std(data_window) * len(data_window) * 0.1
        energy_term = np.exp(-E_destroy / (E_available + 1e-10))

        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 for mutual destruction (scales as 1/δ³)"""
        if len(data) < 10:
            return 10.0

        curvature = np.abs(np.gradient(np.gradient(data)))
        delta = 1.0 / (np.mean(curvature[-10:]) + 0.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 (Stable) vs Black Hole (Crash) event"""
        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}

        psi_result = self.compute_red_star_potential(window)

        # Simplified decision: Use relative threshold based on observed Psi_Red range
        # Calibrated for real market: most time is stable with occasional risk periods
        if psi_result['Psi_Red'] > 0.0015:  # Near median
            prediction = 'RED_STAR'
            confidence = min(psi_result['Psi_Red'] / 0.01, 1.0)
        else:
            prediction = 'BLACK_HOLE'
            confidence = 1.0 - min(psi_result['Psi_Red'] / 0.01, 1.0)

        self._update_element_activations(window, psi_result, 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"""
        features = self._extract_features(window)

        self.element_activations[0] = 1.0 - np.abs(features['control_parameter']) / 2.0
        self.element_activations[1] = 1.0 - np.clip(np.std(window) / 2.0, 0, 1)
        self.element_activations[2] = 1.0 - psi_result['H_final'] / 2.0
        self.element_activations[3] = features['feedback_strength']
        self.element_activations[4] = features['feedback_strength'] * 0.8
        self.element_activations[5] = 1.0 - features['prompt_energy']
        self.element_activations[6] = 1.0 - features['anomaly_score'] / 3.0
        self.element_activations[7] = psi_result['energy_term']
        self.element_activations[8] = psi_result['P_esc']
        self.element_activations[9] = features['anomaly_score'] / 3.0
        self.element_activations[10] = psi_result['P_esc']
        self.element_activations[11] = 1.0 - psi_result['P_esc']
        self.element_activations[12] = psi_result['energy_term']
        self.element_activations[13] = features['prompt_energy']
        self.element_activations[14] = 1.0 if prediction == 'RED_STAR' else 0.0
        self.element_activations[15] = psi_result['Psi_Red']

        self.element_activations = np.clip(self.element_activations, 0.0, 1.0)

    def analyze_market(self, ticker: str = "^GSPC", period: str = "2y",
                      plot_results: bool = True) -> dict:
        """
        Main analysis pipeline: Fetch data → Predict → Visualize
        """
        print(f"\n🛸 CCT Market Crash Predictor")
        print(f"Trained Seed θ*: {self.theta:.4f}")
        print(f"Market: {ticker} | Period: {period}")
        print("=" * 70)

        # Fetch real market data
        df, prices, dates = self.fetch_market_data(ticker, period)

        print(f"\n🔍 Analyzing {len(prices)} data points with CCT Classifier...")

        # Run predictions
        predictions = []
        red_star_indices = []
        black_hole_indices = []

        for idx in range(self.WINDOW, len(prices)):
            result = self.predict_event(prices, 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)

            if idx % 100 == 0:
                print(f"  Processed {idx}/{len(prices)} | "
                      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
        
        # Debug: Show Psi_Red statistics
        psi_values = [p['Psi_Red'] for p in self.prediction_log]
        print(f"\n📈 Psi_Red Statistics:")
        print(f"  Min: {np.min(psi_values):.6f}")
        print(f"  Max: {np.max(psi_values):.6f}")
        print(f"  Mean: {np.mean(psi_values):.6f}")
        print(f"  Median: {np.median(psi_values):.6f}")

        print("=" * 70)
        print(f"📊 CCT ANALYSIS RESULTS:")
        print(f"  Total Predictions: {total_predictions}")
        print(f"  Red Star (Stable) Events: {len(red_star_indices)} ({red_star_pct:.1f}%)")
        print(f"  Black Hole (Crash Risk) Events: {len(black_hole_indices)} ({black_hole_pct:.1f}%)")
        print(f"  Average Confidence: {np.mean([p['confidence'] for p in self.prediction_log]):.3f}")

        # Identify significant crash risk periods (weekly breakdown)
        if black_hole_indices:
            print(f"\n⚠️  HIGH CRASH RISK PERIODS (Weekly Breakdown):")
            crash_periods = self._cluster_events_weekly(black_hole_indices, dates, prices)
            for i, (start, end, severity, black_hole_count, total_days) in enumerate(crash_periods, 1):
                print(f"  {i}. {start.strftime('%Y-%m-%d')} to {end.strftime('%Y-%m-%d')} "
                      f"(Severity: {severity:.1%} | Risk Days: {black_hole_count}/{total_days})")

        # Plot results
        if plot_results:
            self._plot_results(prices, dates, predictions, 
                             red_star_indices, black_hole_indices, df)

        return {
            'data': df,
            'prices': prices,
            'dates': dates,
            '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 _cluster_events_weekly(self, indices: list, dates, prices: np.ndarray) -> list:
        """Clusters Black Hole events into weekly periods for detailed breakdown"""
        if not indices:
            return []

        # Convert indices to dates
        event_dates = [dates[i] for i in indices]
        event_prices = [prices[i] for i in indices]

        # Group by ISO week
        weeks = {}
        for i, (evt_date, evt_price) in enumerate(zip(event_dates, event_prices)):
            # Use ISO calendar: (year, week, weekday)
            iso_year, iso_week, _ = evt_date.isocalendar()
            week_key = (iso_year, iso_week)

            if week_key not in weeks:
                weeks[week_key] = {
                    'dates': [],
                    'prices': [],
                    'all_indices': []
                }
            weeks[week_key]['dates'].append(evt_date)
            weeks[week_key]['prices'].append(evt_price)
            weeks[week_key]['all_indices'].append(indices[i])

        # Build weekly periods
        clusters = []
        for week_key in sorted(weeks.keys()):
            week_data = weeks[week_key]
            start = min(week_data['dates'])
            end = max(week_data['dates'])
            price_slice = week_data['prices']

            # Severity: price drop within the week
            if len(price_slice) > 1 and max(price_slice) > 0:
                severity = 1.0 - (min(price_slice) / max(price_slice))
            else:
                severity = 0.0

            # Count total trading days in that week (from original data)
            iso_year, iso_week = week_key
            # Find first and last trading day of that week
            week_trading_days = []
            for i in range(len(dates)):
                d_iso = dates[i].isocalendar()
                if d_iso[0] == iso_year and d_iso[1] == iso_week:
                    week_trading_days.append(i)

            total_days = len(week_trading_days) if week_trading_days else 5
            black_hole_count = len(week_data['all_indices'])

            clusters.append((start, end, severity, black_hole_count, total_days))

        return clusters

    def _cluster_events(self, indices: list, dates, prices: np.ndarray,
                       threshold: int = 5) -> list:
        """Clusters nearby events into significant periods"""
        if not indices:
            return []

        clusters = []
        cluster_start = indices[0]
        cluster_end = indices[0]

        for idx in indices[1:]:
            if idx - cluster_end <= threshold:
                cluster_end = idx
            else:
                price_slice = prices[cluster_start:cluster_end+1]
                if len(price_slice) > 0 and np.max(price_slice) > 0:
                    severity = np.min(price_slice) / np.max(price_slice) - 1.0
                    severity = abs(severity)
                else:
                    severity = 0.0
                clusters.append((dates[cluster_start], dates[cluster_end],
                               severity))
                cluster_start = idx
                cluster_end = idx

        # Last cluster
        price_slice = prices[cluster_start:cluster_end+1]
        if len(price_slice) > 0 and np.max(price_slice) > 0:
            severity = np.min(price_slice) / np.max(price_slice) - 1.0
            severity = abs(severity)
        else:
            severity = 0.0
        clusters.append((dates[cluster_start], dates[cluster_end], severity))

        return clusters

    def _plot_results(self, prices, dates, predictions, 
                     red_star_idx, black_hole_idx, df):
        """Visualizes CCT analysis on real market data"""
        fig, axs = plt.subplots(3, 1, figsize=(16, 12))

        # Plot 1: Price Chart with Event Markers
        ax = axs[0]
        ax.plot(dates, prices, linewidth=1.5, label=f'{df.name if hasattr(df, "name") else "Market"} Price', 
               color='steelblue')
        if red_star_idx:
            ax.scatter([dates[i] for i in red_star_idx], 
                      [prices[i] for i in red_star_idx],
                      c='green', s=30, label='Red Star (Stable)', 
                      zorder=5, marker='*', alpha=0.7)
        if black_hole_idx:
            ax.scatter([dates[i] for i in black_hole_idx], 
                      [prices[i] for i in black_hole_idx],
                      c='red', s=30, label='Black Hole (Crash Risk)', 
                      zorder=5, marker='v', alpha=0.7)
        ax.set_title("CCT Market Analysis: Price with Event Classification", 
                    fontsize=14, fontweight='bold')
        ax.set_xlabel("Date")
        ax.set_ylabel("Price")
        ax.legend(loc='upper left', fontsize=9)
        ax.grid(True, alpha=0.3)

        # Plot 2: Red Star Potential & Entropy
        ax = axs[1]
        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))
        plot_dates = [dates[i] for i in time_steps]

        ax.plot(plot_dates, psi_values, label='Ψ_Red (Stability Potential)',
               color='red', linewidth=2)
        ax.plot(plot_dates, h_values, label='H(T) (Market Entropy)',
               color='orange', linewidth=2, alpha=0.7)
        ax.axhline(y=self.P_CRITICAL, color='red', linestyle=':',
                  label=f'P_critical ({self.P_CRITICAL})', alpha=0.5)
        ax.axhline(y=self.H_COLLAPSE, color='green', linestyle=':',
                  label=f'H_collapse ({self.H_COLLAPSE})', alpha=0.5)
        ax.set_title("CCT Stability Indicators", fontsize=14, fontweight='bold')
        ax.set_xlabel("Date")
        ax.set_ylabel("Value")
        ax.legend(fontsize=9, loc='upper left')
        ax.grid(True, alpha=0.3)

        # Plot 3: 16-Element Semantic State
        ax = axs[2]
        elements = np.arange(16)
        colors = ['green' if v > 0.5 else 'red' if v < 0.3 else 'orange' 
                 for v in self.element_activations]
        ax.bar(elements, self.element_activations, color=colors, alpha=0.7, 
              edgecolor='black', linewidth=0.5)
        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 (Current Market Condition)", 
                    fontsize=14, fontweight='bold')
        ax.set_ylabel("Activation")
        ax.set_ylim(0, 1.1)
        ax.grid(True, alpha=0.3, axis='y')

        plt.tight_layout()
        plt.savefig('cct_market_analysis.png', dpi=150, bbox_inches='tight')
        print(f"\n📊 Analysis plot saved to: cct_market_analysis.png")
        plt.close()


if __name__ == "__main__":
    # Initialize CCT Market Crash Predictor
    # Calibrated thresholds for real market data
    predictor = CCT_MarketCrash_Predictor(
        trained_theta=0.7342,
        collapse_threshold=0.5,      # Increased from 0.3 for real market volatility
        critical_escape_prob=0.3,     # Decreased from 0.5 for better sensitivity
        window_size=30,               # Smaller window for more responsive detection
        prediction_horizon=10
    )

    # Analyze S&P 500 for crash patterns
    results = predictor.analyze_market(
        ticker="^GSPC",  # S&P 500
        period="2y",
        plot_results=True
    )
