import os
import sys
import torch
import torch.nn as nn
import torch.nn.functional as F  # Add this line
import torchvision
import torchvision.transforms as transforms

sys.path.append("..")
from tools import show

seed = 4321
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device("mps")

num_epochs = 1
batch_size = 64

transform = transforms.Compose([transforms.ToTensor()])

train_dataset = torchvision.datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(root="./data", train=False, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)


class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 8, 5, 1, 0)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(8, 1, 5, 1, 0)
        self.fc1 = nn.Linear(1 * 4 * 4, 10)

    def forward(self, x):
        x = self.forward_unsuper(x)
        x = self.pool(x)
        x = self.pool(self.conv2(x))
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        return x

    def forward_unsuper(self, x):
        weight = self.conv1.weight.reshape(self.conv1.weight.shape[0], -1)
        weight = weight.permute(1, 0)
        mean = torch.mean(weight, dim=0)
        weight = weight - mean
        sum = torch.sum(torch.abs(weight), dim=0)
        weight = weight / sum
        weight = weight.permute(1, 0)
        weight = weight.reshape(self.conv1.weight.shape)
        x = torch.conv2d(x, weight)
        return x

    def printFector(self, x, label, dir=""):
        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), dir + "/input_image.png", Contrast=[0, 1.0])
        # show.DumpTensorToLog(x, "input_image.log")
        x = self.conv1(x)
        w = self.conv1.weight
        show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), dir + "/conv1_weight.png", Contrast=[-1.0, 1.0])
        # show.DumpTensorToLog(w, "conv1_weight.log")

        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]), dir + "/conv1_output.png", Contrast=[-1.0, 1.0])
        # show.DumpTensorToLog(x, "conv1_output.png")

        x = self.pool(x)
        x = self.conv2(x)
        w = self.conv2.weight
        show.DumpTensorToImage(
            w.view(-1, w.shape[2], w.shape[3]).cpu(), dir + "/conv2_weight.png", Contrast=[-1.0, 1.0]
        )
        show.DumpTensorToImage(
            x.view(-1, x.shape[2], x.shape[3]).cpu(), dir + "/conv2_output.png", Contrast=[-1.0, 1.0]
        )
        x = self.pool(x)
        show.DumpTensorToImage(x.view(-1, x.shape[2], x.shape[3]).cpu(), dir + "/pool_output.png", Contrast=[-1.0, 1.0])
        pool_shape = x.shape
        x = x.view(x.shape[0], -1)
        x = self.fc1(x)
        show.DumpTensorToImage(
            self.fc1.weight.view(-1, pool_shape[2], pool_shape[3]), dir + "/fc_weight.png", Contrast=[-1.0, 1.0]
        )
        show.DumpTensorToImage(x.view(-1).cpu(), dir + "/fc_output.png")

        criterion = nn.CrossEntropyLoss()
        loss = criterion(x, label)
        loss.backward()

        if self.conv1.weight.requires_grad:
            w = self.conv1.weight.grad
            show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), dir + "/conv1_weight_grad.png")
        if self.conv2.weight.requires_grad:
            w = self.conv2.weight.grad
            show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), dir + "/conv2_weight_grad.png")
        if self.fc1.weight.requires_grad:
            show.DumpTensorToImage(
                self.fc1.weight.grad.view(-1, pool_shape[2], pool_shape[3]), dir + "/fc_weight_grad.png"
            )


model = ConvNet().to(device)
model.train()

# Train the model unsuper
epochs = 20
model.conv1.weight.requires_grad = True
model.conv2.weight.requires_grad = False
model.fc1.weight.requires_grad = False
n_total_steps = len(train_loader)
for epoch in range(epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        outputs = model.forward_unsuper(images)

        outputs = outputs.permute(0, 2, 3, 1)  # 64 8 24 24 -> 64 24 24 8
        sample = outputs.reshape(-1, outputs.shape[3])  # -> 36864 8
        abs = torch.abs(sample).detach()
        max, max_index = torch.max(abs, dim=1)
        mean = torch.mean(abs, dim=1)
        mean = torch.expand_copy(mean.reshape(-1, 1), sample.shape)
        max = torch.expand_copy(max.reshape(-1, 1), sample.shape)

        all = range(0, sample.shape[0])
        ratio_max = abs / mean
        ratio_nor = (max - abs) / max
        ratio_nor[all, max_index] = ratio_max[all, max_index].clone()
        ratio_nor = torch.where(torch.isnan(ratio_nor), 1.0, ratio_nor)
        label = sample * ratio_nor

        loss = F.l1_loss(sample, label)
        model.conv1.weight.grad = None
        loss.backward()

        model.conv1.weight.data = model.conv1.weight.data - model.conv1.weight.grad * 10

        if (i + 1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.8f}")

w = model.conv1.weight.grad
show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]).cpu(), "conv1_weight_grad.png")
w = model.conv1.weight.data
show.DumpTensorToImage(w.view(-1, w.shape[2], w.shape[3]), "conv1_weight_update.png", Contrast=[-1.0, 1.0])

# loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
# images, labels = next(iter(loader))
# images = images.to(device)

# Train the model
model.conv1.weight.requires_grad = False
model.conv2.weight.requires_grad = True
model.fc1.weight.requires_grad = True
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.2)
n_total_steps = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if (i + 1) % 100 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{n_total_steps}], Loss: {loss.item():.4f}")

print("Finished Training")

test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False)
test_loader = iter(test_loader)
images, labels = next(test_loader)
images = images.to(device)
labels = labels.to(device)
model.printFector(images, labels, "dump1")

images, labels = next(test_loader)
images = images.to(device)
labels = labels.to(device)
model.printFector(images, labels, "dump2")

# Test the model
with torch.no_grad():
    n_correct = 0
    n_samples = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)

        # max returns (value ,index)
        _, predicted = torch.max(outputs.data, 1)
        n_samples += labels.size(0)
        n_correct += (predicted == labels).sum().item()

    acc = 100.0 * n_correct / n_samples
    print(f"Accuracy of the network on the 10000 test images: {acc} %")