#!/usr/bin/env python3
"""
Entropy-Triggered Adaptive Depth (ETAD) MLP for CIFAR-10
- N-layer MLP with per-layer exit heads
- Learns entropy target profile per layer
- Skips computation when entropy drops below threshold
- Single file: train + test
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
from collections import defaultdict

# ============================================================================
# CONFIGURATION
# ============================================================================
class Config:
    # Architecture
    INPUT_DIM = 32 * 32 * 3  # 3072
    HIDDEN_DIM = 100
    NUM_HIDDEN_LAYERS = 10
    OUTPUT_DIM = 10
    
    # Training
    BATCH_SIZE = 100
    EPOCHS = 300
    LR = 0.001
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # Entropy thresholds (learnable per layer)
    ENTROPY_TARGET_START = 0 #3.0  # L0 target
    ENTROPY_TARGET_END = 5    # Final layer target
    MIN_DEPTH = 2               # Minimum layers before allowing exit
    
    # Loss weights
    LAMBDA_ENTROPY = 0.1        # Entropy target loss weight
    LAMBDA_DEPTH = 0.01         # Penalize excessive depth
    
    # Logging
    LOG_INTERVAL = 100
    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 batch_entropy(logits, eps=1e-8):
    """Mean entropy across batch"""
    return compute_entropy(logits, eps=eps).mean()

# ============================================================================
# MODEL: ETAD MLP WITH PER-LAYER EXIT HEADS
# ============================================================================
class ETADMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        # Input projection
        self.input_proj = nn.Linear(config.INPUT_DIM, config.HIDDEN_DIM)
        
        # Hidden layers
        self.hidden_layers = nn.ModuleList([
            nn.Linear(config.HIDDEN_DIM, config.HIDDEN_DIM)
            for _ in range(config.NUM_HIDDEN_LAYERS)
        ])
        
        # Per-layer exit heads (predict at any layer)
        self.exit_heads = nn.ModuleList([
            nn.Linear(config.HIDDEN_DIM, config.OUTPUT_DIM)
            for _ in range(config.NUM_HIDDEN_LAYERS + 1)  # +1 for after input_proj
        ])
        
        # Learnable entropy thresholds per layer
        # Initialized to decay from START to END

        self.entropy_thresholds = nn.Parameter(
            torch.linspace(config.ENTROPY_TARGET_START, config.ENTROPY_TARGET_END, config.NUM_HIDDEN_LAYERS + 1)
        )
        
        # Layer usage counters (for tracking, not trainable)
        self.layer_exit_counts = defaultdict(int)
        
    def forward(self, x, return_all=False, force_depth=None):
        """
        Args:
            x: Input tensor [B, 3072]
            return_all: If True, return all layer outputs (for analysis)
            force_depth: If set, exit at this layer regardless of entropy
            
        Returns:
            logits: Final predictions
            exit_layers: Which layer each sample exited at
            entropies: Entropy at each layer (if return_all)
        """
        batch_size = x.size(0)
        
        # Track exit layer for each sample (-1 = not exited yet)
        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_proj(x))
        logits = self.exit_heads[0](h)
        entropy = compute_entropy(logits)  # [B]
        all_entropies.append(entropy.clone())
        all_logits.append(logits.clone())
        
        # Check exit condition for layer 0
        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
        if final_logits is None:
            final_logits = logits.clone()
        final_logits = torch.where(
            exited.unsqueeze(1),
            logits,
            final_logits
        )
        
        # Hidden layers
        for l, layer in enumerate(self.hidden_layers):
            # Skip computation for samples that already exited
            active_mask = exit_layers == -1
            if not active_mask.any():
                break
                
            h = F.relu(layer(h))
            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:
                # Don't exit before MIN_DEPTH
                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 for exited samples
            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
        )
        
        # Update exit counts (for analysis)
        if self.training:
            for idx in exit_layers.cpu().numpy():
                self.layer_exit_counts[idx] += 1
        
        if return_all:
            return final_logits, exit_layers, all_entropies, all_logits
        
        return final_logits, exit_layers
    
    def get_entropy_profile(self):
        """Return current entropy thresholds (the learned profile)"""
        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 for new evaluation period"""
        self.layer_exit_counts.clear()

