import sys, os, shutil
import matplotlib.pyplot as plt
from matplotlib import collections
from matplotlib.lines import Line2D
import numpy as np
from scipy import stats
import pmx
import re

def get_significance( exp, predA, predB, n, alpha=0.05, bootnum=1000, bReturnDiff=False ):
    #rmseA = np.sqrt( np.mean(np.power(exp-predA,2.0)))
    #rmseB = np.sqrt( np.mean(np.power(exp-predB,2.0)))
    errorA = np.power(exp-predA,2.0)
    errorB = np.power(exp-predB,2.0)    
    diff = []
    for boot in range(0,bootnum):
        rand = np.random.choice(n, n, replace=True)
        foo_rmseA = np.sqrt( np.mean( errorA[rand] ))
        foo_rmseB = np.sqrt( np.mean( errorB[rand] ))
        foo_diff = foo_rmseA-foo_rmseB
        diff.append(foo_diff)
        
    diff = np.sort(diff)
    pos = np.shape( numpy.where( diff > 0.0 ) )[1]/float(bootnum)
    neg = np.shape( numpy.where( diff < 0.0 ) )[1]/float(bootnum)
    bSignif = False
    # more than 0.975 values are negative
    if pos < alpha/2.0:
        bSignif = True
    # more than 0.975 values are positive
    elif neg < alpha/2.0:
        bSignif = True 
    if bReturnDiff:
        return(bSignif,diff)
    else:
        return(bSignif)


def compare_datasets_arrays( valsArr1,valsArr2,errsArr1=[],errsArr2=[], func='aue', bootnum=1000,
                           keys=[],bScaleToKcal=False,defaultError=0.0):

    # maybe convert to array
    if isinstance(valsArr1, dict) and isinstance(valsArr2, dict):
        valsArr1,errsArr1,valsArr2,errsArr2,names = construct_arrays( valsArr1, valsArr2, keys,defaultError=defaultError )    
    #print(errsArr1)
    
    ######## calc value #######
    if 'aue' in func:
        val = calc_aue( valsArr1, valsArr2 )
    elif 'rmse' in func:
        val = calc_rmse( valsArr1, valsArr2 )
    elif 'pearson' in func:
        val = calc_pearson( valsArr1, valsArr2 )      
    elif 'kendall' in func:
        val = calc_kendall( valsArr1, valsArr2 )                
            
    # bootstrap CI
    err = []
    n = np.shape(valsArr1)[0]
    # errsArr
    if len(errsArr1)==0:
        errsArr1 = np.zeros(n)
    if len(errsArr2)==0:
        errsArr2 = np.zeros(n)
        
    for boot in range(0,bootnum):
        d1 = []
        for r in range(0,n):
            d1.append(np.random.normal(valsArr1[r],errsArr1[r],size=1)[0])       
        d1 = np.array(d1)
        d2 = []
        for r in range(0,n):
            d2.append(np.random.normal(valsArr2[r],errsArr2[r],size=1)[0])   
        d2 = np.array(d2)
        rand = np.random.choice(n, n, replace=True)
        
        fooval = 0.0
        if 'aue' in func:
            fooval = calc_aue( d1[rand], d2[rand] )
        elif 'rmse' in func:
            fooval = calc_rmse( d1[rand], d2[rand] )
        elif 'pearson' in func:
            fooval = calc_pearson( d1[rand], d2[rand] )            
        elif 'kendall' in func:
            fooval = calc_kendall( d1[rand], d2[rand] )                    
        err.append( fooval )
        
    err = np.sort(err)
    low = err[int(bootnum*0.025)]
    high = err[int(bootnum*0.975)]
    err = np.sqrt( np.var(err,ddof=0) )     
    
    scale = 1.0
    if bScaleToKcal==True and (func=='aue' or func=='rmse'):
        scale = 1.0/4.184
    
    return([np.round(val*scale,2),np.round(low*scale,2),np.round(high*scale,2)])   

def calc_aue( arr1, arr2 ):
    val = np.mean(np.abs(arr1-arr2))
    return(val)

