"""
=============================================================
MMA-CCT LIGHT: Memory-Efficient Neutrino Detector
CCT-ODE Framework - Minimal Implementation

Target: <500MB GPU memory, runs on laptop CPU
=============================================================
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import math
import numpy as np
from typing import Dict, Optional
from dataclasses import dataclass

# =============================================================================
# CONFIGURATION (Memory-Efficient)
# =============================================================================

@dataclass
class Config:
    # Model dimensions (small)
    d_model: int = 64           # Down from 256
    n_heads: int = 4            # Down from 8
    n_layers: int = 2           # Down from 6
    max_hits: int = 50          # Down from 200
    n_features: int = 12        
    n_classes: int = 2
    
    # Data
    train_samples: int = 5000   # Smaller dataset
    val_samples: int = 1000
    test_samples: int = 1000
    signal_ratio: float = 0.1   # 10% signal
    
    # Training
    batch_size: int = 32        # Smaller batch
    epochs: int = 5
    lr: float = 1e-3            # Higher LR for faster small model
    rare_weight: float = 3.0


CONFIG = Config()

# =============================================================================
# SIMPLE NEUTRINO SIMULATOR
# =============================================================================

class SimpleNeutrinoSimulator:
    """Minimal neutrino event simulator"""
    
    def __init__(self):
        self.signal_count = 0
        self.bg_count = 0
        
    def simulate(self, is_signal: bool) -> Dict:
        """Generate a single event"""
        
        if is_signal:
            self.signal_count += 1
            # Signal: more hits, higher energy
            num_hits = np.random.poisson(40) + 5
            energy = np.random.uniform(1.0, 50.0)  # GeV
            
            # Direction (random)
            theta = np.random.uniform(0, np.pi)
            phi = np.random.uniform(0, 2 * np.pi)
            direction = [
                np.sin(theta) * np.cos(phi),
                np.sin(theta) * np.sin(phi),
                np.cos(theta)
            ]
            
            # Flavor (random)
            flavor = np.random.randint(0, 3)
            
        else:
            self.bg_count += 1
            # Background: fewer hits, lower energy
            num_hits = np.random.poisson(10) + 2
            energy = np.random.uniform(0.01, 0.5)
            
            theta = np.random.uniform(0, np.pi)
            phi = np.random.uniform(0, 2 * np.pi)
            direction = [
                np.sin(theta) * np.cos(phi),
                np.sin(theta) * np.sin(phi),
                np.cos(theta)
            ]
            flavor = -1
        
        # Generate hits (simplified: just x, y, z positions)
        hits = np.random.randn(num_hits, 3) * 50  # 50m spread
        
        # Features extracted from hits
        if num_hits > 0:
            features = [
                np.log10(energy + 1e-6) / 3.0,  # log energy
                direction[0], direction[1], direction[2],  # direction
                num_hits / 50.0,  # normalized hit count
                hits[:, 0].mean() / 500.0,  # centroid x
                hits[:, 1].mean() / 500.0,  # centroid y
                hits[:, 2].mean() / 500.0,  # centroid z
                hits[:, 0].std() / 50.0,    # spread x
                hits[:, 1].std() / 50.0,    # spread y
                hits[:, 2].std() / 50.0,    # spread z
                flavor / 3.0 if flavor >= 0 else -1  # flavor
            ]
        else:
            features = [0.0] * 12
        
        # Normalize hits to [-1, 1]
        hits = hits / 500.0
        
        # Pad to max_hits
        if len(hits) < CONFIG.max_hits:
            padding = np.zeros((CONFIG.max_hits - len(hits), 3))
            hits = np.vstack([hits, padding])
        else:
            hits = hits[:CONFIG.max_hits]
        
        return {
            'hits': torch.FloatTensor(hits),
            'features': torch.FloatTensor(features),
            'label': torch.tensor(1 if is_signal else 0, dtype=torch.long),
            'num_hits': num_hits,
            'energy': energy
        }


class TinyDataset(Dataset):
    """Memory-efficient dataset"""
    
    def __init__(self, num_samples, signal_ratio, simulator):
        self.num_samples = num_samples
        self.simulator = simulator
        
        # Generate all data upfront (small dataset)
        num_signal = int(num_samples * signal_ratio)
        num_bg = num_samples - num_signal
        
        self.data = []
        
        # Generate background
        for _ in range(num_bg):
            self.data.append(simulator.simulate(is_signal=False))
        
        # Generate signal
        for _ in range(num_signal):
            self.data.append(simulator.simulate(is_signal=True))
        
        # Shuffle
        np.random.shuffle(self.data)
    
    def __len__(self):
        return self.num_samples
    
    def __getitem__(self, idx):
        return self.data[idx]


# =============================================================================
# LIGHTWEIGHT CCT-ODE MODEL
# =============================================================================

class TinyAttention(nn.Module):
    """Minimal attention mechanism with sensitivity tracking"""
    
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        
        # Single combined projection (memory efficient)
        self.proj = nn.Linear(d_model, d_model * 3)  # Q, K, V combined
        
    def forward(self, x, mask=None, track_entropy=False):
        """
        x: [batch, seq, d_model]
        Returns: [batch, seq, d_model], entropy
        """
        B, L, D = x.shape
        
        # Combined projection
        qkv = self.proj(x)
        qkv = qkv.reshape(B, L, 3, self.n_heads, self.d_k)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, B, heads, L, d_k]
        
        Q, K, V = qkv[0], qkv[1], qkv[2]
        
        # Attention
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        attn = F.softmax(scores, dim=-1)
        
        # Compute entropy for CCT-ODE tracking
        if track_entropy:
            prob = F.softmax(scores, dim=-1)
            entropy = -torch.sum(prob * torch.log(prob + 1e-10), dim=-1).mean()
        else:
            entropy = 0.0
        
        # Apply attention
        out = torch.matmul(attn, V)
        out = out.transpose(1, 2).reshape(B, L, D)
        
        return out, entropy


class TinyCCTLayer(nn.Module):
    """Single CCT-ODE layer (minimal)"""
    
    def __init__(self, d_model, n_heads, dropout=0.1):
        super().__init__()
        self.attn = TinyAttention(d_model, n_heads)
        self.ff = nn.Sequential(
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        
        self.layer_entropy = 0.0
        
    def forward(self, x, mask=None):
        out, entropy = self.attn(self.norm1(x), mask, track_entropy=True)
        self.layer_entropy = entropy.item()
        x = x + out
        
        x = x + self.ff(self.norm2(x))
        return x


class LightNeutrinoCCT(nn.Module):
    """
    Memory-efficient CCT-ODE Neutrino Detector.
    
    Architecture:
    - Hit embedding → Small Transformer (2 layers)
    - Feature MLP
    - Simple fusion + Classification
    """
    
    def __init__(self, config):
        super().__init__()
        self.config = config
        
        # ========== HIT BRANCH (Small Transformer) ==========
        self.hit_embed = nn.Sequential(
            nn.Linear(3, config.d_model),  # x, y, z only
            nn.LayerNorm(config.d_model),
            nn.GELU()
        )
        
        self.hit_pos = nn.Parameter(torch.randn(1, config.max_hits, config.d_model) * 0.02)
        
        self.hit_layers = nn.ModuleList([
            TinyCCTLayer(config.d_model, config.n_heads)
            for _ in range(config.n_layers)
        ])
        
        # ========== FEATURE BRANCH (MLP) ==========
        self.feature_net = nn.Sequential(
            nn.Linear(config.n_features, config.d_model),
            nn.LayerNorm(config.d_model),
            nn.GELU(),
            nn.Linear(config.d_model, config.d_model)
        )
        
        # ========== FUSION + CLASSIFICATION ==========
        self.fusion = nn.Linear(config.d_model * 2, config.d_model)
        
        self.classifier = nn.Sequential(
            nn.LayerNorm(config.d_model),
            nn.Linear(config.d_model, config.d_model // 2),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(config.d_model // 2, config.n_classes)
        )
        
        # ========== RARE EVENT DETECTOR ==========
        self.rare_detector = nn.Sequential(
            nn.Linear(config.d_model, 16),
            nn.GELU(),
            nn.Linear(16, 1),
            nn.Sigmoid()
        )
        
        # Entropy tracking (CCT-ODE)
        self.layer_entropies = []
        
    def forward(self, hits, features, track_entropy=True):
        """
        hits: [batch, max_hits, 3]
        features: [batch, n_features]
        """
        B = hits.shape[0]
        
        # ========== HIT BRANCH ==========
        # Embed hits
        hit_x = self.hit_embed(hits) + self.hit_pos
        hit_mask = (hits.abs().sum(dim=-1) > 1e-5).float()  # [B, max_hits]
        
        # Pass through transformer layers
        for layer in self.hit_layers:
            hit_x = layer(hit_x, mask=hit_mask.unsqueeze(1).unsqueeze(2))
        
        # Simple mean pooling (no complex attention pooling)
        hit_mask = hit_mask.unsqueeze(-1)
        hit_pooled = (hit_x * hit_mask).sum(dim=1) / (hit_mask.sum(dim=1) + 1e-6)
        
        # ========== FEATURE BRANCH ==========
        feat_rep = self.feature_net(features)
        
        # ========== FUSION ==========
        joint = torch.cat([hit_pooled, feat_rep], dim=-1)
        fused = self.fusion(joint)
        fused = F.gelu(fused)
        
        # ========== CLASSIFICATION ==========
        logits = self.classifier(fused)
        
        # ========== RARE EVENT ==========
        rare_prob = self.rare_detector(fused)
        
        # Track entropy
        if track_entropy:
            self.layer_entropies = [l.layer_entropy for l in self.hit_layers]
        
        return {
            'logits': logits,
            'rare_prob': rare_prob,
            'hit_repr': hit_pooled,
            'feat_repr': feat_rep,
            'fused': fused,
            'entropies': self.layer_entropies if track_entropy else []
        }
    
    def reset(self):
        self.layer_entropies = []


# =============================================================================
# CCT-ODE LOSS (Simplified)
# =============================================================================

class CCTLoss(nn.Module):
    """Conditional Collapse Theory loss - minimal version"""
    
    def __init__(self, rare_weight=3.0):
        super().__init__()
        self.rare_weight = rare_weight
        self.ce = nn.CrossEntropyLoss()
        
    def forward(self, outputs, labels):
        # Cross entropy
        loss_ce = self.ce(outputs['logits'], labels)
        
        # Rare event bonus
        signal_mask = (labels == 1).float()
        loss_rare = -self.rare_weight * (outputs['rare_prob'].squeeze() * signal_mask).mean()
        
        # Entropy bonus (encourage low entropy = high collapse)
        entropy_penalty = sum(outputs.get('entropies', []))
        
        total = loss_ce + loss_rare + 0.01 * entropy_penalty
        
        return total, {
            'ce': loss_ce.item(),
            'rare': loss_rare.item(),
            'entropy': entropy_penalty,
            'total': total.item()
        }


# =============================================================================
# TRAINING FUNCTION (Memory-Efficient)
# =============================================================================

def train_light_cct(config, device='cpu'):
    """Train the light CCT model"""
    
    print("=" * 50)
    print("MMA-CCT LIGHT - Memory Efficient Training")
    print("=" * 50)
    
    # Create simulator and datasets
    print("\n[1/4] Creating datasets...")
    sim = SimpleNeutrinoSimulator()
    
    train_ds = TinyDataset(config.train_samples, config.signal_ratio, sim)
    val_ds = TinyDataset(config.val_samples, config.signal_ratio, SimpleNeutrinoSimulator())
    test_ds = TinyDataset(config.test_samples, config.signal_ratio, SimpleNeutrinoSimulator())
    
    train_loader = DataLoader(train_ds, batch_size=config.batch_size, shuffle=True)
    val_loader = DataLoader(val_ds, batch_size=config.batch_size)
    test_loader = DataLoader(test_ds, batch_size=config.batch_size)
    
    print(f"  Train: {len(train_ds)} samples")
    print(f"  Val: {len(val_ds)} samples")
    print(f"  Test: {len(test_ds)} samples")
    
    # Create model
    print("\n[2/4] Creating model...")
    model = LightNeutrinoCCT(config).to(device)
    
    num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"  Parameters: {num_params:,}")
    
    # Estimate memory
    memory_bytes = num_params * 4  # float32
    memory_mb = memory_bytes / (1024 ** 2)
    print(f"  Est. memory: {memory_mb:.1f} MB (weights only)")
    
    # Training setup
    criterion = CCTLoss(rare_weight=config.rare_weight)
    optimizer = optim.AdamW(model.parameters(), lr=config.lr, weight_decay=0.01)
    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.epochs)
    
    # Training loop
    print("\n[3/4] Training...")
    print("-" * 50)
    
    best_val_loss = float('inf')
    
    for epoch in range(config.epochs):
        model.train()
        epoch_loss = 0
        epoch_correct = 0
        epoch_total = 0
        
        for batch in train_loader:
            hits = batch['hits'].to(device)
            features = batch['features'].to(device)
            labels = batch['label'].to(device)
            
            optimizer.zero_grad()
            outputs = model(hits, features)
            loss, _ = criterion(outputs, labels)
            
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            
            epoch_loss += loss.item()
            preds = outputs['logits'].argmax(1)
            epoch_correct += (preds == labels).sum().item()
            epoch_total += labels.size(0)
            
            model.reset()
        
        scheduler.step()
        
        train_loss = epoch_loss / len(train_loader)
        train_acc = epoch_correct / epoch_total
        
        # Validation
        model.eval()
        val_loss = 0
        val_correct = 0
        val_total = 0
        
        with torch.no_grad():
            for batch in val_loader:
                hits = batch['hits'].to(device)
                features = batch['features'].to(device)
                labels = batch['label'].to(device)
                
                outputs = model(hits, features, track_entropy=False)
                loss, _ = criterion(outputs, labels)
                
                val_loss += loss.item()
                preds = outputs['logits'].argmax(1)
                val_correct += (preds == labels).sum().item()
                val_total += labels.size(0)
                
                model.reset()
        
        val_loss = val_loss / len(val_loader)
        val_acc = val_correct / val_total
        
        print(f"Epoch {epoch+1:2d}/{config.epochs} | "
              f"Train: {train_loss:.3f}/{train_acc:.3f} | "
              f"Val: {val_loss:.3f}/{val_acc:.3f}")
        
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'light_cct_best.pt')
    
    print("-" * 50)
    
    # Test evaluation
    print("\n[4/4] Evaluating...")
    model.load_state_dict(torch.load('light_cct_best.pt'))
    model.eval()
    
    all_preds, all_labels, all_rare = [], [], []
    
    with torch.no_grad():
        for batch in test_loader:
            hits = batch['hits'].to(device)
            features = batch['features'].to(device)
            labels = batch['label'].to(device)
            
            outputs = model(hits, features, track_entropy=False)
            preds = outputs['logits'].argmax(1)
            
            all_preds.extend(preds.cpu().tolist())
            all_labels.extend(labels.cpu().tolist())
            all_rare.extend(outputs['rare_prob'].reshape(-1).cpu().tolist())
            
            model.reset()
    
    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)
    all_rare = np.array(all_rare)
    
    # Metrics
    accuracy = (all_preds == all_labels).mean()
    tp = ((all_preds == 1) & (all_labels == 1)).sum()
    tn = ((all_preds == 0) & (all_labels == 0)).sum()
    fp = ((all_preds == 1) & (all_labels == 0)).sum()
    fn = ((all_preds == 0) & (all_labels == 1)).sum()
    
    precision = tp / (tp + fp + 1e-6)
    recall = tp / (tp + fn + 1e-6)
    f1 = 2 * precision * recall / (precision + recall + 1e-6)
    
    try:
        from sklearn.metrics import roc_auc_score
        auc = roc_auc_score(all_labels, all_rare)
    except:
        auc = 0.5
    
    print("\n" + "=" * 50)
    print("RESULTS")
    print("=" * 50)
    print(f"Accuracy:  {accuracy:.4f}")
    print(f"Precision: {precision:.4f}")
    print(f"Recall:    {recall:.4f}")
    print(f"F1 Score:  {f1:.4f}")
    print(f"ROC-AUC:   {auc:.4f}")
    print(f"\nConfusion Matrix:")
    print(f"  TP: {tp:4d}  FN: {fn:4d}")
    print(f"  FP: {fp:4d}  TN: {tn:4d}")
    print("=" * 50)
    
    return model


# =============================================================================
# MAIN
# =============================================================================

if __name__ == "__main__":
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"Using device: {device}")
    
    model = train_light_cct(CONFIG, device)
    
    print("\nModel saved to: light_cct_best.pt")