# Copyright 2016 Allen Institute for Brain Science
# This file is part of Allen SDK.
#
# Allen SDK is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
# Allen SDK is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Allen SDK.  If not, see <http://www.gnu.org/licenses/>.

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import numpy as np
import pandas as pd
import os, logging
from numpy import arange

def plot_drifting_grating_traces(dg, save_dir):
    '''saves figures with a Ori X TF grid of mean resposes'''
    logging.info("Plotting Ori and TF mean response for all cells")
    
    blank =  dg.sweep_response[dg.stim_table.temporal_frequency==0]          
    for nc in range(dg.numbercells):
        if np.mod(nc,20)==0:
            logging.info("Cell #%s", str(nc))
        xtime = np.arange(-1*dg.interlength/dg.acquisition_rate, (dg.sweeplength+dg.interlength)/dg.acquisition_rate, 1/dg.acquisition_rate)
        plt.figure(nc, figsize=(20,16))
        vmax=0
        vmin=0
        try:
            blank_p =  blank[str(nc)].mean() + (blank[str(nc)].std()/len(blank[str(nc)]))
            blank_n =  blank[str(nc)].mean() - (blank[str(nc)].std()/len(blank[str(nc)]))
        except:
            blank_p = blank.iloc[:,nc].apply(np.mean) + (blank.iloc[:,nc].apply(np.std) / blank.iloc[:,nc].apply(len))
            blank_n = blank.iloc[:,nc].apply(np.mean) - (blank.iloc[:,nc].apply(np.std) / blank.iloc[:,nc].apply(len))
        for ori in dg.orivals:
            ori_pt = np.where(dg.orivals == ori)[0][0]
            for tf in dg.tfvals[1:]:
                tf_pt = np.where(dg.tfvals == tf)[0][0]
                sp_pt = (5*ori_pt)+tf_pt
                subset_response = dg.sweep_response[(dg.stim_table.temporal_frequency==tf)&(dg.stim_table.orientation==ori)]               
                try:
                    subset_response_p =  subset_response[str(nc)].mean() + (subset_response[str(nc)][:-1].std()/len(subset_response[str(nc)]))
                    subset_response_n =  subset_response[str(nc)].mean() - (subset_response[str(nc)][:-1].std()/len(subset_response[str(nc)]))
                except:
                    subset_response_p =  subset_response.iloc[:,nc].apply(np.mean) + (subset_response.iloc[:,nc].apply(np.std) / subset_response.iloc[:,nc].apply(len))
                    subset_response_n =  subset_response.iloc[:,nc].apply(np.mean) - (subset_response.iloc[:,nc].apply(np.std) / subset_response.iloc[:,nc].apply(len))
                ax = plt.subplot(8,5,sp_pt)
                while len(xtime)>len(subset_response[str(nc)].mean()):
                    xtime = np.delete(xtime, -1)
                try:
                    ax.fill_between(xtime, subset_response_p, subset_response_n, color='b', alpha=0.5)
                except:
                    pass
                try:
                    ax.fill_between(xtime, blank_p, blank_n, color='k', alpha=0.5)
                except:
                    pass
                try:
                    ax.plot(xtime, subset_response[str(nc)].mean(), color='b', lw=2)
                except:
                    pass
                ax.plot(xtime, subset_response[str(nc)].mean(), color='b', lw=2)
                #TODO: remove the [:119]  and [:-1] and the try/except                 
                ax.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
                ax.axvspan(0, dg.sweeplength/dg.acquisition_rate ,ymin=0, ymax=1, facecolor='gray', alpha=0.3)
                ax.set_xlim(-1,3)
                ax.set_xticks(range(-1,4))
                ax.yaxis.set_major_locator(MaxNLocator(4))                 
                vmax = np.where(np.amax(subset_response_p)>vmax, np.amax(subset_response_p), vmax)
                vmin = np.where(np.amin(subset_response_n)<vmin, np.amin(subset_response_n), vmin)                 

                if ori_pt<7:
                    ax.set_xticks([])
                else:
                    ax.set_xlabel("Time (s)", fontsize=20)
                if tf_pt>1:
                    ax.set_yticks([])
                else:
                    ax.set_ylabel(str(dg.orivals[ori_pt]), fontsize=24)
                if ori_pt==0:
                    ax.set_title(str(dg.tfvals[tf_pt]), fontsize=24)
                
        for i in range(1,sp_pt+1):
            ax = plt.subplot(8,5,i)
            ax.set_ylim(vmin, vmax)
        plt.tick_params(labelsize=16)
        plt.tight_layout()
        plt.suptitle("Cell " + str(nc+1), fontsize=20)
        plt.subplots_adjust(top=0.9)
        filename = 'Traces DG Cell_'+str(nc+1)+'.png'
        fullfilename = os.path.join(save_dir, filename) 
        plt.savefig(fullfilename)   
        plt.close()

def plot_ns_traces(nsa, save_dir):
    logging.info("Plotting Natural Scene traces for each cell")
    xtime = np.arange(-1*nsa.interlength/nsa.acquisition_rate, (nsa.sweeplength+nsa.interlength)/nsa.acquisition_rate, 1/nsa.acquisition_rate)
    blank = nsa.sweep_response[nsa.stim_table.frame==-1]       
    for nc in range(nsa.numbercells):
        if np.mod(nc,20)==0:
            logging.info("Cell #%s", str(nc))
        vmax=0
        vmin=0
        blank_p =  blank[str(nc)].mean() + (blank[str(nc)].std()/len(blank[str(nc)]))
        blank_n =  blank[str(nc)].mean() - (blank[str(nc)].std()/len(blank[str(nc)])) 
        plt.figure(nc, figsize=(30,25))
        for ns in range(nsa.number_scenes-1):
            subset_response = nsa.sweep_response[nsa.stim_table.frame==ns]               
            subset_response_p =  subset_response[str(nc)].mean() + (subset_response[str(nc)][:].std()/len(subset_response[str(nc)]))
            subset_response_n =  subset_response[str(nc)].mean() - (subset_response[str(nc)][:].std()/len(subset_response[str(nc)]))
            ax = plt.subplot(10,12,ns+1)
            try:
                ax.fill_between(xtime, subset_response_p, subset_response_n, color='b', alpha=0.5)
            except:
                xtime = xtime[:-1]
                ax.fill_between(xtime, subset_response_p, subset_response_n, color='b', alpha=0.5)
            ax.fill_between(xtime, blank_p, blank_n, color='k', alpha=0.5)                 
            ax.plot(xtime, subset_response[str(nc)].mean(), color='b', lw=2)                 
            ax.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
            ax.axvspan(0, nsa.sweeplength/nsa.acquisition_rate ,ymin=0, ymax=1, facecolor='gray', alpha=0.3)
            ax.yaxis.set_major_locator(MaxNLocator(4))                 
            vmax = np.where(np.amax(subset_response_p)>vmax, np.amax(subset_response_p), vmax)
            vmin = np.where(np.amin(subset_response_n)<vmin, np.amin(subset_response_n), vmin)     
            if ns<108:
                ax.set_xticks([])
            if np.mod(ns,12):
                ax.set_yticks([])
        for i in range(1,nsa.number_scenes):
            ax = plt.subplot(10,12,i)
            ax.set_ylim(vmin, vmax)
        plt.tight_layout()
        plt.suptitle("Cell " + str(nc+1), fontsize=20)
        plt.subplots_adjust(top=0.9)
        filename = 'NS Traces Cell_'+str(nc+1)+'.png'
        fullfilename = os.path.join(save_dir, filename) 
        plt.savefig(fullfilename)   
        plt.close()                

