import os
import time

import torch
import torch.nn as nn
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.optim as optim

from tqdm import tqdm


def get_mean_std(loader):
  '''
  Calculates mean and std of input images.

  Args:
    loader (torch.DataLoader): Loader with images

  Returns:
    mean (torch.Tensor): Mean of images in loader
    std (torch.Tensor): Standard deviation of images in loader
  '''
  channels_sum, channels_squared_sum, num_batches = 0, 0, 0

  for data, _ in loader:
    channels_sum += torch.mean(data, dim=[0,2,3])  # mean across [no. of examples, height, width]
    channels_squared_sum += torch.mean(data**2, dim=[0,2,3])  # squared mean across [no. of examples, height, width]
    num_batches += 1

    mean = channels_sum/num_batches
    std = (channels_squared_sum/(num_batches-mean**2))**0.5

    return mean, std


class Net(nn.Module):
    '''
    model definition
    '''

    def __init__(self):
        super(Net, self).__init__()

        self.layer1 = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=5),
            nn.ReLU(),
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(32, 32, kernel_size=5, bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),
            nn.Dropout2d(0.25),
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3),
            nn.ReLU(),
        )
        self.layer4 = nn.Sequential(
            nn.Conv2d(64, 64, kernel_size=3, bias=False),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d((2, 2)),
            nn.Dropout2d(0.25),
            nn.Flatten(),
        )
        self.layer5 = nn.Sequential(
            nn.Linear(576, 256, bias=False),
            nn.BatchNorm1d(256),
            nn.ReLU(),
        )
        self.layer6 = nn.Sequential(
            nn.Linear(256, 128, bias=False),
            nn.BatchNorm1d(128),
            nn.ReLU(),
        )
        self.layer7 = nn.Sequential(
            nn.Linear(128, 84, bias=False),
            nn.BatchNorm1d(84),
            nn.ReLU(),
            nn.Dropout(0.25),
        )
        self.layer8 = nn.Sequential(
            nn.Linear(84, 10),
            nn.LogSoftmax(dim=1),
        )

    def forward(self, x):
        x = transforms.Normalize(mean, std)(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        x = self.layer6(x)
        x = self.layer7(x)
        x = self.layer8(x)

        return x


# downloads and loads MNIST train set
transform = transforms.Compose([transforms.ToTensor(), transforms.RandomAffine(degrees=10, translate=(0.1,0.1))])
train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True, pin_memory=True)

# downloads and loads MNIST test set
val_data = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
val_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=False, pin_memory=True)

# uses GPU if available
if torch.cuda.is_available():
  dev = "cuda:0"
else:
  dev = "cpu"

device = torch.device(dev)

# gets mean and std of dataset
mean, std = get_mean_std(train_loader)


def run_model():
    # defines parameters
    model = Net().to(device=device)
    optimizer = optim.Adam(model.parameters(), lr=0.1)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.2, patience=2)
    criterion = nn.NLLLoss()

    # iterates through epochs
    for epoch in range(30):
        print(f"\nEpoch {epoch+1}/{30}.")

        # train loop
        model.train()

        total_train_loss = 0
        total_correct = 0

        for i, (images, labels) in enumerate(tqdm(train_loader)):
            images = images.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()

            outputs = model(images)

            loss = criterion(outputs, labels)
            total_train_loss += loss.item()

            loss.backward()
            optimizer.step()

            # Calculates train accuracy
            outputs_probs = nn.functional.softmax(
                outputs, dim=1)  # gets probabilities
            for idx, preds in enumerate(outputs_probs):
                # if label with max probability matches true label
                if labels[idx] == torch.argmax(preds.data):
                    total_correct += 1

        train_loss = total_train_loss/(i+1)
        train_accuracy = total_correct/len(train_data)

        print(f"Train set:- Loss: {train_loss}, Accuracy: {train_accuracy}.")

        # saves model state
        if not os.path.exists("./saved_models"):
            os.mkdir("./saved_models")
        torch.save(model.state_dict(), f"./saved_models/mnist-cnn-{time.time()}.pt")

        # val loop
        model.eval()

        total_val_loss = 0
        total_correct = 0

        with torch.no_grad():
            for i, (images, labels) in enumerate(tqdm(val_loader)):
                images = images.to(device)
                labels = labels.to(device)

                outputs = model(images)

                loss = criterion(outputs, labels)
                total_val_loss += loss.item()

                outputs_probs = nn.functional.softmax(outputs, dim=1)
                for idx, preds in enumerate(outputs_probs):
                    if labels[idx] == torch.argmax(preds.data):
                        total_correct += 1

        val_loss = total_val_loss/(i+1)
        val_accuracy = total_correct/len(val_data)

        print(f"Val set:- Loss: {val_loss}, Accuracy: {val_accuracy}.")

        # adjusts lr
        scheduler.step(val_loss)