# Path: src/helper.py
# def validate_function_key(slot_key):
# def get_output_folder():
# def get_models_folder():
# def get_dataset_folder():
# def get_home_folder():
# def pickle_object(data, name, directory=get_output_folder()):
# def generate_img_key(image):
# def get_tensor(image):
# def get_output_vector(presses):
# def get_key_state_from_vector(vector):
# 
# Path: src/keylog.py
# class Keyboard(enum.Enum):
# class KeyLog:
#     def __init__(self, logging_delay=0.3):
#     def start(self):
#     def on_press(self, key):
#     def on_release(self, key):
#     def record(self):
#     def save_to_file(self, file_name="log.json"):
#     def _get_key_value(key):
# 
# Path: src/agents/train_valid_data.py
# def get_mario_train_valid_loader(batch_size,
#                                  augment,
#                                  random_seed,
#                                  valid_size=0.2,
#                                  shuffle=True,
#                                  history=1,
#                                  num_workers=1,
#                                  pin_memory=False):
#     """
#     Utility function for loading and returning train and valid
#     multi-process iterators over the MarioKart dataset. A sample
#     9x9 grid of the images can be optionally displayed.
#     If using CUDA, num_workers should be set to 1 and pin_memory to True.
#     Params
#     ------
#     - data_dir: path directory to the dataset.
#     - batch_size: how many samples per batch to load.
#     - augment: whether to apply the data augmentation scheme
#       mentioned in the paper. Only applied on the train split.
#     - random_seed: fix seed for reproducibility.
#     - valid_size: percentage split of the training set used for
#       the validation set. Should be a float in the range [0, 1].
#     - shuffle: whether to shuffle the train/validation indices.
#     - show_sample: plot 9x9 sample grid of the dataset.
#     - num_workers: number of subprocesses to use when loading the dataset.
#     - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
#       True if using GPU.
#     Returns
#     -------
#     - train_loader: training set iterator.
#     - valid_loader: validation set iterator.
#     """
#     error_msg = "[!] valid_size should be in the range [0, 1]."
#     assert ((valid_size >= 0) and (valid_size <= 1)), error_msg
# 
#     normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
#                                      std=[0.5, 0.5, 0.5])
# 
#     # define transforms
#     valid_transform = transforms.Compose([
#         transforms.ToTensor(),
#         normalize
#     ])
#     if augment:
#         train_transform = transforms.Compose([
#             transforms.RandomCrop(32, padding=4),
#             transforms.ToTensor(),
#             normalize
#         ])
#     else:
#         train_transform = transforms.Compose([
#             transforms.ToTensor(),
#             normalize
#         ])
# 
#     # load the dataset
#     train_dataset = MarioKartDataset(transform=train_transform, history=history)
#     valid_dataset = MarioKartDataset(transform=valid_transform, history=history)
# 
#     num_train = len(train_dataset)
#     indices = list(range(num_train))
#     split = int(np.floor(valid_size * num_train))
# 
#     if shuffle:
#         np.random.seed(random_seed)
#         np.random.shuffle(indices)
# 
#     train_idx, valid_idx = indices[split:], indices[:split]
# 
#     train_sampler = SubsetRandomSampler(train_idx)
#     valid_sampler = SubsetRandomSampler(valid_idx)
# 
#     train_loader = torch.utils.data.DataLoader(train_dataset,
#                                                batch_size=batch_size, sampler=train_sampler,
#                                                num_workers=num_workers, pin_memory=pin_memory)
# 
#     valid_loader = torch.utils.data.DataLoader(valid_dataset,
#                                                batch_size=batch_size, sampler=valid_sampler,
#                                                num_workers=num_workers, pin_memory=pin_memory)
#     return train_loader, valid_loader


import logging
import os
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from src import helper, keylog
from src.agents.train_valid_data import get_mario_train_valid_loader
from torch.autograd import Variable
                y_label = y_label.cuda()
            nn_label = Variable(y_label)

            # Forward pass
            forward_pass = mkcrnn(x)
            loss = loss_func(forward_pass, nn_label)  # compute loss
            optimizer.zero_grad()  # zero gradients from previous step
            loss.backward()  # compute gradients
            optimizer.step()  # apply backpropagation

            # Log training data
            if step % 50 == 0:
                print('Epoch: ', epoch, 'Step: ', step, '| training loss: %.4f' % loss.data[0])
                valid_loss = 0
                for (valid_x, valid_y) in valid_loader:
                    valid_y_label = valid_y.view(-1, output_vec)
                    if torch.cuda.is_available():
                        valid_y_label = valid_y_label.cuda()
                    valid_nn_label = Variable(valid_y_label)
                    valid_forward_pass = mkcrnn(valid_x)
                    valid_loss_eval = loss_func(valid_forward_pass, valid_nn_label)  # compute validation loss
                    valid_loss += valid_loss_eval.data[0]
                print('Epoch: ', epoch, 'Step: ', step, '| validation loss: %.4f' % valid_loss)
                validation_losses.append(valid_loss)

    # Save model
    torch.save(mkcrnn, os.path.join(helper.get_models_folder(), "mkcrnn_{}_gru.pkl".format(history)))

    # Save validation curve data
    fig_data = [validation_losses, history, num_epochs, batch_size, learning_rate]

