import gc
import json
import os.path as osp

import ipdb
import numpy as np
import torch
import torch.utils.data as data
from datasets.dataset_builder import DATA_LAYERS

FEATURE_SIZES = {
    "rgb_anet_resnet50": 2048,
    "flow_anet_resnet50": 2048,
    "rgb_kinetics_bninception": 1024,
    "flow_kinetics_bninception": 1024,
    "rgb_kinetics_resnet50": 2048,
    "flow_kinetics_resnet50": 2048,
    "flow_nv_kinetics_bninception": 1024,
    "rgb_kinetics_i3d": 2048,
    "flow_kinetics_i3d": 2048,
}


@DATA_LAYERS.register("THUMOS")
@DATA_LAYERS.register("TVSERIES")
@DATA_LAYERS.register("ASSEMBLY101-O")
@DATA_LAYERS.register("EPIC-TENT-O")
class THUMOSDataset(data.Dataset):

    def __init__(self, cfg, mode="train"):
        self.root_path = cfg["root_path"]
        self.mode = mode
        self.training = mode == "train"
        self.window_size = cfg["window_size"]
        self.stride = cfg["stride"]
        data_name = cfg["data_name"]
        self.vids = json.load(open(cfg["video_list_path"]))[data_name][
            mode + "_session_set"
        ]  # [:5]  # list of video names
        self.num_classes = cfg["num_classes"]
        self.inputs = []
        self._load_features(cfg)
        self._init_features()

    def _load_features(self, cfg):
        self.removed = 0
        self.annotation_type = cfg["annotation_type"]
        self.rgb_type = cfg["rgb_type"]
        self.flow_type = cfg["flow_type"]
        self.target_all = {}
        self.rgb_inputs = {}
        self.flow_inputs = {}
        dummy_target = np.zeros((self.window_size - 1, self.num_classes))
        dummy_rgb = np.zeros((self.window_size - 1, FEATURE_SIZES[cfg["rgb_type"]]))
        dummy_flow = np.zeros((self.window_size - 1, FEATURE_SIZES[cfg["flow_type"]]))
        for vid in self.vids:
            try:
                target = np.load(
                    osp.join(self.root_path, self.annotation_type, vid + ".npy")
                )
                rgb = np.load(osp.join(self.root_path, self.rgb_type, vid + ".npy"))
                # checkpoint con rgb al posto del flow
                if cfg["flow_type"] == "flow_anet_resnet50":
                    flow = np.load(
                        osp.join(
                            self.root_path + "/rgb_as_flow", self.rgb_type, vid + ".npy"
                        )
                    )
                    flow = np.zeros(flow.shape)
                else:  # checkpoint con flow
                    my_string = (
                        "assembly_optical_flow_BNInception/" + vid + "/assembling.npy"
                    )
                    flow = np.load(osp.join(self.root_path, self.flow_type, my_string))
                # # concatting dummy target at the front
                # ! LEO  to save train data leave these lines (within the if block) with the comment
                if self.training:
                    self.target_all[vid] = np.concatenate(
                        (dummy_target, target), axis=0
                    )
                    self.rgb_inputs[vid] = np.concatenate((dummy_rgb, rgb), axis=0)
                    self.flow_inputs[vid] = np.concatenate((dummy_flow, flow), axis=0)
                else:
                    self.target_all[vid] = target
                    self.rgb_inputs[vid] = rgb
                    self.flow_inputs[vid] = flow
            except Exception as e:
                # print excpetion
                print("---- Exception in loading video ", e)
                # remove the video from the list if it does not have the required features
                self.vids.remove(vid)
                self.removed += 1
                print("---- Removed video ", vid)
        print("---- Removed videos ", self.removed)

    def _init_features(self):
        # del self.inputs
        # gc.collect()
        self.inputs = []
        # remove 'nusar-2021_action_both_9056-b08a_9056_user_id_2021-02-22_141934'
        if (
            "nusar-2021_action_both_9056-b08a_9056_user_id_2021-02-22_141934"
            in self.vids
        ):
            self.vids.remove(
                "nusar-2021_action_both_9056-b08a_9056_user_id_2021-02-22_141934"
            )
        ipdb.set_trace()
        for vid in self.vids:
            target = self.target_all[vid]
            #!Leo to save train data leave these lines (within the if block) with the comment
            ipdb.set_trace()
            if self.training:
                seed = np.random.randint(self.stride)
                for start, end in zip(
                    range(seed, target.shape[0], self.stride),
                    range(seed + self.window_size, target.shape[0] + 1, self.stride),
                ):
                    self.inputs.append([vid, start, end, target[start:end]])
            else:
                start = 0
                end = target.shape[0]
                self.inputs.append([vid, start, end, target[start:end]])

    def __getitem__(self, index):
        vid, start, end, target = self.inputs[index]
        rgb_input = self.rgb_inputs[vid][start:end]
        flow_input = self.flow_inputs[vid][start:end]
        rgb_input = torch.tensor(rgb_input.astype(np.float32))
        flow_input = torch.tensor(flow_input.astype(np.float32))
        target = torch.tensor(target.astype(np.float32))
        return rgb_input, flow_input, target, vid, start, end  #!  vid added by Leo

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


