## Code modified based on https://github.com/MoonBlvd/bidireaction-trajectory-prediction/blob/main/datasets/ETH_UCY.py

import os
import sys
sys.path.append('./Trajectron-plus-plus')
sys.path.append('./Trajectron-plus-plus/trajectron')
from .trajectron import NodeTypeDataset
import numpy as np
import torch
from torch.utils import data
import dill
import json
import random

def chunks(lst, n):
    for i in range(0, len(lst), n):
        yield lst[i:i + n]


class ETHUCYDataLayer(data.Dataset):

    def __init__(self, args, split):
        self.args = args
        self.split = split
        self.batch_size = args.batch_size

        conf_json = open(args.ETH_CONFIG, 'r')
        hyperparams = json.load(conf_json)

        hyperparams['minimum_history_length'] = self.args.enc_steps-1 if self.split == 'test' else 1
        hyperparams['maximum_history_length'] = self.args.enc_steps-1

        hyperparams['state'] = {'PEDESTRIAN':{'position':['x','y'], 'velocity':['x','y'], 'acceleration':['x','y']}}
        hyperparams['pred_state'] = {'PEDESTRIAN':{'position':['x','y']}}

        args.data_root = args.dataset.lower()

        # File can be generated by using srcipts from Trajectron++ (https://github.com/StanfordASL/Trajectron-plus-plus)
        if split == 'train':
            f = open(os.path.join(args.eth_root, args.data_root, 'train', args.data_root+'_train.pkl'), 'rb')
        elif split == 'val':
            f = open(os.path.join(args.eth_root, args.data_root, 'val', args.data_root+'_val.pkl'), 'rb')
        elif split == 'test':
            f = open(os.path.join(args.eth_root, args.data_root, 'test', args.data_root+'_test.pkl'), 'rb')
        else:
            raise ValueError()

        train_env = dill.load(f, encoding='latin1')

        node_type=train_env.NodeType[0]
        train_env.attention_radius[(node_type, node_type)] = 3.0 #10.0
        augment = False
        if split=='train':
            min_history_timesteps = 1
            augment = True if self.args.augment else False
        else:
            min_history_timesteps = 7
        self.dataset = NodeTypeDataset(train_env, 
                                        node_type, 
                                        hyperparams['state'],
                                        hyperparams['pred_state'],
                                        scene_freq_mult=hyperparams['scene_freq_mult_train'],
                                        node_freq_mult=hyperparams['node_freq_mult_train'],
                                        hyperparams=hyperparams, 
                                        augment=augment, 
                                        min_history_timesteps=min_history_timesteps,
                                        min_future_timesteps=hyperparams['prediction_horizon'],
                                        return_robot=False)

        self.len_dict = {}
        for index in range(len(self.dataset)):
            first_history_index, x_t, y_t, x_st_t, y_st_t,scene_name,timestep = self.dataset.__getitem__(index)
            if first_history_index not in self.len_dict:
                self.len_dict[first_history_index] = []
            self.len_dict[first_history_index].append(index)
        self.shuffle_dataset()

    def shuffle_dataset(self):
        self._init_inputs()

    def _init_inputs(self):
        '''
        shuffle the data based on its length
        '''
        self.inputs = []
        for length in self.len_dict:
            indices = self.len_dict[length]
            random.shuffle(indices)
            self.inputs.extend(list(chunks(self.len_dict[length], self.batch_size)))

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

    def __getitem__(self, index):
        indices = self.inputs[index]

        ret = {
            'input_x': [],
            'input_x_st': [],
            'target_y': [],
            'target_y_st': [],
            'first_history_index':[],
            'scene_name': [],
            'timestep': [],
        }
        
        for idx in indices:
            this_ret = self.getitem_one(idx)
            ret['input_x'].append(this_ret['input_x'])
            ret['input_x_st'].append(this_ret['input_x_st'])
            ret['target_y'].append(torch.as_tensor(this_ret['target_y']).type(torch.FloatTensor))
            ret['first_history_index'].append(torch.as_tensor(this_ret['first_history_index']).type(torch.LongTensor))
            ret['scene_name'].append(this_ret['scene_name'])
            ret['timestep'].append(this_ret['timestep'])
            

        ret['input_x'] = torch.stack(ret['input_x'])
        ret['input_x_st'] = torch.stack(ret['input_x_st'])
        ret['target_y'] = torch.stack(ret['target_y'])
        
        ret['first_history_index'] = torch.stack(ret['first_history_index'])
        # to locate image
        ret['scene_name'] = ret['scene_name']
        ret['timestep'] = ret['timestep']

        return ret

    def getitem_one(self, index):
        first_history_index, x_t, y_t, x_st_t, y_st_t, scene_name, timestep = self.dataset.__getitem__(index)
        ret = {}
        all_t = torch.cat((x_t[:,:2], y_t),dim=0)
        y_t = self.get_target(all_t, 0, self.args.enc_steps, self.args.enc_steps, self.args.dec_steps)
        ret['first_history_index'] = first_history_index
        ret['input_x'] = x_t
        ret['input_x_st'] = x_st_t
        ret['target_y'] = y_t
        ret['target_y_st'] = y_st_t
        ret['scene_name'] = scene_name
        ret['timestep'] = timestep
        return ret

    def get_target(self, session, start, end, observe_length, predict_length):
        '''
        Prepare the target for loss

        ''' 
        target = np.zeros((observe_length, predict_length, session.shape[-1]))
        for i, target_start in enumerate(range(start, end)):
            '''the target of time t is the change of bbox/ego motion at times [t+1,...,t+5}'''
            target_start = target_start + 1
            try:
                target[i,:,:] = np.asarray(session[target_start:target_start+predict_length,:] - 
                                           session[target_start-1:target_start,:])
            except:
                print("segment start: ", start)
                print("sample start: ", target_start)
                print("segment end: ", end)
                print(session.shape)
                raise ValueError()
        return target


