import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
import os
import sys
import ephys.rasters
from ephys import core, events, rasters
import pandas as pd

"""
Explanation: Initial visualizations of your data
This notebook is meant to let you quickly look at your post-manually sorted data
End of explanation
"""


block_path = '/home/brad/cube/Ice/kperks/B1112/klusta/phy041316/Pen01_Lft_AP300_ML700__Site04_Z3300__B1112_cat_P01_S04_1/'
block_path = '/mnt/cube/btheilma/emily/B1171/klusta/phy103116/Pen02_Lft_AP2500_ML1200__Site14_Z2570__B1171_cat_P02_S14_1/'
block_path = '/mnt/cube/btheilma/emily/B1171/klusta/phy103116/Pen02_Lft_AP2500_ML1200__Site15_Z2780__B1171_cat_P02_S15_1/'
block_path = '/mnt/cube/btheilma/emily/B604/klusta/phy103116/Pen01_Lft_AP2300_ML1250__Site13_Z2641__B604_cat_P01_S13_1/'
block_path = '/mnt/cube/btheilma/emily/B604/klusta/phy103116/Pen01_Lft_AP2300_ML1250__Site15_Z2780__B604_cat_P01_S15_1/'
block_path = '/home/brad/emily/P01S18/'
#block_path = '/mnt/cube/btheilma/emily/B604/klusta/phy103116/Pen01_Lft_AP2300_ML1250__Site16_Z2865__B604_cat_P01_S16_1/'
spikes     = core.load_spikes(block_path)
clusters   = core.load_clusters(block_path)
trials     = events.load_trials(block_path)
fs         = core.load_fs(block_path)
rec        = 0 

st = trials[trials['stimulus']=='T3S3D40']
print(st)
Activetrials = trials[pd.notnull(trials['response'])]


"""
Explanation: Load the data
Set 'block_path' to the path of the folder containing your post-manually sorted data
End of explanation
"""


from importlib import reload
reload(rasters)
good_clusters = clusters[clusters['quality']=='Good']
good_cluster_ids = good_clusters['cluster'].unique()
print(len(good_cluster_ids))

MUA_clusters = clusters[clusters['quality']=='MUA']
good_cluster_ids=MUA_clusters['cluster'].unique()

t=np.unique(trials['stimulus'].values)
print(t)
len(t)
stims = trials['stimulus'].unique()
print(stims[4])

# We have 16 stims, so let's go with a 4x4 subplot layout
subplot_xy = [6, 4]
figsize = (28,20)
period = [-2, 2]

# Make a figure directory
raster_dir = os.path.join(block_path, 'rasters')
if not os.path.exists(raster_dir):
    os.mkdir(raster_dir)

for clusterID in good_cluster_ids:
    print(clusterID)
    fig = rasters.plot_unit_raster_emily(spikes, trials, clusterID, period, rec, fs, subplot_xy, figsize)
    raster_fname = os.path.join(raster_dir, 'unit{}_raster.png'.format(str(clusterID)))
    fig.savefig(raster_fname)
    plt.close(fig)

"""
Explanation: Plot rasters of each trial from each Good unit for all stims
End of explanation
"""


reload(rasters)
clusterID = 9
stims = ['E_scaled_burung', 'C_scaled_burung', 'B_scaled_burung', 'H_scaled_burung']
stims = ['T3S3D-1', 'T3S3D40', 'T3S8D-1', 'T3S8D40', 'T40S40D-1', 'T40S40D3', 'T40S70D-1', 'T40S70D3']
trials_to_plot = trials[trials['stimulus'].isin(stims)]

subplot_xy=[4,2]
figsize=(28, 20)
period=[-2, 2]
# Make a figure directory
raster_dir = os.path.join(block_path, 'rasters')
if not os.path.exists(raster_dir):
    os.mkdir(raster_dir)

fig = rasters.plot_unit_raster_emily(spikes, trials_to_plot, clusterID, period, rec, fs, subplot_xy, figsize)
raster_fname = os.path.join(raster_dir, 'unit{}_raster_specific_stims.png'.format(str(clusterID)))
fig.savefig(raster_fname)
plt.close(fig)

"""
Explanation: See a cell / stim combo that looked interesting?  Plot it here!
End of explanation
"""


mua_clusters = clusters[clusters['quality']=='MUA']
mua_cluster_ids = good_clusters['cluster'].unique()

