#!/usr/bin/env python

import numpy as np
import re
#import pudb
import matplotlib.gridspec as gridspec
from matplotlib import rc
from collections import OrderedDict
import argparse
import string
import ast
import os

#import scipy.optimize
from scipy.odr import ODR, Model, Data, RealData
from scipy.optimize import curve_fit
from scipy.integrate import quad

#############################################################################################
# function to save a figure
#############################################################################################

def save(path, ext='png', close=True, verbose=True):
    """Save a figure from pyplot.
    Parameters
    ----------
    path : string
        The path (and filename, without the extension) to save the
        figure to.
    ext : string (default='png')
        The file extension. This must be supported by the active
        matplotlib backend (see matplotlib.backends module).  Most
        backends support 'png', 'pdf', 'ps', 'eps', and 'svg'.
    close : boolean (default=True)
        Whether to close the figure after saving.  If you want to save
        the figure multiple times (e.g., to multiple formats), you
        should NOT close it in between saves or you will have to
        re-plot it.
    verbose : boolean (default=True)
        Whether to print information about when and where the image
        has been saved.
    """
    
    # Extract the directory and filename from the given path
    directory = os.path.split(path)[0]
    filename = "%s.%s" % (os.path.split(path)[1], ext)
    if directory == '':
        directory = '.'

    # If the directory does not exist, create it
    if not os.path.exists(directory):
        os.makedirs(directory)

    # The final path to save to
    savepath = os.path.join(directory, filename)

    if verbose:
        print(("Saving figure to '%s'..." % savepath), end=' ')

    # Actually save the figure
    plt.savefig(savepath)
    
    # Close it
    if close:
        plt.close()

    if verbose:
            print("Done")



##################################################################################################################################
# function extracting data
##################################################################################################################################

# routine that extracts all data from a top file
def extractData(fname,ynorm=1.0) :
	with open(fname) as f :
		lines = f.readlines()
	matching = False
	observable = None
	data = OrderedDict()
	for line in lines :
		if matching and re.match("^$",line) != None :
			matching = False
			data[observable] = np.array(data[observable])
		if matching :
			# collapse blanks
			line=' '.join(line.split())
			line=re.sub("D","E",line)
			try :
				datarow = list(map(float,line.split(' ')))
			except ValueError as e :
				print("\t{}".format(e))
				print("\t{}".format(line))
			try:
				data[observable].append([datarow[0], datarow[1], ynorm*datarow[2], ynorm*datarow[3]])
			except:
                        	continue
		match = re.match("^# ([^ ]*).*$",line)
		if match != None :
			matching = True
			observable = (line.split(" index",1)[0]).replace("# ", "")
			data[observable]=[]
	return data

def extractDataSlim(fname,pattern,ynorm=1.0) :
	return filterData(extractData(fname, ynorm), pattern)

##################################################################################################################################
# function filtering data, will keep only keys matching a pattern
##################################################################################################################################
def filterData(data, pattern) :
	compiledPattern = re.compile(pattern)
	keysToKeep = []
	for key in list(data.keys()) :
		if compiledPattern.match(key) != None :
			keysToKeep += [key]
	return {key: data[key] for key in keysToKeep}

##################################################################################################################################
# function to handle data
##################################################################################################################################

# routine that rebin data extract from top file
def rebinDataWeightedAverage(rr,newsize=5) :
    x0, x1, ym, errym = [], [], [], []
    for i in range(len(rr)/newsize):
        x0=np.append(x0, [rr[i*newsize, 0]])
        x1=np.append(x1, [rr[(i+1)*newsize-1, 1]])
        y=0
        erry=0
        for j in range(newsize):
            if rr[j+i*newsize,2]!=0:
                y=y+rr[j+i*newsize,2]/rr[j+i*newsize,3]**2
                erry=erry+1/rr[j+i*newsize,3]**2
        if y!=0:
            y=y/erry
            erry=np.sqrt(1/erry)
        ym=np.append(ym, [y])
        errym=np.append(errym, [erry])     
    r=np.c_[x0,x1,ym, errym]
    return r

