# The script plots the neural and behavioural data contained in a .h5 file of the Diomedi et al., 2023 dataset
# in the form of a raster plot
# authors: Stefano Diomedi
# date: 10/2022

import numpy as np
import matplotlib.pyplot as plt
from supportFunctions import get_all_data_from_level_h5_rec
import tkinter as tk
from tkinter import filedialog
import os
import h5py

# Defining the custom colormap
clmap = np.array([
    [0, 0, 0.6667],
    [0, 0, 1.0000],
    [0, 0.3333, 1.0000],
    [0, 0.6667, 1.0000],
    [0, 1.0000, 1.0000],
    [0.3333, 1.0000, 0.6667],
    [0.6667, 1.0000, 0.3333],
    [1.0000, 1.0000, 0],
    [1.0000, 0.6667, 0],
    [1.0000, 0.3333, 0]
])

# Get the current directory
current_dir = os.getcwd()

# Path to open the file dialog window in the data branch
path_folder_data = os.path.abspath(os.path.join(current_dir, "..", ".."))
path_folder_data = os.path.join(path_folder_data, "data")

# Select H5 dataset file
root = tk.Tk()
root.withdraw()
filename = filedialog.askopenfilename(initialdir=path_folder_data)


group_name = '/DATA'  # it can be 'unit_XXX', 'unit_XXX/condition_YY' or 'unit_XXX/condition_YY/trial_ZZ'

markers = []
all_strings_mk = []
spikes = []
all_strings_sp = []
markers, all_strings_mk, spikes, all_strings_sp = get_all_data_from_level_h5_rec.get_all_data_from_level_h5_rec(filename, group_name, markers, all_strings_mk, spikes, all_strings_sp)

str2readmk = '/DATA/unit_001/condition_01/trial_01/event_markers'

with h5py.File(filename, 'r') as f:
    marker_labels = f[str2readmk].attrs['Marker labels']

count = 1
fig, ax = plt.subplots()
first_iteration = True   
for spike_data, marker_data in zip(spikes, markers):

    ax.scatter(spike_data, np.ones_like(spike_data) * count, marker='|', color='k')
    sz = 30

    for mk in range(marker_data.size):
        cc = clmap[mk]

        if first_iteration:
            ax.scatter(marker_data[mk], np.ones_like(marker_data[mk]) * count,
                       s=sz, color=cc, marker='o', edgecolors=cc, label=marker_labels[mk])
        else:
            ax.scatter(marker_data[mk], np.ones_like(marker_data[mk]) * count,
                       s=sz, color=cc, marker='o', edgecolors=cc)


    count = count + 1
    first_iteration= False

ax.set_ylim([0.5, count ])
ax.set_xlabel('Time (ms)')
ax.set_ylabel('Trial #')
ax.legend()
ax.set_title('Raster Plot')
plt.show()