#!/usr/bin/env python3
"""
Deep Entropy Red/Black Market Crash Predictor
- 30-layer MLP with per-layer exit heads (ETAD architecture)
- Integrates CCT Red Star/Black Hole crash prediction theory
- Uses multi-scale semantic entropy computation
- PyTorch implementation with adaptive depth prediction
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import pandas as pd
import yfinance as yf
import warnings
from collections import defaultdict
from datetime import datetime, timedelta

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

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.stats import linregress

# ============================================================================
# CONFIGURATION
# ============================================================================
class Config:
    # Architecture
    INPUT_DIM = 50  # Window size for market data
    HIDDEN_DIM = 128
    NUM_HIDDEN_LAYERS = 30  # 30 hidden layers as requested
    OUTPUT_DIM = 2  # Red Star (0) vs Black Hole (1)

    # Training
    BATCH_SIZE = 64
    EPOCHS = 20
    LR = 0.001
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'

    # Entropy thresholds (learnable per layer)
    ENTROPY_TARGET_START = 0.5  # Initial entropy target
    ENTROPY_TARGET_END = 0.1    # Final layer target (more certain)
    MIN_DEPTH = 5               # Minimum layers before allowing exit

    # Loss weights
    LAMBDA_ENTROPY = 0.1
    LAMBDA_DEPTH = 0.05

    # CCT Parameters
    THETA = 0.7342  # Trained ML Seed
    H_COLLAPSE = 0.3
    WINDOW_SIZE = 50
    PREDICTION_HORIZON = 10

    # Logging
    LOG_INTERVAL = 50
    SEED = 42

# ============================================================================
# ENTROPY UTILS
# ============================================================================
def compute_entropy(logits, dim=-1, eps=1e-8):
    """Compute Shannon entropy of probability distribution"""
    probs = F.softmax(logits, dim=dim)
    return -torch.sum(probs * torch.log(probs + eps), dim=dim)

def compute_semantic_entropy_numpy(data_window):
    """
    Multi-scale semantic entropy computation (NumPy version for data preprocessing)
    Combines distributional, spectral, and predictability entropy
    """
    data_window = data_window[~np.isnan(data_window)]
    
    if len(data_window) < 10:
        return 1.0
    
    # 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))
    
    # 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))
    
    # Predictability Entropy (AR model residuals)
    if len(data_window) > 20:
        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 batch_entropy(logits, eps=1e-8):
    """Mean entropy across batch"""
    return compute_entropy(logits, eps=eps).mean()

# ============================================================================
# MARKET DATA FEATURE ENGINEERING
# ============================================================================
class MarketFeatureExtractor:
    """Extracts CCT-based features from market data"""
    
    def __init__(self, window_size=50):
        self.window_size = window_size
        
    def extract_features(self, prices, current_idx):
        """
        Extract feature vector from price data at given index
        Returns feature array of shape (window_size,) or (window_size, n_features)
        """
        start_idx = max(0, current_idx - self.window_size + 1)
        window = prices[start_idx:current_idx + 1]
        
        if len(window) < self.window_size:
            # Pad with first value if insufficient data
            padding = np.full(self.window_size - len(window), window[0])
            window = np.concatenate([padding, window])
        
        return window
    
    def extract_multi_features(self, prices, current_idx):
        """Extract multiple features per timestep"""
        start_idx = max(0, current_idx - self.window_size + 1)
        end_idx = current_idx + 1
        window = prices[start_idx:end_idx]
        
        # Price features
        returns = np.diff(window) / (window[:-1] + 1e-10)
        returns = np.append(returns, 0)  # Pad to same length
        
        # Volatility (rolling std)
        vol = np.array([np.std(window[max(0, i-5):i+1]) for i in range(len(window))])
        vol = vol / (window + 1e-10)
        
        # Momentum
        momentum = np.array([
            (window[i] / window[max(0, i-5)] - 1) if i >= 5 else 0 
            for i in range(len(window))
        ])
        
        # Entropy per window
        entropy_vals = np.array([
            compute_semantic_entropy_numpy(window[max(0, i-10):i+1]) 
            for i in range(len(window))
        ])
        
        # Stack features: (window_size, 5)
        features = np.column_stack([
            window / (window[0] + 1e-10),  # Normalized price
            returns,
            vol,
            momentum,
            entropy_vals
        ])
        
        return features.astype(np.float32)

# ============================================================================
# DATASET
# ============================================================================
class MarketCrashDataset:
    """Creates training dataset from historical market data"""
    
    def __init__(self, ticker='BTC-USD', period='5y', window_size=50, horizon=10):
        self.ticker = ticker
        self.period = period
        self.window_size = window_size
        self.horizon = horizon
        
        # Fetch and prepare data
        self.df = self._fetch_data()
        self.prices = self.df['Close'].values
        self.feature_extractor = MarketFeatureExtractor(window_size)
        
        # Generate labels based on future returns
        self.labels = self._generate_labels()
        
    def _fetch_data(self):
        """Fetch market data"""
        print(f"📊 Fetching {self.ticker} data ({self.period})...")
        stock = yf.Ticker(self.ticker)
        df = stock.history(period=self.period)

        if df.empty:
            raise ValueError(f"No data fetched for {self.ticker}")

        print(f"✓ Retrieved {len(df)} data points")
        return df
    
    def _generate_labels(self):
        """
        Generate crash labels based on future price movement
        0 = Red Star (stable/recovering), 1 = Black Hole (crash)
        """
        labels = np.zeros(len(self.prices), dtype=np.int64)
        
        for i in range(len(self.prices) - self.horizon):
            future_return = (self.prices[i + self.horizon] - self.prices[i]) / (self.prices[i] + 1e-10)
            
            # Crash if future drop > 10%
            if future_return < -0.10:
                labels[i] = 1  # Black Hole
            else:
                labels[i] = 0  # Red Star
        
        return labels
    
    def create_samples(self, use_multi_features=True):
        """Create (features, labels) samples"""
        samples = []
        
        for i in range(self.window_size, len(self.prices) - self.horizon):
            if use_multi_features:
                features = self.feature_extractor.extract_multi_features(self.prices, i)
                # Flatten to 1D: (window_size * n_features,)
                features_flat = features.flatten()
            else:
                features_flat = self.feature_extractor.extract_features(self.prices, i)
            
            label = self.labels[i]
            samples.append((features_flat, label))
        
        return samples

# ============================================================================
# MODEL: ETAD MLP WITH CCT INTEGRATION
# ============================================================================
class DeepEntropyRedBlackMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        # Input projection
        self.input_proj = nn.Linear(config.INPUT_DIM, config.HIDDEN_DIM)
        self.input_norm = nn.LayerNorm(config.HIDDEN_DIM)
        
        # 30 Hidden layers with residual connections
        self.hidden_layers = nn.ModuleList([
            nn.Sequential(
                nn.Linear(config.HIDDEN_DIM, config.HIDDEN_DIM),
                nn.LayerNorm(config.HIDDEN_DIM),
                nn.ReLU(),
                nn.Dropout(0.1)
            )
            for _ in range(config.NUM_HIDDEN_LAYERS)
        ])
        
        # Per-layer exit heads (predict at any layer)
        self.exit_heads = nn.ModuleList([
            nn.Sequential(
                nn.Linear(config.HIDDEN_DIM, config.HIDDEN_DIM // 2),
                nn.ReLU(),
                nn.Linear(config.HIDDEN_DIM // 2, config.OUTPUT_DIM)
            )
            for _ in range(config.NUM_HIDDEN_LAYERS + 1)
        ])
        
        # Learnable entropy thresholds per layer
        self.entropy_thresholds = nn.Parameter(
            torch.linspace(config.ENTROPY_TARGET_START, config.ENTROPY_TARGET_END, config.NUM_HIDDEN_LAYERS + 1)
        )
        
        # CCT integration: Learnable Red Star potential head
        self.red_star_head = nn.Sequential(
            nn.Linear(config.HIDDEN_DIM, config.HIDDEN_DIM // 2),
            nn.ReLU(),
            nn.Linear(config.HIDDEN_DIM // 2, 1),
            nn.Sigmoid()
        )
        
        # Layer usage counters
        self.layer_exit_counts = defaultdict(int)
        
    def forward(self, x, return_all=False, force_depth=None):
        """
        Args:
            x: Input tensor [B, INPUT_DIM]
            return_all: If True, return all layer outputs
            force_depth: If set, exit at this layer regardless of entropy
            
        Returns:
            logits: Final predictions [B, 2]
            exit_layers: Which layer each sample exited at [B]
            red_star_potential: Ψ_Red values [B, 1]
        """
        batch_size = x.size(0)
        
        # Track exit layer for each sample
        exit_layers = torch.full((batch_size,), -1, dtype=torch.long, device=x.device)
        final_logits = None
        all_entropies = []
        all_logits = []
        
        # Layer 0: After input projection
        h = F.relu(self.input_norm(self.input_proj(x)))
        logits = self.exit_heads[0](h)
        entropy = compute_entropy(logits)
        all_entropies.append(entropy.clone())
        all_logits.append(logits.clone())
        
        # Check exit condition
        if force_depth is None:
            can_exit = 0 >= self.config.MIN_DEPTH
            should_exit = (entropy <= self.entropy_thresholds[0]) & can_exit & (exit_layers == -1)
        else:
            should_exit = torch.zeros(batch_size, dtype=torch.bool, device=x.device) if 0 < force_depth else torch.ones(batch_size, dtype=torch.bool, device=x.device)
        
        exited = should_exit & (exit_layers == -1)
        exit_layers[exited] = 0
        final_logits = logits.clone()
        final_logits = torch.where(exited.unsqueeze(1), logits, final_logits)
        
        # 30 Hidden layers with adaptive exit
        for l, layer in enumerate(self.hidden_layers):
            active_mask = exit_layers == -1
            if not active_mask.any():
                break
            
            # Residual connection
            h_prev = h
            h = layer(h)
            h = h + h_prev  # Residual
            
            logits = self.exit_heads[l + 1](h)
            entropy = compute_entropy(logits)
            all_entropies.append(entropy.clone())
            all_logits.append(logits.clone())
            
            # Check exit condition
            layer_idx = l + 1
            if force_depth is None:
                can_exit = layer_idx >= self.config.MIN_DEPTH
                should_exit = (entropy <= self.entropy_thresholds[layer_idx]) & can_exit
            else:
                should_exit = torch.zeros_like(entropy, dtype=torch.bool) if layer_idx < force_depth else torch.ones_like(entropy, dtype=torch.bool)
            
            exited = should_exit & active_mask
            exit_layers[exited] = layer_idx
            
            # Update final logits
            final_logits = torch.where(exited.unsqueeze(1), logits, final_logits)
        
        # Force exit remaining samples at final layer
        remaining = exit_layers == -1
        exit_layers[remaining] = self.config.NUM_HIDDEN_LAYERS
        final_logits = torch.where(remaining.unsqueeze(1), logits, final_logits)
        
        # Compute Red Star Potential (CCT integration)
        red_star_potential = self.red_star_head(h)
        
        # Update exit counts
        if self.training:
            for idx in exit_layers.cpu().numpy():
                self.layer_exit_counts[idx] += 1
        
        if return_all:
            return final_logits, exit_layers, red_star_potential, all_entropies, all_logits
        
        return final_logits, exit_layers, red_star_potential
    
    def get_entropy_profile(self):
        """Return current entropy thresholds"""
        return self.entropy_thresholds.detach().cpu().numpy()
    
    def get_exit_distribution(self):
        """Return distribution of exit layers"""
        total = sum(self.layer_exit_counts.values())
        if total == 0:
            return {}
        return {k: v / total for k, v in sorted(self.layer_exit_counts.items())}
    
    def reset_exit_counts(self):
        """Reset exit counters"""
        self.layer_exit_counts.clear()

# ============================================================================
# LOSS FUNCTIONS
# ============================================================================
class ETADLoss(nn.Module):
    def __init__(self, config, class_weights=None):
        super().__init__()
        self.config = config
        self.ce = nn.CrossEntropyLoss(weight=class_weights)
        
    def forward(self, logits, targets, exit_layers, red_star_potential=None, entropies=None):
        """
        Args:
            logits: Final predictions [B, 2]
            targets: Ground truth [B] (0=Red Star, 1=Black Hole)
            exit_layers: Which layer each sample exited at [B]
            red_star_potential: Ψ_Red values [B, 1]
            entropies: List of entropy tensors per layer
        """
        # 1. Task loss (cross-entropy)
        task_loss = self.ce(logits, targets)
        
        # 2. Entropy target loss
        entropy_loss = 0
        if entropies is not None:
            for l, entropy_l in enumerate(entropies):
                exited_mask = (exit_layers == l)
                if exited_mask.any():
                    target_entropy = self.config.ENTROPY_TARGET_START * \
                                    np.exp(-l / 5) + self.config.ENTROPY_TARGET_END
                    entropy_diff = (entropy_l[exited_mask] - target_entropy).abs()
                    entropy_loss += entropy_diff.mean()
            entropy_loss /= len(entropies)
        
        # 3. Depth penalty
        avg_depth = exit_layers.float().mean()
        depth_penalty = avg_depth / self.config.NUM_HIDDEN_LAYERS
        
        # 4. Red Star potential consistency loss (CCT integration)
        # Encourage Ψ_Red to be high for Red Star (0) and low for Black Hole (1)
        psi_loss = 0
        if red_star_potential is not None:
            # For Red Star samples (target=0), Ψ_Red should be high
            # For Black Hole samples (target=1), Ψ_Red should be low
            red_mask = (targets == 0)
            black_mask = (targets == 1)
            
            if red_mask.any():
                psi_loss += F.mse_loss(red_star_potential[red_mask], torch.ones_like(red_star_potential[red_mask]))
            if black_mask.any():
                psi_loss += F.mse_loss(red_star_potential[black_mask], torch.zeros_like(red_star_potential[black_mask]))
            psi_loss /= 2
        
        # Total loss
        total_loss = (
            task_loss +
            self.config.LAMBDA_ENTROPY * entropy_loss +
            self.config.LAMBDA_DEPTH * depth_penalty +
            0.2 * psi_loss  # CCT consistency weight
        )
        
        return total_loss, {
            'task_loss': task_loss.item(),
            'entropy_loss': entropy_loss.item() if entropies is not None else 0,
            'depth_penalty': depth_penalty.item(),
            'avg_depth': avg_depth.item(),
            'psi_loss': psi_loss.item() if isinstance(psi_loss, torch.Tensor) else 0
        }

# ============================================================================
# TRAINING & EVALUATION
# ============================================================================
def prepare_dataset(config, ticker='BTC-USD', period='5y'):
    """Prepare training dataset from market data"""
    dataset = MarketCrashDataset(ticker=ticker, period=period, 
                                window_size=config.WINDOW_SIZE, 
                                horizon=config.PREDICTION_HORIZON)
    
    samples = dataset.create_samples(use_multi_features=True)
    
    # Update input dim based on actual features
    if len(samples) > 0:
        config.INPUT_DIM = len(samples[0][0])
    
    # Split train/test
    split_idx = int(0.8 * len(samples))
    train_samples = samples[:split_idx]
    test_samples = samples[split_idx:]
    
    print(f"📊 Dataset: {len(samples)} samples ({len(train_samples)} train, {len(test_samples)} test)")
    print(f"   Red Star: {sum(1 for _, l in samples if l == 0)} | Black Hole: {sum(1 for _, l in samples if l == 1)}")
    
    return train_samples, test_samples, dataset

def create_data_loader(samples, batch_size, shuffle=True):
    """Create PyTorch data loader from samples"""
    def collate_fn(batch):
        features = torch.tensor([s[0] for s in batch], dtype=torch.float32)
        labels = torch.tensor([s[1] for s in batch], dtype=torch.long)
        return features, labels

    from torch.utils.data import DataLoader, TensorDataset
    features = torch.tensor([s[0] for s in samples], dtype=torch.float32)
    labels = torch.tensor([s[1] for s in samples], dtype=torch.long)
    dataset = TensorDataset(features, labels)
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)

def compute_class_weights(samples):
    """Compute class weights to handle imbalanced data (soft balancing)"""
    n_red = sum(1 for _, l in samples if l == 0)
    n_black = sum(1 for _, l in samples if l == 1)
    total = n_red + n_black
    
    # Soft balancing: sqrt instead of full ratio to avoid extreme weights
    # Increased Black Hole weight for better detection
    w_red = np.sqrt(total / (2.0 * n_red))
    w_black = np.sqrt(total / (2.0 * n_black)) * 1.5  # Boost Black Hole by 50%
    
    print(f"⚖️  Class Weights: Red Star={w_red:.3f}, Black Hole={w_black:.3f}")
    
    return torch.tensor([w_red, w_black], dtype=torch.float32)

def balance_classes_with_kmeans(samples, n_subclasses=None, verbose=True):
    """
    Balance imbalanced classes by splitting dominant class into KMeans subclasses.
    
    If Red Star has 1000 samples and Black Hole has 100, split Red into 9 subclasses
    of ~100 samples each, creating 10 balanced classes.
    """
    from sklearn.cluster import KMeans
    
    features = np.array([s[0] for s in samples])
    labels = np.array([s[1] for s in samples])
    
    # Find unique classes and their counts
    unique, counts = np.unique(labels, return_counts=True)
    class_counts = dict(zip(unique, counts))
    
    # Identify minority class size (target per class)
    minority_class = min(class_counts, key=class_counts.get)
    minority_count = class_counts[minority_class]
    
    if verbose:
        print(f"\n⚖️  Original class distribution:")
        for cls, cnt in sorted(class_counts.items()):
            print(f"   Class {cls}: {cnt} samples")
        print(f"   Minority class: {minority_class} ({minority_count} samples)")
    
    # For each dominant class, split into subclasses using KMeans
    balanced_features = []
    balanced_labels = []
    
    # Keep minority class as-is
    minority_mask = labels == minority_class
    balanced_features.append(features[minority_mask])
    balanced_labels.append(labels[minority_mask])
    
    # Split each dominant class
    for cls in unique:
        if cls == minority_class:
            continue
        
        cls_mask = labels == cls
        cls_features = features[cls_mask]
        
        # Number of subclasses = round up to get ~minority_count per subclass
        n_sub = max(1, int(np.ceil(len(cls_features) / minority_count)))
        if n_subclasses is not None:
            n_sub = n_subclasses
        
        if verbose:
            print(f"   Splitting class {cls} ({len(cls_features)} samples) into {n_sub} subclasses")
        
        # KMeans clustering
        kmeans = KMeans(n_clusters=n_sub, random_state=42, n_init=10)
        cluster_labels = kmeans.fit_predict(cls_features)
        
        # Assign new class labels
        # Start new class numbering after the original classes
        new_base = max(unique) + 1
        for sub in range(n_sub):
            sub_mask = cluster_labels == sub
            balanced_features.append(cls_features[sub_mask])
            balanced_labels.append(np.full(np.sum(sub_mask), new_base + sub))
    
    # Combine all
    all_features = np.vstack(balanced_features)
    all_labels = np.concatenate(balanced_labels)
    
    # Remap labels to start from 0
    unique_labels = np.unique(all_labels)
    label_map = {old: new for new, old in enumerate(unique_labels)}
    all_labels = np.array([label_map[l] for l in all_labels])
    
    # Track Black Hole class (original class 1)
    bh_class_mapped = label_map.get(1, 1)  # Default to 1 if not found
    
    # Create new samples
    balanced_samples = list(zip(all_features, all_labels))
    
    # Shuffle
    np.random.seed(42)
    indices = np.random.permutation(len(balanced_samples))
    balanced_samples = [balanced_samples[i] for i in indices]
    
    n_classes = len(np.unique(all_labels))
    samples_per_class = len(balanced_samples) // n_classes
    
    if verbose:
        new_counts = dict(zip(*np.unique(all_labels, return_counts=True)))
        print(f"\n📊 Balanced class distribution ({n_classes} classes):")
        for cls, cnt in sorted(new_counts.items()):
            print(f"   Class {cls}: {cnt} samples")
        print(f"   Total: {len(balanced_samples)} samples")
    
    return balanced_samples, n_classes, bh_class_mapped

def train_epoch(model, loader, optimizer, loss_fn, config, epoch):
    """Train one epoch"""
    model.train()
    total_loss = 0
    correct = 0
    total = 0
    metrics = defaultdict(float)
    
    for batch_idx, (data, target) in enumerate(loader):
        data, target = data.to(config.DEVICE), target.to(config.DEVICE)
        
        optimizer.zero_grad()
        
        # Forward with entropy tracking
        logits, exit_layers, psi_red, entropies, _ = model(data, return_all=True)
        
        # Compute loss
        loss, batch_metrics = loss_fn(logits, target, exit_layers, psi_red, entropies)
        
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        
        # Track metrics
        total_loss += loss.item()
        pred = logits.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)
        
        for k, v in batch_metrics.items():
            metrics[k] += v
        
        if batch_idx % config.LOG_INTERVAL == 0:
            print(f'Epoch {epoch} [{batch_idx}/{len(loader)}] '
                  f'Loss: {loss.item():.4f} Acc: {100.*correct/total:.2f}% '
                  f'AvgDepth: {batch_metrics["avg_depth"]:.2f}')
    
    model.reset_exit_counts()
    
    return {
        'loss': total_loss / len(loader),
        'accuracy': 100. * correct / total,
        **{k: v / len(loader) for k, v in metrics.items()}
    }

@torch.no_grad()
def evaluate(model, loader, loss_fn, config, bh_threshold=0.25, n_original_classes=2):
    """Evaluate model with Black Hole detection threshold"""
    model.eval()
    total_loss = 0
    correct = 0
    total = 0
    metrics = defaultdict(float)
    exit_distribution = defaultdict(int)
    accuracy_per_layer = defaultdict(lambda: {'correct': 0, 'total': 0})
    all_predictions = []
    all_targets = []
    
    # Per-class metrics (mapping back to original 2 classes)
    # Classes 0..n_original_classes-1 are original, rest are KMeans subclasses
    # For 2-class: class 0 = Red Star (or subclasses), class 1 = Black Hole
    true_pos = 0  # Black Hole correctly predicted
    false_pos = 0  # Red Star incorrectly predicted as Black Hole
    true_neg = 0  # Red Star correctly predicted
    false_neg = 0  # Black Hole incorrectly predicted as Red Star
    
    for data, target in loader:
        data, target = data.to(config.DEVICE), target.to(config.DEVICE)

        logits, exit_layers, psi_red, entropies, all_logits = model(data, return_all=True)

        # For test data with 2-class targets, compute loss only on first 2 output logits
        if config.OUTPUT_DIM > 2 and target.max() < 2:
            # Only use Black Hole (class 1) vs combined Red Star for loss
            loss, batch_metrics = loss_fn(logits[:, :2], target, exit_layers, psi_red, entropies)
        else:
            loss, batch_metrics = loss_fn(logits, target, exit_layers, psi_red, entropies)

        total_loss += loss.item()
        probs = F.softmax(logits, dim=1)
        
        # If model has more than 2 output classes (KMeans subclasses),
        # aggregate: Black Hole = mapped class, Red Star = all other classes
        if config.OUTPUT_DIM > 2:
            bh_class = getattr(config, 'BH_CLASS', 1)
            # Black Hole = specific class, all others = Red Star
            bh_prob = probs[:, bh_class]  # Black Hole probability
            pred = (bh_prob >= bh_threshold).long()
        else:
            # Standard 2-class
            bh_probs = probs[:, 1]
            pred = (bh_probs >= bh_threshold).long()
            
        correct += pred.eq(target).sum().item()
        total += target.size(0)

        all_predictions.extend(pred.cpu().numpy())
        all_targets.extend(target.cpu().numpy())
        
        # Per-class metrics
        for p, t in zip(pred.cpu().numpy(), target.cpu().numpy()):
            if t == 1 and p == 1:
                true_pos += 1
            elif t == 0 and p == 1:
                false_pos += 1
            elif t == 0 and p == 0:
                true_neg += 1
            elif t == 1 and p == 0:
                false_neg += 1

        # Track exit distribution
        for idx in exit_layers.cpu().numpy():
            exit_distribution[idx] += 1

        # Track accuracy per exit layer
        for i, (l, p, t) in enumerate(zip(exit_layers, pred, target)):
            layer_idx = l.item()
            accuracy_per_layer[layer_idx]['total'] += 1
            if p.item() == t.item():
                accuracy_per_layer[layer_idx]['correct'] += 1

        for k, v in batch_metrics.items():
            metrics[k] += v
    
    model.reset_exit_counts()
    
    # Compute per-layer accuracy
    layer_accuracy = {}
    for l, stats in accuracy_per_layer.items():
        if stats['total'] > 0:
            layer_accuracy[l] = 100. * stats['correct'] / stats['total']
    
    # Compute precision, recall, F1 for Black Hole class
    precision_bh = true_pos / (true_pos + false_pos) if (true_pos + false_pos) > 0 else 0.0
    recall_bh = true_pos / (true_pos + false_neg) if (true_pos + false_neg) > 0 else 0.0
    f1_bh = 2 * precision_bh * recall_bh / (precision_bh + recall_bh) if (precision_bh + recall_bh) > 0 else 0.0
    
    return {
        'loss': total_loss / len(loader),
        'accuracy': 100. * correct / total,
        'exit_distribution': {k: v / total for k, v in sorted(exit_distribution.items())},
        'layer_accuracy': layer_accuracy,
        'predictions': np.array(all_predictions),
        'targets': np.array(all_targets),
        'true_pos': true_pos,
        'false_pos': false_pos,
        'true_neg': true_neg,
        'false_neg': false_neg,
        'precision_bh': precision_bh,
        'recall_bh': recall_bh,
        'f1_bh': f1_bh,
        **{k: v / len(loader) for k, v in metrics.items()}
    }

# ============================================================================
# VISUALIZATION
# ============================================================================
def plot_results(train_metrics_history, test_metrics_history, final_test_metrics, model, dataset):
    """Create comprehensive visualization"""
    fig, axs = plt.subplots(3, 2, figsize=(18, 14))
    
    # 1. Training/Testing Loss & Accuracy
    ax = axs[0, 0]
    ax.plot(train_metrics_history['loss'], label='Train Loss', color='blue')
    ax.plot(test_metrics_history['loss'], label='Test Loss', color='red')
    ax.set_title('Loss Over Epochs', fontsize=14, fontweight='bold')
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Loss')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    ax = axs[0, 1]
    ax.plot(train_metrics_history['accuracy'], label='Train Accuracy', color='blue')
    ax.plot(test_metrics_history['accuracy'], label='Test Accuracy', color='red')
    ax.set_title('Accuracy Over Epochs', fontsize=14, fontweight='bold')
    ax.set_xlabel('Epoch')
    ax.set_ylabel('Accuracy (%)')
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    # 2. Exit Layer Distribution
    ax = axs[1, 0]
    exit_dist = final_test_metrics['exit_distribution']
    layers = list(exit_dist.keys())
    percentages = [exit_dist[l] * 100 for l in layers]
    ax.bar(layers, percentages, color='steelblue', alpha=0.7)
    ax.set_title('Exit Layer Distribution', fontsize=14, fontweight='bold')
    ax.set_xlabel('Layer')
    ax.set_ylabel('Percentage (%)')
    ax.grid(True, alpha=0.3)
    
    # 3. Layer-wise Accuracy
    ax = axs[1, 1]
    layer_acc = final_test_metrics['layer_accuracy']
    if layer_acc:
        sorted_layers = sorted(layer_acc.keys())
        accuracies = [layer_acc[l] for l in sorted_layers]
        ax.plot(sorted_layers, accuracies, marker='o', color='red', linewidth=2)
        ax.set_title('Accuracy by Exit Layer', fontsize=14, fontweight='bold')
        ax.set_xlabel('Layer')
        ax.set_ylabel('Accuracy (%)')
        ax.grid(True, alpha=0.3)
    
    # 4. Entropy Profile
    ax = axs[2, 0]
    entropy_profile = model.get_entropy_profile()
    ax.plot(range(len(entropy_profile)), entropy_profile, marker='s', color='orange', linewidth=2)
    ax.set_title('Learned Entropy Threshold Profile', fontsize=14, fontweight='bold')
    ax.set_xlabel('Layer')
    ax.set_ylabel('Entropy Threshold')
    ax.grid(True, alpha=0.3)
    
    # 5. Price Chart with Predictions
    ax = axs[2, 1]
    prices = dataset.prices
    dates = dataset.df.index
    
    # Test set starts at 80% of data
    split_idx = int(0.8 * (len(prices) - Config.WINDOW_SIZE - Config.PREDICTION_HORIZON))
    test_start = Config.WINDOW_SIZE + split_idx
    
    # Align test predictions with dates
    n_test = len(final_test_metrics['predictions'])
    pred_indices = range(test_start, test_start + n_test)
    pred_dates = [dates[i] for i in pred_indices]
    
    # Plot price for test period
    ax.plot(pred_dates, prices[test_start:test_start+n_test], 
            linewidth=1.5, label='Price', color='steelblue')
    
    # Mark Black Hole predictions
    black_hole_mask = (final_test_metrics['predictions'] == 1)
    if black_hole_mask.any():
        bh_indices = np.array(list(pred_indices))[black_hole_mask]
        ax.scatter([dates[i] for i in bh_indices],
                  [prices[i] for i in bh_indices],
                  c='red', s=30, marker='v', alpha=0.7, label='Black Hole Predicted')
    
    # Mark Red Star predictions
    red_star_mask = (final_test_metrics['predictions'] == 0)
    if red_star_mask.any():
        rs_indices = np.array(list(pred_indices))[red_star_mask]
        ax.scatter([dates[i] for i in rs_indices],
                  [prices[i] for i in rs_indices],
                  c='green', s=30, marker='*', alpha=0.7, label='Red Star Predicted')
    
    ax.set_title('Market Price with Crash Predictions', fontsize=14, fontweight='bold')
    ax.set_xlabel('Date')
    ax.set_ylabel('Price')
    ax.legend(fontsize=9)
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('deep_entropy_redblack_results.png', dpi=150, bbox_inches='tight')
    print(f"\n📊 Results saved to deep_entropy_redblack_results.png")
    plt.close()

# ============================================================================
# MAIN
# ============================================================================
def main():
    # Set seed
    torch.manual_seed(Config.SEED)
    np.random.seed(Config.SEED)
    
    print(f"🚀 Deep Entropy Red/Black Market Crash Predictor")
    print(f"   Device: {Config.DEVICE}")
    print(f"   Layers: {Config.NUM_HIDDEN_LAYERS} hidden + exit heads")
    print(f"   Min exit depth: {Config.MIN_DEPTH}")
    print(f"   Input dim: {Config.INPUT_DIM} (will be updated)")
    print()
    
    # Prepare dataset
    train_samples, test_samples, dataset = prepare_dataset(Config, ticker='BTC-USD', period='5y')

    print(f"   Actual input dim: {Config.INPUT_DIM}")
    print()

    # Balance classes with KMeans
    print("⚖️  Balancing classes with KMeans...")
    train_samples_balanced, n_classes, bh_class_mapped = balance_classes_with_kmeans(train_samples)
    
    # Update config for new number of classes
    Config.OUTPUT_DIM = n_classes
    Config.BH_CLASS = bh_class_mapped
    print(f"   Updated OUTPUT_DIM: {n_classes}")
    print(f"   Black Hole class (mapped): {bh_class_mapped}")
    print()

    # Create data loaders
    train_loader = create_data_loader(train_samples_balanced, Config.BATCH_SIZE, shuffle=True)
    test_loader = create_data_loader(test_samples, Config.BATCH_SIZE, shuffle=False)

    # Initialize model
    model = DeepEntropyRedBlackMLP(Config).to(Config.DEVICE)
    loss_fn = ETADLoss(Config)

    optimizer = optim.Adam(model.parameters(), lr=Config.LR, weight_decay=1e-5)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=Config.EPOCHS, eta_min=1e-6)
    
    # Training loop
    print("📈 Training...")
    print("=" * 80)
    
    best_acc = 0
    train_metrics_history = {'loss': [], 'accuracy': []}
    test_metrics_history = {'loss': [], 'accuracy': []}
    
    for epoch in range(1, Config.EPOCHS + 1):
        # Train
        train_metrics = train_epoch(model, train_loader, optimizer, loss_fn, Config, epoch)
        
        # Evaluate
        test_metrics = evaluate(model, test_loader, loss_fn, Config)
        
        # Store metrics
        train_metrics_history['loss'].append(train_metrics['loss'])
        train_metrics_history['accuracy'].append(train_metrics['accuracy'])
        test_metrics_history['loss'].append(test_metrics['loss'])
        test_metrics_history['accuracy'].append(test_metrics['accuracy'])
        
        # Print summary
        print(f"\n{'='*80}")
        print(f"EPOCH {epoch:2d} SUMMARY")
        print(f"{'='*80}")
        print(f"Train Loss: {train_metrics['loss']:.4f} | Acc: {train_metrics['accuracy']:.2f}%")
        print(f"Test  Loss: {test_metrics['loss']:.4f} | Acc: {test_metrics['accuracy']:.2f}%")
        print(f"Avg Depth: {test_metrics['avg_depth']:.2f} / {Config.NUM_HIDDEN_LAYERS}")
        print(f"Entropy Profile: {model.get_entropy_profile().round(3)}")
        print(f"Exit Distribution: {test_metrics['exit_distribution']}")
        
        # Save best
        if test_metrics['accuracy'] > best_acc:
            best_acc = test_metrics['accuracy']
            torch.save({
                'model': model.state_dict(),
                'config': Config,
                'entropy_thresholds': model.get_entropy_profile(),
                'epoch': epoch,
                'best_acc': best_acc
            }, 'deep_entropy_redblack_best.pt')
            print(f"   ✓ Saved best model (Acc: {best_acc:.2f}%)")
        
        scheduler.step()
    
    print(f"\n{'='*80}")
    print(f"✅ Training Complete! Best Test Acc: {best_acc:.2f}%")
    print(f"{'='*80}")
    
    # Final analysis
    print("\n📊 FINAL ANALYSIS")
    print("=" * 80)
    
    # Show exit layer distribution
    exit_dist = test_metrics['exit_distribution']
    print(f"Exit Layer Distribution:")
    for layer, pct in exit_dist.items():
        bar = '█' * int(pct * 50)
        print(f"  Layer {layer:2d}: {pct*100:5.1f}% {bar}")
    
    # Show accuracy vs depth
    print(f"\nAccuracy by Exit Layer:")
    for layer, acc in sorted(test_metrics['layer_accuracy'].items()):
        print(f"  Layer {layer:2d}: {acc:.2f}%")
    
    # Compute efficiency gain
    avg_depth = test_metrics['avg_depth']
    efficiency = 1 - (avg_depth / Config.NUM_HIDDEN_LAYERS)
    print(f"\n⚡ Efficiency Gain: {efficiency*100:.1f}% fewer layers on average")
    print(f"   (Avg depth: {avg_depth:.2f} vs max {Config.NUM_HIDDEN_LAYERS})")
    
    # Red Star vs Black Hole statistics
    predictions = test_metrics['predictions']
    targets = test_metrics['targets']
    red_star_count = np.sum(predictions == 0)
    black_hole_count = np.sum(predictions == 1)
    actual_red = np.sum(targets == 0)
    actual_black = np.sum(targets == 1)
    
    print(f"\n🔴 Red Star Predictions: {red_star_count} ({red_star_count/len(predictions)*100:.1f}%) [Actual: {actual_red}]")
    print(f"⚫ Black Hole Predictions: {black_hole_count} ({black_hole_count/len(predictions)*100:.1f}%) [Actual: {actual_black}]")
    
    # Confusion matrix
    print(f"\n📋 Confusion Matrix:")
    print(f"                    Predicted")
    print(f"                  RedStar  BlackHole")
    print(f"  Actual RedStar   {test_metrics['true_neg']:4d}     {test_metrics['false_pos']:4d}")
    print(f"  Actual BlackHole {test_metrics['false_neg']:4d}     {test_metrics['true_pos']:4d}")
    
    print(f"\n🎯 Black Hole Detection:")
    print(f"   Precision: {test_metrics['precision_bh']:.3f}")
    print(f"   Recall:    {test_metrics['recall_bh']:.3f}")
    print(f"   F1 Score:  {test_metrics['f1_bh']:.3f}")
    
    # Plot results
    print("\n📊 Generating visualizations...")
    plot_results(train_metrics_history, test_metrics_history, test_metrics, model, dataset)
    
    print("\n✅ Done!")

if __name__ == '__main__':
    main()