import numpy as np

import torch
from torch.utils.data import Dataset, Sampler
from torch_geometric.data import Data

from rorlkit.torch.data_management import rosbag_data


def normalize_image(image):
    assert image.dtype == np.uint8
    return np.float64(image) / 255.0


class ImageDataset(Dataset):

    def __init__(self, images, should_normalize=True):
        super().__init__()
        self.dataset = images
        self.dataset_len = len(self.dataset)
        assert should_normalize == (images.dtype == np.uint8)
        self.should_normalize = should_normalize

    def __len__(self):
        return self.dataset_len

    def __getitem__(self, idxs):
        samples = self.dataset[idxs, :]
        if self.should_normalize:
            samples = normalize_image(samples)
        return np.float32(samples)


class InfiniteRandomSampler(Sampler):

    def __init__(self, data_source):
        self.data_source = data_source
        self.iter = iter(torch.randperm(len(self.data_source)).tolist())

    def __iter__(self):
        return self

    def __next__(self):
        try:
            idx = next(self.iter)
        except StopIteration:
            self.iter = iter(torch.randperm(len(self.data_source)).tolist())
            idx = next(self.iter)
        return idx

    def __len__(self):
        return 2 ** 62


class InfiniteSequenceSampler(Sampler):
    def __init__(self, data_source, batch_size):
        super().__init__(data_source)
        self.data_source = data_source
        self.sequence_length = batch_size
        self.starts_iter = iter(torch.randperm(len(self.data_source) - self.sequence_length).tolist())
        start_idx = next(self.starts_iter)
        self.current_iter = iter(torch.arange(start_idx, start_idx + self.sequence_length).tolist())
        self.iter_idx = 0

    def __iter__(self):
        return self

    def __next__(self):
        try:
            idx = next(self.current_iter)
            self.iter_idx += 1
            if self.iter_idx >= self.sequence_length:
                start_idx = next(self.starts_iter)
                self.current_iter = iter(torch.arange(start_idx, start_idx + self.sequence_length).tolist())
                self.iter_idx = 0
        except StopIteration:
            self.starts_iter = iter(torch.randperm(len(self.data_source) - self.sequence_length).tolist())
            start_idx = next(self.starts_iter)
            self.current_iter = iter(torch.arange(start_idx, start_idx + self.sequence_length).tolist())
            idx = next(self.current_iter)
        return idx

    def __len__(self):
        return 2 ** 62


class InfiniteWeightedRandomSampler(Sampler):

    def __init__(self, data_source, weights):
        assert len(data_source) == len(weights)
        assert len(weights.shape) == 1
        self.data_source = data_source
        # Always use CPU
        self._weights = torch.from_numpy(weights)
        self.iter = self._create_iterator()

    def update_weights(self, weights):
        self._weights = weights
        self.iter = self._create_iterator()

    def _create_iterator(self):
        return iter(
            torch.multinomial(
                self._weights, len(self._weights), replacement=True
            ).tolist()
        )

    def __iter__(self):
        return self

    def __next__(self):
        try:
            idx = next(self.iter)
        except StopIteration:
            self.iter = self._create_iterator()
            idx = next(self.iter)
        return idx

    def __len__(self):
        return 2 ** 62


class MotionGraphDataset(Dataset):
    def __init__(self, bag_files, topics):
        super(MotionGraphDataset, self).__init__()
        self._n_bag = len(bag_files)
        self._db = []
        self._edge_index = torch.tensor([[0, 0], [1, 2]], dtype=torch.long)
        self.pg = rosbag_data.PosesGetter(bag_files, topics)
        self._build_db()

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

    def __getitem__(self, item):
        graph_curr, graph_next = self._db[item]
        return graph_curr, graph_next

    def _build_db(self):
        """Build the motion graph database as list[list[Data]].

        :return: None
        """
        for record_idx in range(self._n_bag):
            msg_dict = self.pg.msg_dict_lists[record_idx]
            n_unit = msg_dict[self.pg.unit_key]
            # Omit some samples if they were not belong to the demonstration
            idx_list = np.arange(0, n_unit-1)
            for unit_idx in idx_list:
                graphs = self.get_motion_graphs(record_idx, unit_idx)
                self._db.append(graphs)
        print('Created database of length {}'.format(self.__len__()))

    def get_motion_graphs(self, record_idx, unit_idx, shift=None):
        """Get two consecutive graphs from the database.

        :param record_idx: int The index of the bag file.
        :param unit_idx: int Starting index of the unit in the record. Each unit forms a graph.
        :param shift: array[9] Shift value for children features, only used for experiment
        :return: list[Data] A list for 2 (the current and the next) motion graphs.
        """
        graphs = []
        # 2 is for getting the current and the next frames of the motion
        for i in range(2):
            _idx = unit_idx + i
            # Get root node features
            # TODO In a upper body fashion, we do not consider the movement of the base
            base_pose_feature = rosbag_data.get_identity_feature()
            # Get children features
            left_hand_feature, right_hand_feature = self.pg.get_unit_features(record_idx, _idx)
            if shift is not None:
                left_hand_feature += shift
                right_hand_feature += shift
            # Create node features
            x = torch.tensor([base_pose_feature, left_hand_feature, right_hand_feature], dtype=torch.float)
            # Create graph
            data_i = Data(x=x, edge_index=self._edge_index)
            graphs.append(data_i)
        return graphs

    def get_motion_features(self, record_idx, unit_idx):
        left_hand_feature, right_hand_feature = self.pg.get_unit_features(record_idx, unit_idx)
        return left_hand_feature, right_hand_feature

    def feature_array_to_graph(self, feature_array):
        base_pose_feature = rosbag_data.get_identity_feature()
        # Get children features
        left_hand_feature = feature_array[1]
        right_hand_feature = feature_array[2]
        # Create node features
        x = torch.tensor([base_pose_feature, left_hand_feature, right_hand_feature], dtype=torch.float)
        # Create graph
        graph = Data(x=x, edge_index=self._edge_index)
        return graph
