import sys
from torch.utils.data import Dataset, DataLoader
import os
import os.path as osp
import glob
import numpy as np
import random
import cv2
import pickle as pkl
import json
import h5py
import torch
import matplotlib.pyplot as plt
from lib.utils.misc import process_dataset_for_video

class MPIINFDataset(Dataset):
    def __init__(self, config, is_train=True):
        self.frame_interval = config.DATA.FRAME_INTERVAL
        # for mpi dataset, we convert its order to match with h36m
        self.mpi2h36m = [10, 9, 8, 11, 12, 13, 4, 3, 2, 5, 6, 7, 1, 14, 15, 16, 0]
        # randomization will lead to inferior performance
        self.scale_path = "../data/mpi_train_scales.pkl" if config.USE_GT else "../data/mpi_train_scales_pre.pkl"
        self.use_same_norm_2d = config.DATA.USE_SAME_NORM_2D
        self.use_same_norm_3d = config.DATA.USE_SAME_NORM_3D
        self.is_train = is_train
        self.data_path = config.DATA.TRAIN_PATH if self.is_train else config.DATA.VALID_PATH
        self.head_root_distance = 1 / config.TRAIN.CAMERA_SKELETON_DISTANCE
        # whether to use dataset adapted from k[MaÌ]inetics
        self.use_gt = config.USE_GT
        self.use_ideal_scale = config.DATA.USE_IDEAL_SCALE
        self.min_diff_dist = config.DATA.MIN_DIFF_DIST
        self.use_scaler = config.TRAIN.USE_SCALER
        self.bound_azim = float(config.TRAIN.BOUND_AZIM) # y axis rotation  
        self.bound_elev = float(config.TRAIN.BOUND_ELEV)
        self._load_data_set()

    def _load_data_set(self):
        if self.is_train:
            print('start loading mpiinf {} data.'.format("train" if self.is_train else "test"))
        key = "joint_2d_gt" if self.use_gt else "joint_2d_pre"
        fp = h5py.File(self.data_path, "r")
        self.kp2ds = np.array(fp[key])[:, self.mpi2h36m, :2]
        self.kp2ds[:, :, 0] = (self.kp2ds[..., 0] - 1024.0) / 1024.0
        self.kp2ds[:, :, 1] = (self.kp2ds[..., 1] - 1024.0) / 1024.0
        # self.kp2ds = np.maximum(np.minimum(self.kp2ds, 1.0), -1.0)
        # locate root at the origin 
        self.kp2ds = self.kp2ds - self.kp2ds[:, 13:14]
        self.kp2ds[:, 13] = 1e-5
        # imagenames will be used to sample frames 
        self.imagenames = [name.decode() for name in fp['imagename'][:]]
        if 'seqname' not in fp.keys():
            # first we close the already opened (read-only) h5 
            fp.close()
            print("Process corresponding dataset...")
            process_dataset_for_video(self.data_path, is_mpi=True)
            fp = h5py.File(self.data_path, "r")
        self.sequence_lens = np.array(fp['seqlen'])
        self.sequence_names = [name.decode() for name in fp['seqname'][:]]
        self.indices_in_seq = np.array(fp['index_in_seq'])

        # normlize again so that the mean distance of head and root is 1/c
        if not self.use_same_norm_2d:
            factor_gt = self.head_root_distance / (np.tile(np.linalg.norm(self.kp2ds[:, -1] - self.kp2ds[:, 13], axis=1).reshape(-1, 1, 1), (1, 17, 2)) + 1e-8)
        else:
            factor_gt = self.head_root_distance / np.linalg.norm(self.kp2ds[:, -1] - self.kp2ds[:, 13], axis=1).mean()
        self.kp2ds = self.kp2ds * factor_gt 

        self.kp3ds = np.array(fp['joint_3d_gt'])[:, self.mpi2h36m, :3] / 1000.0
        # factor_3d = np.linalg.norm(self.kp3ds[:, -1] - self.kp3ds[:, 13], axis=1).mean())
        factor_filename = "../data/mpi_{}_factor_3d.pkl".format("train" if self.is_train else "test")
        if not self.use_same_norm_3d:
            factor_3d = (np.tile(np.linalg.norm(self.kp3ds[:, -1] - self.kp3ds[:, 13], axis=1).reshape(-1, 1, 1), (1, 17, 3)) + 1e-8)
            print(factor_3d.shape)
            with open(factor_filename, "wb") as f:
                pkl.dump(factor_3d, f)
        
        if osp.exists(self.scale_path):
            with open(self.scale_path, "rb") as f:
                self.scales = pkl.load(f)['scale']
        else:
            if self.use_scaler:
                pass
                # raise Warning("You haven't generated the computed scales, if you don't need to observe the scale error during training, \njust ignore this warning because it won't affect training.")
            self.scales = None

        if self.use_ideal_scale: 
            # scales computed from projection of 3d 
            with open("../data/mpi_{}_scales.pkl".format("train" if self.is_train else "valid"), "rb") as f:
                scales = pkl.load(f)
            self.kp2ds = self.kp2ds * scales

        fp.close()
        print('finished load mpiinf {} data, total {} samples'.format("train" if self.is_train else "test", \
            self.kp2ds.shape[0]))

        # generate the rotation factors 
        num_examples = self.kp2ds.shape[0]
        rotation_y = (2 * np.random.random_sample((num_examples, 1)) - 1) * self.bound_azim 
        rotation_x = (2 * np.random.random_sample((num_examples, 1)) - 1) * self.bound_elev 
        rotation_z = (2 * np.random.random_sample((num_examples, 1)) - 1) * self.bound_elev / 2
        rotation_1 = np.concatenate((rotation_y, rotation_x, rotation_z), axis=1)
        rotation_2 = rotation_1.copy()
        rotation_2[:, 0] = rotation_2[:, 0] + np.pi
        self.rotation = np.concatenate((rotation_1, rotation_2), axis=0)
        np.random.shuffle(self.rotation)
        self.rotation = torch.from_numpy(self.rotation).float()

        self.kp2ds = torch.from_numpy(self.kp2ds).float()
        self.kp3ds = torch.from_numpy(self.kp3ds).float()
        if self.scales is not None:
            self.scales = torch.from_numpy(self.scales).float()

    def get_seqnames(self):
        return self.sequence_names

    def __len__(self):
        return self.kp2ds.shape[0]

    def __getitem__(self, index):
        seq_len = self.sequence_lens[index]
        index_in_seq = self.indices_in_seq[index]
        kps_3d = self.kp3ds[index]
        rot = self.rotation[index]
        if not self.is_train: 
            kps_2d = self.kp2ds[index]
            # don't use
            diff1 = diff2 = self.kp2ds[index]
        else:
            kps_2d = self.kp2ds[index]
            if self.frame_interval + index < seq_len:
                diff1_index = index + self.frame_interval
            else: 
                diff1_index = index - self.frame_interval
            diff1 = self.kp2ds[diff1_index]

            diff_dist = np.random.randint(-index_in_seq, seq_len - index_in_seq)
            while abs(diff_dist) < self.min_diff_dist: 
                diff_dist = np.random.randint(-index_in_seq, seq_len - index_in_seq)
            diff2_index = index + diff_dist 
            diff2 = self.kp2ds[diff2_index]

        rot = self.rotation[index]
        # for valdiation, simply ignore scale
        if self.scales is None or not self.is_train: 
            scale = 0
        else:
            scale = self.scales[index]
        return kps_2d, kps_3d, rot, diff1, diff2, scale

