import matplotlib.pyplot as plt
import numpy as np

import alf.io
from brainbox.singlecell import peths
from brainbox.examples.plot_all_peths import get_session_path
from v1_protocol_utilities.example_plotting import *

from oneibl.one import ONE # to download data

one = ONE()
#eid = one.search(subject='KS004', date=['2019-09-25'], task_protocol='ephysChoiceWorld')[0]
eid = one.search(lab='wittenlab', date='2019-08-04')
datasets = one.load(eid, download_only=True)
#ses_path = get_session_path(datasets)
ses_path = datasets[0].local_path.parent #local path where the data has been downloaded

spikes = alf.io.load_object(ses_path, 'spikes')
trials = alf.io.load_object(ses_path, 'trials')

# check which neurons are responsive
#are_neurons_responsive(spike_times,spike_clusters,stimulus_intervals=None,spontaneous_period=None,p_value_threshold=.05):
# spontaenous period is just 1 interval! on and off time
responsive = are_neurons_responsive(spikes.times, spikes.clusters, np.vstack((trials.stimOn_times, trials.stimOn_times + 0.5)).T, np.array([0, trials.stimOn_times[0]-1]).reshape(-1), 0.001)
#peths(spike_times, spike_clusters, cluster_ids, align_times, pre_time=0.2,
#          post_time=0.5, bin_size=0.025, smoothing=0.025, return_fr=True):

responsive_neuron = np.unique(spikes.clusters)[responsive] 

peth, bs = peths(spikes.times, spikes.clusters, responsive_neuron, trials.stimOn_times)

plt.plot(peth.tscale, peth.means.T)

#for m in np.arange(peth.means.shape[0]):
#    plt.fill_between(peth.tscale,
#                     peth.means[m, :].T - peth.stds[m, :].T / 20,
#                     peth.means[m, :].T + peth.stds[m, :].T / 20,
#                     alpha=0.2, edgecolor='#1B2ACC', facecolor='#089FFF',
#                     linewidth=4, linestyle='dashdot', antialiased=True)
