# From https://github.com/brentyi/multimodalfilter/blob/master/crossmodal/tasks/_push.py

import argparse
import sys
from typing import Any, Dict, List, NamedTuple, Tuple

from PIL.Image import NONE

import fannypack
import fannypack as fp
import numpy as np
import torch

from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from robustness.visual_robust import visual_robustness
from robustness.timeseries_robust import timeseries_robustness

dataset_urls = {
    # Mujoco URLs
    "gentle_push_10.hdf5": "https://drive.google.com/file/d/1qmBCfsAGu8eew-CQFmV1svodl9VJa6fX/view?usp=sharing",
    "gentle_push_100.hdf5": "https://drive.google.com/file/d/1PmqQy5myNXSei56upMy3mXKu5Lk7Fr_g/view?usp=sharing",
    "gentle_push_300.hdf5": "https://drive.google.com/file/d/18dr1z0N__yFiP_DAKxy-Hs9Vy_AsaW6Q/view?usp=sharing",
    "gentle_push_1000.hdf5": "https://drive.google.com/file/d/1JTgmq1KPRK9HYi8BgvljKg5MPqT_N4cR/view?usp=sharing",
    # Real data (kloss_dataset=True)
    "kloss_train0.hdf5": "https://drive.google.com/file/d/1nk4BO0rcVTKw22vYq6biewiwAFUPevM1/view?usp=sharing",
    "kloss_train1.hdf5": "https://drive.google.com/file/d/1gBWoB2PCrgYlLjuDJQm6BFAf_xwMqRxa/view?usp=sharing",
    "kloss_train2.hdf5": "https://drive.google.com/file/d/15W2zj52bSITxIRVRi7ajehAmz14RU33M/view?usp=sharing",
    "kloss_train3.hdf5": "https://drive.google.com/file/d/1WhRFu4SDlIYKnLYLyDdgOQYjP20JOTLE/view?usp=sharing",
    "kloss_train4.hdf5": "https://drive.google.com/file/d/1-ur_hzyBvd1_QCLTamaO8eWJ7rXii7y4/view?usp=sharing",
    "kloss_train5.hdf5": "https://drive.google.com/file/d/1ni8vEy4c1cmCKP2ZlWfXqLo7a4sdRFwe/view?usp=sharing",
    "kloss_val.hdf5": "https://drive.google.com/file/d/1-CRocf7I4mTLBp7Tjo7-D-QvkwcGZkNo/view?usp=sharing",
}


class TrajectoryNumpy(NamedTuple):
    """Named tuple containing states, observations, and controls represented in NumPy."""

    states: Any
    observations: Any
    controls: Any