# Path: src/helper.py
# def validate_function_key(slot_key):
# def get_output_folder():
# def get_models_folder():
# def get_dataset_folder():
# def get_home_folder():
# def pickle_object(data, name, directory=get_output_folder()):
# def generate_img_key(image):
# def get_tensor(image):
# def get_output_vector(presses):
# def get_key_state_from_vector(vector):
# 
# Path: src/keylog.py
# class Keyboard(enum.Enum):
# class KeyLog:
#     def __init__(self, logging_delay=0.3):
#     def start(self):
#     def on_press(self, key):
#     def on_release(self, key):
#     def record(self):
#     def save_to_file(self, file_name="log.json"):
#     def _get_key_value(key):
# 
# Path: src/agents/train_valid_data.py
# def get_mario_train_valid_loader(batch_size,
#                                  augment,
#                                  random_seed,
#                                  valid_size=0.2,
#                                  shuffle=True,
#                                  history=1,
#                                  num_workers=1,
#                                  pin_memory=False):
#     """
#     Utility function for loading and returning train and valid
#     multi-process iterators over the MarioKart dataset. A sample
#     9x9 grid of the images can be optionally displayed.
#     If using CUDA, num_workers should be set to 1 and pin_memory to True.
#     Params
#     ------
#     - data_dir: path directory to the dataset.
#     - batch_size: how many samples per batch to load.
#     - augment: whether to apply the data augmentation scheme
#       mentioned in the paper. Only applied on the train split.
#     - random_seed: fix seed for reproducibility.
#     - valid_size: percentage split of the training set used for
#       the validation set. Should be a float in the range [0, 1].
#     - shuffle: whether to shuffle the train/validation indices.
#     - show_sample: plot 9x9 sample grid of the dataset.
#     - num_workers: number of subprocesses to use when loading the dataset.
#     - pin_memory: whether to copy tensors into CUDA pinned memory. Set it to
#       True if using GPU.
#     Returns
#     -------
#     - train_loader: training set iterator.
#     - valid_loader: validation set iterator.
#     """
#     error_msg = "[!] valid_size should be in the range [0, 1]."
#     assert ((valid_size >= 0) and (valid_size <= 1)), error_msg
# 
#     normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
#                                      std=[0.5, 0.5, 0.5])
# 
#     # define transforms
#     valid_transform = transforms.Compose([
#         transforms.ToTensor(),
#         normalize
#     ])
#     if augment:
#         train_transform = transforms.Compose([
#             transforms.RandomCrop(32, padding=4),
#             transforms.ToTensor(),
#             normalize
#         ])
#     else:
#         train_transform = transforms.Compose([
#             transforms.ToTensor(),
#             normalize
#         ])
# 
#     # load the dataset
#     train_dataset = MarioKartDataset(transform=train_transform, history=history)
#     valid_dataset = MarioKartDataset(transform=valid_transform, history=history)
# 
#     num_train = len(train_dataset)
#     indices = list(range(num_train))
#     split = int(np.floor(valid_size * num_train))
# 
#     if shuffle:
#         np.random.seed(random_seed)
#         np.random.shuffle(indices)
# 
#     train_idx, valid_idx = indices[split:], indices[:split]
# 
#     train_sampler = SubsetRandomSampler(train_idx)
#     valid_sampler = SubsetRandomSampler(valid_idx)
# 
#     train_loader = torch.utils.data.DataLoader(train_dataset,
#                                                batch_size=batch_size, sampler=train_sampler,
#                                                num_workers=num_workers, pin_memory=pin_memory)
# 
#     valid_loader = torch.utils.data.DataLoader(valid_dataset,
#                                                batch_size=batch_size, sampler=valid_sampler,
#                                                num_workers=num_workers, pin_memory=pin_memory)
#     return train_loader, valid_loader


# Path: src/agents/mk_crnn_gru_train.py
import logging
import os
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from src import helper, keylog
from src.agents.train_valid_data import get_mario_train_valid_loader
from torch.autograd import Variable

                y_label = y_label.cuda()
            nn_label = Variable(y_label)

            # Forward pass
            forward_pass = mkcrnn(x)
            loss = loss_func(forward_pass, nn_label)  # compute loss
            optimizer.zero_grad()  # zero gradients from previous step
            loss.backward()  # compute gradients
            optimizer.step()  # apply backpropagation

            # Log training data
            if step % 50 == 0:
                print('Epoch: ', epoch, 'Step: ', step, '| training loss: %.4f' % loss.data[0])
                valid_loss = 0
                for (valid_x, valid_y) in valid_loader:
                    valid_y_label = valid_y.view(-1, output_vec)
                    if torch.cuda.is_available():
                        valid_y_label = valid_y_label.cuda()
                    valid_nn_label = Variable(valid_y_label)
                    valid_forward_pass = mkcrnn(valid_x)
                    valid_loss_eval = loss_func(valid_forward_pass, valid_nn_label)  # compute validation loss
                    valid_loss += valid_loss_eval.data[0]
                print('Epoch: ', epoch, 'Step: ', step, '| validation loss: %.4f' % valid_loss)
                validation_losses.append(valid_loss)

    # Save model
    torch.save(mkcrnn, os.path.join(helper.get_models_folder(), "mkcrnn_{}_gru.pkl".format(history)))

    # Save validation curve data
    fig_data = [validation_losses, history, num_epochs, batch_size, learning_rate]

    helper.pickle_object(fig_data, "mkcrnn_training_data_{}_gru".format(history))