from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np

from oneibl.one import ONE
import alf.io as ioalf
import ibllib.plots as iblplt

from brainbox.processing import bincount2D

T_BIN = 0.01

# get the data from flatiron and the current folder
one = ONE()
eid = one.search(subject='ZM_1150', date='2019-05-07', number=1)
D = one.load(eid[0], clobber=False, download_only=True)
session_path = Path(D.local_path[0]).parent

# load objects
spikes = ioalf.load_object(session_path, 'spikes')
clusters = ioalf.load_object(session_path, 'clusters')
channels = ioalf.load_object(session_path, 'channels')
trials = ioalf.load_object(session_path, '_ibl_trials')

# compute raster map as a function of cluster number
R, times, clusters = bincount2D(spikes['times'], spikes['clusters'], T_BIN)

# plot raster map
plt.imshow(R, aspect='auto', cmap='binary', vmax=T_BIN / 0.001 / 4,
           extent=np.r_[times[[0, -1]], clusters[[0, -1]]], origin='lower')
# plot trial start and reward time
reward = trials['feedback_times'][trials['feedbackType'] == 1]
iblplt.vertical_lines(trials['intervals'][:, 0], ymin=0, ymax=clusters[-1],
                      color='k', linewidth=0.5, label='trial starts')
iblplt.vertical_lines(reward, ymin=0, ymax=clusters[-1], color='m', linewidth=0.5,
                      label='valve openings')
plt.xlabel('Time (s)')
plt.ylabel('Cluster #')
plt.legend()
