import pandas as pd
import numpy as np
import scipy
from scipy import stats
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import seaborn as sns

##################################################################################################################################
#         Functions description:

#   LoadData(file_name, sheet_name): load data from Excel file

#   CalculatePrefIndex(data):  calculate preference index and total egg-laying rate per replica

#   stat_PI_different_from0(data2): run statistics to determine if preferences of each pools are significantly different from 0

#   stat_compare_pools(data, dataname, p1, p2): statistics to compare distributions of pref indexes from pools p1 & p2

#   Graph_PI_Swarm(PI, PIpvalues): make bee swarm graph of preference indexes

#   Graph_Dose_Response(data): make graph of egg-laying rate for two pools, with varying conditions (x axis)

#   AnalyzePI(file_name, sheet_name): combines all functions to load, analyze and plot pref index data

#   AnalyzeDoseResponse(file_name, sheet_name): combines functions to load and plot egg-laying rate data



##################################################################################################################################

def LoadData(file_name, sheet_name): # Load data from specified sheet from an Excel file
                # for PI, each pool = 2 columns (left and right egg-laying substrate)
                # for dose response, each pool = 1 column
                
                # first row of table must contain name of pool
                              
    file_path='/Users/cavey/Desktop/' #change this
    
    excel_file_path=file_path+file_name
    data = pd.read_excel(excel_file_path, sheet_name=sheet_name)
    return data


##################################################################################################################################

# calculate Preference Index for each pool (pairs of columns)

def CalculatePrefIndex(data):
    
    nrows=len(data.index)
    ncols=len(data.columns)
    
    # Make dataframes for PI (pref index) and TOT (total egg laying rate on the two sides)
    # all PI pools in 1 column, index gives name of pool for each value
    
    PInamelist=[]
    TOTnamelist=[]
    vallist=[]
    picol=0 #num of PI pair
    
    for col in range(0,ncols,2):
        for i in range(nrows): 
            PInamelist+=[picol] 
            TOTnamelist+=[picol]
            vallist+=[0]                     
        picol+=1
        
    #make dataframe containing PI    
    preftable = pd.DataFrame({'pool': PInamelist, 'PIval':vallist}) #1st column is name of pool for each replica, 
                                                                    #2nd column will contain PI values
    #make dataframe containing TOT egg laying rate         
    tottable = pd.DataFrame({'pool': TOTnamelist, 'TOTval':vallist}) #1st column is name of pool for each replica
                                                                     #2nd column will contain TOTeggs values
    
    # make a dataframe for raw data excluding replicas where toteggs<mineggs => will be used for statistics
    data2=pd.DataFrame(float("nan"), columns = data.columns.values, index=data.index.values)
          
    # set minimum num of eggs required per chamber to calculate PI     
    mineggs = 5 #i.e.: exclude chambers with less than 5 eggs total
        
    j=0
    k=0
    for col in range(0,ncols-1,2):
        k=0 # count for copying data to data2 for each pool
        
        for i in range(nrows):
            # check if row is empty (nan), i.e. when unequal numbers of replica per pools
            m=data.iloc[i,col]
            if m != m: # not equal to itself <=> is a nan
                pref=float("nan")
                toteggs=float("nan")
                
            else:    
                toteggs=data.iloc[i,col]+data.iloc[i,col+1] #calculate total num of eggs for each replica
                
                if int(toteggs)>mineggs:
                    #calculate PI (pref for 1st column):
                    pref=(data.iloc[i,col]-data.iloc[i,col+1])/toteggs
                    
                    # copy data to data2 (for PIpvalues statistics)
                    data2.iloc[k, col]=data.iloc[i,col]
                    data2.iloc[k, col+1]=data.iloc[i,col+1]
                    k+=1
                        
                else: # toteggs is smaller than threshold
                    pref=float("nan")
                                      
            preftable.loc[j,'PIval']=pref
            tottable.loc[j,'TOTval']=toteggs
            j+=1

                
    #remove nans:  
    prefNoNans = preftable[preftable.PIval == preftable.PIval] #keeps only rows where PIval is equal to itself (i.e.: it is not a nan)
    totNoNans = tottable[tottable.TOTval == tottable.TOTval] #idem for TOTval
        
        
    return prefNoNans, totNoNans, data2   
     
##################################################################################################################################  

#                         STATISTICS

##################################################################################################################################  


