print(f'Can I can use GPU now? -- {torch.cuda.is_available()}')

rand_tensor = torch.rand(5,2)
simple_model = nn.Sequential(nn.Linear(2,10), nn.ReLU(), nn.Linear(10,1))
print(f'input is on {rand_tensor.device}')
print(f'model parameters are on {[param.device for param in simple_model.parameters()]}')
print(f'output is on {simple_model(rand_tensor).device}')

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Move rand_tensor and model onto the GPU device
rand_tensor = rand_tensor.to(device)
simple_model = simple_model.to(device)

print(f'input is on {rand_tensor.device}')
print(f'model parameters are on {[param.device for param in simple_model.parameters()]}')
print(f'output is on {simple_model(rand_tensor).device}')

# Download MNIST and transformation
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True)

testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=128, shuffle=False)

# Let's first train the FC model.
criterion = nn.CrossEntropyLoss()

start = time.time()
max_epoch = 4
def train(model, trainloader, criterion, optimizer, max_epoch):
    model.train()
    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f'Training with {total_params} trainable parameters')
    for epoch in range(max_epoch):
        running_loss = 0.0
        correct = 0
        total = 0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
        epoch_loss = running_loss / len(trainloader)
        epoch_acc = 100 * correct / total
        print(f'Epoch {epoch+1}, Loss: {epoch_loss:.4f}, Train Accuracy: {epoch_acc:.2f}%')
        # Evaluate on test set after each epoch
        test_acc = test(model, testloader)
        print(f'Test Accuracy after Epoch {epoch+1}: {test_acc:.2f}%')

