import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import copy

# ------------------------------
# 1. MLP with hook‑overcoming module
# ------------------------------
class HookOvercomingMLP(nn.Module):
    def __init__(self, input_dim=784, hidden_dims=[256,128], output_dim=10, 
                 hook_feature_dim=32, escape_strength=0.1):
        super().__init__()
        self.escape_strength = escape_strength

        # Main classifier layers
        self.fc1 = nn.Linear(input_dim, hidden_dims[0])
        self.fc2 = nn.Linear(hidden_dims[0], hidden_dims[1])
        self.fc3 = nn.Linear(hidden_dims[1], output_dim)
        self.dropout = nn.Dropout(0.2)

        # Hook detection module – reads hidden state & gradient norm, outputs escape vector
        # The escape vector will be added to fc1 weights
        self.hook_detector = nn.Sequential(
            nn.Linear(hidden_dims[0] + 1, hook_feature_dim),  # hidden mean + grad_norm
            nn.ReLU(),
            nn.Linear(hook_feature_dim, hidden_dims[0] * input_dim)  # same size as fc1.weight
        )
        # Keep a buffer to store the last computed escape vector
        self.register_buffer('last_escape', torch.zeros(hidden_dims[0], input_dim))

    def forward(self, x, detect_hook=False, grad_norm=None):
        x = x.view(x.size(0), -1)
        h1 = F.relu(self.fc1(x))
        h2 = F.relu(self.fc2(h1))
        h2 = self.dropout(h2)
        out = self.fc3(h2)

        if detect_hook and grad_norm is not None:
            # Build hook detection features: mean of hidden activity + gradient norm
            h1_mean = h1.mean(dim=1).mean(dim=0).detach()  # scalar
            feat = torch.cat([h1_mean.view(1), grad_norm.view(1)]).unsqueeze(0)  # [1,2]
            # Expand to batch 1, but our hook_detector expects [1, hidden_dim+1]
            # Actually we designed it with hidden_dim input, let's fix:
            # We'll just use h1_mean and grad_norm as two scalars, map to hidden+1 -> hook_feature_dim
            # Simpler: concat the mean of h1 (flattened) with grad_norm
            # But that would be huge. Instead, use the first principal component?
            # Let's use: mean activation of each neuron (already reduced across batch) + grad_norm
            # Actual implementation: mean of h1 across batch and features gives a scalar, plus grad_norm
            # Instead use histogram? Keep simple: concatenate the mean vector of h1 (size hidden_dim0)
            # That would be 256 + 1 -> too big. Let's reduce by taking mean of h1 across batch -> 256-dim
            # But that's large. Let's use only the first 10 neurons' means for efficiency.
            h1_sub = h1.mean(dim=0)[:10]  # shape [10]
            feat = torch.cat([h1_sub, grad_norm.view(1)])  # [11]
            # Pass through hook_detector (first layer expects 11 inputs, output weight_flat)
            escape_flat = self.hook_detector(feat.unsqueeze(0))  # [1, fc1_weight_flat]
            escape = escape_flat.view(self.fc1.weight.shape)     # [out_features, in_features]
            self.last_escape = escape.detach()
            return out, escape
        else:
            return out

    def apply_escape(self):
        """Add the last computed escape vector to fc1 weights (in‑place)."""
        with torch.no_grad():
            self.fc1.weight.data += self.escape_strength * self.last_escape

    def zero_escape(self):
        self.last_escape.zero_()


# ------------------------------
# 2. Training functions
# ------------------------------
def compute_gradient_norm(model):
    """Compute L2 norm of gradients of fc1 layer (first linear layer)."""
    if model.fc1.weight.grad is not None:
        return torch.norm(model.fc1.weight.grad).item()
    return 0.0

def train_epoch(model, loader, optimizer, device, hook_detection_interval=200, 
                grad_norm_threshold=0.01, patience=5):
    model.train()
    total_loss = 0
    correct = 0
    total = 0

    # Simple stall tracking
    recent_grad_norms = []
    stall_counter = 0

    for batch_idx, (data, target) in enumerate(tqdm(loader, desc="Training")):
        data, target = data.to(device), target.to(device)

        # Forward + loss
        output = model(data)
        loss_ce = F.cross_entropy(output, target)

        # --- Hook detection & escape ---
        grad_norm = compute_gradient_norm(model)
        recent_grad_norms.append(grad_norm)
        if len(recent_grad_norms) > 10:
            recent_grad_norms.pop(0)

        # Condition: gradient norm is very small for multiple steps (local minimum)
        is_stuck = (len(recent_grad_norms) >= patience and 
                    all(gn < grad_norm_threshold for gn in recent_grad_norms[-patience:]))

        # Also detect plateau by loss not decreasing? We'll keep simple.

        if is_stuck and (batch_idx % hook_detection_interval == 0):
            # Get escape vector from the detector
            with torch.enable_grad():
                # Need gradients flowing through hook_detector:
                # We recompute forward with detect_hook=True and get escape
                out2, escape = model(data, detect_hook=True, grad_norm=torch.tensor(grad_norm).to(device))
                # Apply escape to weights directly (in-place)
                model.apply_escape()
                # Compute loss again after escape (to see improvement)
                loss_after = F.cross_entropy(out2, target)
                # Loss term that encourages escape to lower loss: we want loss_after < loss_ce
                recovery_loss = F.relu(loss_ce - loss_after + 0.05)  # positive if escape didn't help
                total_loss_term = loss_ce + 0.1 * recovery_loss
                # Backward through both classification and recovery loss
                optimizer.zero_grad()
                total_loss_term.backward()
                optimizer.step()
                # Reset escape buffer for next time
                model.zero_escape()
                # Continue to next batch (skip default backward)
                continue

        # Normal training step
        optimizer.zero_grad()
        loss_ce.backward()
        optimizer.step()

        total_loss += loss_ce.item()
        pred = output.argmax(dim=1, keepdim=True)
        correct += pred.eq(target.view_as(pred)).sum().item()
        total += target.size(0)

    return total_loss / len(loader), correct / total

