"""
Leaps of Intelligence: ResNet18 → MLP via Lottery Ticket Hypothesis
===================================================================
"""

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 time

# ============================================================
# Configuration
# ============================================================
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
BATCH_SIZE = 128
EPOCHS_CONV = 1
EPOCHS_MLP = 2
PRUNE_RATIOS = [0.5, 0.7, 0.8, 0.9]
MNIST_PATH = './data'

# ============================================================
# Modified ResNet18 for MNIST (28x28)
# ============================================================
class ModifiedResNet18(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        
        self.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        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)
        
    def _make_layer(self, in_ch, out_ch, blocks, stride):
        layers = []
        layers.append(BasicBlock(in_ch, out_ch, stride))
        for _ in range(1, blocks):
            layers.append(BasicBlock(out_ch, out_ch, 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)
        return self.fc(x)

    def get_param_count(self):
        return sum(p.numel() for p in self.parameters())


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


# ============================================================
# Leaps MLP: Direct extraction from pruned ResNet
# ============================================================
class LeapsExtractor:
    """
    Extract the 'leaps of intelligence' from a trained/pruned ResNet.
    The key insight: each layer-to-layer connection can be expressed as 
    a linear projection, even if originally a convolution.
    """
    
    @staticmethod
    def extract_fc_layers(resnet):
        """Extract the FC layer weights"""
        return resnet.fc.weight.data.clone()
    
    @staticmethod
    def extract_conv_as_fc(resnet, layer_name, in_spatial=28):
        """
        Convert conv weights to FC by unfolding spatial dimensions.
        This preserves the learned 'leaps' as direct connections.
        """
        if layer_name == 'conv1':
            w = resnet.conv1.weight.data  # (64, 1, 3, 3)
            # For 1x1 conv equivalent, use average pooled weights
            w_pooled = w.mean(dim=(2, 3))  # (64, 1) - spatial invariance
            return w_pooled
        elif layer_name.startswith('layer'):
            # Get weights from the last conv in the layer
            layer = getattr(resnet, layer_name)
            last_conv = None
            for m in layer.modules():
                if isinstance(m, nn.Conv2d):
                    last_conv = m
            if last_conv:
                w = last_conv.weight.data
                # Average over spatial for FC equivalence
                return w.mean(dim=(2, 3))  # (out_ch, in_ch)
        return None


class LeapsMLP(nn.Module):
    """
    MLP that uses the leap patterns learned by ResNet.
    Architecture matches the layer-to-layer flow: 
    Input → [Leap1] → [Leap2] → [Leap3] → Output
    """
    def __init__(self, input_size=784, hidden1=512, hidden2=256, output=10):
        super().__init__()
        # Three "leap" layers - each is a direct semantic connection
        self.leap1 = nn.Linear(input_size, hidden1, bias=False)
        self.leap2 = nn.Linear(hidden1, hidden2, bias=False)
        self.leap3 = nn.Linear(hidden2, output, bias=False)
        
        # Non-linear leaps (skip connections)
        self.leap1_b = nn.Linear(input_size, hidden2, bias=False)  # Direct skip
        
    def forward(self, x):
        x = x.view(x.size(0), -1)
        
        # Leap 1: Direct mapping with bypass
        h1 = F.relu(self.leap1(x))
        skip = F.relu(self.leap1_b(x))
        
        # Leap 2: Semantic abstraction
        h2 = F.relu(self.leap2(h1))
        
        # Leap 3: Output projection
        out = self.leap3(h2 + skip * 0.3)  # Combine with skip
        return out
    
    def inject_leaps(self, resnet, prune_ratio):
        """Inject learned leap patterns from pruned ResNet"""
        with torch.no_grad():
            # Layer 1 leap: from flattened image
            # Convert first conv to FC-like by using the learned patterns
            conv_w = resnet.conv1.weight.data  # (64, 1, 3, 3)
            # Apply magnitude-based selection to find "leap" connections
            leap_w = self._extract_top_leaps(conv_w, self.leap1.out_features)
            self.leap1.weight.data[:leap_w.size(0), :leap_w.size(1)] = leap_w
            
            # Layer 2 leap: from layer4 output to hidden
            fc_w = resnet.fc.weight.data  # (10, 512)
            out_features = min(self.leap3.weight.size(0), fc_w.size(0))
            in_features = min(self.leap3.weight.size(1), fc_w.size(1))
            self.leap3.weight.data[:out_features, :in_features] = fc_w[:out_features, :in_features]
            
    def _extract_top_leaps(self, conv_weights, target_out):
        """Extract the highest-magnitude connections as leaps"""
        # Flatten and take top connections
        flat = conv_weights.abs().flatten()
        k = min(target_out * 100, flat.numel())
        if k <= 0:
            return conv_weights.new_zeros((0, conv_weights[0].numel()))
        if k >= flat.numel():
            return conv_weights.view(conv_weights.size(0), -1)[:target_out, :]

        _, indices = flat.topk(k)
        mask = torch.zeros_like(flat)
        mask[indices] = 1
        mask = mask.view_as(conv_weights)
        return (conv_weights * mask).view(conv_weights.size(0), -1)[:target_out, :]


# ============================================================
# Random MLP for comparison
# ============================================================
class RandomMLP(nn.Module):
    def __init__(self, input_size=784, hidden1=512, hidden2=256, output=10):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden1)
        self.fc2 = nn.Linear(hidden1, hidden2)
        self.fc3 = nn.Linear(hidden2, output)
        
    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)