def plot_sg_traces(sg, save_dir):
    logging.info("Plotting Static Grating traces for each cell")
    xtime = np.arange(-1*sg.interlength/sg.acquisition_rate, (sg.sweeplength+sg.interlength)/sg.acquisition_rate, 1/sg.acquisition_rate)     
    blank = sg.sweep_response[sg.stim_table.spatial_frequency==0]
    for nc in range(sg.numbercells):
        if np.mod(nc,20)==0:
            logging.info("Cell #%s", str(nc))
        vmax=0
        vmin=0
        blank_p =  blank[str(nc)].mean() + (blank[str(nc)].std()/len(blank[str(nc)]))
        blank_n =  blank[str(nc)].mean() - (blank[str(nc)].std()/len(blank[str(nc)]))
        while len(xtime)>len(blank_p):
            xtime = np.delete(xtime, -1)  
        plt.figure(nc, figsize=(30,30))
        ph_dict = {0:0, 0.25:6, 0.5:77, 0.75:83}
        for ori in sg.orivals:
            ori_pt = np.where(sg.orivals == ori)[0][0]
            for sf in sg.sfvals[1:]:
                sf_pt = np.where(sg.sfvals == sf)[0][0]
                for phase in sg.phasevals:
                    ph_pt = ph_dict[phase]
                    subplotnum = sf_pt + (ori_pt*11) + ph_pt
                    subset_response = sg.sweep_response[(sg.stim_table.spatial_frequency==sf)&(sg.stim_table.orientation==ori)&(sg.stim_table.phase==phase)]
                    subset_response_p =  subset_response[str(nc)].mean() + (subset_response[str(nc)][:].std()/len(subset_response[str(nc)]))
                    subset_response_n =  subset_response[str(nc)].mean() - (subset_response[str(nc)][:].std()/len(subset_response[str(nc)]))
                    ax = plt.subplot(13,11,subplotnum)
                    ax.fill_between(xtime, subset_response_p, subset_response_n, color='b', alpha=0.5)
                    ax.fill_between(xtime, blank_p, blank_n, color='k', alpha=0.5)                 
                    ax.plot(xtime, subset_response[str(nc)].mean(), color='b', lw=2)                 
                    ax.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
                    ax.axvspan(0, sg.sweeplength/sg.acquisition_rate ,ymin=0, ymax=1, facecolor='gray', alpha=0.3)
                    ax.yaxis.set_major_locator(MaxNLocator(4))                 
                    vmax = np.where(np.amax(subset_response_p)>vmax, np.amax(subset_response_p), vmax)
                    vmin = np.where(np.amin(subset_response_n)<vmin, np.amin(subset_response_n), vmin)    
                    if np.mod(subplotnum,11)!=1:
                        ax.set_yticks([])
                    else:
                        ax.set_ylabel(ori, fontsize=20)
                    if subplotnum<133:
                        ax.set_xticks([])
                    if subplotnum<12:
                        ax.set_title(sf, fontsize=20)
                    if subplotnum==3:
                        ax.set_title("Phase 0.0", fontsize=20)
                    if subplotnum==9:
                        ax.set_title("Phase 0.25", fontsize=20)
                    if subplotnum==80:
                        ax.set_title("Phase 0.5", fontsize=20)
                    if subplotnum==86:
                        ax.set_title("Phase 0.75", fontsize=20)
        for i in range(1,144):
            ax = plt.subplot(13,11,i)
            ax.set_ylim(vmin, vmax)
        plt.tight_layout()
        plt.suptitle("Cell " + str(nc+1), fontsize=20)
        plt.subplots_adjust(top=0.9)
        filename = 'SG Traces Cell_'+str(nc+1)+'.png'
        fullfilename = os.path.join(save_dir, filename) 
        plt.savefig(fullfilename)   
        plt.close()                

def plot_lsn_traces(lsn, save_dir):
    logging.info("Plotting LSN traces for all cells")
    xtime = np.arange(-lsn.interlength/lsn.acquisition_rate, 
                       (lsn.interlength + lsn.sweeplength)/lsn.acquisition_rate, 
                       1.0/lsn.acquisition_rate)

    for nc in range(lsn.numbercells):
        if np.mod(nc,20)==0:
            logging.info("Cell #%s", str(nc))
            
        plt.figure(nc, figsize=(24,20))
        vmax=0
        vmin=0
        one_cell = lsn.sweep_response[str(nc)]
        
        for yp in range(16):
            for xp in range(28):
                sp_pt = (yp*28)+xp+1
                on_frame = np.where(lsn.LSN[:,yp,xp]==255)[0]
                off_frame = np.where(lsn.LSN[:,yp,xp]==0)[0]
                subset_on = one_cell[lsn.stim_table.frame.isin(on_frame)]
                subset_off = one_cell[lsn.stim_table.frame.isin(off_frame)]

                subset_on_mean = subset_on.mean()
                subset_off_mean = subset_off.mean()

                ax = plt.subplot(16,28,sp_pt)
                ax.plot(xtime, subset_on_mean, color='r', lw=2)
                ax.plot(xtime, subset_off_mean, color='b', lw=2)
                ax.axvspan(0, lsn.sweeplength/lsn.acquisition_rate ,ymin=0, ymax=1, facecolor='gray', alpha=0.3)
                vmax = np.where(np.amax(subset_on_mean)>vmax, np.amax(subset_on_mean), vmax)
                vmax = np.where(np.amax(subset_off_mean)>vmax, np.amax(subset_off_mean), vmax)
                vmin = np.where(np.amin(subset_on_mean)<vmin, np.amin(subset_on_mean), vmin)
                vmin = np.where(np.amin(subset_off_mean)<vmin, np.amin(subset_off_mean), vmin)
                ax.set_xticks([])
                ax.set_yticks([])
                
        for i in range(1,sp_pt+1):
            ax = plt.subplot(16,28,i)
            ax.set_ylim(vmin, vmax)
        
        plt.tight_layout()
        plt.suptitle("Cell " + str(nc+1), fontsize=20)
        plt.subplots_adjust(top=0.9)
        filename = 'Traces LSN Cell_'+str(nc+1)+'.png'
        fullfilename = os.path.join(save_dir, filename) 
        plt.savefig(fullfilename)   
        plt.close()          


