'''
    Author: Yeshwanth Bethi
    Email: Y.Bethi@westernsydney.edu.au
    Date: 29/01/2021
    Description: 
'''

import h5py
import numpy as np
import matplotlib.pyplot as plt
from math import floor as floor
from mpl_toolkits.mplot3d import Axes3D
import random


#Load the data
dataset = h5py.File('labelled_ebssa.h5', 'r')

#Select the recording file. To check the list of files print(dataset.keys())
recording = 'ISS_2_td'
recording_data = dataset[recording]

# Check the number of detected objects in the file
nObj = np.asarray(recording_data['nObj'],dtype=np.int)[0,0]

# If there are any objects plot the trajectories.
if nObj > 0:
    # 3D plotting using Matplotlib
    fig = plt.figure()
    ax = plt.axes(projection='3d')

    # Find the first and last time stamp of the events recording
    initial_time = np.asarray(recording_data['TD']['ts'],dtype=np.int64)[0,0]
    final_time = np.asarray(recording_data['TD']['ts'],dtype=np.int64)[0,-1]

    # Find the number of frames per recording. Each frame = 1e4 uSecs
    nFrames = floor((final_time - initial_time)/10000)

    # Load the X,Y,Ts,Ids of the objects. Some of the arrays have to be converted into Integer arrays
    xCoords = np.asarray(recording_data['Obj']['x'])[0,:]
    yCoords = np.asarray(recording_data['Obj']['y'])[0,:]
    ids = np.asarray(recording_data['Obj']['id'],dtype=np.int32)[0,:]
    ts = np.asarray(recording_data['Obj']['ts'],dtype=np.int64)[0,:]

    # Make arrays to use for plotting.
    xCoords_plot =  np.zeros((nFrames,nObj))
    yCoords_plot = np.zeros((nFrames,nObj))

    # Fill the missing frames with NaN values so that they dont show up in the plot
    xCoords_plot.fill(np.nan)
    yCoords_plot.fill(np.nan)


    # For all objects do:
    for i in range(nObj):

        # Find the frame number corresponding to each timestamp of the object detection.
        xCoords_plot[np.floor((ts[ids==i]-initial_time)/1e4).astype(np.int),i] = xCoords[ids==i]
        yCoords_plot[np.floor((ts[ids==i]-initial_time)/1e4).astype(np.int), i] = yCoords[ids == i]

        # Plot the specific object in 3D
        ax.plot3D(xCoords_plot[:,i], yCoords_plot[:,i], np.arange(nFrames), c=(random.random(),random.random(),random.random()))



ax.set_xlabel('$X$', fontsize=10)
ax.set_ylabel('$Y$',fontsize=10)
ax.set_zlabel('$T$', fontsize=10, rotation = 0)

plt.show()

# To see all the recording names.
print(dataset.keys())