# importing project dependencies
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cv2 as cv
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils




# method to preprocess the dataset

def preprocess_dataset(csv_path, selected_gestures):
    """Preprocesses the CSV dataset to extract required data + apply scaling
    Args: 
        csv_path::string- Path to the csv file 
        selected_gestures::list- List of integers specifying selected labels
    Returns:
        images::np.ndarray- Array of images
        labels::np.ndarray- Array of labels
    """
    # reading dataframe
    df = pd.read_csv(csv_path)

    # selecting 4 gestures for pause, play, next and previous
    mask = df['label'].isin(selected_gestures)
    df = df[mask]

    # separating labels and image data
    images = df.drop(columns=['label']).to_numpy()     
    labels = df.label.to_numpy()

    # renaming labels
    for g, l in zip(selected_gestures, range(4)):
        labels = np.where(labels==g, l, labels)

    # reshaping image data
    images = images.reshape(images.shape[0], 28, 28, 1) # m, n_W, n_H, n_C

    # normalizing images
    images = images / 255

    return images, labels


selected_gestures = [0, 1, 21, 22] # indexes for selected gestures
train_path = 'sign_mnist_train/sign_mnist_train.csv'
test_path = 'sign_mnist_test/sign_mnist_test.csv'

train_x, train_y = preprocess_dataset(train_path, selected_gestures)
test_x, test_y = preprocess_dataset(test_path, selected_gestures)





# creating dataset

class GestureDataset(Dataset):
    def __init__(self, images, labels, transforms):
        self.images = images
        self.labels = labels
        self.transforms = transforms

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        if torch.is_tensor(idx):
            idx = idx.tolist()
        if self.transforms:
            image = self.transforms(self.images[idx])
        else:
            image = self.images[idx]
        label = self.labels[idx]
        return {'image': image, 'label': label}

# training dataset
transform =  transforms.ToTensor()
train_dataset = GestureDataset(train_x, train_y, transform)

# validation dataset
test_dataset = GestureDataset(test_x, test_y, transform)






# dataloaders (to input data to the model)

train_loader = DataLoader(train_dataset, 
                            batch_size=64,
                            shuffle=True,
                            num_workers=0, # DO NOT CHANGE NUM WORKERS (Windows specific PyTorch bug)
                            )
test_loader = DataLoader(test_dataset, 
                            batch_size=128,
                            shuffle=False,
                            num_workers=0, # DO NOT CHANGE NUM WORKERS
                            )




# model
class GestureModel(nn.Module):
    """Neural net for recognizing hand gestures
    Input dims: m x 1 x 28 x 28
    Output dims: m x 4
    """
    def __init__(self):
        super().__init__()
        # input: m x 1 x 28 x 28
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(16),
        ) # m x 16 x 28 x 28
        self.conv2 = nn.Sequential(
            # nn.Dropout(p=0.2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2),
        ) # m x 32 x 14 x 14
        self.conv3 = nn.Sequential(
            # nn.Dropout(p=0.2),
            nn.Conv2d(32, 64, kernel_size=3),
            nn.ReLU(),
            nn.BatchNorm2d(64),
            nn.MaxPool2d(2),
        ) # m x 64 x 6 x 6
        self.conv4 = nn.Sequential(
            # nn.Dropout(p=0.2),
            nn.Conv2d(64, 128, kernel_size=3),
            nn.ReLU(),
            nn.BatchNorm2d(128),
            nn.MaxPool2d(2),
        ) # m x 128 x 2 x 2
        self.classifier = nn.Sequential(
            nn.Flatten(), # m x 128*2*2
            nn.Dropout(p=0.2),
            nn.Linear(128*2*2, 4),
            nn.Softmax(dim=1),
        ) # m x 4

    def forward(self, input):
        x = self.conv1(input)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        output = self.classifier(x)
        return output

    def accuracy(self, model_output, y):
        """Calculates accuracy on a batch's predictions.
        Args:
            model_output::torch.Tensor::shape(m x 4)- Output generated by the model on a batch of inputs
            y::torch.Tensor::shape(m)- Original labels for the images
        Returns:
            accuracy::int- Accuracy for the batch
        """
        y_hat = torch.argmax(model_output, axis=1)
        accuracy = torch.sum(y_hat == y).item()/y.shape[0]
        return accuracy

    def training_step(self, batch):
        """Single training step on a batch of data."""
        images = batch[0]
        labels = batch[1]
        output = self(images)
        batch_loss = nn.CrossEntropyLoss()(output, labels)
        return batch_loss

    def validation_step(self, batch):
        """Single validation step on a batch of data."""
        images = batch[0]
        labels = batch[1]
        output = self(images)
        batch_loss = nn.CrossEntropyLoss()(output, labels)
        batch_acc = self.accuracy(output, labels)
        return torch.Tensor([batch_loss, batch_acc])

    def epoch_end(self, epoch, result):
        print("Epoch: {} || LR: {:.5f} || Training Loss: {:.5f} || Validation Loss: {:.5f} || Validation Accuracy: {:.5f}"
                .format(epoch, result[3], result[2], result[0], result[1]))