# ============================================================================
# LOSS FUNCTIONS
# ============================================================================
class ETADLoss(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.ce = nn.CrossEntropyLoss()
        
    def forward(self, logits, targets, exit_layers, entropies=None):
        """
        Args:
            logits: Final predictions [B, 10]
            targets: Ground truth [B]
            exit_layers: Which layer each sample exited at [B]
            entropies: List of entropy tensors per layer (optional)
        """
        # 1. Task loss (cross-entropy)
        task_loss = self.ce(logits, targets)
        
        # 2. Entropy target loss (encourage entropy to match thresholds at exit)
        entropy_loss = 0
        if entropies is not None:
            for l, entropy_l in enumerate(entropies):
                # Only penalize layers where samples actually exited
                exited_mask = (exit_layers == l)
                if exited_mask.any():
                    target_entropy = self.config.ENTROPY_TARGET_START * \
                                    np.exp(-l / 3) + self.config.ENTROPY_TARGET_END
                    # Penalize if entropy at exit is too far from target
                    entropy_diff = (entropy_l[exited_mask] - target_entropy).abs()
                    entropy_loss += entropy_diff.mean()
            
            entropy_loss /= len(entropies)
        
        # 3. Depth penalty (encourage early exit when possible)
        avg_depth = exit_layers.float().mean()
        depth_penalty = avg_depth / self.config.NUM_HIDDEN_LAYERS
        
        # Total loss
        total_loss = (
            task_loss +
            self.config.LAMBDA_ENTROPY * entropy_loss +
            self.config.LAMBDA_DEPTH * depth_penalty
        )
        
        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()
        }

# ============================================================================
# TRAINING & EVALUATION
# ============================================================================
def get_cifar10_loaders(config):
    """Load CIFAR-10 with proper preprocessing"""
    transform_train = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), 
                           (0.2023, 0.1994, 0.2010)),
        transforms.Lambda(lambda x: x.view(-1))  # Flatten to 3072
    ])
    
    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), 
                           (0.2023, 0.1994, 0.2010)),
        transforms.Lambda(lambda x: x.view(-1))
    ])
    
    train_dataset = datasets.CIFAR10(
        root='../data', train=True, download=True, transform=transform_train
    )
    test_dataset = datasets.CIFAR10(
        root='../data', train=False, download=True, transform=transform_test
    )
    
    train_loader = DataLoader(train_dataset, batch_size=config.BATCH_SIZE, 
                             shuffle=True, num_workers=2)
    test_loader = DataLoader(test_dataset, batch_size=config.BATCH_SIZE, 
                            shuffle=False, num_workers=2)
    
    return train_loader, test_loader

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)

        for _ in range(7):        
            optimizer.zero_grad()
            
            # Forward with entropy tracking
            logits, exit_layers, entropies, _ = model(data, return_all=True)
            
            # Compute loss
            loss, batch_metrics = loss_fn(logits, target, exit_layers, entropies)
            
            loss.backward()
            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}')
    
    # Reset exit counts for next epoch
    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):
    """Evaluate model"""
    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})
    
    for data, target in loader:
        data, target = data.to(config.DEVICE), target.to(config.DEVICE)
        
        logits, exit_layers, entropies, all_logits = model(data, return_all=True)
        
        loss, batch_metrics = loss_fn(logits, target, exit_layers, entropies)
        
        total_loss += loss.item()
        pred = logits.argmax(dim=1)
        correct += pred.eq(target).sum().item()
        total += target.size(0)
        
        # 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']
    
    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,
        **{k: v / len(loader) for k, v in metrics.items()}
    }

# ============================================================================
# MAIN
# ============================================================================
def main():
    # Set seed
    torch.manual_seed(Config.SEED)
    np.random.seed(Config.SEED)
    
    print(f"🚀 ETAD MLP for CIFAR-10")
    print(f"   Device: {Config.DEVICE}")
    print(f"   Layers: {Config.NUM_HIDDEN_LAYERS} hidden + exit heads")
    print(f"   Min exit depth: {Config.MIN_DEPTH}")
    print()
    
    # Load data
    train_loader, test_loader = get_cifar10_loaders(Config)
    
    # Initialize model
    model = ETADMLP(Config).to(Config.DEVICE)
    optimizer = optim.Adam(model.parameters(), lr=Config.LR)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.5)
    loss_fn = ETADLoss(Config)
    
    # Training loop
    print("📈 Training...")
    print("=" * 80)
    
    best_acc = 0
    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)
        
        # 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(2)}")
        print(f"Exit Distribution: {test_metrics['exit_distribution']}")
        print(f"Layer Accuracy: {test_metrics['layer_accuracy']}")
        
        # Save best
        if test_metrics['accuracy'] > best_acc:
            best_acc = test_metrics['accuracy']
            torch.save({
                'model': model.state_dict(),
                'entropy_thresholds': model.get_entropy_profile(),
                'epoch': epoch,
            }, 'etad_mlp_best.pt')
        
        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 which layers are actually used
    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})")

if __name__ == '__main__':
    main()