def calc_rmse( arr1, arr2 ):
    val = np.sqrt( np.mean( np.power(arr1-arr2,2.0) ) )
    return(val)

def calc_pearson( arr1, arr2 ):
    val = np.corrcoef(arr1,arr2)[0,1]
    return(val)

def calc_kendall( arr1, arr2 ):
    val, p_value = stats.kendalltau(arr1, arr2)
    return(val)

###### plot ########
###########################################################################
##### adapted from https://github.com/andrewcharlesjones/plottify #########
###########################################################################

def autosize(fig=None, figsize=None):

    ## Take current figure if no figure provided
    if fig is None:
        fig = plt.gcf()

    if figsize is None:
        ## Get size of figure
        figsize = fig.get_size_inches()
    else:
        ## Set size of figure
        fig.set_size_inches(figsize)

    scale_figsize = 1.0/fig.get_axes()[0].get_gridspec().ncols

    ## Set sizes of objects
    ##      -- Note that this function currently adjusts sizes only
    ##         according to the horizontal width fo the figure.

    ## Font sizes
    fontsize_labels = figsize[0] * scale_figsize * 3.5
    fontsize_ticks = figsize[0] * scale_figsize * 2.75
    
    ## text size
   # plt.rcParams['font.size'] = int(fontsize_labels)
    #print(fontsize_labels)

    ## Scatter point size
    scatter_size = ( figsize[0]*scale_figsize * 0.95) ** 2.5

    ## Line width
    linewidth = figsize[0]*scale_figsize*0.1

    ## Spine width
    spine_width = 0.10 * figsize[0]*scale_figsize

    ## Tick length and width
    tick_length = spine_width*2
    tick_width = spine_width

    ## Change all of these sizes
    axes = fig.get_axes()
    for ax in axes:

        ## Set label font sizes
        for item in [ax.title, ax.xaxis.label, ax.yaxis.label]:
            item.set_fontsize(fontsize_labels)

        ## Set tick font sizes
        for item in ax.get_xticklabels() + ax.get_yticklabels():
            item.set_fontsize(fontsize_ticks)

        ## Set line widths
        plot_objs = [child for child in ax.get_children() if isinstance(child, Line2D)]
        for plot_obj in plot_objs:
            plot_obj.set_linewidth(linewidth)

        ## Set scatter point sizes
        plot_objs = [
            child
            for child in ax.get_children()
            if isinstance(child, collections.PathCollection)
        ]
        for plot_obj in plot_objs:
            plot_obj.set_sizes([scatter_size])

        ## Set spine widths
        for spine in ax.spines.values():
            spine.set_linewidth(spine_width)

        ## Set tick width and length
        ax.tick_params(width=tick_width, length=tick_length)

    ## Set tight layout
    plt.tight_layout()


def construct_arrays( valX, valY, keys=[], defaultError=0.0 ):

    if len(keys)==0:
        keys= list(valX.keys())
    
    vals1 = []
    vals2 = []
    errs1 = []
    errs2 = []
    names = []
    for key in keys:
        if key not in valY.keys():
            continue
        
        names.append(key)
        dataX = valX[key]
        dataY = valY[key]
        
        if np.isnan(dataX[0]) or (np.isnan(dataY[0])):
            continue

        val1 = 0.0
        err1 = defaultError

        if hasattr(dataX, "__len__"):
            val1 = dataX[0]
            if len(dataX)>1:
                if not np.isnan(dataX[-1]):
                    if dataX[-1]>0.0:
                        err1 = dataX[-1]
        else:
            val1 = dataX

        val2 = 0.0
        err2 = defaultError
        if hasattr(dataY, "__len__"):
            val2 = dataY[0]
            if len(dataY)>1:
                if not np.isnan(dataY[-1]):
                    if dataY[-1]>0.0:
                        err2 = dataY[-1]
        else:
            val2 = dataY

        vals1.append(val1)
        vals2.append(val2)
        errs1.append(err1)
        errs2.append(err2)
        
    vals1 = np.array(vals1)
    vals2 = np.array(vals2)  
    errs1 = np.array(errs1)
    errs2 = np.array(errs2)  
    
    return(vals1,errs1,vals2,errs2,names)
    
    