def rebinData(rr,newsize=5) :
    x0, x1, ym, errym = [], [], [], []
    for i in range(int(len(rr)/newsize)):
        x0=np.append(x0, [rr[i*newsize, 0]])
        x1=np.append(x1, [rr[(i+1)*newsize-1, 1]])
        y=0
        erry=0
        for j in range(newsize):
            y=y+rr[j+i*newsize,2]
            erry=erry+rr[j+i*newsize,3]**2
        y=y/newsize
        erry=np.sqrt(erry)/newsize
        ym=np.append(ym, [y])
        errym=np.append(errym, [erry])     
    r=np.c_[x0,x1,ym, errym]
    return r


# routine that attaches 2 data samples
def attachData(data1, data2, xmax=2600) :
    x0, x1, ym, errym = [], [], [], []
    for i in range(len(data1)):
        x0=np.append(x0, data1[i,0])
        x1=np.append(x1, data1[i,1])
        ym=np.append(ym, data1[i,2])
        errym=np.append(errym, data1[i,3])
    for i in range(len(data2)):
        if(data2[i,0]<xmax):
            x0=np.append(x0, data2[i,0])
            x1=np.append(x1, data2[i,1])
            ym=np.append(ym, data2[i,2])
            errym=np.append(errym, data2[i,3])  
    r=np.c_[x0,x1,ym, errym]
    return r
                             
    
def smearAllData(dataIn,resolution):
    dataOut=[]
    deltaBin=-dataIn[0,0]+dataIn[0,1]
    area1=0
    area2=0
    numBin=int(round(resolution/deltaBin))
    x1=[]
    x2=[]
    yArray=[]
    yErrArray=[]
    for i in range(len(dataIn)):
        area1=area1+dataIn[i,2]
        x=(dataIn[i,0]+dataIn[i,1])/2.
        x1=np.append(x1, dataIn[i,0])
        x2=np.append(x2, dataIn[i,1])
        y=0
        erry=0
        for j in range(len(dataIn)):
            #if(j>=0)and(j<len(dataIn)):
            xx=(dataIn[j,0]+dataIn[j,1])/2.
            yy=dataIn[j,2]
            erryy=dataIn[j,3]
            gauss=np.exp(-((x-xx)/resolution)**2/2.)
            y=y+yy*gauss
            erry=erry+(erryy*gauss)**2
        erry=erry**0.5/len(dataIn)
        area2=area2+y
        yArray=np.append(yArray,y)
        yErrArray=np.append(yErrArray,erry)
    yArray=yArray/area2*area1
    yErrArray=yErrArray/area2*area1
    return np.c_[x1,x2,yArray,yErrArray]


def getMeanInRange(dataIn,xmin,xmax):
    deltaBin=-dataIn[0,0]+dataIn[0,1]
    area=0
    mean=0
    errorMean=0
    for i in range(len(dataIn)):
        x=(dataIn[i,0]+dataIn[i,1])/2.
        if(x>=xmin)and(x<=xmax):
            mean=mean+dataIn[i,2]*x
            errorMean=errorMean+(dataIn[i,3]*x)**2
            area=area+dataIn[i,2]
    mean=mean/area
    errorMean=(errorMean)**0.5/area
    return [mean, errorMean]

def getMaxBin(dataIn,size):
    area=0
    errorArea=0
    x=0
    for i in range(len(dataIn)):
        xi=(dataIn[i,0]+dataIn[i,1])/2.
        areai=0
        for j in range(len(dataIn)):
          xj=(dataIn[j,0]+dataIn[j,1])/2.
          if(abs(xi-xj)<size/2):
              areai=areai+dataIn[j,2]
        if(areai>area):
            area=areai
            x=xi
    return x
    
            
def truncate(f, n):
    '''Truncates/pads a float f to n decimal places without rounding'''
    s = '{}'.format(f)
    if 'e' in s or 'E' in s:
        return '{0:.{1}f}'.format(f, n)
    i, p, d = s.partition('.')
    return '.'.join([i, (d+'0'*n)[:n]])
          