def _plot_3sa(dg, nm1, nm3, save_dir):
    logging.info("Plotting for all cell")
    for nc in range(dg.numbercells):
        if np.mod(nc,20)==0:
            logging.info("Cell #%s", str(nc))
        plt.figure(nc, figsize=(20,20))
        ax1 = plt.subplot2grid((6,6), (0,0), colspan=4) #full trace
        ax2 = plt.subplot2grid((6,6), (0,4))            #histogram of F
        ax3 = plt.subplot2grid((6,6), (1,0), colspan=4) #running speed
        ax4 = plt.subplot2grid((6,6), (1,4))
        ax5 = plt.subplot2grid((6,6), (1,5), sharex=ax4, sharey=ax4)
        ax6 = plt.subplot2grid((6,6), (2,0), colspan=2)
        ax7 = plt.subplot2grid((6,6), (2,2), colspan=2)
        ax8 = plt.subplot2grid((6,6), (2,4), colspan=2)
        ax9 = plt.subplot2grid((6,6), (3,0))
        ax10 = plt.subplot2grid((6,6), (4,0), colspan=3)
        ax11 = plt.subplot2grid((6,6), (5,0), colspan=3)
        ax12 = plt.subplot2grid((6,6), (4,3), colspan=3)
        ax13 = plt.subplot2grid((6,6), (5,3), colspan=3)

        xtime = np.arange(0,np.size(dg.celltraces,1),1.)
        xtime /= dg.acquisition_rate
        dif = np.ediff1d(dg.stim_table.start.values, to_begin=8000, to_end=8000)
        test = np.argwhere(dif>5000)
        ax1.plot(xtime, dg.celltraces[nc,:])        
        for i in range(len(test)-1): 
            ax1.axvspan(xmin=(dg.stim_table.start.iloc[test[i]].values/dg.acquisition_rate), xmax=(dg.stim_table.end.iloc[test[i+1]-1].values/dg.acquisition_rate), color='gray', alpha=0.3)
        ax1.axvspan(xmin=nm1.stim_table.start.min()/nm1.acquisition_rate, xmax=((nm1.stim_table.start.max()+nm1.sweeplength)/nm1.acquisition_rate), color='red', alpha=0.3)
        dif = np.ediff1d(nm3.stim_table.start.values, to_begin=8000, to_end=8000)
        test = np.argwhere(dif>5000)
        for i in range(len(test)-1): 
            ax1.axvspan(xmin=(nm3.stim_table.start.iloc[test[i]].values/nm3.acquisition_rate), xmax=((nm3.stim_table.end.iloc[test[i+1]-1].values+nm3.sweeplength)/nm3.acquisition_rate), color='blue', alpha=0.3)
        ax1.set_xlabel("Time (s)", fontsize=20)
        ax1.set_ylabel("Fluorescence", fontsize=20)
        
        ax2.hist(dg.celltraces[nc,:], bins=70)
        ax2.set_yscale('log')
        ax2.set_xlabel("Fluorescence", fontsize=20)
        ax2.set_ylabel("Count", fontsize=20)
        
        xtime = np.arange(0, np.size(dg.dxcm), 1.)
        xtime /= dg.acquisition_rate
        ax3.plot(xtime, dg.dxcm, color='k')
        ax3.set_xlabel("Time (s)", fontsize=20)
        ax3.set_xlabel("Speed (cm/s)", fontsize=20)
        
        smax = nm1.binned_cells_sp[nc,np.argmax(nm1.binned_cells_sp[nc,:,0]),0] + nm1.binned_cells_sp[nc,np.argmax(nm1.binned_cells_sp[nc,:,0]),1]
        vmax = nm1.binned_cells_vis[nc,np.argmax(nm1.binned_cells_vis[nc,:,0]),0] + nm1.binned_cells_vis[nc,np.argmax(nm1.binned_cells_vis[nc,:,0]),1]
        rmax = np.where(smax>vmax, smax, vmax)
        smin = nm1.binned_cells_sp[nc,np.argmin(nm1.binned_cells_sp[nc,:,0]),0] - nm1.binned_cells_sp[nc,np.argmin(nm1.binned_cells_sp[nc,:,0]),1]
        vmin = nm1.binned_cells_vis[nc,np.argmin(nm1.binned_cells_vis[nc,:,0]),0] - nm1.binned_cells_vis[nc,np.argmin(nm1.binned_cells_vis[nc,:,0]),1]
        rmin = np.where(smin<vmin, smin, vmin)      
        
        ax4.errorbar(nm1.binned_dx_sp[:,0], nm1.binned_cells_sp[nc,:,0], yerr=nm1.binned_cells_sp[nc,:,1], fmt='.', color='k')
        ax4.set_ylim(rmin, rmax)        
        ax4.set_xlabel("Speed (cm/s)", fontsize=20)
        ax4.set_ylabel("DF/F", fontsize=20)
        ax4.set_title("Spontaneous", fontsize=20)
        
        ax5.errorbar(nm1.binned_dx_vis[:,0], nm1.binned_cells_vis[nc,:,0], yerr=nm1.binned_cells_vis[nc,:,1], fmt='.')
        ax5.set_ylim(rmin, rmax)        
        ax5.set_xlabel("Speed (cm/s)", fontsize=20)
        ax5.set_ylabel("DF/F", fontsize=20)
        ax5.set_title("Visual Stimuli", fontsize=20)
        
        peakori = dg.peak.ori_dg[nc]
        peaktf = dg.peak.tf_dg[nc]
        ax6.errorbar(dg.orivals, dg.response[:,peaktf,nc,0], yerr=dg.response[:,peaktf,nc,1], fmt='bo-', lw=2)
        ax6.fill_between(dg.orivals, np.repeat(dg.response[0,0,nc,0]+dg.response[0,0,nc,1], dg.number_ori), np.repeat(dg.response[0,0,nc,0]-dg.response[0,0,nc,1], dg.number_ori), color='gray', alpha=0.5)
        ax6.axhline(y=dg.response[0,0,nc,0], ls='--', color='k')
        ax6.annotate(str(dg.tfvals[peaktf]) + " Hz", xy=(0,0.9), xycoords='axes fraction', fontsize=14)
        ax6.set_xticks(dg.orivals)
        ax7.set_xlim(-10,325)
        ax6.set_xlabel("Direction (d)", fontsize=20)
        ax6.set_ylabel("Mean DF/F (%)", fontsize=20)
        ax6.yaxis.set_major_locator(MaxNLocator(6))
        
        ax7.errorbar(range(5), dg.response[peakori,1:,nc,0], yerr=dg.response[peakori,1:,nc,1],  fmt='bo-', lw=2)
        ax7.fill_between(range(5), np.repeat(dg.response[0,0,nc,0]+dg.response[0,0,nc,1], 5), np.repeat(dg.response[0,0,nc,0]-dg.response[0,0,nc,1], 5), color='gray', alpha=0.5)
        ax7.axhline(y=dg.response[0,0,nc,0], ls='--', color='k', lw=2)
        ax7.annotate(str(dg.orivals[peakori]) + " Deg", xy=(0,0.9), xycoords='axes fraction', fontsize=14)
        ax7.set_xlim(-0.2, 4.2)
        ax7.set_xticks(range(5))
        ax7.set_xticklabels(dg.tfvals[1:])
        ax7.set_xlabel("Temporal frequency (Hz)", fontsize=20)
        
        subset = dg.sweep_response[(dg.stim_table.orientation==dg.orivals[peakori])&(dg.stim_table.temporal_frequency==dg.tfvals[peaktf])]
        xtime = np.arange(-1*dg.interlength/dg.acquisition_rate, (dg.sweeplength+dg.interlength)/dg.acquisition_rate, 1/dg.acquisition_rate)
        while len(xtime)>len(subset[str(nc)].mean()):
            xtime = np.delete(xtime, -1)
        for index, row in subset.iterrows():
            ax8.plot(xtime, subset[str(nc)][index], lw=2)   
        ax8.set_xlim(-1,3)
        ax8.annotate(str(dg.orivals[peakori])+" Deg / "+str(dg.tfvals[peaktf])+" Hz", xy=(0,0.9), xycoords='axes fraction', fontsize=14)
        ax8.set_xlabel("Time (s)", fontsize=20)
        ax8.set_ylabel("DF/F (%)", fontsize=20)
        ax8.axvspan(0, dg.sweeplength/dg.acquisition_rate ,ymin=0, ymax=1, facecolor='gray', alpha=0.3)
        ax8.yaxis.set_major_locator(MaxNLocator(6))
        ax8.set_title("Trial responses to prefered ori/tf", fontsize=20)   
        
        im = ax9.imshow(dg.response[:,1:,nc,0], cmap='gray', interpolation='none')
        ax9.set_ylabel("Direction (d)", fontsize=20)
        ax9.set_yticks(range(8))
        ax9.set_yticklabels(dg.orivals)
        ax9.set_xlabel("Temporal frequency (Hz)", fontsize=20)
        ax9.set_xticks(range(5))
        ax9.set_xticklabels(dg.tfvals[1:])        
        cbar = plt.colorbar(im, ax=ax9)
        cbar.ax.set_ylabel('DF/F (%)', fontsize=8)
        for t in cbar.ax.get_yticklabels():
            t.set_fontsize(8)
        
        xtime = np.arange(0, nm1.sweeplength/nm1.acquisition_rate, 1/nm1.acquisition_rate)
        while len(xtime)>len(nm1.sweep_response[str(nc)].mean()):
            xtime = np.delete(xtime, -1)        
        for index, row in nm1.sweep_response.iterrows():
            ax10.plot(xtime, nm1.sweep_response[str(nc)][index], lw=2)
        ax10.set_xlabel("Time (s)", fontsize=20)
        ax10.set_ylabel("DF/F", fontsize=20)
        ax10.set_title("Natural Movie 1", fontsize=20, color='red')
        
        temp = np.empty((len(nm1.stim_table), nm1.sweeplength))
        for i in range(len(nm1.stim_table)):
            temp[i,:] = nm1.sweep_response[str(nc)].iloc[i]
        ax11.imshow(temp, cmap='gray', interpolation='none', aspect=40)
        ax11.set_ylabel("Trials", fontsize=20)
        ax11.set_xticks([])
        
        xtime = np.arange(0, nm3.sweeplength/nm3.acquisition_rate, 1/nm3.acquisition_rate)
        while len(xtime)>len(nm3.sweep_response[str(nc)].mean()):
            xtime = np.delete(xtime, -1)        
        for index, row in nm3.sweep_response.iterrows():
            ax12.plot(xtime, nm3.sweep_response[str(nc)][index], lw=2)
        ax12.set_xlabel("Time (s)", fontsize=20)
        ax12.set_ylabel("DF/F", fontsize=20)
        ax12.set_title("Natural Movie Long", fontsize=20, color='blue')
        
        temp = np.empty((len(nm3.stim_table), nm3.sweeplength))
        for i in range(len(nm3.stim_table)):
            temp[i,:] = nm3.sweep_response[str(nc)].iloc[i]
        ax13.imshow(temp, cmap='gray', interpolation='none', aspect=100)
        ax13.set_ylabel("Trials", fontsize=20)
        ax13.set_xticks([])
        
        plt.tick_params(labelsize=16)    
        plt.tight_layout()
        filename = 'Cell_'+str(nc+1)+'_3SA.png'
        fullfilename = os.path.join(save_dir, filename) 
        plt.savefig(fullfilename)   
        plt.close()      

