# PyTorch
import torch
from torch.utils.data import IterableDataset, DataLoader
from donkeycar.utils import train_test_split
from donkeycar.parts.tub_v2 import Tub
from torchvision import transforms
from typing import List, Any
from donkeycar.pipeline.types import TubRecord, TubDataset
from donkeycar.pipeline.sequence import TubSequence
import pytorch_lightning as pl


def get_default_transform(for_video=False, for_inference=False, resize=True):
    """
    Creates a default transform to work with torchvision models

    Video transform:
    All pre-trained models expect input images normalized in the same way, 
    i.e. mini-batches of 3-channel RGB videos of shape (3 x T x H x W), 
    where H and W are expected to be 112, and T is a number of video frames 
    in a clip. The images have to be loaded in to a range of [0, 1] and 
    then normalized using mean = [0.43216, 0.394666, 0.37645] and 
    std = [0.22803, 0.22145, 0.216989].
    """

    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    input_size = (224, 224)

    if for_video:
        mean = [0.43216, 0.394666, 0.37645]
        std = [0.22803, 0.22145, 0.216989]
        input_size = (112, 112)

    transform_items = [
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ]

    if resize:
        transform_items.insert(0, transforms.Resize(input_size))

    return transforms.Compose(transform_items)


class TorchTubDataset(IterableDataset):
    '''
    Loads the dataset, and creates a train/test split.
    '''

    def __init__(self, config, records: List[TubRecord], transform=None):
        """Create a PyTorch Tub Dataset

        Args:
            config (object): the configuration information
            records (List[TubRecord]): a list of tub records
            transform (function, optional): a transform to apply to the data
        """
        self.config = config

        # Handle the transforms
        if transform:
            self.transform = transform
        else:
            self.transform = get_default_transform()

        self.sequence = TubSequence(records)
        self.pipeline = self._create_pipeline()
        self.len = len(records)

    def _create_pipeline(self):
        """ This can be overridden if more complicated pipelines are
            required """

        def y_transform(record: TubRecord):
            angle: float = record.underlying['user/angle']
            throttle: float = record.underlying['user/throttle']
            predictions = torch.tensor([angle, throttle], dtype=torch.float)

            # Normalize to be between [0, 1]
            # angle and throttle are originally between [-1, 1]
            predictions = (predictions + 1) / 2
            return predictions

        def x_transform(record: TubRecord):
            # Loads the result of Image.open()
            img_arr = record.image(cached=True, as_nparray=False)
            return self.transform(img_arr)

        # Build pipeline using the transformations
        pipeline = self.sequence.build_pipeline(x_transform=x_transform,
                                                y_transform=y_transform)
        return pipeline

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

    def __iter__(self):
        return iter(self.pipeline)

class TorchTubDataModule(pl.LightningDataModule):

    def __init__(self, config: Any, tub_paths: List[str], transform=None):
        """Create a PyTorch Lightning Data Module to contain all data loading logic

        Args:
            config (object): the configuration information
            tub_paths (List[str]): a list of paths to the tubs to use (minimum size of 1).
                                   Each tub path corresponds to another training run.
            transform (function, optional): a transform to apply to the data
        """
        super().__init__()

        self.config = config
        self.tub_paths = tub_paths

        # Handle the transforms
        if transform:
            self.transform = transform
        else:
            self.transform = get_default_transform()

        self.tubs: List[Tub] = [Tub(tub_path, read_only=True)
                                for tub_path in self.tub_paths]
        self.records: List[TubRecord] = []

    def setup(self, stage=None):
        """Load all the tub data and set up the datasets.

        Args:
            stage ([string], optional): setup expects a string arg stage. 
                                        It is used to separate setup logic for trainer.fit 
                                        and trainer.test. Defaults to None.
        """
        # Loop through all the different tubs and load all the records for each of them
        for tub in self.tubs:
            for underlying in tub:
                record = TubRecord(self.config, tub.base_path,
                                   underlying=underlying)
                self.records.append(record)

        train_records, val_records = train_test_split(
            self.records, test_size=(1. - self.config.TRAIN_TEST_SPLIT))

        assert len(val_records) > 0, "Not enough validation data. Add more data"

        self.train_dataset = TorchTubDataset(
            self.config, train_records, transform=self.transform)
        self.val_dataset = TorchTubDataset(
            self.config, val_records, transform=self.transform)

    def train_dataloader(self):
        # The number of workers are set to 0 to avoid errors on Macs and Windows
        # See: https://github.com/rusty1s/pytorch_geometric/issues/366#issuecomment-498022534
        return DataLoader(self.train_dataset, batch_size=self.config.BATCH_SIZE, num_workers=0)

    def val_dataloader(self):
        # The number of workers are set to 0 to avoid errors on Macs and Windows
        # See: https://github.com/rusty1s/pytorch_geometric/issues/366#issuecomment-498022534
        return DataLoader(self.val_dataset, batch_size=self.config.BATCH_SIZE, num_workers=0)

[((43, 11, 43, 46), 'torchvision.transforms.Compose', 'transforms.Compose', ({(43, 30, 43, 45): 'transform_items'}, {}), '(transform_items)', False, 'from torchvision import transforms\n'), ((36, 8, 36, 29), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', False, 'from torchvision import transforms\n'), ((37, 8, 37, 48), 'torchvision.transforms.Normalize', 'transforms.Normalize', (), '', False, 'from torchvision import transforms\n'), ((67, 24, 67, 44), 'donkeycar.pipeline.sequence.TubSequence', 'TubSequence', ({(67, 36, 67, 43): 'records'}, {}), '(records)', False, 'from donkeycar.pipeline.sequence import TubSequence\n'), ((142, 37, 143, 72), 'donkeycar.utils.train_test_split', 'train_test_split', (), '', False, 'from donkeycar.utils import train_test_split\n'), ((155, 15, 155, 95), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import IterableDataset, DataLoader\n'), ((160, 15, 160, 93), 'torch.utils.data.DataLoader', 'DataLoader', (), '', False, 'from torch.utils.data import IterableDataset, DataLoader\n'), ((41, 34, 41, 63), 'torchvision.transforms.Resize', 'transforms.Resize', ({(41, 52, 41, 62): 'input_size'}, {}), '(input_size)', False, 'from torchvision import transforms\n'), ((78, 26, 78, 76), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((123, 32, 123, 61), 'donkeycar.parts.tub_v2.Tub', 'Tub', (), '', False, 'from donkeycar.parts.tub_v2 import Tub\n'), ((138, 25, 139, 57), 'donkeycar.pipeline.types.TubRecord', 'TubRecord', (), '', False, 'from donkeycar.pipeline.types import TubRecord, TubDataset\n')]