def rebinDataXmin(rr,xmin=0,newsize=2) :
    x0, x1, ym, errym = [], [], [], []
    count=0
    for j in range(len(rr)):
        if(rr[j,1]<=xmin):
            x0=np.append(x0, rr[j,0])
            x1=np.append(x1, rr[j,1])
            ym=np.append(ym, rr[j,2])
            errym=np.append(errym, rr[j,3])
        else:
            count=count+1
            if(count==1):
                x0=np.append(x0, rr[j,0])
                x1=np.append(x1, rr[j,1])
                ym=np.append(ym, rr[j,2]*(rr[j,1]-rr[j,0]))
                errym=np.append(errym, rr[j,3]**2*(rr[j,1]-rr[j,0]))
            else:
                x1[-1] = rr[j,1] #update xmax
                ym[-1]+=rr[j,2]*(rr[j,1]-rr[j,0])
                errym[-1]+=rr[j,3]**2*(rr[j,1]-rr[j,0])**2
            if(count==newsize or j==len(rr)-1):
                ym[-1]/=(x1[-1]-x0[-1])
                errym[-1]=np.sqrt(errym[-1])/(x1[-1]-x0[-1])
                if(count==newsize): count=0
    #merge last 2 bins
    if(count>0):
        ym[-2]=(ym[-2]*(x1[-2]-x0[-2])+ym[-1]*(x1[-1]-x0[-1]))/(x1[-1]-x0[-2])
        errym[-2]=np.sqrt((errym[-2]*(x1[-2]-x0[-2]))**2+(errym[-1]*(x1[-1]-x0[-1]))**2)/(x1[-1]-x0[-2])
        x1[-2]=x1[-1]        
        x0=x0[:-1]
        x1=x1[:-1]
        ym=ym[:-1]
        errym=errym[:-1]
        
               
    r=np.c_[x0,x1,ym, errym]
    return r

def rebinDatav2(x_old,y_old,ex_old,ey_old, newsize=5):
    x0_old = x_old-ex_old
    x1_old = x_old+ex_old
    x0, x1, ym, errym = np.array([]), np.array([]), np.array([]), np.array([])
    for i in range(int(len(x0_old)/newsize)):
        x0=np.append(x0, [x0_old[i*newsize]])
        x1=np.append(x1, [x1_old[(i+1)*newsize-1]])
        y=0
        erry=0
        for j in range(newsize):
            y=y+y_old[j+i*newsize]
            erry=erry+ey_old[j+i*newsize]**2
        y=y/newsize
        erry=np.sqrt(erry)/newsize
        ym=np.append(ym, [y])
        errym=np.append(errym, [erry])     
    x    = (x1+x0)/2
    errx = (x1-x0)/2
    return x, ym, errx, errym
         
# def rebinDataXmin(rr,xmin=0,newsize=2) :
#     x0, x1, ym, errym = [], [], [], []
#     count=0
#     for j in range(len(rr)-1,-1,-1):
#         if(rr[j,0] >= xmin):
#             count = count +1
#             if(count==1):
#                 x0=np.append(rr[j,0],x0)
#                 x1=np.append(rr[j,1],x1)
#                 ym=np.append(rr[j,2]*(rr[j,1]-rr[j,0]),ym)
#                 errym=np.append(rr[j,3]**2*(rr[j,1]-rr[j,0]), errym)
#             else:
#                 x0[0] = rr[j,0]
#                 ym[0] = ym[0] + rr[j,2]*(rr[j,1]-rr[j,0])
#                 errym[0] = errym[0] + rr[j,3]**2*(rr[j,1]-rr[j,0])**2
#             if(count==newsize or j==0):
#                 ym[0] = ym[0]/(x1[0]-x0[0])
#                 errym[0] = np.sqrt(errym[0])/(x1[0]-x0[0])
#                 count =0
#         else:
#             if(count >0):
#                 ym[0] = ym[0]/(x1[0]-x0[0])
#                 errym[0] = np.sqrt(errym[0])/(x1[0]-x0[0])
#                 count =0
#             else:
#                 x0=np.append(rr[j,0],x0)
#                 x1=np.append(rr[j,1],x1)
#                 ym=np.append(rr[j,2],ym)
#                 errym=np.append(rr[j,3], errym)
#     r=np.c_[x0,x1,ym, errym]
#     return r