def _plot_3sc(lsn, nm1, nm2, save_dir):
    logging.info("Plotting for all cells")
    for nc in range(lsn.numbercells):
        if np.mod(nc,20)==0:
            logging.info("Cell #%s", str(nc))
                         
        plt.figure(nc, figsize=(20,20)) 
        ax1 = plt.subplot2grid((6,6), (0,0), colspan=4) #full trace
        ax2 = plt.subplot2grid((6,6), (0,4))            #histogram of F
        ax11 = plt.subplot2grid((6,6), (1,0), colspan=4)
        ax12 = plt.subplot2grid((6,6), (1,4))
        ax13 = plt.subplot2grid((6,6), (1,5), sharex=ax12, sharey=ax12)
        ax3 = plt.subplot2grid((6,6), (2,0), colspan=3)  #movie 1
        ax4 = plt.subplot2grid((6,6), (3,0), colspan=3)
        ax5 = plt.subplot2grid((6,6), (2,3), colspan=3) #movie 2
        ax6 = plt.subplot2grid((6,6), (3,3), colspan=3)
        ax7 = plt.subplot2grid((6,6), (4,0), colspan=3) #receptive fields
        ax8 = plt.subplot2grid((6,6), (4,3), colspan=3)
        ax9 = plt.subplot2grid((6,6), (5,0), colspan=3)
        ax10 = plt.subplot2grid((6,6), (5,3), colspan=3)
        
        xtime = np.arange(0,np.size(lsn.celltraces,1),1.)
        xtime /= lsn.acquisition_rate
        dif = np.ediff1d(lsn.stim_table.start.values, to_begin=8000, to_end=8000)
        test = np.argwhere(dif>5000)
        ax1.plot(xtime, lsn.celltraces[nc,:])        
        for i in range(len(test)-1): 
            ax1.axvspan(xmin=(lsn.stim_table.start.iloc[test[i]].values/lsn.acquisition_rate), xmax=(lsn.stim_table.end.iloc[test[i+1]-1].values/lsn.acquisition_rate), color='gray', alpha=0.3)
        ax1.axvspan(nm1.stim_table.start.min()/nm1.acquisition_rate, nm1.stim_table.end.max()/nm1.acquisition_rate, ymin=0, ymax=1, color='red', alpha=0.3)
        ax1.axvspan(nm2.stim_table.start.min()/nm2.acquisition_rate, nm2.stim_table.end.max()/nm2.acquisition_rate, ymin=0, ymax=1, color='green', alpha=0.3)
        ax1.set_xlabel("Time (s)", fontsize=20)
        ax1.set_ylabel("Fluorescence", fontsize=20)
        
        ax2.hist(lsn.celltraces[nc,:], bins=70)
        ax2.set_yscale('log')
        ax2.set_xlabel("Fluorescence", fontsize=20)
        ax2.set_ylabel("Count", fontsize=20)
        
        xtime = np.arange(0, np.size(lsn.dxcm), 1.)
        xtime /= lsn.acquisition_rate
        ax11.plot(xtime, lsn.dxcm, color='k')
        ax11.set_xlabel("Time (s)", fontsize=20)
        ax11.set_xlabel("Speed (cm/s)", fontsize=20)
        
        smax = nm1.binned_cells_sp[nc,np.argmax(nm1.binned_cells_sp[nc,:,0]),0] + nm1.binned_cells_sp[nc,np.argmax(nm1.binned_cells_sp[nc,:,0]),1]
        vmax = nm1.binned_cells_vis[nc,np.argmax(nm1.binned_cells_vis[nc,:,0]),0] + nm1.binned_cells_vis[nc,np.argmax(nm1.binned_cells_vis[nc,:,0]),1]
        rmax = np.where(smax>vmax, smax, vmax)
        smin = nm1.binned_cells_sp[nc,np.argmin(nm1.binned_cells_sp[nc,:,0]),0] - nm1.binned_cells_sp[nc,np.argmin(nm1.binned_cells_sp[nc,:,0]),1]
        vmin = nm1.binned_cells_vis[nc,np.argmin(nm1.binned_cells_vis[nc,:,0]),0] - nm1.binned_cells_vis[nc,np.argmin(nm1.binned_cells_vis[nc,:,0]),1]
        rmin = np.where(smin<vmin, smin, vmin)        
        
        ax12.errorbar(nm1.binned_dx_sp[:,0], nm1.binned_cells_sp[nc,:,0], yerr=nm1.binned_cells_sp[nc,:,1], fmt='.', color='k')
        ax12.set_ylim(rmin, rmax)        
        ax12.set_xlabel("Speed (cm/s)", fontsize=20)
        ax12.set_ylabel("DF/F", fontsize=20)
        ax12.set_title("Spontaneous", fontsize=20)
        
        ax13.errorbar(nm1.binned_dx_vis[:,0], nm1.binned_cells_vis[nc,:,0], yerr=nm1.binned_cells_vis[nc,:,1], fmt='.')
        ax13.set_ylim(rmin, rmax)        
        ax13.set_xlabel("Speed (cm/s)", fontsize=20)
        ax13.set_ylabel("DF/F", fontsize=20)
        ax13.set_title("Visual Stimuli", fontsize=20)
        
        xtime = np.arange(0, nm1.sweeplength/nm1.acquisition_rate, 1/nm1.acquisition_rate)
        for index, row in nm1.sweep_response.iterrows():
            ax3.plot(xtime, nm1.sweep_response[str(nc)][index], lw=2)
        ax3.set_xlabel("Time (s)", fontsize=20)
        ax3.set_ylabel("DF/F", fontsize=20)
        ax3.set_title("Natural Movie 1", fontsize=20, color='red')
        
        temp = np.empty((len(nm1.stim_table), nm1.sweeplength))
        for i in range(len(nm1.stim_table)):
            temp[i,:] = nm1.sweep_response[str(nc)].iloc[i]
        ax4.imshow(temp, cmap='gray', interpolation='none', aspect=40)
        ax4.set_ylabel("Trials", fontsize=20)
        ax4.set_xticks([])
        
        xtime = np.arange(0, nm2.sweeplength/nm2.acquisition_rate, 1/nm2.acquisition_rate)
        for index, row in nm2.sweep_response.iterrows():
            ax5.plot(xtime, nm2.sweep_response[str(nc)][index], lw=2)
        ax5.set_xlabel("Time (s)", fontsize=20)
        ax5.set_ylabel("DF/F", fontsize=20)
        ax5.set_title("Natural Movie 2", fontsize=20, color='green')
        
        temp = np.empty((len(nm2.stim_table), nm2.sweeplength))
        for i in range(len(nm2.stim_table)):
            temp[i,:] = nm2.sweep_response[str(nc)].iloc[i]
        ax6.imshow(temp, cmap='gray', interpolation='none', aspect=40)
        ax6.set_ylabel("Trials", fontsize=20)
        ax6.set_xticks([])
        
        vMax = np.where(np.amax(lsn.receptive_field[:,:,nc,0])>np.amax(lsn.receptive_field[:,:,nc,1]), np.amax(lsn.receptive_field[:,:,nc,0]), np.amax(lsn.receptive_field[:,:,nc,1]))
        vMin = np.where(np.amin(lsn.receptive_field[:,:,nc,0])<np.amin(lsn.receptive_field[:,:,nc,1]), np.amin(lsn.receptive_field[:,:,nc,0]), np.amin(lsn.receptive_field[:,:,nc,1]))

        imon = ax7.imshow(lsn.receptive_field[:,:,nc,0], cmap='gray', interpolation='None', vmin=vMin, vmax=vMax)
        ax7.set_title("ON", fontsize=20)
        ax7.set_xticks([])
        ax7.set_yticks([])
        cbar = plt.colorbar(imon, ax=ax7, fraction=0.046, pad=0.04)
        cbar.ax.set_ylabel('DF/F (%)', fontsize=10)
        for t in cbar.ax.get_yticklabels():
            t.set_fontsize(8)
        
        imoff = ax8.imshow(lsn.receptive_field[:,:,nc,1], cmap='gray', interpolation='None', vmin=vMin, vmax=vMax)
        ax8.set_title("OFF", fontsize=20)
        ax8.set_xticks([])
        ax8.set_yticks([])
        cbar = plt.colorbar(imoff, ax=ax8, fraction=0.046, pad=0.04)
        cbar.ax.set_ylabel('DF/F (%)', fontsize=10)
        for t in cbar.ax.get_yticklabels():
            t.set_fontsize(8)
        
        zon = (lsn.receptive_field[:,:,nc,0]-np.mean(lsn.receptive_field[:,:,nc,0]))/np.std(lsn.receptive_field[:,:,nc,0])
        zon = np.where(abs(zon)>2, zon, 0)
        Vmax_on = np.where(abs(np.amax(zon))>abs(np.amin(zon)), np.amax(zon), -1*np.amin(zon))
        zoff = (lsn.receptive_field[:,:,nc,1]-np.mean(lsn.receptive_field[:,:,nc,1]))/np.std(lsn.receptive_field[:,:,nc,1])
        zoff = np.where(abs(zoff)>2, zoff, 0)
        Vmax_off = np.where(abs(np.amax(zoff))>abs(np.amin(zoff)), np.amax(zoff), -1*np.amin(zoff))
        Vmax = np.where(Vmax_on>Vmax_off, Vmax_on, Vmax_off)
        imzon = ax9.imshow(zon, cmap='RdBu_r', interpolation='none', vmin=-1*Vmax, vmax=Vmax)
        ax9.set_title("On Z-score", fontsize=20)
        ax9.set_xticks([])
        ax9.set_yticks([])
        cbar = plt.colorbar(imzon, ax=ax9, fraction=0.046, pad=0.04)
        
        zoff = (lsn.receptive_field[:,:,nc,1]-np.mean(lsn.receptive_field[:,:,nc,1]))/np.std(lsn.receptive_field[:,:,nc,1])
        zoff = np.where(abs(zoff)>2, zoff, 0)
        imzoff = ax10.imshow(zoff, cmap='RdBu', interpolation='none', vmin=-1*Vmax, vmax=Vmax)
        ax10.set_title("Off Z-score", fontsize=20)
        ax10.set_xticks([])
        ax10.set_yticks([])
        cbar = plt.colorbar(imzoff, ax=ax10, fraction=0.046, pad=0.04)
        
        plt.tick_params(labelsize=16)    
        plt.tight_layout()
        filename = 'Cell_'+str(nc+1)+'_3SC.png'
        fullfilename = os.path.join(save_dir, filename) 
        plt.savefig(fullfilename)   
        plt.close()