class PushTask():
    """Dataset definition and model registry for pushing task."""

    @classmethod
    def add_dataset_arguments(cls, parser: argparse.ArgumentParser) -> None:
        """Add dataset options to an argument parser.

        Args:
            parser (argparse.ArgumentParser): Parser to add arguments to.
        """
        parser.add_argument("--no_vision", action="store_true")
        parser.add_argument("--no_proprioception", action="store_true")
        parser.add_argument("--no_haptics", action="store_true")
        parser.add_argument("--image_blackout_ratio", type=float, default=0.0)
        parser.add_argument("--sequential_image_rate", type=int, default=1)
        parser.add_argument("--kloss_dataset", action="store_true")

    @classmethod
    def get_dataset_args(cls, args: argparse.Namespace) -> Dict[str, Any]:
        """Get a dataset_args dictionary from a parsed set of arguments.

        Args:
            args (argparse.Namespace): Parsed arguments.
        """
        dataset_args = {
            "use_vision": not args.no_vision,
            "use_proprioception": not args.no_proprioception,
            "use_haptics": not args.no_haptics,
            "image_blackout_ratio": args.image_blackout_ratio,
            "sequential_image_rate": args.sequential_image_rate,
            "kloss_dataset": args.kloss_dataset,
        }
        return dataset_args

    @classmethod
    def get_dataloader(
        cls, subsequence_length: int, modalities=None, test_multimodal_only=False,
        batch_size=32, drop_last=True, test_noises=range(10), train_sampler=None, **dataset_args
    ):
        # Load trajectories into memory
        train_trajectories = cls.get_train_trajectories(**dataset_args)
        val_trajectories = cls.get_eval_trajectories(**dataset_args)
        test_trajectories = cls.get_test_trajectories(
            modalities if not test_multimodal_only else ['multimodal'],
            noises=test_noises, **dataset_args)
        train_loader = DataLoader(
            SubsequenceDataset(train_trajectories,
                               subsequence_length, modalities),
            batch_size=batch_size,
            sampler=train_sampler,
            shuffle=True if train_sampler is None else False,
            drop_last=drop_last,
        )
        val_loader = DataLoader(
            SubsequenceDataset(
                val_trajectories, subsequence_length, modalities),
            batch_size=batch_size,
            shuffle=True,
        )
        test_loader = dict()
        for modality, trajectories in test_trajectories.items():
            test_loader[modality] = []
            for traj in trajectories:
                test_loader[modality].append(DataLoader(
                    SubsequenceDataset(traj, subsequence_length, modalities),
                    batch_size=batch_size,
                    shuffle=False,
                )
                )
        return train_loader, val_loader, test_loader

    @classmethod
    def get_train_trajectories(
        cls, **dataset_args
    ) -> List[TrajectoryNumpy]:

        kloss_dataset = (
            dataset_args["kloss_dataset"] if "kloss_dataset" in dataset_args else False
        )
        if kloss_dataset:
            return _load_trajectories(
                "kloss_train0.hdf5",
                "kloss_train1.hdf5",
                "kloss_train2.hdf5",
                "kloss_train3.hdf5",
                "kloss_train4.hdf5",
                "kloss_train5.hdf5",
                **dataset_args,
            )
        else:
            return _load_trajectories("gentle_push_1000.hdf5", **dataset_args)

    @classmethod
    def get_eval_trajectories(
        cls, **dataset_args
    ) -> List[TrajectoryNumpy]:

        kloss_dataset = (
            dataset_args["kloss_dataset"] if "kloss_dataset" in dataset_args else False
        )
        if kloss_dataset:
            return _load_trajectories(("kloss_val.hdf5", 50), **dataset_args)
        else:
            return _load_trajectories("gentle_push_10.hdf5", **dataset_args)

    @classmethod
    def get_test_trajectories(
        cls, modalities, noises=range(10), **dataset_args
    ):

        kloss_dataset = (
            dataset_args["kloss_dataset"] if "kloss_dataset" in dataset_args else False
        )
        if kloss_dataset:
            raise Exception('No test dataset for kloss')
        else:
            trajectories = dict()
            if modalities == None or 'image' in modalities:
                trajectories['image'] = []
                for i in noises:
                    trajectories['image'].append(_load_trajectories(
                        "gentle_push_300.hdf5", visual_noise=i/10, **dataset_args))
            if modalities == None or 'gripper_pos' in modalities:
                trajectories['proprio'] = []
                for i in noises:
                    trajectories['proprio'].append(_load_trajectories(
                        "gentle_push_300.hdf5", prop_noise=i/10, **dataset_args))
            if modalities == None or 'gripper_sensors' in modalities:
                trajectories['haptics'] = []
                for i in noises:
                    trajectories['haptics'].append(_load_trajectories(
                        "gentle_push_300.hdf5", haptics_noise=i/10, **dataset_args))
            if modalities == None or 'controls' in modalities:
                trajectories['controls'] = []
                for i in noises:
                    trajectories['controls'].append(_load_trajectories(
                        "gentle_push_300.hdf5", controls_noise=i/10, **dataset_args))
            if modalities == None or 'multimodal' in modalities:
                trajectories['multimodal'] = []
                for i in noises:
                    trajectories['multimodal'].append(_load_trajectories(
                        "gentle_push_300.hdf5", multimodal_noise=i/10, **dataset_args))
            return trajectories