# Determine if Pref Indexes are significantly different from O for each pool (i.e: significant preference for one of the two sides)
# => compare num eggs on the two sides for each replica :MANN WHITNEY SIGNED RANK TEST (=paired difference test)

def stat_PI_different_from0(data2):    
    
    nrows=len(data2.index)
    ncols=len(data2.columns)
        
    strcounts=data2.count() #list of num of data points for each pool
                
    col=0
    pvals=[]

    for i in range(int(ncols/2)):
        count=strcounts[i*2]
        
        if count>1 : # do stats only if more than 1 replica in the pool
             
            res = scipy.stats.wilcoxon(data2.iloc[0:count-1, col],data2.iloc[0:count-1, col+1]) #calculate statistic
            pvals.append(res[1]) #store pvalue in list pvals
                  
        else:
            #print('no data')
            pvals.append(1)
                
        col+=2
            
    PIpvalues=pd.Series(pvals)
    return PIpvalues

##################################################################################################################################
 
# compare distributions of PI or TOT for pool p1 vs pool p2 (enter pool nums, 0-based)

def stat_compare_pools(data, dataname, p1, p2):
    gp = data.groupby('pool', sort=False) 
    dic=gp.groups
    res = scipy.stats.mannwhitneyu(data.loc[dic[p1], dataname], data.loc[dic[p2], dataname])
    print('Mann Whitney U test for pool'+str(p1)+' vs pool'+str(p2)+': p='+str(res[1]))
    
    

##################################################################################################################################

#                                   GRAPHS

##################################################################################################################################  

# Make Bee Swarm of PI, filled circles for pools with a significant difference, empty circles for pools with non-significant preference

def Graph_PI_Swarm(PI, PIpvalues):
    
    #calculate mean & stdev PI for each pool:
    avg = PI.groupby('pool', sort=False)['PIval'].mean()
    std = PI.groupby('pool', sort=False)['PIval'].std()
 
    #make color palette:
        #user has to enter list of desired colors here:
    poolcolors=['blue', 'red', 'green', 'magenta'] #complete such that list length must be at least = number of pools
    
    if len(poolcolors)<len(avg):
        print('Warning: user needs to enter color list, assigning auto colors')
        poolcolors=['grey']*len(avg)
            
    colors=[] # will contain colors of filled circles
    edgecolorlist=[] # will contain colors of circle outlines
               
    i=0
    for pval in PIpvalues:
        if pval<0.05:    
            colors.append(poolcolors[i]) # filled circles
            edgecolorlist.append(poolcolors[i])
        else:   
            colors.append('white') # white (=empty) circles
            edgecolorlist.append(poolcolors[i])      
        i+=1
    
    #### Make Bee Swarm plot:
            
    # graph settings:
    dotsize=6
    whiskerwidth=0.4
    
    ## draw bee swarm:
    g = sns.swarmplot(x="pool", y="PIval", palette=sns.color_palette(colors), size=dotsize, data=PI, linewidth=0, zorder=1) 
    
    
    #### Draw empty cicles for pools with no preference (PIpvalue >0.05)
    
    #record xy positions of dots to add circle outlines
    artists = g.get_children()
    listoffsets=[]
    offsets = []
    for a in artists:
        if type(a) is matplotlib.collections.PathCollection:
            offsets = a.get_offsets() 
            listoffsets.append(offsets) #store list of xy positions of each dot, for each pool
        
    #draw circle outlines if PIpvalue>0.05
    j=0
    for pvals in PIpvalues:
        xycoords=listoffsets[j]
        if pvals>=0.05:
            for i in xycoords:
                plt.scatter(i[0], i[1], marker='o', color='white', edgecolor=edgecolorlist[j], linewidths=1, s=dotsize*5, zorder=2)    
        j+=1
    
    
    #### add shaded bars for avg and whiskers for std
    for i in range(len(avg)):
        
        #colored bar:
        left, bottom, width, height = (i- whiskerwidth/2, 0,  whiskerwidth, avg[i])
        rect=mpatches.Rectangle((left,bottom),width,height, alpha=0.2, facecolor=poolcolors[i])
        plt.gca().add_patch(rect)
        
        #horizontal line for avg
        plt.plot([i- whiskerwidth/2, i+ whiskerwidth/2], [avg[i],avg[i]], linewidth=2, color='black', zorder=5) 
        
        #whiskers:
        plt.plot([i, i], [avg[i]-std[i],avg[i]+std[i]], linewidth=2, color='black', zorder=3) 
        plt.plot([i- whiskerwidth/3, i+ whiskerwidth/3], [avg[i]+std[i],avg[i]+std[i]], linewidth=2, color='black', zorder=3)
        plt.plot([i- whiskerwidth/3, i+ whiskerwidth/3], [avg[i]-std[i],avg[i]-std[i]], linewidth=2, color='black', zorder=3)
            
                              
    #horizontal zero line: 
    plt.plot([-1, plt.axis()[1]], [0, 0], linewidth=1, color='black')
    g.set(ylabel="Pref Index")
    g.set(xlabel=" ")
    g.set(xticks=[])
    sns.despine(fig=None, ax=None, top=True, right=True, left=False, bottom=True, offset=None, trim=False)
    plt.ylim(-1.1,1.1)
    plt.show()
   
 
 