def _plot_3sb(sg, nm1, ns, save_dir):
    logging.info("Plotting for all cells")
    for nc in range(sg.numbercells):
        if np.mod(nc,20)==0:
            logging.info("Cell #%s", str(nc))
        plt.figure(nc, figsize=(20,24)) 
        ax1 = plt.subplot2grid((6,6), (0,0), colspan=4) #full trace
        ax2 = plt.subplot2grid((6,6), (0,4))            #histogram of F
        ax14 = plt.subplot2grid((6,6), (1,4))           #speed tuning
        ax16 = plt.subplot2grid((6,6), (1,5), sharex=ax14, sharey=ax14)
        ax15 = plt.subplot2grid((6,6), (1,0), colspan=4)
        ax3 = plt.subplot2grid((6,6),(2,0), colspan=2)  #Ori tuning
        ax4 = plt.subplot2grid((6,6),(2,2), colspan=2)  #sf tuning
        ax13 = plt.subplot2grid((6,6),(2,4), colspan=2)  #response at peak
        ax5 = plt.subplot2grid((6,6), (3,0))
        ax6 = plt.subplot2grid((6,6), (3,1))
        ax7 = plt.subplot2grid((6,6), (3,2))
        ax8 = plt.subplot2grid((6,6), (3,3))
        ax9 = plt.subplot2grid((6,6),(4,0), colspan=3)  #movie
        ax10 = plt.subplot2grid((6,6),(5,0), colspan=3)
        ax11 = plt.subplot2grid((6,6),(4,3), colspan=2) #natural scenes
        ax12 = plt.subplot2grid((6,6),(5,3), colspan=2)
        
        
        xtime = np.arange(0,np.size(sg.celltraces,1),1.)
        xtime /= sg.acquisition_rate
        dif = np.ediff1d(sg.stim_table.start.values, to_begin=8000, to_end=8000)
        test = np.argwhere(dif>5000)
        ax1.plot(xtime, sg.celltraces[nc,:])        
        for i in range(len(test)-1): 
            ax1.axvspan(xmin=(sg.stim_table.start.iloc[test[i]].values/sg.acquisition_rate), xmax=(sg.stim_table.end.iloc[test[i+1]-1].values/sg.acquisition_rate), color='gray', alpha=0.3)
        ax1.axvspan(nm1.stim_table.start.min()/nm1.acquisition_rate, nm1.stim_table.end.max()/nm1.acquisition_rate, ymin=0, ymax=1, color='red', alpha=0.3)
        dif = np.ediff1d(ns.stim_table.start.values, to_begin=8000, to_end=8000)
        test = np.argwhere(dif>5000)
        for i in range(len(test)-1): 
            ax1.axvspan(xmin=(ns.stim_table.start.iloc[test[i]].values/ns.acquisition_rate), xmax=(ns.stim_table.end.iloc[test[i+1]-1].values/ns.acquisition_rate), color='blue', alpha=0.3)
        ax1.set_xlabel("Time (s)", fontsize=20)
        ax1.set_ylabel("Fluorescence", fontsize=20)
        
        ax2.hist(sg.celltraces[nc,:], bins=70)
        ax2.set_yscale('log')
        ax2.set_xlabel("Fluorescence", fontsize=20)
        ax2.set_ylabel("Count", fontsize=20)
        
        xtime = np.arange(0, np.size(sg.dxcm), 1.)
        xtime /= sg.acquisition_rate
        ax15.plot(xtime, sg.dxcm, color='k')
        ax15.set_xlabel("Time (s)", fontsize=20)
        ax15.set_xlabel("Speed (cm/s)", fontsize=20)
        
        peakori = sg.peak.ori_sg[nc]
        peaksf = sg.peak.sf_sg[nc]
        ax3.errorbar(sg.orivals, sg.response[:,peaksf,0,nc,0], yerr=sg.response[:,peaksf,0,nc,1], color='blue', fmt='o-', lw=2)
        ax3.errorbar(sg.orivals, sg.response[:,peaksf,1,nc,0], yerr=sg.response[:,peaksf,1,nc,1], color='cornflowerblue', fmt='o-', lw=2)
        ax3.errorbar(sg.orivals, sg.response[:,peaksf,2,nc,0], yerr=sg.response[:,peaksf,2,nc,1], color='steelblue', fmt='o-', lw=2)
        ax3.errorbar(sg.orivals, sg.response[:,peaksf,3,nc,0], yerr=sg.response[:,peaksf,3,nc,1], color='lightskyblue', fmt='o-', lw=2)
        ax3.fill_between(sg.orivals, np.repeat(sg.response[0,0,0,nc,0]+sg.response[0,0,0,nc,1], sg.number_ori), np.repeat(sg.response[0,0,0,nc,0]-sg.response[0,0,0,nc,1], sg.number_ori), color='gray', alpha=0.5)
        ax3.axhline(y=sg.response[0,0,0,nc,0], ls='--', color='k', lw=2)
        ax3.set_xlim(-10, 160)
        ax3.set_xticks(sg.orivals)
        ax3.set_xlabel("Orientation (d)", fontsize=20)
        ax3.set_ylabel("DF/F (%)", fontsize=20)
        
        ax4.errorbar(range(5), sg.response[peakori,1:,0,nc,0], yerr=sg.response[peakori,1:,0,nc,1], color='blue', fmt='o-', lw=2)
        ax4.errorbar(range(5), sg.response[peakori,1:,1,nc,0], yerr=sg.response[peakori,1:,1,nc,1], color='cornflowerblue', fmt='o-', lw=2)
        ax4.errorbar(range(5), sg.response[peakori,1:,2,nc,0], yerr=sg.response[peakori,1:,2,nc,1], color='steelblue', fmt='o-', lw=2)
        ax4.errorbar(range(5), sg.response[peakori,1:,3,nc,0], yerr=sg.response[peakori,1:,3,nc,1], color='lightskyblue', fmt='o-', lw=2)
        ax4.fill_between(range(5), np.repeat(sg.response[0,0,0,nc,0]+sg.response[0,0,0,nc,1], 5), np.repeat(sg.response[0,0,0,nc,0]-sg.response[0,0,0,nc,1], 5), color='gray', alpha=0.5)
        ax4.axhline(y=sg.response[0,0,0,nc,0], ls='--', color='k', lw=2)
        ax4.set_xlim(-0.2, 4.2)
        ax4.set_xticks(range(5))
        ax4.set_xticklabels(sg.sfvals[1:])
        ax4.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        
        xtime = np.arange(-1*sg.interlength/sg.acquisition_rate, (sg.sweeplength+sg.interlength)/sg.acquisition_rate, 1/sg.acquisition_rate)
        peakori = sg.peak.ori_sg[nc]
        peaksf = sg.peak.sf_sg[nc]
        peakphase = sg.peak.phase_sg[nc]
        subset = sg.sweep_response[(sg.stim_table.orientation==sg.orivals[peakori])&(sg.stim_table.spatial_frequency==sg.sfvals[peaksf])&(sg.stim_table.phase==sg.phasevals[peakphase])]
        subset_p =  subset[str(nc)].mean() + (subset[str(nc)].std()/np.sqrt(len(subset[str(nc)])))
        subset_n =  subset[str(nc)].mean() - (subset[str(nc)].std()/np.sqrt(len(subset[str(nc)])))
        try:
            ax13.fill_between(xtime, subset_p, subset_n, color='b', alpha=0.5)
        except:
            xtime = xtime[:-1]
            ax13.fill_between(xtime, subset_p, subset_n, color='b', alpha=0.5)
        blank = sg.sweep_response[(sg.stim_table.orientation==0)&(sg.stim_table.spatial_frequency==0)&(sg.stim_table.phase==0)]
        blank_p = blank[str(nc)].mean() + (blank[str(nc)].std()/np.sqrt(len(blank[str(nc)])))
        blank_n = blank[str(nc)].mean() - (blank[str(nc)].std()/np.sqrt(len(blank[str(nc)])))        
        ax13.fill_between(xtime, blank_p, blank_n, color='gray', alpha=0.5)        
        ax13.plot(xtime, subset[str(nc)].mean(), color='b', lw=2)        
        ax13.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)
        ax13.axvspan(0, sg.sweeplength/sg.acquisition_rate ,ymin=0, ymax=1, facecolor='gray', alpha=0.3)        
        ax13.yaxis.set_major_locator(MaxNLocator(4)) 
        ax13.set_xlabel("Time (s)", fontsize=20)
        ax13.set_ylabel("DF/F (%)", fontsize=20)
        
        Vmax = sg.response[:,1:,:,nc,0].max()
        ax5.imshow(sg.response[:,1:,0,nc,0], cmap='gray', interpolation='none', vmin=0, vmax=Vmax)
        ax5.set_ylabel("Orientation (d)", fontsize=20)
        ax5.set_yticks(range(6))
        ax5.set_yticklabels(sg.orivals)
        ax5.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        ax5.set_xticks(range(5))
        ax5.set_xticklabels(sg.sfvals[1:])
        ax5.set_title("Phase 0.0", color='blue', fontsize=20)
        
        ax6.imshow(sg.response[:,1:,1,nc,0], cmap='gray', interpolation='none', vmin=0, vmax=Vmax)
        ax6.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        ax6.set_xticks(range(5))
        ax6.set_xticklabels(sg.sfvals[1:])
        ax6.set_yticks(range(6))
        ax6.set_yticklabels(sg.orivals)
        ax6.set_title("Phase 0.25", color='cornflowerblue', fontsize=20)
        
        ax7.imshow(sg.response[:,1:,2,nc,0], cmap='gray', interpolation='none', vmin=0, vmax=Vmax)
        ax7.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        ax7.set_xticks(range(5))
        ax7.set_xticklabels(sg.sfvals[1:])
        ax7.set_yticks(range(6))
        ax7.set_yticklabels(sg.orivals)
        ax7.set_title("Phase 0.5", color='steelblue', fontsize=20)
        
        ax8.imshow(sg.response[:,1:,3,nc,0], cmap='gray', interpolation='none', vmin=0, vmax=Vmax)
        ax8.set_xlabel("Spatial frequency (cpd)", fontsize=20)
        ax8.set_xticks(range(5))
        ax8.set_xticklabels(sg.sfvals[1:])
        ax8.set_yticks(range(6))
        ax8.set_yticklabels(sg.orivals)
        ax8.set_title("Phase 0.75", color='lightskyblue', fontsize=20)
        
        xtime = np.arange(0, nm1.sweeplength/nm1.acquisition_rate, 1/nm1.acquisition_rate)
        while len(xtime)>nm1.sweeplength:
            xtime = np.delete(xtime, -1)            
        for index, row in nm1.sweep_response.iterrows():
            ax9.plot(xtime, nm1.sweep_response[str(nc)][index], lw=2)
        ax9.set_xlabel("Time (s)", fontsize=20)
        ax9.set_ylabel("DF/F", fontsize=20)
        ax9.set_title("Natural Movie 1", fontsize=20, color='red')
        
        temp = np.empty((len(nm1.stim_table), nm1.sweeplength))
        for i in range(len(nm1.stim_table)):
            temp[i,:] = nm1.sweep_response[str(nc)].iloc[i]
        ax10.imshow(temp, cmap='gray', interpolation='none', aspect=40)
        ax10.set_ylabel("Trials", fontsize=20)
        ax10.set_xticks([])
        
        temp = np.copy(ns.response[1:,nc,:2])
        scene_response = pd.DataFrame(temp, columns=('response','error'))
        scene_response = scene_response.sort(columns='response', ascending=False)
        ax11.errorbar(range(ns.number_scenes-1), scene_response.response, yerr=scene_response.error, fmt='o', color='k')
        ax11.fill_between(range(ns.number_scenes-1), np.repeat(ns.response[0,nc,0]+ns.response[0,nc,1], ns.number_scenes-1), np.repeat(ns.response[0,nc,0]-ns.response[0,nc,1], ns.number_scenes-1), color='gray', alpha=0.3)
        ax11.axhline(y=ns.response[0,nc,0], ls='--', lw=2, color='k')
        ax11.set_xlim(-2,120)
        ax11.set_title("Natural Scenes", fontsize=20, color='blue')
        ax11.set_xlabel("Scene", fontsize=20)
        ax11.set_ylabel("DF/F (%)", fontsize=20)
        
        xtime = np.arange(-1*ns.interlength/ns.acquisition_rate, (ns.sweeplength+ns.interlength)/ns.acquisition_rate, 1/ns.acquisition_rate)
        nsp = np.argmax(ns.response[1:,nc,0])
        subset_response = ns.sweep_response[ns.stim_table.frame==nsp]               
        subset_response_p = subset_response[str(nc)].mean() + (subset_response[str(nc)][:].std()/np.sqrt(len(subset_response[str(nc)])))
        subset_response_n = subset_response[str(nc)].mean() - (subset_response[str(nc)][:].std()/np.sqrt(len(subset_response[str(nc)])))
        try:        
            ax12.fill_between(xtime, subset_response_p, subset_response_n, color='b', alpha=0.5)
        except:
            xtime = xtime[:-1]
            ax12.fill_between(xtime, subset_response_p, subset_response_n, color='b', alpha=0.5)
        blank = ns.sweep_response[ns.stim_table.frame==-1]
        blank_p = blank[str(nc)].mean() + (blank[str(nc)].std()/np.sqrt(len(blank[str(nc)])))
        blank_n = blank[str(nc)].mean() - (blank[str(nc)].std()/np.sqrt(len(blank[str(nc)])))        
        ax12.fill_between(xtime, blank_p, blank_n, color='gray', alpha=0.5)          
        ax12.plot(xtime, subset_response[str(nc)].mean(), color='b', lw=2)
        ax12.plot(xtime, blank[str(nc)].mean(), color='k', lw=2)                 
        ax12.axvspan(0, ns.sweeplength/ns.acquisition_rate ,ymin=0, ymax=1, facecolor='gray', alpha=0.3)
        ax12.yaxis.set_major_locator(MaxNLocator(4)) 
        ax12.set_xlabel("Time (s)", fontsize=20)
        ax12.set_ylabel("DF/F (%)", fontsize=20)
        
        
        plt.tick_params(labelsize=16)    
        plt.tight_layout()
        filename = 'Cell_'+str(nc+1)+'_3SB.png'
        fullfilename = os.path.join(save_dir, filename) 
        plt.savefig(fullfilename)   
        plt.close()