@DATA_LAYERS.register("THUMOS_ANTICIPATION")
@DATA_LAYERS.register("TVSERIES_ANTICIPATION")
class THUMOSDataset(data.Dataset):

    def __init__(self, cfg, mode="train"):
        self.root_path = cfg["root_path"]
        self.mode = mode
        self.training = mode == "train"
        self.window_size = cfg["window_size"]
        self.stride = cfg["stride"]
        self.anticipation_length = cfg["anticipation_length"]
        data_name = cfg["data_name"].split("_")[0]
        self.vids = json.load(open(cfg["video_list_path"]))[data_name][
            mode + "_session_set"
        ]  # list of video names
        self.num_classes = cfg["num_classes"]
        self.inputs = []
        self._load_features(cfg)
        self._init_features()

    def _load_features(self, cfg):
        self.annotation_type = cfg["annotation_type"]
        self.rgb_type = cfg["rgb_type"]
        self.flow_type = cfg["flow_type"]
        self.target_all = {}
        self.rgb_inputs = {}
        self.flow_inputs = {}
        dummy_target = np.zeros((self.window_size - 1, self.num_classes))
        dummy_rgb = np.zeros((self.window_size - 1, FEATURE_SIZES[cfg["rgb_type"]]))
        dummy_flow = np.zeros((self.window_size - 1, FEATURE_SIZES[cfg["flow_type"]]))
        for vid in self.vids:
            target = np.load(
                osp.join(self.root_path, self.annotation_type, vid + ".npy")
            )
            rgb = np.load(osp.join(self.root_path, self.rgb_type, vid + ".npy"))
            flow = np.load(osp.join(self.root_path, self.flow_type, vid + ".npy"))
            if self.training:
                self.target_all[vid] = np.concatenate((dummy_target, target), axis=0)
                self.rgb_inputs[vid] = np.concatenate((dummy_rgb, rgb), axis=0)
                self.flow_inputs[vid] = np.concatenate((dummy_flow, flow), axis=0)
            else:
                self.target_all[vid] = target
                self.rgb_inputs[vid] = rgb
                self.flow_inputs[vid] = flow

    def _init_features(self):
        del self.inputs
        gc.collect()
        self.inputs = []

        for vid in self.vids:
            target = self.target_all[vid]
            if self.training:
                seed = np.random.randint(self.stride)
                for start, end in zip(
                    range(seed, target.shape[0], self.stride),
                    range(
                        seed + self.window_size,
                        target.shape[0] - self.anticipation_length,
                        self.stride,
                    ),
                ):
                    self.inputs.append(
                        [
                            vid,
                            start,
                            end,
                            target[start:end],
                            target[end : end + self.anticipation_length],
                        ]
                    )
            else:
                start = 0
                end = target.shape[0] - self.anticipation_length
                ant_target = []
                for s in range(0, target.shape[0] - self.anticipation_length):
                    ant_target.append(target[s : s + self.anticipation_length])

                self.inputs.append(
                    [vid, start, end, target[start:end], np.array(ant_target)]
                )

    def __getitem__(self, index):
        vid, start, end, target, ant_target = self.inputs[index]
        rgb_input = self.rgb_inputs[vid][start:end]
        flow_input = self.flow_inputs[vid][start:end]
        rgb_input = torch.tensor(rgb_input.astype(np.float32))
        flow_input = torch.tensor(flow_input.astype(np.float32))
        target = torch.tensor(target.astype(np.float32))
        ant_target = torch.tensor(ant_target.astype(np.float32))
        return rgb_input, flow_input, target, ant_target

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


@DATA_LAYERS.register("FINEACTION")
class FINEACTIONDataset(data.Dataset):

    def __init__(self, cfg, mode="train"):
        self.root_path = cfg["root_path"]
        self.mode = mode
        self.training = mode == "train"
        self.window_size = cfg["window_size"]
        self.stride = cfg["stride"]
        data_name = cfg["data_name"]
        self.vids = json.load(open(cfg["video_list_path"]))[data_name][
            mode + "_session_set"
        ]  # list of video names
        self.num_classes = cfg["num_classes"]
        self.inputs = []
        self._load_features(cfg)
        self._init_features()

    def _load_features(self, cfg):
        self.annotation_type = cfg["annotation_type"]
        self.rgb_type = cfg["rgb_type"]
        self.flow_type = cfg["flow_type"]

    def _init_features(self, seed=0):
        # self.inputs = []
        del self.inputs
        gc.collect()
        self.inputs = []
        for vid in self.vids:
            target = np.load(
                osp.join(self.root_path, self.annotation_type, vid + ".npy")
            )
            if self.training:
                seed = np.random.randint(self.stride)
                for start, end in zip(
                    range(seed, target.shape[0], self.stride),
                    range(seed + self.window_size, target.shape[0] + 1, self.stride),
                ):
                    self.inputs.append([vid, start, end])
            else:
                start = 0
                end = target.shape[0]
                self.inputs.append([vid, start, end])

    def __getitem__(self, index):
        vid, start, end = self.inputs[index]
        rgb_input = np.load(
            osp.join(self.root_path, self.rgb_type, vid + ".npy"), mmap_mode="r"
        )[start:end]
        flow_input = np.load(
            osp.join(self.root_path, self.flow_type, vid + ".npy"), mmap_mode="r"
        )[start:end]
        target = np.load(
            osp.join(self.root_path, self.annotation_type, vid + ".npy"), mmap_mode="r"
        )[start:end]
        rgb_input = torch.tensor(rgb_input.astype(np.float32))
        flow_input = torch.tensor(flow_input.astype(np.float32))
        target = torch.tensor(target.astype(np.float32))
        return rgb_input, flow_input, target

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