def evaluate(model, loader, device):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
            total += target.size(0)
    return correct / total

# ------------------------------
# 3. Data loading and main loop
# ------------------------------
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # Data preparation
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])
    train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    test_dataset = datasets.MNIST('./data', train=False, transform=transform)

    train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)
    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)

    # Models: one with hook overcoming, one baseline (standard MLP)
    model_with_hook = HookOvercomingMLP().to(device)
    baseline_model = HookOvercomingMLP().to(device)  # same architecture but will not use hook detection
    optimizer_hook = optim.Adam(model_with_hook.parameters(), lr=0.001)
    optimizer_base = optim.Adam(baseline_model.parameters(), lr=0.001)

    epochs = 10
    train_losses_hook, train_accs_hook, test_accs_hook = [], [], []
    train_losses_base, train_accs_base, test_accs_base = [], [], []

    print("\n--- Training Hook‑Aware MLP ---")
    for epoch in range(1, epochs+1):
        train_loss, train_acc = train_epoch(model_with_hook, train_loader, optimizer_hook, device)
        test_acc = evaluate(model_with_hook, test_loader, device)
        train_losses_hook.append(train_loss)
        train_accs_hook.append(train_acc)
        test_accs_hook.append(test_acc)
        print(f"Epoch {epoch:2d} | Loss: {train_loss:.4f} | Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

    print("\n--- Training Baseline MLP ---")
    for epoch in range(1, epochs+1):
        # Baseline training (no hook detection) – we reuse train_epoch but disable the hook logic by skipping condition
        # To reuse the same code but without escapes, we can temporarily set gradient norm threshold to very low
        # Or just write a simple loop. For clarity, we'll use a plain training loop without hook detection.
        model_with_hook.train()  # reuse the same class but we won't trigger escapes
        total_loss = 0
        correct = 0
        total = 0
        for data, target in tqdm(train_loader, desc="Baseline training"):
            data, target = data.to(device), target.to(device)
            optimizer_base.zero_grad()
            output = baseline_model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()
            optimizer_base.step()
            total_loss += loss.item()
            pred = output.argmax(dim=1)
            correct += pred.eq(target).sum().item()
            total += target.size(0)
        train_acc = correct / total
        test_acc = evaluate(baseline_model, test_loader, device)
        train_losses_base.append(total_loss/len(train_loader))
        train_accs_base.append(train_acc)
        test_accs_base.append(test_acc)
        print(f"Epoch {epoch:2d} | Loss: {total_loss/len(train_loader):.4f} | Train Acc: {train_acc:.4f} | Test Acc: {test_acc:.4f}")

    # Final comparison
    print("\n===== Final Results =====")
    print(f"Hook-Aware MLP  final test accuracy: {test_accs_hook[-1]*100:.2f}%")
    print(f"Baseline MLP    final test accuracy: {test_accs_base[-1]*100:.2f}%")
    print(f"Improvement: {(test_accs_hook[-1] - test_accs_base[-1])*100:.2f}%")

    # Optional plot
    plt.figure(figsize=(10,5))
    plt.subplot(1,2,1)
    plt.plot(range(1,epochs+1), train_accs_hook, label='Hook-Aware Train')
    plt.plot(range(1,epochs+1), train_accs_base, label='Baseline Train')
    plt.plot(range(1,epochs+1), test_accs_hook, '--', label='Hook-Aware Test')
    plt.plot(range(1,epochs+1), test_accs_base, '--', label='Baseline Test')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.title('Accuracy Comparison')
    plt.subplot(1,2,2)
    plt.plot(range(1,epochs+1), train_losses_hook, label='Hook-Aware')
    plt.plot(range(1,epochs+1), train_losses_base, label='Baseline')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.title('Training Loss')
    plt.tight_layout()
    plt.savefig('hook_comparison.png')
    plt.show()

if __name__ == "__main__":
    main()