def test(model, testloader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100 * correct / total
    return accuracy

ourfc = OurFC()
optimizer = optim.Adam(ourfc.parameters(), lr=0.001)
train(ourfc, trainloader, criterion, optimizer, max_epoch)
final_test_acc = test(ourfc, testloader)
print(f'Final Accuracy on test set: {final_test_acc:.2f}%')
end = time.time()
print(f'Finished Training after {end-start} s ')

# Let's then train the OurCNN model.
start = time.time()
ourcnn = OurCNN()
optimizer = optim.Adam(ourcnn.parameters(), lr=0.001)
train(ourcnn, trainloader, criterion, optimizer, max_epoch)
test(ourcnn, testloader)
end = time.time()
print(f'Finished Training after {end-start} s ')

ourfc = OurFC()
total_params = sum(p.numel() for p in ourfc.parameters())
print(f'OurFC has a total of {total_params} parameters')

ourcnn = OurCNN()
total_params = sum(p.numel() for p in ourcnn.parameters())
print(f'OurCNN has a total of {total_params} parameters')

classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

# Create the transform, load/download CIFAR10 train and test datasets with transform
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

# Define trainloader and testloader
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(trainset, batch_size=9, shuffle=True, num_workers=2)

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

# Code to display images
batch_idx, (images, targets) = next(enumerate(train_loader))
fig, ax = plt.subplots(3,3,figsize = (9,9))
for i in range(3):
    for j in range(3):
        image = images[i*3+j].permute(1,2,0)
        image = image/2 + 0.5
        ax[i,j].imshow(image)
        ax[i,j].set_axis_off()
        ax[i,j].set_title(f'{classes[targets[i*3+j]]}')
fig.show()

# Train neural network.
start = time.time()
max_epoch = 4
def train(net, train_loader, criterion, optimizer, epoch):
    net.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for i, data in enumerate(train_loader, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    print(f'Epoch {epoch}, Loss: {running_loss/len(train_loader):.4f}, '
          f'Train Accuracy: {100 * correct / total:.2f}%')

def test(net, criterion, test_loader, epoch):
    net.eval()
    correct = 0
    total = 0
    test_loss = 0.0
    predictions = []
    with torch.no_grad():
        for data in test_loader:
            images, labels = data
            outputs = net(images)
            loss = criterion(outputs, labels)
            test_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
            predictions.extend(predicted.tolist())
    accuracy = 100 * correct / total
    print(f'Epoch {epoch}, Test Loss: {test_loss/len(test_loader):.4f}, '
          f'Test Accuracy: {accuracy:.2f}%')
    return {'prediction': predictions, 'accuracy': accuracy}

net = CNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
for epoch in range(1, max_epoch + 1):
    train(net, train_loader, criterion, optimizer, epoch)
    test(net, criterion, test_loader, epoch)
output = test(net, criterion, test_loader, epoch)
end = time.time()
print(f'Finished Training after {end-start} s ')

total_images = 5
predictions = output['prediction']
targets = torch.tensor(testset.targets)
misclassified = {1: None, 3: None, 5: None, 7: None, 9: None}
class_indices = {1: 'car', 3: 'cat', 5: 'dog', 7: 'horse', 9: 'truck'}

# Find misclassified images
print("Searching for misclassified images...")
for idx, (pred, true) in enumerate(zip(predictions, targets)):
    if true.item() in misclassified and pred != true.item() and misclassified[true.item()] is None:
        misclassified[true.item()] = (idx, pred, true.item())
        print(f"Found misclassified image for class {class_indices[true.item()]}: "
              f"Index={idx}, Predicted={classes[pred]}, True={class_indices[true.item()]}")

# Check if all classes have misclassified images
missing_classes = [class_id for class_id, value in misclassified.items() if value is None]
if missing_classes:
    print(f"Warning: Could not find misclassified images for classes: "
          f"{[class_indices[c] for c in missing_classes]}")

# Plot misclassified images
fig, ax = plt.subplots(1, 5, figsize=(15, 3))
for i, class_id in enumerate([1, 3, 5, 7, 9]):
    if misclassified[class_id] is not None:
        idx, pred_label, true_label = misclassified[class_id]
        image, _ = testset[idx]
        image = image.permute(1, 2, 0).detach().cpu()
        image = image / 2 + 0.5  # Denormalize
        if image.min() < 0 or image.max() > 1:
            print(f"Warning: Image at index {idx} has invalid range: min={image.min()}, max={image.max()}")
            image = torch.clamp(image, 0, 1)
        ax[i].imshow(image)
        ax[i].set_title(f'True: {class_indices[true_label]}
Pred: {classes[pred_label]}')
        ax[i].set_axis_off()
    else:
        print(f"No misclassified image found for class {class_indices[class_id]}")
        ax[i].set_axis_off()
        ax[i].set_title(f'No misclassified
{class_indices[class_id]}')
plt.tight_layout()
plt.show()

def calculate_accuracy(model, data_loader):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in data_loader:
            images = images.to(device)
            labels = labels.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    return correct / total

test_accuracy = calculate_accuracy(resnet18, test_loader)
print(f'Test Accuracy: {100 * test_accuracy:.2f}%')

def test(model, criterion, data_loader, epoch=None):
    model.eval()
    correct = 0
    total = 0
    test_loss = 0.0
    with torch.no_grad():
        for images, labels in data_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)
            test_loss += loss.item()
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = 100 * correct / total
    if epoch is not None:
        print(f'Epoch {epoch}, Test Loss: {test_loss/len(data_loader):.4f}, Test Accuracy: {accuracy:.2f}%')
    else:
        print(f'Test Loss: {test_loss/len(data_loader):.4f}, Test Accuracy: {accuracy:.2f}%')
    return accuracy

# Define train function
def train(model, train_loader, criterion, optimizer, epoch):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    print(f'Epoch {epoch}, Train Loss: {running_loss/len(data_loader):.4f}, '
          f'Train Accuracy: {100 * correct / total:.2f}%')

# Directly train the whole model.
start = time.time()
# Modify the final layer to match CIFAR-10's 10 classes
resnet18.fc = nn.Linear(resnet18.fc.in_features, 10).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(resnet18.parameters(), lr=0.001)
max_epoch = 5

for epoch in range(1, max_epoch + 1):
    train(resnet18, train_loader, criterion, optimizer, epoch)
    test(resnet18, criterion, test_loader, epoch)
test(resnet18, criterion, test_loader)
end = time.time()
print(f'Finished Training after {end-start} s ')

# Train the model!
start = time.time()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, resnet18.parameters()), lr=0.001)
max_epoch = 5

for epoch in range(1, max_epoch + 1):
    train(resnet18, train_loader, criterion, optimizer, epoch)
    test(resnet18, criterion, test_loader, epoch)
test(resnet18, criterion, test_loader)
end = time.time()
print(f'Finished Training after {end-start} s ')