def plot_running_a(dg, nm1, nm3, save_dir):
    logging.info("Plotting running data summary")
    nc = -1
    plt.figure(1, figsize=(10,8))
    ax1 = plt.subplot2grid((4,4),(0,0), colspan=3)
    ax2 = plt.subplot2grid((4,4),(0,3))
    ax3 = plt.subplot2grid((4,4),(1,0))
    ax4 = plt.subplot2grid((4,4),(1,1))
    ax5 = plt.subplot2grid((4,4),(1,2))
    ax6 = plt.subplot2grid((4,4), (2,0), colspan=2)
    ax7 = plt.subplot2grid((4,4), (3,0), colspan=2)
    ax8 = plt.subplot2grid((4,4), (2,2), colspan=2)
    ax9 = plt.subplot2grid((4,4), (3,2), colspan=2)
    
    xtime = np.arange(0, np.size(dg.dxcm), 1.)
    xtime /= dg.acquisition_rate
    dif = np.ediff1d(dg.stim_table.start.values, to_begin=8000, to_end=8000)
    test = np.argwhere(dif>5000)
    ax1.plot(xtime, dg.dxcm, color='k')        
    for i in range(len(test)-1): 
        ax1.axvspan(xmin=(dg.stim_table.start.iloc[test[i]].values/dg.acquisition_rate), xmax=(dg.stim_table.end.iloc[test[i+1]-1].values/dg.acquisition_rate), color='gray', alpha=0.3)
    ax1.axvspan(xmin=nm1.stim_table.start.min()/nm1.acquisition_rate, xmax=((nm1.stim_table.start.max()+nm1.sweeplength)/nm1.acquisition_rate), color='red', alpha=0.3)
    dif = np.ediff1d(nm3.stim_table.start.values, to_begin=8000, to_end=8000)
    test = np.argwhere(dif>5000)
    for i in range(len(test)-1): 
        ax1.axvspan(xmin=(nm3.stim_table.start.iloc[test[i]].values/nm3.acquisition_rate), xmax=((nm3.stim_table.end.iloc[test[i+1]-1].values+nm3.sweeplength)/nm3.acquisition_rate), color='blue', alpha=0.3)
    ax1.set_xlabel("Time (s)", fontsize=20)
    ax1.set_ylabel("Speed (cm/s)", fontsize=20)
    
    dx = dg.dxcm[np.logical_not(np.isnan(dg.dxcm))]
    ax2.hist(dx, bins=80, range=(-20,100), color='gray')
    ax2.set_xlabel("Speed (cm/s)", fontsize=20)
    
    run_peak = np.where(dg.response[:,1:,nc,0] == np.nanmax(dg.response[:,1:,nc,0]))
    peakori = run_peak[0][0]
    peaktf = run_peak[1][0]+1
    
    ax3.errorbar(dg.orivals, dg.response[:,peaktf,nc,0], yerr=dg.response[:,peaktf,nc,1], fmt='b.-', lw=2)
    ax3.fill_between(dg.orivals, np.repeat(dg.response[0,0,nc,0]+dg.response[0,0,nc,1], dg.number_ori), np.repeat(dg.response[0,0,nc,0]-dg.response[0,0,nc,1], dg.number_ori), color='gray', alpha=0.5)
    ax3.axhline(y=dg.response[0,0,nc,0], ls='--', color='k')
    ax3.annotate(str(dg.tfvals[peaktf]) + " Hz", xy=(0,0.9), xycoords='axes fraction', fontsize=14)
    ax3.set_xtick = (dg.orivals)
    ax3.set_xlabel("Direction (deg)", fontsize=20)
    ax3.set_ylabel("Speed (cm/s)", fontsize=20)
    ax3.yaxis.set_major_locator(MaxNLocator(6))
    
    ax4.errorbar(dg.tfvals[1:], dg.response[peakori,1:,nc,0], yerr=dg.response[peakori,1:,nc,1], fmt='b.-', lw=2)
    ax4.fill_between(dg.tfvals[1:], np.repeat(dg.response[0,0,nc,0]+dg.response[0,0,nc,1], dg.number_tf-1), np.repeat(dg.response[0,0,nc,0]-dg.response[0,0,nc,1], dg.number_tf-1), color='gray', alpha=0.5)
    ax4.axhline(y=dg.response[0,0,nc,0], ls='--', color='k')  
    ax4.annotate(str(dg.orivals[peakori]) + " Deg", xy=(0,0.9), xycoords='axes fraction', fontsize=14)
    ax4.set_xticks = (dg.tfvals[1:])
    ax4.set_xlabel("Temporal frequency (Hz)", fontsize=20)
    ax4.yaxis.set_major_locator(MaxNLocator(6))
    
    im = ax5.imshow(dg.response[:,1:,nc,0], cmap='gray', interpolation='none')
    ax5.set_ylabel("Direction", fontsize=16)
    ax5.set_xlabel("TF", fontsize=16)
    ax5.set_yticks(range(dg.number_ori))
    ax5.set_yticklabels(list(dg.orivals.astype(int).astype(str)))
    ax5.set_xticks(range(dg.number_tf-1))
    ax5.set_xticklabels(list(dg.tfvals[1:].astype(int).astype(str)))
    cbar = plt.colorbar(im, ax=ax5)
    cbar.ax.set_ylabel('Speed (cm/s)', fontsize=8)
    for t in cbar.ax.get_yticklabels():
        t.set_fontsize(8)
    
    xtime = np.arange(0, nm1.sweeplength/nm1.acquisition_rate, 1/nm1.acquisition_rate)
    while len(xtime)>len(nm1.sweep_response['dx'].mean()):
        xtime = np.delete(xtime, -1)        
    for index, row in nm1.sweep_response.iterrows():
        ax6.plot(xtime, nm1.sweep_response['dx'][index], lw=2)
    ax6.set_xlabel("Time (s)", fontsize=20)
    ax6.set_ylabel("DF/F", fontsize=20)
    ax6.set_title("Natural Movie 1", fontsize=20, color='red')
    
    temp = np.empty((len(nm1.stim_table), nm1.sweeplength))
    for i in range(len(nm1.stim_table)):
        temp[i,:] = nm1.sweep_response['dx'].iloc[i][:,0]
    ax7.imshow(temp, cmap='gray', interpolation='none', aspect=40)
    ax7.set_ylabel("Trials", fontsize=20)
    ax7.set_xticks([])
    
    xtime = np.arange(0, nm3.sweeplength/nm3.acquisition_rate, 1/nm3.acquisition_rate)
    while len(xtime)>len(nm3.sweep_response['dx'].mean()):
        xtime = np.delete(xtime, -1)        
    for index, row in nm3.sweep_response.iterrows():
        ax8.plot(xtime, nm3.sweep_response['dx'][index], lw=2)
    ax8.set_xlabel("Time (s)", fontsize=20)
    ax8.set_ylabel("DF/F", fontsize=20)
    ax8.set_title("Natural Movie Long", fontsize=20, color='blue')
    
    temp = np.empty((len(nm3.stim_table), nm3.sweeplength))
    for i in range(len(nm3.stim_table)):
        temp[i,:] = nm3.sweep_response['dx'].iloc[i][:,0]
    ax9.imshow(temp, cmap='gray', interpolation='none', aspect=100)
    ax9.set_ylabel("Trials", fontsize=20)
    ax9.set_xticks([])

    plt.tick_params(labelsize=16)
    plt.tight_layout()
    plt.suptitle("Running Summary", fontsize=20)
    plt.subplots_adjust(top=0.9)
    filename = 'Running Summary.png'
    fullfilename = os.path.join(save_dir, filename) 
    plt.savefig(fullfilename)   
    plt.close()              