# ============================================================
# Pruning
# ============================================================
def magnitude_prune(model, prune_ratio):
    """Prune lowest magnitude weights"""
    masks = {}
    for name, param in model.named_parameters():
        if 'weight' in name and param.dim() >= 2:
            flat = param.data.abs().flatten()
            k = int(flat.numel() * (1 - prune_ratio))
            threshold = flat.kthvalue(k)[0] if k > 0 else 0
            mask = (param.data.abs() > threshold).float()
            masks[name] = mask
            param.data *= mask
    return model, masks


def reset_weights(model):
    """Reset to small random (lottery ticket style)"""
    for m in model.modules():
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)


# ============================================================
# Training
# ============================================================
def train(model, loader, epochs, device, name="Model", lr=0.001):
    optimizer = optim.Adam(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    
    model.train()
    for epoch in range(epochs):
        correct, total, loss_sum = 0, 0, 0
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            out = model(data)
            loss = criterion(out, target)
            loss.backward()
            optimizer.step()
            
            loss_sum += loss.item()
            pred = out.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)
        
        acc = 100 * correct / total
        print(f"  {name} Ep {epoch+1}/{epochs} - Loss: {loss_sum/len(loader):.4f}, Acc: {acc:.2f}%")
    
    return model


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


# ============================================================
# MAIN
# ============================================================
def main():
    print("="*60)
    print("LEAPS OF INTELLIGENCE: ResNet18 → MLP via Lottery Ticket")
    print("="*60)
    print(f"Device: {DEVICE}")
    
    # Load 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)
    
    print(f"Train: {len(train_set)}, Test: {len(test_set)}\n")
    
    # Train ResNet18 (scaffold)
    print("1. Training ResNet18 (scaffold)...")
    resnet = ModifiedResNet18().to(DEVICE)
    print(f"   Parameters: {resnet.get_param_count():,}")
    
    start = time.time()
    resnet = train(resnet, train_loader, EPOCHS_CONV, DEVICE, "ResNet18")
    resnet_time = time.time() - start
    
    resnet_acc = test(resnet, test_loader, DEVICE, "ResNet18")
    print(f"   Time: {resnet_time:.1f}s\n")
    
    # Pruning & extraction experiment
    print("2. Finding 'Leaps of Intelligence'...\n")
    results = []
    
    for prune_ratio in PRUNE_RATIOS:
        print(f"  === Prune ratio: {prune_ratio*100:.0f}% ===")
        
        # Retrain with pruning (lottery ticket style)
        resnet_p = ModifiedResNet18().to(DEVICE)
        train(resnet_p, train_loader, 5, DEVICE, f"Pruned-{prune_ratio*100:.0f}%")
        
        # Prune
        _, mask = magnitude_prune(resnet_p, prune_ratio)
        remaining = sum(m.sum().item() for m in mask.values())
        total_p = sum(m.numel() for m in mask.values())
        print(f"  Remaining params: {remaining:,.0f} ({remaining/total_p*100:.1f}%)")
        
        # Create LeapsMLP and inject patterns
        leaps_mlp = LeapsMLP(input_size=784, hidden1=512, hidden2=256, output=10).to(DEVICE)
        leaps_mlp.inject_leaps(resnet_p, prune_ratio)
        
        # Train LeapsMLP
        print(f"  Training LeapsMLP...")
        leaps_mlp = train(leaps_mlp, train_loader, EPOCHS_MLP, DEVICE, "LeapsMLP")
        leaps_acc = test(leaps_mlp, test_loader, DEVICE, "LeapsMLP")
        
        # Compare with random MLP (same architecture)
        random_mlp = RandomMLP(input_size=784, hidden1=512, hidden2=256, output=10).to(DEVICE)
        random_mlp = train(random_mlp, train_loader, EPOCHS_MLP, DEVICE, "RandomMLP")
        random_acc = test(random_mlp, test_loader, DEVICE, "RandomMLP")
        
        results.append({
            'prune': prune_ratio,
            'resnet': resnet_acc,
            'leaps': leaps_acc,
            'random': random_acc,
            'remaining': remaining
        })
        print()
    
    # Summary
    print("="*60)
    print("RESULTS SUMMARY")
    print("="*60)
    print(f"{'Prune%':>8} | {'ResNet18':>10} | {'LeapsMLP':>10} | {'RandomMLP':>10} | {'Params':>10}")
    print("-"*60)
    for r in results:
        print(f"{r['prune']*100:>7.0f}% | {r['resnet']:>10.2f}% | {r['leaps']:>10.2f}% | {r['random']:>10.2f}% | {r['remaining']:>10,.0f}")
    print("-"*60)
    print("\nIf LeapsMLP > RandomMLP: 'leaps of intelligence' were extracted!")
    
    return results


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