from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from oneibl.one import ONE
import alf.io
import scipy.stats
plt.ion()

def scatter_raster(spikes):
 
    '''
    Create a scatter plot, time vs depth for each spike
    colored by cluster id; including vertical lines
    for stimulus type boundary times

    Note that interval should be at most 10**6 else
    the plot is too memory expensive

    :param spike: spike = alf.io.load_object(alf_path, 'spikes')
    :type spike: dict
    :type restrict: [int, int]
    :param restrict: array of clusters to be plotted
    :rtype: plot
    '''
    
    downsample_factor = 20

    uclusters = np.unique(spikes['clusters'])
    cols = ['c','b','g','y','k','r','m']
    cols_cat = (cols*int(len(uclusters)/len(cols)+10))[:len(uclusters)]
    col_dict = dict(zip(uclusters, cols_cat))

    # downsample 
    z = spikes['clusters'][::downsample_factor]
    x = spikes['times'][::downsample_factor]
    y = spikes['depths'][::downsample_factor]

    cols_int =[col_dict[x] for x in z]

    plt.scatter(x, y, marker='o', s=0.01, c = cols_int)

    plt.ylabel('depth [um]')
    plt.xlabel('time [sec]')
    plt.title('downsample factor: %s' %downsample_factor)  

if __name__ == '__main__':

    one = ONE()
    eid = one.search(subject='ZM_2407', date='2019-11-05', number=3)
    D = one.load(eid[0], clobber=False, download_only=True)
    alf_path = Path(D.local_path[0]).parent

    spikes = alf.io.load_object(alf_path, 'spikes')
    scatter_raster(spikes)