def _load_trajectories(
    *input_files,
    use_vision: bool = True,
    use_proprioception: bool = True,
    use_haptics: bool = True,
    vision_interval: int = 10,
    image_blackout_ratio: float = 0.0,
    sequential_image_rate: int = 1,
    start_timestep: int = 0,
    kloss_dataset: bool = False,
    visual_noise: int = 0,
    prop_noise: int = 0,
    haptics_noise: int = 0,
    controls_noise: int = 0,
    multimodal_noise: int = 0,
) -> List[TrajectoryNumpy]:
    """Loads a list of trajectories from a set of input files, where each trajectory is
    a tuple containing...
        states: an (T, state_dim) array of state vectors
        observations: a key->(T, *) dict of observations
        controls: an (T, control_dim) array of control vectors

    Each input can either be a string or a (string, int) tuple, where int indicates the
    maximum number of trajectories to import.


    Args:
        *input_files: Trajectory inputs. Should be members of
            `crossmodal.push_data.dataset_urls.keys()`.

    Keyword Args:
        use_vision (bool, optional): Set to False to zero out camera inputs.
        vision_interval (int, optional): # of times each camera image is duplicated. For
            emulating a slow image rate.
        use_proprioception (bool, optional): Set to False to zero out kinematics data.
        use_haptics (bool, optional): Set to False to zero out F/T sensors.
        image_blackout_ratio (float, optional): Dropout probabiliity for camera inputs.
            0.0 = no dropout, 1.0 = all images dropped out.
        sequential_image_rate (int, optional): If value is `N`, we only send 1 image
            frame ever `N` timesteps. All others are zeroed out.
        start_timestep (int, optional): If value is `N`, we skip the first `N` timesteps
            of each trajectory.

    Returns:
        List[TrajectoryNumpy]: list of trajectories.
    """
    trajectories = []

    assert 1 > image_blackout_ratio >= 0
    assert image_blackout_ratio == 0 or sequential_image_rate == 1

    for name in input_files:
        max_trajectory_count = sys.maxsize
        if type(name) == tuple:
            name, max_trajectory_count = name
        assert type(max_trajectory_count) == int

        # Load trajectories file into memory, all at once
        with fannypack.data.TrajectoriesFile(
            fannypack.data.cached_drive_file(name, dataset_urls[name])
        ) as f:
            raw_trajectories = list(f)

        # Iterate over each trajectory
        for raw_trajectory_index, raw_trajectory in enumerate(raw_trajectories):
            if raw_trajectory_index >= max_trajectory_count:
                break

            if kloss_dataset:
                timesteps = len(raw_trajectory["pos"])
            else:
                timesteps = len(raw_trajectory["object-state"])

            # State is just (x, y)
            state_dim = 2
            states = np.full((timesteps, state_dim), np.nan, dtype=np.float32)

            if kloss_dataset:
                states[:, 0] = raw_trajectory["pos"][:, 0]
                states[:, 1] = raw_trajectory["pos"][:, 2]
            else:
                states[:, :2] = raw_trajectory["Cylinder0_pos"][:, :2]  # x, y

            # Pull out observations
            ## This is currently consisted of:
            ## > gripper_pos: end effector position
            ## > gripper_sensors: F/T, contact sensors
            ## > image: camera image

            observations = {}

            if kloss_dataset:
                observations["gripper_pos"] = raw_trajectory["tip"]
            else:
                observations["gripper_pos"] = raw_trajectory["eef_pos"]
            if prop_noise != 0:
                observations["gripper_pos"] = timeseries_robustness(
                    [observations["gripper_pos"]], noise_level=prop_noise, struct_drop=False)[0]
            assert observations["gripper_pos"].shape == (timesteps, 3)

            if kloss_dataset:
                observations["gripper_sensors"] = np.zeros(
                    (timesteps, 7), dtype=np.float32)
                observations["gripper_sensors"][:,
                                                :3] = raw_trajectory["force"]
                observations["gripper_sensors"][:,
                                                6] = raw_trajectory["contact"]
            else:
                observations["gripper_sensors"] = np.concatenate(
                    (
                        raw_trajectory["force"],
                        raw_trajectory["contact"][:, np.newaxis],
                    ),
                    axis=1,
                )
            if haptics_noise != 0:
                observations["gripper_sensors"] = timeseries_robustness(
                    [observations["gripper_sensors"]], noise_level=haptics_noise, struct_drop=False)[0]
            assert observations["gripper_sensors"].shape[1] == 7

            # Zero out proprioception or haptics if unused
            if not use_proprioception:
                observations["gripper_pos"][:] = 0
            if not use_haptics:
                observations["gripper_sensors"][:] = 0

            # Get image
            if kloss_dataset:
                observations["image"] = np.mean(
                    raw_trajectory["image"], axis=-1)
            else:
                observations["image"] = raw_trajectory["image"].copy()
            if visual_noise != 0:
                observations["image"] = np.array(visual_robustness(
                    observations["image"], noise_level=visual_noise))
            assert observations["image"].shape == (timesteps, 32, 32)

            # Mask image observations based on dataset args
            image_mask: np.ndarray
            if not use_vision:
                # Use a zero mask
                image_mask = np.zeros((timesteps, 1, 1), dtype=np.float32)
            elif image_blackout_ratio == 0.0:
                # Apply sequential rate
                image_mask = np.zeros((timesteps, 1, 1), dtype=np.float32)
                image_mask[::sequential_image_rate, 0, 0] = 1.0
            else:
                # Apply blackout rate
                image_mask = (
                    (np.random.uniform(size=(timesteps,)) > image_blackout_ratio)
                    .astype(np.float32)
                    .reshape((timesteps, 1, 1))
                )
            observations["image"] *= image_mask

            # Pull out controls
            ## This is currently consisted of:
            ## > previous end effector position
            ## > end effector position delta
            ## > binary contact reading
            if kloss_dataset:
                eef_positions = raw_trajectory["tip"]
            else:
                eef_positions = raw_trajectory["eef_pos"]
            eef_positions_shifted = np.roll(eef_positions, shift=1, axis=0)
            eef_positions_shifted[0] = eef_positions[0]

            # Force controls to be of type float32
            # NOTE: In numpy 1.20+, this can be done
            # through dtype kwarg to concatenate
            controls = np.empty((timesteps, 7), dtype=np.float32)
            np.concatenate(
                [
                    eef_positions_shifted,
                    eef_positions - eef_positions_shifted,
                    raw_trajectory["contact"][
                        :, np.newaxis
                    ],  # "contact" key same for both kloss and normal dataset
                ],
                axis=1,
                out=controls
            )
            if controls_noise != 0:
                controls = timeseries_robustness(
                    [controls], noise_level=controls_noise, struct_drop=False)[0]

            if multimodal_noise != 0:
                tmp = timeseries_robustness([observations["image"], observations["gripper_pos"],
                                            observations["gripper_sensors"], controls], noise_level=multimodal_noise, rand_drop=False)
                observations["image"] = tmp[0]
                observations["gripper_pos"] = tmp[1]
                observations["gripper_sensors"] = tmp[2]
                controls = tmp[3]

            # Normalize data
            if kloss_dataset:
                observations["gripper_pos"] -= np.array(
                    [[-0.00360131, 0.0, 0.00022349]],
                    dtype=np.float32,
                )
                observations["gripper_pos"] /= np.array(
                    [[0.07005621, 1.0, 0.06883541]],
                    dtype=np.float32,
                )
                observations["gripper_sensors"] -= np.array(
                    [
                        [
                            3.04424347e-02,
                            1.61328610e-02,
                            -2.47517393e-04,
                            0.00000000e00,
                            0.00000000e00,
                            0.00000000e00,
                            6.25842857e-01,
                        ]
                    ],
                    dtype=np.float32,
                )
                observations["gripper_sensors"] /= np.array(
                    [[2.09539968, 2.0681382, 0.00373115, 1.0, 1.0, 1.0, 0.48390451]],
                    dtype=np.float32,
                )
                states -= np.array(
                    [[-0.00279736, -0.00027878]],
                    dtype=np.float32,
                )
                states /= np.array(
                    [[0.06409658, 0.06649422]],
                    dtype=np.float32,
                )
                controls -= np.array(
                    [
                        [
                            -3.55868486e-03,
                            0.00000000e00,
                            2.34369027e-04,
                            -4.26185595e-05,
                            0.00000000e00,
                            -1.08724583e-05,
                            6.25842857e-01,
                        ]
                    ],
                    dtype=np.float32,
                )
                controls /= np.array(
                    [
                        [
                            0.0693582,
                            1.0,
                            0.06810329,
                            0.01176415,
                            1.0,
                            0.0115694,
                            0.48390451,
                        ]
                    ],
                    dtype=np.float32,
                )

            else:
                observations["gripper_pos"] -= np.array(
                    [[0.46806443, -0.0017836, 0.88028437]],
                    dtype=np.float32,
                )
                observations["gripper_pos"] /= np.array(
                    [[0.02410769, 0.02341035, 0.04018243]],
                    dtype=np.float32,
                )
                observations["gripper_sensors"] -= np.array(
                    [
                        [
                            4.9182904e-01,
                            4.5039989e-02,
                            -3.2791464e00,
                            -3.3874984e-03,
                            1.1552566e-02,
                            -8.4817986e-04,
                            2.1303751e-01,
                        ]
                    ],
                    dtype=np.float32,
                )
                observations["gripper_sensors"] /= np.array(
                    [
                        [
                            1.6152629,
                            1.666905,
                            1.9186896,
                            0.14219016,
                            0.14232528,
                            0.01675198,
                            0.40950698,
                        ]
                    ],
                    dtype=np.float32,
                )
                states -= np.array(
                    [[0.4970164, -0.00916641]],
                    dtype=np.float32,
                )
                states /= np.array(
                    [[0.0572766, 0.06118315]],
                    dtype=np.float32,
                )
                controls -= np.array(
                    [
                        [
                            4.6594709e-01,
                            -2.5247163e-03,
                            8.8094306e-01,
                            1.2939950e-04,
                            -5.4364675e-05,
                            -6.1112235e-04,
                            2.2041667e-01,
                        ]
                    ],
                    dtype=np.float32,
                )
                controls /= np.array(
                    [
                        [
                            0.02239027,
                            0.02356066,
                            0.0405312,
                            0.00054858,
                            0.0005754,
                            0.00046352,
                            0.41451886,
                        ]
                    ],
                    dtype=np.float32,
                )

            trajectories.append(
                TrajectoryNumpy(
                    states[start_timestep:],
                    fannypack.utils.SliceWrapper(observations)[
                        start_timestep:],
                    controls[start_timestep:],
                )
            )

            # Reduce memory usage
            raw_trajectories[raw_trajectory_index] = None
            del raw_trajectory

    ## Uncomment this line to generate the lines required to normalize data
    # _print_normalization(trajectories)

    return trajectories