reload(rasters)
mua_clusters = clusters[clusters['quality']=='MUA']
mua_cluster_ids = mua_clusters['cluster'].unique()
print(len(good_cluster_ids))
# We have 16 stims, so let's go with a 4x4 subplot layout
subplot_xy = [6,4]
figsize = (22,20)
period = [-2, 2]

# Make a figure directory
raster_dir = os.path.join(block_path, 'rasters')
if not os.path.exists(raster_dir):
    os.mkdir(raster_dir)

for clusterID in mua_cluster_ids:
    print(clusterID)
    fig = rasters.plot_unit_raster_emily(spikes, trials, clusterID, period, rec, fs, subplot_xy, figsize)
    raster_fname = os.path.join(raster_dir, 'unit{}_MUAraster.png'.format(str(clusterID)))
    fig.savefig(raster_fname)
    plt.close(fig)

from ephys.spiketrains import get_spiketrain

def plot_trial_raster_emily(spikes, trials, clusters, trialID, 
                            stim, period, fs, ax=None, **kwargs):

    stim_trials = trials[trials['stimulus']==stim]
    stim_recs = stim_trials['recording'].values
    ntrials = len(stim_trials)

    stim_starts = stim_trials['time_samples'].values
    stim_ends = stim_trials['stimulus_end'].values

    stim_start = stim_starts[trialID]
    stim_end = stim_ends[trialID]
    stim_end_seconds = (stim_end - stim_start)/fs
    srec = stim_recs[trialID]

    clusterIDs = clusters['cluster'].values
    window = [period[0], stim_end_seconds+period[1]]
    raster_data = []
    for clu in clusterIDs:
        sptrain = get_spiketrain(srec, stim_start, clu, spikes, window, fs)
        raster_data.append(sptrain)
    rasters.do_raster(raster_data, window, [0, stim_end_seconds], ax, **kwargs) 


%matplotlib inline
cluster_group = ['Good', 'MUA']
if cluster_group != None:
    mask = np.ones(len(clusters.index)) < 0
    for grp in cluster_group:
        mask = np.logical_or(mask, clusters['quality'] == grp)
goodMUAclusters = clusters[mask]
fig = plt.figure()
ax = plt.gca()
plot_trial_raster_emily(spikes, Activetrials, goodMUAclusters, 6, 'T40S40D3', [-2, 2], 0, fs, ax=ax, spike_color='b')
plot_trial_raster_emily(spikes, Activetrials, goodMUAclusters, 6, 'T40S70D3', [-2, 2], 0, fs, ax=ax, spike_color='y')

plt.show()

fig = plt.figure()
ax = plt.gca()
plot_trial_raster_emily(spikes, Activetrials, goodMUAclusters, 6, 'T40S70D3', [-2, 2], 0, fs, ax=ax)
plt.show()

"""
Explanation: Plot Average, Gaussian Smoothed PSTHs for each Unit
Plot rasters of each MUA unit
End of explanation
"""


from ephys.spiketrains import get_spiketrain
from ephys import clust

def plot_trial_raster_emily(spikes, trials, clusters, trialID, 
                            stim, period, fs, ax=None, **kwargs):

    stim_trials = trials[trials['stimulus']==stim]
    stim_recs = stim_trials['recording'].values
    ntrials = len(stim_trials)
    print('Ntrials: {}'.format(ntrials))

    stim_starts = stim_trials['time_samples'].values
    stim_ends = stim_trials['stimulus_end'].values

    stim_start = stim_starts[trialID]
    stim_end = stim_ends[trialID]
    stim_end_seconds = (stim_end - stim_start)/fs
    srec = stim_recs[trialID]

    clusterIDs = clusters['cluster'].values
    window = [period[0], stim_end_seconds+period[1]]
    raster_data = []
    for clu in clusterIDs:
        sptrain = get_spiketrain(srec, stim_start, clu, spikes, window, fs)
        raster_data.append(sptrain)
    rasters.do_raster(raster_data, window, [0, 2.0, 2.5, stim_end_seconds], ax, **kwargs) 


correctTrials = trials[trials['correct']==True]
incorrectTrials = trials[trials['correct']==False]

widenarrow_threshold = 0.000230
cluster_group = ['Good', 'MUA']
clusters_list = clusters[clusters.quality.isin(cluster_group)]['cluster'].unique()
(wide, narrow) = clust.get_wide_narrow(block_path, clusters_list, widenarrow_threshold)

