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

# 1. Define a simple architecture for MNIST
class MNISTNet(nn.Module):
    def __init__(self):
        super(MNISTNet, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

# 2. Setup Data and Train Model briefly
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

train_loader = DataLoader(datasets.MNIST('../data', train=True, download=True, transform=transform), batch_size=64, shuffle=True)
test_loader = DataLoader(datasets.MNIST('../data', train=False, transform=transform), batch_size=1000, shuffle=False)

model = MNISTNet().to(device)
optimizer = optim.Adam(model.parameters(), lr=0.003)

print("Training baseline model...")
model.train()
for epoch in range(2):  # Quick training for demonstration
    for data, target in 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()

# 3. Compute Baseline Probabilities on Test Set
model.eval()
baseline_logits = []
with torch.no_grad():
    for data, _ in test_loader:
        data = data.to(device)
        baseline_logits.append(model(data))
baseline_logits = torch.cat(baseline_logits, dim=0)
p_S = F.softmax(baseline_logits, dim=-1)

# 4. Compute the Informational Pull-Coefficient (Gravity) per Neuron
# Instead of single weights (which are too granular), we remove hidden units in fc2
print("\nMeasuring Informational Gravity (Pull-Coefficient) per hidden neuron in fc2...")
pull_coefficients = []

# KL Divergence helper function
def kl_divergence(p, q, eps=1e-10):
    return torch.sum(p * (torch.log(p + eps) - torch.log(q + eps)), dim=-1).mean().item()

with torch.no_grad():
    # Loop over the 64 neurons of the fc2 layer
    for neuron_idx in range(64):
        # Save the original weight and bias data so we can restore them
        orig_weight = model.fc2.weight.data[neuron_idx].clone()
        orig_bias = model.fc2.bias.data[neuron_idx].clone()
        
        # "Remove" the component by zeroing out its incoming impacts
        model.fc2.weight.data[neuron_idx] = 0.0
        model.fc2.bias.data[neuron_idx] = 0.0
        
        # Evaluate perturbed model
        perturbed_logits = []
        for data, _ in test_loader:
            data = data.to(device)
            perturbed_logits.append(model(data))
        perturbed_logits = torch.cat(perturbed_logits, dim=0)
        p_remove = F.softmax(perturbed_logits, dim=-1)
        
        # Calculate Pull Coefficient G_info
        g_info = kl_divergence(p_S, p_remove)
        pull_coefficients.append((neuron_idx, g_info))
        
        # Restore the component
        model.fc2.weight.data[neuron_idx] = orig_weight
        model.fc2.bias.data[neuron_idx] = orig_bias

# 5. Classify Essentiality based on Framework Thresholds
print("\n--- Structural Analysis Results ---")
for idx, g in sorted(pull_coefficients, key=lambda x: x[1], reverse=True)[:10]:
    # Mapping to your paper's classification boundaries (Scale adjusted for standard KL values)
    if g > 0.1:
        role = "LOAD-BEARING (Essential)"
    elif g > 0.01:
        role = "ENHANCING"
    else:
        role = "DECORATIVE"
    print(f"Neuron {idx:02d} | Pull-Coefficient (G_info): {g:.6f} | Structural Classification: {role}")