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

from perforatedai import globals_perforatedai as GPA
from perforatedai import utils_perforatedai as UPA


class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        return F.log_softmax(self.fc2(x), dim=1)


def train(model, device, loader, optimizer):
    model.train()
    correct = 0
    for data, target in loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        out = model(data)
        loss = F.nll_loss(out, target)
        loss.backward()
        optimizer.step()
        pred = out.argmax(dim=1)
        correct += pred.eq(target).sum().item()
    return 100.0 * correct / len(loader.dataset)


def validate(model, device, loader):
    model.eval()
    correct = 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()
    return 100.0 * correct / len(loader.dataset)


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST("./data", train=True, download=True, transform=transform),
        batch_size=64, shuffle=True
    )

    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST("./data", train=False, transform=transform),
        batch_size=1000, shuffle=False
    )

    # ===== PAI SETTINGS =====
    GPA.pc.set_testing_dendrite_capacity(False)
    GPA.pc.set_verbose(False)
    GPA.pc.set_n_epochs_to_switch(15)

    model = Net().to(device)
    model = UPA.perforate_model(model)

    GPA.pai_tracker.set_optimizer(optim.Adam)
    optimizer, _ = GPA.pai_tracker.setup_optimizer(
        model, {"params": model.parameters(), "lr": 0.001}, None
    )

    epoch = -1
    while True:
        epoch += 1
        print(f"\nEpoch {epoch}")

        train_acc = train(model, device, train_loader, optimizer)
        val_acc = validate(model, device, test_loader)

        print(f"Train Accuracy: {train_acc:.2f}%")
        print(f"Validation Accuracy: {val_acc:.2f}%")

        model, restructured, training_complete = GPA.pai_tracker.add_validation_score(
            val_acc, model
        )

        if restructured:
            print("🔁 Model restructured → reinitializing optimizer")
            optimizer, _ = GPA.pai_tracker.setup_optimizer(
                model, {"params": model.parameters(), "lr": 0.001}, None
            )

        if training_complete:
            print("\n✅ Training complete — PAI graph generated")
            break


if __name__ == "__main__":
    main()