import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm
import copy

# ----------------------------------------------------------------------
# Stable division operator (replaces /)
# ----------------------------------------------------------------------
def stable_div(a, b, eps=1e-8):
    """a / b but with eps to avoid division by zero."""
    return a / (b + eps)

# ----------------------------------------------------------------------
# Simple CNN for CIFAR-10
# ----------------------------------------------------------------------
class SimpleCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool2d(2, 2)
        self.drop = nn.Dropout(0.25)
        self.fc = nn.Linear(128 * 4 * 4, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = self.drop(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

# ----------------------------------------------------------------------
# Training loop with bug‑fixing via stable division of updates
# ----------------------------------------------------------------------
def train_with_bug_fixing(model, device, train_loader, optimizer, epochs=20, eps=1e-8):
    model.train()
    # Store previous gradients mapped by parameter id
    prev_grads = {id(p): torch.zeros_like(p) for p in model.parameters() if p.requires_grad}

    for epoch in range(1, epochs+1):
        total_loss = 0
        correct = 0
        for batch_idx, (data, target) in enumerate(tqdm(train_loader, desc=f'Epoch {epoch}')):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = F.cross_entropy(output, target)
            loss.backward()

            # ----------------------------------------------------------
            # Bug fixing: compute gradient ratios and adjust learning rates
            # ----------------------------------------------------------
            with torch.no_grad():
                for param in model.parameters():
                    if param.grad is None:
                        continue
                    param_id = id(param)
                    if param_id not in prev_grads:
                        prev_grads[param_id] = torch.zeros_like(param.grad)
                        continue

                    current_grad = param.grad.clone()
                    prev_grad = prev_grads[param_id]

                    # Compute ratio |current / prev| using stable division
                    ratio = stable_div(torch.abs(current_grad), torch.abs(prev_grad) + eps, eps)
                    # Clamp ratio to [0.1, 10] for stability
                    ratio = torch.clamp(ratio, 0.1, 10.0)

                    # Adjust the learning rate for this parameter
                    for param_group in optimizer.param_groups:
                        param_ids = [id(p) for p in param_group['params']]
                        if param_id in param_ids:
                            mean_ratio = ratio.mean().item()
                            param_group['lr'] *= (mean_ratio ** 0.5)
                            param_group['lr'] = max(1e-5, min(1e-2, param_group['lr']))
                            break  # param belongs to only one group

                    # Store current grad for next iteration
                    prev_grads[param_id] = current_grad

            # Perform update with adjusted LRs
            optimizer.step()

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

        avg_loss = total_loss / len(train_loader.dataset)
        accuracy = 100. * correct / len(train_loader.dataset)
        print(f'Epoch {epoch}: Loss {avg_loss:.4f}, Accuracy {accuracy:.2f}%')
        print(f'Current LRs: {[pg["lr"] for pg in optimizer.param_groups]}')

    return model

# ----------------------------------------------------------------------
# Test function
# ----------------------------------------------------------------------
def test(model, device, test_loader):
    model.eval()
    correct = 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, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    accuracy = 100. * correct / len(test_loader.dataset)
    print(f'Test accuracy: {accuracy:.2f}%')
    return accuracy

# ----------------------------------------------------------------------
# Main
# ----------------------------------------------------------------------
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using {device}")

    # CIFAR-10 data
    transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])
    test_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
    testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=test_transform)
    train_loader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
    test_loader = torch.utils.data.DataLoader(testset, batch_size=100, shuffle=False, num_workers=2)

    model = SimpleCNN().to(device)
    # Use SGD with a base learning rate; we will adjust it dynamically
    optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

    print("Training with bug‑fixing via stable division of updates...")
    train_with_bug_fixing(model, device, train_loader, optimizer, epochs=20)
    test(model, device, test_loader)

if __name__ == '__main__':
    main()