wideClus = clusters[clusters.cluster.isin(wide)]
narrowClus = clusters[clusters.cluster.isin(narrow)]
stims = correctTrials['stimulus'].unique()

%matplotlib inline
import seaborn as sns
sns.set_context('poster')
sns.set_style('white')
my_trial_nums = [0, 2]

for stim in stims:
    print(stim)
   
    # plot Wide Correct / Wide Incorrect
    f, axarr = plt.subplots(2,2, figsize=(28, 12))
    #Target
    plot_trial_raster_emily(spikes, correctTrials, wideClus, my_trial_nums[0], stim, [-1, 1], fs, ax= axarr[0,0])
    axarr[0,0].set_ylabel('Wide Cluster')
    axarr[0,0].set_title(stim + ' correct')
    #SD
    plot_trial_raster_emily(spikes, incorrectTrials, wideClus, my_trial_nums[0], stim, [-1, 1], fs, ax= axarr[0,1])
    axarr[0,1].set_title(stim + ' incorrect')
    
    plot_trial_raster_emily(spikes, correctTrials, wideClus, my_trial_nums[1], stim, [-1, 1], fs, ax= axarr[1,0])
    axarr[1,0].set_ylabel('Wide Cluster')
    axarr[1,0].set_title(stim + ' correct')
    #SD
    plot_trial_raster_emily(spikes, incorrectTrials, wideClus, my_trial_nums[1], stim, [-1, 1], fs, ax= axarr[1,1])
    axarr[1,1].set_title(stim + ' incorrect')
    
    plt.savefig('/home/brad/figures/B604_P01S15_raster_wide_CvsX_2trials_{}.png'.format(stim))
    
    f, axarr = plt.subplots(2,2, figsize=(28, 12))
    #Target
    plot_trial_raster_emily(spikes, correctTrials, narrowClus, my_trial_nums[0], stim, [-1, 1], fs, ax= axarr[0,0])
    axarr[0,0].set_ylabel('Narrow Cluster')
    axarr[0,0].set_title(stim + ' correct')
    #SD
    plot_trial_raster_emily(spikes, incorrectTrials, narrowClus, my_trial_nums[0], stim, [-1, 1], fs, ax= axarr[0,1])
    axarr[0,1].set_title(stim + ' incorrect')
    
    plot_trial_raster_emily(spikes, correctTrials, narrowClus, my_trial_nums[1], stim, [-1, 1], fs, ax= axarr[1,0])
    axarr[1,0].set_ylabel('Narrow Cluster')
    axarr[1,0].set_title(stim + ' correct')
    #SD
    plot_trial_raster_emily(spikes, incorrectTrials, narrowClus, my_trial_nums[1], stim, [-1, 1], fs, ax= axarr[1,1])
    axarr[1,1].set_title(stim + ' incorrect')
    
    plt.savefig('/home/brad/figures/B604_P01S15_raster_narrow_CvsX_2trials_{}.png'.format(stim))
    # plot Narrow Correct / Narrow Incorrect

for stim in ['T40S40D3']:
    f, axarr = plt.subplots(2,2, figsize=(28, 12))
    #Target
    plot_trial_raster_emily(spikes, correctTrials, narrowClus, my_trial_nums[0], stim, [-1, 1], fs, ax= axarr[0,0])
    axarr[0,0].set_ylabel('Narrow Cluster')
    axarr[0,0].set_title(stim + ' correct')
    #SD
    plot_trial_raster_emily(spikes, incorrectTrials, narrowClus, my_trial_nums[0], stim, [-1, 1], fs, ax= axarr[0,1])
    axarr[0,1].set_title(stim + ' incorrect')
    
    plot_trial_raster_emily(spikes, correctTrials, narrowClus, my_trial_nums[1], stim, [-1, 1], fs, ax= axarr[1,0])
    axarr[1,0].set_ylabel('Narrow Cluster')
    axarr[1,0].set_title(stim + ' correct')
    #SD
    plot_trial_raster_emily(spikes, incorrectTrials, narrowClus, my_trial_nums[1], stim, [-1, 1], fs, ax= axarr[1,1])
    axarr[1,1].set_title(stim + ' incorrect')
    
    plt.savefig('/home/brad/figures/B604_P01S15_raster_narrow_CvsX_2trials_{}.png'.format(stim))

"""
Explanation: Wide Narrow Correct Incorrect
End of explanation
"""


