# Path: src/agents/mk_dataset.py
# class MarioKartDataset(Dataset):
#     """Mario Kart dataset."""
# 
#     def __init__(self, log_file=os.path.join(helper.get_dataset_folder(), 'mario_kart', "mario_kart.json"),
#                  transform=None, history=1):
#         """
#         Args:
#             log_file: file path location to the log file containing key press states.
#         """
#         self.images = os.path.join(helper.get_dataset_folder(), 'mario_kart', "images")
#         with open(log_file, 'r') as f:
#             self.data = json.load(f).get('data')
#         self.transform = transform
#         self.history = history
# 
#     def __len__(self):
#         return len(self.data)
# 
#     def __getitem__(self, idx):
#         key_state = self.data[idx]
#         img_count = key_state.get('count')
#         tensors = []
#         for i in range(self.history):
#             img_name = os.path.join(self.images, "{}.png".format(img_count - i))
#             image = cv2.imread(img_name)
#             img_tensor = helper.get_tensor(image)
#             if img_tensor is not None:
#                 tensors.append(img_tensor)
#             else:
#                 logger.warning("Image not found, using last tensor found")
#                 tensors.append(tensors[-1])
#         tensor = torch.stack(tensors)
#         return tensor, helper.get_output_vector(key_state['presses'])


import numpy as np
import torch
from torch.utils.data.sampler import SubsetRandomSampler
from torchvision import transforms
from src.agents.mk_dataset import MarioKartDataset
    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)

# Path: src/agents/mk_dataset.py
# class MarioKartDataset(Dataset):
#     """Mario Kart dataset."""
# 
#     def __init__(self, log_file=os.path.join(helper.get_dataset_folder(), 'mario_kart', "mario_kart.json"),
#                  transform=None, history=1):
#         """
#         Args:
#             log_file: file path location to the log file containing key press states.
#         """
#         self.images = os.path.join(helper.get_dataset_folder(), 'mario_kart', "images")
#         with open(log_file, 'r') as f:
#             self.data = json.load(f).get('data')
#         self.transform = transform
#         self.history = history
# 
#     def __len__(self):
#         return len(self.data)
# 
#     def __getitem__(self, idx):
#         key_state = self.data[idx]
#         img_count = key_state.get('count')
#         tensors = []
#         for i in range(self.history):
#             img_name = os.path.join(self.images, "{}.png".format(img_count - i))
#             image = cv2.imread(img_name)
#             img_tensor = helper.get_tensor(image)
#             if img_tensor is not None:
#                 tensors.append(img_tensor)
#             else:
#                 logger.warning("Image not found, using last tensor found")
#                 tensors.append(tensors[-1])
#         tensor = torch.stack(tensors)
#         return tensor, helper.get_output_vector(key_state['presses'])


# Path: src/agents/train_valid_data.py
import numpy as np
import torch
from torch.utils.data.sampler import SubsetRandomSampler
from torchvision import transforms
from src.agents.mk_dataset import MarioKartDataset

    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)