##################################################################################################################################

#                   EGG-LAYING RATE GRAPH
    
# egg-laying rate dose response graphs, written for two pools only

# x axis coordinates (e.g.: sugar concentration) must be specified in this function
    
    
##################################################################################################################################

def Graph_Dose_Response(data):  

    poolcolors=['blue', 'red']
    
    #make empty dataframe for avg and std:
    df = pd.DataFrame(index=np.arange(0,(len(data.columns)/2)), columns = ['xcoords', 'AVGpool0', 'STDpool0', 'AVGpool1', 'STDpool1'])
    
    #enter list of x axis coordinates (e.g.: concentrations):
    xcoords=[0.0001, 0.001, 0.01, 0.1, 1]
    
    if len(xcoords)!=len(df.index):
        print('Error: incorrect list of x coordinates') 
        
    else:
    
        #calculate avg and std   
        row=0      
        for i in range(0,len(data.columns), 2):
            df.loc[row, 'AVGpool0']=data.iloc[0:,i].mean()
            df.loc[row, 'STDpool0']=data.iloc[0:,i].std()
            row+=1
    
        row=0    
        for i in range(1,len(data.columns), 2):
            df.loc[row, 'AVGpool1']=data.iloc[0:,i].mean() 
            df.loc[row, 'STDpool1']=data.iloc[0:,i].std()
            row+=1
               
        #Make plot
        fig, ax = plt.subplots()
            
        
        x=xcoords
        
        for pool in range(2):
            avgname='AVGpool'+str(pool)
            stdname='STDpool'+str(pool)
            y=df[avgname].to_list()
            err=df[stdname].to_list()
           
            y1=[]
            y2=[]
            for j in range(len(y)): #calculate avg+std and avg-std for filled area
                y1.append(y[j]-err[j])
                y2.append(y[j]+err[j])
                
            #plot avg line and fill area for std:
            ax.plot(xcoords, y, color=poolcolors[pool]) # line
            ax.fill_between(xcoords, y1, y2, facecolor=poolcolors[pool], alpha=0.2) # fill area
                
            #add individual dots:
            ax.scatter(xcoords, y, color=poolcolors[pool]) # dots        
                
        plt.xscale('log', subs=[2, 3, 4, 5, 6, 7, 8, 9])    
        plt.xlabel('%')
        plt.ylabel('number of eggs')
            
        ax.spines['bottom'].set_position(('axes', 0))
                
            
        plt.show()
    
        
 
##################################################################################################################################

#                                      MAIN PROGRAM:
        
##################################################################################################################################
       
def AnalyzePI(file_name, sheet_name):
    
    ###  load data from excel
    data=LoadData(file_name, sheet_name)

    ###  calculate Pref Index
    PI,TOT,data2=CalculatePrefIndex(data) 

    ###  run statistics to determine if preferences are significantly different from 0
    PIpvalues=stat_PI_different_from0(data2) 

    ###  make PI graph
    Graph_PI_Swarm(PI, PIpvalues) 

    ###  compare PI pools by Mann Whitney U test
    p1=0 #enter num of 1st pool
    p2=1 #enter num of 2nd pool
    stat_compare_pools(PI, 'PIval', p1, p2)

##################################################################################################################################
    
def AnalyzeDoseResponse(file_name, sheet_name):
    ###  load data from excel
    data=LoadData(file_name, sheet_name)
    ###  make graph
    Graph_Dose_Response(data)
    
##################################################################################################################################
    
# Examples:
AnalyzePI('Test_Data.xlsx','Sheet1')

AnalyzeDoseResponse('Test_Data.xlsx','Sheet_dose_data')