# From https://github.com/stanford-iprl-lab/torchfilter:

def split_trajectories(
    trajectories: List[TrajectoryNumpy], subsequence_length: int, modalities=None
):
    """Helper for splitting a list of trajectories into a list of overlapping
    subsequences.
    For each trajectory, assuming a subsequence length of 10, this function
    includes in its output overlapping subsequences corresponding to
    timesteps...
    ```
        [0:10], [10:20], [20:30], ...
    ```
    as well as...
    ```
        [5:15], [15:25], [25:30], ...
    ```
    Args:
        trajectories (List[torchfilter.base.TrajectoryNumpy]): List of trajectories.
        subsequence_length (int): # of timesteps per subsequence.
    Returns:
        List[torchfilter.base.TrajectoryNumpy]: List of subsequences.
    """

    subsequences = []

    for traj in trajectories:
        # Chop up each trajectory into overlapping subsequences
        trajectory_length = len(traj.states)
        assert len(fp.utils.SliceWrapper(
            traj.observations)) == trajectory_length
        assert len(fp.utils.SliceWrapper(traj.controls)) == trajectory_length

        # We iterate over two offsets to generate overlapping subsequences
        for offset in (0, subsequence_length // 2):

            def split_fn(x: np.ndarray) -> np.ndarray:
                """Helper: splits arrays of shape `(T, ...)` into `(sections,
                subsequence_length, ...)`, where `sections = orig_length //
                subsequence_length`."""
                # Offset our starting point
                x = x[offset:]

                # Make sure our array is evenly divisible
                sections = len(x) // subsequence_length
                new_length = sections * subsequence_length
                x = x[:new_length]

                # Split & return
                return np.split(x, sections)

            for s, o, c in zip(
                # States are always raw arrays
                split_fn(traj.states),
                # Observations and controls can be dictionaries, so we have to jump
                # through some hoops
                fp.utils.SliceWrapper(
                    fp.utils.SliceWrapper(traj.observations).map(split_fn)
                ),
                fp.utils.SliceWrapper(
                    fp.utils.SliceWrapper(traj.controls).map(split_fn)
                ),
            ):
                # Add to subsequences
                if modalities is None:
                    subsequences.append(
                        (
                            o['gripper_pos'],
                            o['gripper_sensors'],
                            o['image'].reshape([*o['image'].shape, 1]),
                            c,
                            s,
                        )
                    )
                else:
                    mods = []
                    for m in modalities:
                        if m == 'gripper_pos':
                            mods.append(o['gripper_pos'])
                        elif m == 'gripper_sensors':
                            mods.append(o['gripper_sensors'])
                        elif m == 'image':
                            mods.append(o['image'].reshape([*o['image'].shape, 1]))
                        elif m == 'control':
                            mods.append(c)
                    mods.append(s)
                    subsequences.append(tuple(mods))
    return subsequences


class SubsequenceDataset(Dataset):
    """A data preprocessor for producing training subsequences from
    a list of trajectories.
    Thin wrapper around `torchfilter.data.split_trajectories()`.
    Args:
        trajectories (list): list of trajectories, where each is a tuple of
            `(states, observations, controls)`. Each tuple member should be
            either a numpy array or dict of numpy arrays with shape `(T, ...)`.
        subsequence_length (int): # of timesteps per subsequence.
    """

    def __init__(
        self, trajectories: List[TrajectoryNumpy], subsequence_length: int, modalities=None
    ):
        # Split trajectory into overlapping subsequences
        self.subsequences = split_trajectories(
            trajectories, subsequence_length, modalities
        )

    def __getitem__(self, index: int):
        """Get a subsequence from our dataset.
        Args:
            index (int): Subsequence number in our dataset.
        Returns:
            tuple: `(states, observations, controls)` tuple that contains
            data for a single subsequence. Each tuple member should be either a
            numpy array or dict of numpy arrays with shape
            `(subsequence_length, ...)`.
        """
        return self.subsequences[index]

    def __len__(self) -> int:
        """Total number of subsequences in the dataset.
        Returns:
            int: Length of dataset.
        """
        return len(self.subsequences)
