"""
Leaps of Intelligence: ResNet18 → MLP via Lottery Ticket Hypothesis
===================================================================
Goal: Find the "true mechanism" (leaps) hidden in conv layers, extract as MLP.

Hypothesis: Dense networks contain sparse sub-networks (winning tickets) 
that encode the actual solution. These "leaps of intelligence" can be 
extracted and expressed as a direct MLP without conv scaffolding.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as transforms
import numpy as np
from collections import OrderedDict
import time

# ============================================================
# Configuration
# ============================================================
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
BATCH_SIZE = 128
EPOCHS_CONV = 10      # Train ResNet18
EPOCHS_MLP = 20       # Train extracted MLP
PRUNE_RATIOS = [0.5, 0.7, 0.8, 0.9, 0.95]  # Different pruning levels to test
MNIST_PATH = './data'

# ============================================================
# Modified ResNet18 for MNIST (28x28 grayscale)
# ============================================================
class ModifiedResNet18(nn.Module):
    """
    Standard ResNet18 but modified for 28x28 MNIST input.
    Removes initial maxpool (too aggressive for small images).
    """
    def __init__(self, num_classes=10):
        super().__init__()
        
        # Initial conv (matches ResNet first layer)
        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        
        # Residual blocks (instead of layer1-4)
        self.layer1 = self._make_layer(64, 64, 2, stride=1)
        self.layer2 = self._make_layer(64, 128, 2, stride=2)
        self.layer3 = self._make_layer(128, 256, 2, stride=2)
        self.layer4 = self._make_layer(256, 512, 2, stride=2)
        
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, num_classes)
        
        # Store which connections are "alive" (for pruning)
        self.alive_masks = {}
        
    def _make_layer(self, in_channels, out_channels, blocks, stride):
        layers = []
        layers.append(BasicBlock(in_channels, out_channels, stride))
        for _ in range(1, blocks):
            layers.append(BasicBlock(out_channels, out_channels, 1))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x
    
    def get_params_count(self):
        """Count total parameters"""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)


class BasicBlock(nn.Module):
    """Standard ResNet Basic Block"""
    expansion = 1
    
    def __init__(self, in_channels, out_channels, stride):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, 1, 1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1, stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )
    
    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out


# ============================================================
# Sparse MLP ("Leaps Network")
# ============================================================
class LeapsMLP(nn.Module):
    """
    MLP that receives the EXACT pruned connectivity pattern from ResNet18.
    Each "leap" is a direct connection between layers, skipping hierarchy.
    
    Instead of: Input → Conv → Conv → Pool → Conv → ... → FC → Output
    We have:    Input → [Leap Layer 1] → [Leap Layer 2] → ... → Output
    
    The "leaps" are the non-zero weights from the pruned ResNet.
    """
    def __init__(self, original_model):
        super().__init__()
        self._build_from_resnet(original_model)
        
    def _build_from_resnet(self, resnet):
        """Extract the sparse leap connections from ResNet"""
        
        # Flatten all FC layers directly
        fc_layers = []
        
        # Conv1 -> FC (first leap from raw pixels)
        fc_layers.append(('conv1_fc', self._extract_conv_to_fc(resnet.conv1, resnet.bn1)))
        
        # Layer outputs -> FC
        for name, module in [('layer1', resnet.layer1), 
                              ('layer2', resnet.layer2),
                              ('layer3', resnet.layer3),
                              ('layer4', resnet.layer4)]:
            fc_layers.append((f'{name}_fc', self._extract_layer_to_fc(module)))
        
        # Final FC
        fc_layers.append(('fc_final', self._extract_fc(resnet.fc)))
        
        self.layers = nn.ModuleDict(OrderedDict(fc_layers))
        
    def _extract_conv_to_fc(self, conv, bn):
        """Extract weights from conv+bn to FC layer"""
        # Get effective weights: conv weights × bn scales / std
        w = conv.weight.data
        if bn is not None:
            w = w * bn.weight.data.view(64, 1, 1, 1) / (bn.running_var.view(64, 1, 1, 1).sqrt() + 1e-5)
        # Flatten spatial
        w = w.view(64, -1)  # (64, 3*3*1)
        return nn.Linear(w.shape[1], w.shape[0], bias=False)
        
    def _extract_layer_to_fc(self, layer):
        """Extract from layer output to FC"""
        # Use the last conv in the layer
        last_conv = list(layer.modules())[-2] if isinstance(list(layer.modules())[-1], BasicBlock) else None
        for m in layer.modules():
            if isinstance(m, nn.Conv2d):
                last_conv = m
        return nn.Linear(512, 10, bias=False)
    
    def _extract_fc(self, fc):
        """Copy FC layer"""
        new_fc = nn.Linear(fc.in_features, fc.out_features, bias=False)
        new_fc.weight.data = fc.weight.data.clone()
        return new_fc
    
    def forward(self, x):
        x = x.view(x.size(0), -1)  # Flatten
        for name, layer in self.layers.items():
            x = layer(x)
            if name != 'fc_final':
                x = F.relu(x)
        return x


# ============================================================
# Alternative: Learnable Leap Connections
# ============================================================
class LeapConvLayer(nn.Module):
    """
    A single "leap" layer that learns non-local connections.
    Instead of conv (local), we use a learned linear projection
    that connects every input to every output (with sparsity).
    """
    def __init__(self, in_features, out_features, sparsity=0.1):
        super().__init__()
        self.layer = nn.Linear(in_features, out_features, bias=False)
        self.sparsity = sparsity
        self._init_sparse()
        
    def _init_sparse(self):
        """Initialize with sparse random weights (top-K selection)"""
        pass  # Will be pruned later
        
    def apply_leap_mask(self, mask):
        """Apply a learned leap pattern (binary mask)"""
        self.layer.weight.data *= mask
        
    def forward(self, x):
        return self.layer(x)


class IntelligentLeapsNet(nn.Module):
    """
    MLP where each layer learns "leaps of intelligence" - 
    direct semantic connections without conv locality.
    """
    def __init__(self, input_size=784, hidden_sizes=[256, 128], output_size=10):
        super().__init__()
        
        self.leap1 = LeapConvLayer(input_size, hidden_sizes[0])
        self.leap2 = LeapConvLayer(hidden_sizes[0], hidden_sizes[1])
        self.leap3 = LeapConvLayer(hidden_sizes[1], output_size)
        
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.leap1(x))
        x = F.relu(self.leap2(x))
        x = self.leap3(x)
        return x


# ============================================================
# Pruning: Find the Winning Ticket
# ============================================================
def magnitude_prune(model, prune_ratio):
    """
    Prune weights with lowest magnitude (standard magnitude pruning).
    Returns pruned model with mask.
    """
    mask_dict = {}
    
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() > 1:  # Only weight matrices
            # Get threshold for pruning
            flat_weights = param.data.abs().flatten()
            threshold = torch.quantile(flat_weights, prune_ratio)
            
            # Create mask
            mask = (param.data.abs() > threshold).float()
            mask_dict[name] = mask
            
            # Apply mask
            param.data *= mask
            
    return model, mask_dict


def find_winning_ticket(model, train_loader, device, prune_ratio=0.7, epochs=5):
    """
    Iterative magnitude pruning to find winning ticket.
    Train → Prune → Reset → Repeat
    """
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    
    # Initial training
    print(f"  Training initial network...")
    for epoch in range(epochs):
        model.train()
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()
    
    # First prune pass
    print(f"  Applying magnitude pruning ({prune_ratio*100:.0f}%)...")
    model, mask_dict = magnitude_prune(model, prune_ratio)
    
    return model, mask_dict


def extract_leaps_as_mlp(resnet_model, mask_dict):
    """
    Convert pruned ResNet to MLP using extracted leap patterns.
    """
    # Build a direct MLP from conv weights
    layers = []
    
    # Input to first hidden (use conv1 weights)
    conv1_w = resnet_model.conv1.weight.data
    conv1_w_flat = conv1_w.view(conv1_w.size(0), -1)
    
    # Apply mask if available
    if 'conv1.weight' in mask_dict:
        mask = mask_dict['conv1.weight'].view(conv1_w.size(0), -1)
        conv1_w_flat = conv1_w_flat * mask
    
    in_features = conv1_w_flat.shape[1]  # 3*3*1 = 9 for MNIST
    out_features = conv1_w_flat.shape[0]  # 64
    
    layers.append(nn.Linear(in_features, out_features, bias=False))
    layers[-1].weight.data = conv1_w_flat
    
    # Hidden layers (simplified - just map from conv outputs)
    # In reality, we'd need to forward through actual model
    # Here we approximate with a direct mapping
    layers.append(nn.Linear(64, 128, bias=False))
    layers.append(nn.Linear(128, 10, bias=False))
    
    mlp = nn.Sequential(*layers)
    return mlp


# ============================================================
# Train/Test Functions
# ============================================================
def train_model(model, train_loader, epochs, device, lr=0.001, name="Model"):
    """Train any model"""
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        correct = 0
        total = 0
        
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)
        
        acc = 100. * correct / total
        print(f"  {name} Epoch {epoch+1}/{epochs} - Loss: {total_loss/len(train_loader):.4f}, Acc: {acc:.2f}%")
    
    return model


def test_model(model, test_loader, device, name="Model"):
    """Evaluate any model"""
    model.eval()
    correct = 0
    total = 0
    
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)
    
    accuracy = 100. * correct / total
    print(f"  {name} Test Accuracy: {accuracy:.2f}%")
    return accuracy


# ============================================================
# MAIN EXPERIMENT
# ============================================================
def main():
    print("="*60)
    print("LEAPS OF INTELLIGENCE: ResNet18 → MLP via Lottery Ticket")
    print("="*60)
    print(f"Device: {DEVICE}")
    print()
    
    # ============================================================
    # 1. Load MNIST Data
    # ============================================================
    print("1. Loading MNIST data...")
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    
    train_set = torchvision.datasets.MNIST(MNIST_PATH, train=True, download=True, transform=transform)
    test_set = torchvision.datasets.MNIST(MNIST_PATH, train=False, download=True, transform=transform)
    
    train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
    test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False)
    
    print(f"  Train: {len(train_set)} samples, Test: {len(test_set)} samples")
    print()
    
    # ============================================================
    # 2. Train ResNet18 (The Scaffold)
    # ============================================================
    print("2. Training ResNet18 (scaffold network)...")
    print(f"  Parameters: {sum(p.numel() for p in ModifiedResNet18().parameters()):,}")
    
    resnet = ModifiedResNet18().to(DEVICE)
    start_time = time.time()
    resnet = train_model(resnet, train_loader, EPOCHS_CONV, DEVICE, name="ResNet18")
    resnet_time = time.time() - start_time
    
    resnet_acc = test_model(resnet, test_loader, DEVICE, name="ResNet18")
    print(f"  Time: {resnet_time:.2f}s")
    print()
    
    # ============================================================
    # 3. Find Winning Tickets (Leaps) via Pruning
    # ============================================================
    print("3. Finding 'Leaps of Intelligence' via pruning...")
    
    results = []
    
    for prune_ratio in PRUNE_RATIOS:
        print(f"\n  --- Pruning ratio: {prune_ratio*100:.0f}% ---")
        
        # Reset and retrain (lottery ticket style)
        resnet_pruned = ModifiedResNet18().to(DEVICE)
        
        # Quick retrain
        train_model(resnet_pruned, train_loader, 3, DEVICE, name="Pruned")
        
        # Apply pruning
        resnet_pruned, mask = magnitude_prune(resnet_pruned, prune_ratio)
        
        # Count remaining parameters
        remaining = sum(p.sum().item() for p in mask.values())
        total_params = sum(p.numel() for p in resnet_pruned.parameters() if 'weight' in p.name or 'bias' in p.name)
        
        print(f"  Remaining params: {remaining:,.0f} ({remaining/total_params*100:.1f}%)")
        
        # ============================================================
        # 4. Extract Leaps as MLP
        # ============================================================
        print(f"  Extracting leaps as MLP...")
        
        # Simple extraction: map from flattened conv to output
        leaps_mlp = IntelligentLeapsNet(input_size=784, hidden_sizes=[256, 128], output_size=10).to(DEVICE)
        
        # Copy leap patterns from pruned network
        with torch.no_grad():
            # First leap: from flattened image to first hidden
            leaps_mlp.leap1.layer.weight.data[:64] = resnet_pruned.conv1.weight.data.view(64, -1)[:64]
            
            # Second leap: first hidden to second
            for i, (name, param) in enumerate(resnet_pruned.named_parameters()):
                if 'layer4.1.conv2.weight' in name:
                    leaps_mlp.leap2.layer.weight.data[:512] = param.data.view(512, -1)[:512]
                    break
            
            # Final leap
            leaps_mlp.leap3.layer.weight.data[:128] = resnet_pruned.fc.weight.data[:128]
        
        # ============================================================
        # 5. Train the Extracted MLP
        # ============================================================
        print(f"  Training extracted MLP...")
        mlp_time = time.time()
        leaps_mlp = train_model(leaps_mlp, train_loader, EPOCHS_MLP, DEVICE, name="LeapsMLP")
        mlp_time = time.time() - mlp_time
        
        mlp_acc = test_model(leaps_mlp, test_loader, DEVICE, name="LeapsMLP")
        
        # ============================================================
        # 6. Also train a random MLP for comparison
        # ============================================================
        random_mlp = IntelligentLeapsNet(input_size=784, hidden_sizes=[256, 128], output_size=10).to(DEVICE)
        random_mlp = train_model(random_mlp, train_loader, EPOCHS_MLP, DEVICE, name="RandomMLP")
        random_acc = test_model(random_mlp, test_loader, DEVICE, name="RandomMLP")
        
        results.append({
            'prune_ratio': prune_ratio,
            'resnet_acc': resnet_acc,
            'leaps_acc': mlp_acc,
            'random_acc': random_acc,
            'remaining_params': remaining
        })
        
        print(f"  Time: {mlp_time:.2f}s")
    
    # ============================================================
    # 7. Results Summary
    # ============================================================
    print("\n" + "="*60)
    print("RESULTS SUMMARY")
    print("="*60)
    print(f"{'Prune%':>8} | {'ResNet':>10} | {'LeapsMLP':>10} | {'RandomMLP':>10} | {'Params':>10}")
    print("-"*60)
    
    for r in results:
        print(f"{r['prune_ratio']*100:>7.0f}% | {r['resnet_acc']:>10.2f}% | {r['leaps_acc']:>10.2f}% | {r['random_acc']:>10.2f}% | {r['remaining_params']:>10,.0f}")
    
    print("-"*60)
    print("\nHypothesis: LeapsMLP (extracted from pruned ResNet) should outperform")
    print("            RandomMLP (random init) if 'leaps of intelligence' exist.")
    
    return results


if __name__ == "__main__":
    results = main()