
from matplotlib.colors import LogNorm
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy import stats
import glob
import argparse


def parse_fmftime(namestring):
    fn = namestring.split('/')[-1]
    exp_id, CAMN, DATE, TIME = fn.split('_', 3)
    FLY_ID = exp_id + '_' + DATE + '_' + TIME
    fmftime = pd.to_datetime(DATE + TIME)
    return FLY_ID, fmftime, exp_id
 
def boutlength_distribution(df, _bins=np.arange(0,10,100)):
    dist = np.histogram(df['bout_duration'], bins=_bins)[0]
    count = sum(dist)
    return dist/float(count)
    
def group_data(df, groups):
    #df.index = df[args.parameter]
    grouped = df.groupby(groups)
    means = grouped.mean()
    ns = grouped.count()
    sems = grouped.sem()
    #sems = grouped.aggregate(lambda x: np.std(x, ddof=1)/np.sqrt(x.count()))  #HANDLES NAN DATA BETTER THAN GROUPED.SEM()
    #sems = grouped.aggregate(lambda x: st.sem(x, axis=None))
    return means, sems, ns


def row_of_histograms(DATADIR, _HANDLE, _treatments, _fig, row):
    for x in np.arange(0,len(_treatments)):
        
        ax = _fig.add_subplot(3,4, 4*row+x+1)
        
        plt.hist2d(foo.maxAngle.values, foo.laterality.values, bins=50, norm=LogNorm())
        if row == 0:
            ax.set_title(_treatments[x])
        if row == 2:
            ax.set_xlabel('Wing Angle')
        if x == 0:
            ax.set_ylabel(_HANDLE + '\nLaterality score')
        #ax.set_xlim(0,1.6)
        #ax.set_ylim(0,1)
    return



treatments = ['00','11','15','65']
genotypes = ['DB072','DB185','DB213']
colourlist = ['#0033CC','#33CC33','#FFAA00', '#CC3300', '#AAAAAA','#0032FF','r','c','m','y', '#000000', '#333333']
if __name__ == "__main__":

    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--boutsfile', type=str, required=True,
                            help='full path to bouts.pickle, generated by bouts.py')  
    parser.add_argument('--parameter', type=str, required=False, default='Time (s)',
                            help='parameter for which to measure bouts')
    parser.add_argument('--min', type=float, required=False, default=0,
                            help='minimum range of histogram')
    parser.add_argument('--max', type=float, required=False, default=400,
                            help='maximum range of histogram')
    parser.add_argument('--bins', type=float, required=False, default=50,
                            help='number of bins for histogram')
    parser.add_argument('--log_bins', type=bool, required=False, default=False,
                            help='make True to generate a log-transformed set of bins.')
    parser.add_argument('--handle', type=str, required=False, default='',
                            help='give a string to select a subset of the dataset.')
    
    args = parser.parse_args()

    if args.log_bins == True:
        if args.min < 0:
            BINS = np.logspace(0, np.log10(args.max - args.min + 1.0), args.bins) - 1.0 + args.min
        elif args.min == 0:
            BINS = np.logspace(0, np.log10(args.max+1.0), args.bins) - 1.0
        else:
            BINS = np.logspace(np.log10(args.min), np.log10(args.max), args.bins)
    else:
        BINS = np.linspace(args.min, args.max, args.bins)
        
    
    bouts = pd.read_pickle(args.boutsfile)
    bouts = bouts[bouts['bout_duration'] > 0]
    print bouts.bout_duration.max()
    g = bouts.groupby(['g1','g2','g3', 'FlyID'])
    df = pd.DataFrame()
    for h, i in g:
        data = np.cumsum(boutlength_distribution(i, BINS))
        tempdf = pd.DataFrame({args.parameter:BINS[:-1], 'Frequency':data, 'FlyID':h[3], 'group':'-'.join(h[:-1])}) 
        df = pd.concat([df, tempdf], axis=0)
    
    fig = plt.figure(figsize=(22,17))
    sets = list(sorted(set(df.group)))
    stims = ['No stim', '1x1sec','1x5sec','6x5sec']
    genotypes = ['mAL-LexA','LexAop-TNTE','mAL-TNTE']
    figH = len(set(bouts.g1))
    figW = len(set(bouts.g3))
    for num in range(len(sets)):
        GROUP = sets[num]
        data = df[df['group'] == GROUP]
        m, s, n = group_data(data, args.parameter)
        ax = fig.add_subplot(figH,figW, num+1)
        x_values= m.index.values
        y_values = m.Frequency.values#np.cumsum(m.Frequency.values)
        psems = s.Frequency.values
        nsems = -1.0*psems
        top_errbar = tuple(map(sum, zip(psems, y_values)))
        bottom_errbar = tuple(map(sum, zip(nsems, y_values)))
        p = plt.plot(x_values, y_values, linewidth=3, zorder=100,
                        linestyle = '-',
                        color=colourlist[0],
                        label=(GROUP))# + ', n= ' + str(max_n))) 
        
        q = plt.fill_between(x_values, 
                            top_errbar, 
                            bottom_errbar, 
                            alpha=0.15, 
                            linewidth=0,
                            zorder=90,
                            color=colourlist[0],
                            )         
        
        
        LD50 = (x_values[np.argmax(y_values > 0.5)], y_values[np.argmax(y_values > 0.5)])
        plt.axvline(x=LD50[0], ymin=0.0, ymax = LD50[1], linewidth=2, linestyle='--', color='r')
        plt.annotate(str(round(LD50[0],1)) +' sec', xy=(LD50[0],0), xytext=(1.5*LD50[0], 0.25),arrowprops=dict(arrowstyle="->"))
        LD80 = (x_values[np.argmax(y_values > 0.8)], y_values[np.argmax(y_values > 0.8)])
        plt.axvline(x=LD80[0], ymin=0.0, ymax = LD80[1], linewidth=2, linestyle='--', color='#CC20AA')
        plt.annotate(str(round(LD80[0],1)) +' sec', xy=(LD80[0],0), xytext=(1.5*LD80[0], 0.4),arrowprops=dict(arrowstyle="->"))
        ax.set_xlim((np.amin(x_values),np.amax(x_values)))
        if args.log_bins == True:
            ax.set_xscale("symlog")
        ax.set_ylim(0.0,1.0)#1.3*(np.amax(y_values)))
        #ax.set_xlabel(args.parameter, fontsize=22) 
        # + u"\u00B1" + ' SEM', fontsize=22)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.yaxis.set_ticks_position('left')
        ax.xaxis.set_ticks_position('bottom') 
        if num < figW:
            ax.set_title(stims[num], fontsize=16)
        if num >= (figW*(figH-1)):
            ax.set_xlabel(args.parameter, fontsize=16)
        
        if num % figW == 0:
            ax.set_ylabel(genotypes[num/figW]+'\nCumulative bout terminations', fontsize=16)
    
    plt.savefig(args.boutsfile.split('.pickle')[0] + 'cumulative_terminations.svg')
    plt.savefig(args.boutsfile.split('.pickle')[0] + 'cumulative_terminations.png')
    plt.savefig(args.boutsfile.split('.pickle')[0] + 'cumulative_terminations.pdf')
        