def plot_scatter( fig, ax, valX, valY, keys=[],
                 bAue=False, bRmse=False, bPearson=False, bKendall=False, defaultError=0.0,
                 title='', bLabels=False,
                 minval=-999.99, maxval=999.99,
                 kcal1=False, bX=False, bY=False,
                 bDG=True,bkJ=True,bScaleToKcal=False,bTextBottomRight=True,
                 circleSize=50, symbol='s', case=''):
    
    cm = plt.get_cmap('coolwarm')
    scale = 1.0    
    if bScaleToKcal==True:
        bkJ = False
        scale = 1.0/4.184
    
    if len(title)>0:
        plt.title(title,y=1.005)

    #######################################
    ##### scatter plot the data ###########
    #######################################
    vals1 = []
    vals2 = []
    vals1,errs1,vals2,errs2,names = construct_arrays( valX, valY, keys=keys, defaultError=defaultError )
    vals1 *= scale
    errs1 *= scale
    vals2 *= scale
    errs2 *= scale
    
    clr = np.abs(vals1-vals2)
    clr = cm(clr/10.0/scale)
    plt.scatter(vals1,vals2,color=clr,s=50,marker=symbol,zorder=2,alpha=0.8)
    plt.errorbar(vals1,vals2,xerr=errs1,yerr=errs2,fmt='none',ecolor='gray',
                 zorder=1,mew=1, capsize=1, elinewidth=1)
    numpoints = np.shape(vals1)[0]

    ######################################
    ###### ranges ########################
    ######################################
    if minval==-999.99:
        minval = np.min( [vals1,vals2] ) - 0.2*np.abs(np.min( [vals1,vals2] ))
    if maxval==999.99:
        maxval = np.max( [vals1,vals2] ) + 0.2*np.abs(np.max( [vals1,vals2] ))
    midval = maxval - (maxval-minval)*0.5
    deltaval = (maxval-minval)*0.1
    plt.xlim(minval,maxval)
    plt.ylim(minval,maxval) 
    
    ######################################
    ###### axes ########################
    ######################################
    ax.set_axisbelow(True)
    # yaxis
    ax.yaxis.set_ticks_position('left')
    #y = [-20,-10,0,10,20]    
#    plt.yticks(y,y,fontsize=16)
    if bY==False:
        ax.yaxis.set_ticklabels([])
    else:
        ylabeltext = r'$\Delta\Delta$G$_{calc}$'
        if bDG==True:
            ylabeltext = r'$\Delta$G$_{calc}$'            
        if bkJ==True:
            ylabeltext = '{0}, kJ/mol'.format(ylabeltext)
        else:
            ylabeltext = '{0}, kcal/mol'.format(ylabeltext)
        plt.ylabel(ylabeltext,rotation=90)            
    # xaxis
    ax.xaxis.set_ticks_position('bottom')
    #ax.tick_params(axis='x',labelbottom='off')
    #x = [-20,-10,0,10,20]    