def get_default_device():
    """Picks the trainig device-- GPU if available, else CPU.
    """
    if torch.cuda.is_available():   # checks if a cuda device is available
        return torch.device('cuda') # sets the default device as the available CUDA device
    else:
        return torch.device('cpu')  # if no CUDA device found, sets CPU as the default device

def to_device(data, device):
    """Move tensor(s) to chosen device
    """
    if isinstance(data, (list,tuple)): # asserts if the data is a list/tuple 
        return [to_device(x, device) for x in data]
    return data.to(device, non_blocking=True)

class DeviceDataLoader():
    """Wrap a dataloader to move data to the default device.
    """
    def __init__(self, dataloader, device):
        self.dl = dataloader
        self.device = device
        
    def __iter__(self):
        for batch in self.dl: 
            image, label = batch['image'], batch['label']
            yield to_device([image.float(), label], self.device)

    def __len__(self):
        return len(self.dl)

# setting default device
device = get_default_device()

# creating device dataloaders
train_ddl = DeviceDataLoader(train_loader, device)
test_ddl = DeviceDataLoader(test_loader, device)

# moving model to device
model = GestureModel()
model.to(device, non_blocking=True)





# validation epoch
@torch.no_grad()
def evaluate(model, validation_ddl):
    """Calculates validation loss and validation accuracy over an epoch of validation/test data"""
    model.eval()    # sets the model to evaluation mode
    output = torch.zeros(2)
    total_batches = 0
    for batch in validation_ddl:
        total_batches += 1
        output += model.validation_step(batch)
    output = output / total_batches
    return output

def get_lr(optimizer):
    """Gets the learning rate of the optimizer.
    """
    for param_group in optimizer.param_groups:
        return param_group['lr']


def fit_one_cycle(epochs, max_lr, model, train_dataloader, val_dataloader, 
                  weight_decay = 0, grad_clip = None, opt_func = torch.optim.Adam):
    
    torch.cuda.empty_cache()    # clears cache in CUDA device
    history = []    # declares an empty list to store result for each epoch
    
    # sets up optimizer with weight decay
    optimizer = opt_func(model.parameters(), max_lr, weight_decay = weight_decay)
    
    # sets up one-cycle learning rate scheduler
    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr, epochs = epochs, steps_per_epoch = len(train_dataloader))
    
    for epoch in range(epochs): 
        model.train()    # initiate training phase
        train_losses = []
        lrs = 0
        for batch in train_dataloader:    # cycles through each batch of the training data
            loss = model.training_step(batch)    
            train_losses.append(loss)
            loss.backward()
            
            # perfomrs gradient clipping
            if grad_clip: 
                nn.utils.clip_grad_value_(model.parameters(), grad_clip)
            
            optimizer.step() # updates parameters based on gradients
            optimizer.zero_grad() # resets gradient values
            
            # records & updates learning rate
            lrs = get_lr(optimizer)
            scheduler.step()
        
        # initaites validation phase
        result = evaluate(model, val_dataloader).tolist()
        result.append(torch.stack(train_losses).mean().item())
        result.append(lrs)
        model.epoch_end(epoch, result)
        history.append(result)
    return history


# training the model
history = fit_one_cycle(epochs=10, max_lr=0.001, model=model, 
              train_dataloader=train_ddl, val_dataloader=test_ddl, 
              weight_decay=0.01, grad_clip=0.001,
              opt_func=torch.optim.Adam)

# saving the trained model
torch.save(model.state_dict(), 'hand_gesture_model.pth')