#    plt.xticks(x,x,fontsize=16)
    if bX==False:
        ax.xaxis.set_ticklabels([])
    else:
        xlabeltext = r'$\Delta\Delta$G$_{exp}$'
        if bDG==True:
            xlabeltext = r'$\Delta$G$_{exp}$'            
        if bkJ==True:
            xlabeltext = '{0}, kJ/mol'.format(xlabeltext)
        else:
            xlabeltext = '{0}, kcal/mol'.format(xlabeltext)
        plt.xlabel(xlabeltext,rotation=0)      
    ax.xaxis.labelpad = 0    
    
    #########################################
    ######## lines ##########################
    #########################################    
    # 1kcal range
    xkcal = [minval,maxval]#[-100,100]
    ykcal1 = [minval-4.184*scale,maxval-4.184*scale]
    ykcal2 = [minval+4.184*scale,maxval+4.184*scale]
    plt.fill_between(xkcal,ykcal1,ykcal2,color='gray',alpha=0.1,zorder=1)
    
    # diagonal
    xx = [minval,maxval]#[-100,100]
    plt.plot(xx,xx, '--',color='black',alpha=0.35,linewidth=3.0,zorder=1)
    
    # regression
    m, b = np.polyfit(vals1,vals2,1)
    plt.plot(xx, m*np.asarray(xx) + b, '-',color='black',linewidth=1.0,alpha=0.5)    
    

    ##################################
    ####### text #####################
    ##################################    
    scale_figsize = 1.0/ax.get_gridspec().ncols
    fontsize = int(fig.get_size_inches()[0]*scale_figsize*3.0)
    kcal = 1.0#/4.184
    i=0
    if bAue==True:
        aue = compare_datasets_arrays( valsArr1=vals1,valsArr2=vals2,errsArr1=errs1,errsArr2=errs2,func='aue' )
        string1 = r'AUE = ${0}_{{{1}}}^{{{2}}}$'.format(np.round(aue[0]*kcal,1),np.round(aue[1]*kcal,1),np.round(aue[2]*kcal,1))
        t = plt.text(minval+0.5*deltaval, maxval-(0.5+0.75*i)*deltaval, string1,fontsize=fontsize)
        i+=1
    if bRmse==True:
        rmse = compare_datasets_arrays( valsArr1=vals1,valsArr2=vals2,errsArr1=errs1,errsArr2=errs2,func='rmse' )
        string2 = r'RMSE = ${0}_{{{1}}}^{{{2}}}$'.format(np.round(rmse[0]*kcal,1),np.round(rmse[1]*kcal,1),np.round(rmse[2]*kcal,1))
        t = plt.text(minval+0.5*deltaval, maxval-(0.5+0.75*i)*deltaval, string2,fontsize=fontsize)
        i+=1        
    if bPearson==True:
        cor = compare_datasets_arrays( valsArr1=vals1,valsArr2=vals2,errsArr1=errs1,errsArr2=errs2,func='pearson' )
        string3 = r'$\rho$ = ${0}_{{{1}}}^{{{2}}}$'.format(np.round(cor[0],2),np.round(cor[1],2),np.round(cor[2],2))
        t = plt.text(minval+0.5*deltaval, maxval-(0.5+0.75*i)*deltaval, string3,fontsize=fontsize)     
        i+=1        
    if bKendall==True:
        tau = compare_datasets_arrays( valsArr1=vals1,valsArr2=vals2,errsArr1=errs1,errsArr2=errs2,func='kendall' )
        string3 = r'$\tau$ = ${0}_{{{1}}}^{{{2}}}$'.format(np.round(tau[0],2),np.round(tau[1],2),np.round(tau[2],2))
        t = plt.text(minval+0.5*deltaval, maxval-(0.5+0.75*i)*deltaval, string3,fontsize=fontsize)   
        i+=1        
    
    # values
    stringVal = r'Values: '+str(numpoints)
    # regression 
    stringR = r'y = '+str(round(m,2))+'x'+' + '+str(round(b,2))
    if bTextBottomRight==True:
        t = plt.text(midval+2*deltaval, midval-4*deltaval, stringVal,fontsize=fontsize)    
        t = plt.text(midval+0.5*deltaval, midval-4.5*deltaval, stringR,fontsize=fontsize)
    else:
        t = plt.text(midval+2*deltaval, midval+4*deltaval, stringVal,fontsize=fontsize)    
        t = plt.text(midval+0.5*deltaval, midval+4.5*deltaval, stringR,fontsize=fontsize)

    ##################################
    ####### datapoint labels #####################
    ##################################    
    fontsize = int(fig.get_size_inches()[0]*0.6)
    if bLabels==True:
        for v1,v2,n in zip(vals1,vals2,names):
            t = n
            plt.text(v1,v2,t,fontsize=fontsize)

