#!/usr/bin/python3

import sys
import argparse
from collections import deque
import itertools
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.colors import colorConverter
from scipy import optimize
from scipy.interpolate import griddata

fit_pol0 = lambda p, x: p[0]
fit_pol1 = lambda p, x: p[0]+p[1]*x
fit_pol2 = lambda p, x: p[0]+(p[1]+p[2]*x)*x
fit_pol3 = lambda p, x: p[0]+(p[1]+(p[2]+p[3]*x)*x)*x
fit_pol4 = lambda p, x: p[0]+(p[1]+(p[2]+(p[3]+p[4]*x)*x)*x)*x

fit_gauss    = lambda p, x: p[0]*np.exp( -((x-p[1])**2.0)/(2.0*p[2]**2.0) )
fit_gauss_p0 = lambda p, x: p[0]*np.exp( -((x-p[1])**2.0)/(2.0*p[2]**2.0) ) + p[3]

def fit(function,params,xd,yd,ylog):
    errfunc  = lambda p, x, y: (function(p, x) - y)
    terrfunc = lambda p, x, y: (function(p, x) - y)
    if ylog:
        terrfunc = lambda p, x, y: (pow(10,function(p, x)) - pow(10,y))
    # fit using xd,yd
    p1,succes=optimize.leastsq(errfunc,params,args=(xd,yd))
    # get error of fit
    errval = sum(errfunc(p1,xd,yd)**2)/(len(xd)-1)
    # get error of untransformed estimates
    terrval = sum(terrfunc(p1,xd,yd)**2)/(len(xd)-1)
        
    return (errval,terrval,p1)

def bin_error(x, y, nbins):
    n, _ = np.histogram(x, bins=nbins)
    sy, _ = np.histogram(x, bins=nbins, weights=y)
    sy2, _ = np.histogram(x, bins=nbins, weights=y*y)
    mean = sy / n
    std = np.sqrt(sy2/n - mean*mean)
    epos = (_[1:] + _[:-1])/2
    return mean,std,epos
        
def bin_by(x, y, nbins=30):
    """
    Bin x by y.
    Returns the binned "x" values and the left edges of the bins
    """
    bins = np.linspace(y.min(), y.max(), nbins+1)
    # To avoid extra bin for the max value
    bins[-1] += 1 

    indicies = np.digitize(y, bins)

    output = []
    for i in range(1, len(bins)):
        output.append(x[indicies==i])

    # Just return the left edges of the bins
    d=np.diff(bins)
    bins = bins[:-1]
    midp = bins+0.5*d

    return output, midp

def pastel(colour, weight=2.4):
    """ Convert colour into a nice pastel shade"""
    rgb = np.asarray(colorConverter.to_rgb(colour))
    # scale colour
    maxc = max(rgb)
    if maxc < 1.0 and maxc > 0:
        # scale colour
        scale = 1.0 / maxc
        rgb = rgb * scale
    # now decrease saturation
    total = sum(rgb)
    slack = 0
    for x in rgb:
        slack += 1.0 - x

    # want to increase weight from total to weight
    # pick x s.t.  slack * x == weight - total
    # x = (weight - total) / slack
    x = (weight - total) / slack

    rgb = [c + (x * (1.0-c)) for c in rgb]

    return rgb


def get_colours2(n):
    """ Return n pastel colours. """
    base = np.asarray([[1,0,0], [0,1,0], [0,0,1]])

    if n <= 3:
        return base[0:n]

    # how many new colours to we need to insert between
    # red and green and between green and blue?
    needed = (((n - 3) + 1) / 2, 1+(n - 3) / 2)

    colours = []
    for start in (0, 1):
        for x in np.linspace(0, 1, needed[start]+2.5):
            capp=(base[start] * (1.0 - x)) + (base[start+1] * x)
            colours.append((base[start] * (1.0 - x)) +
                           (base[start+1] * x))

##    return [pastel(c,2.4) for c in colours[0:n]]
    print( colours)
    return [c for c in colours[0:n]]

def evaluateCondition( condition, codeVars, codeCond, vardict, variables, elems ):
    # first set variables
    if codeVars is None:
        vardict={}
        sourceVars=''
        for x in variables: #range(0,len(variables)):
            sourceVars = sourceVars + x+'=elems[\''+x+'\'];'
            vardict[x]=0
        codeVars = compile(sourceVars, '<string>', 'exec')
    
    vardict['condValue']=False
    vardict['elems']=elems
    exec(codeVars,None,vardict)
    
    # now check condition
    if codeCond is None:
        codeCond = compile('condValue=('+condition+')', '<string>', 'exec')
    exec(codeCond,None,vardict)
    return vardict['condValue'],vardict,codeVars,codeCond 

    
def getVariables(filename,usehdr,sep,prefixtag,starttag):
    '''
    Input:
     filename          column wise text file, lines with '#' are interpreted as comments and ignored
    '''
    # header list separator
    if sep=='': sep=None

    
    # header
    if starttag=='': starttag=None
    hdrfound = False
    variables = []
    foundstart = False # flag if start tag is found
    searchtag = True   # flag if start tag is given
    hasprefix = prefixtag is not None
    if starttag is None:
        foundstart = True
        searchtag = False
    
    with open(filename) as fh:
        for line in fh:
            line = line.strip().rstrip('\n')
            if len(line)==0: continue
            # if start tag is not found (only the case if start tag is actually given)
            if not foundstart:
                if line.startswith(starttag):
                    foundstart = True;
                    if usehdr:
                        hdrfound = True
                        variables = line.split()[1:]
                continue
            # get header if requested
            if usehdr and line[0]=='#' and not hdrfound:
                hdrfound=True
                if len(line)>1:
                    variables=line[1:].split();

            # skip commented or empty lines
            if len(line)==0 or line[0] in '#':
                continue

            if hasprefix:
                if line.startswith(prefixtag):
                    line = line.replace(prefixtag,'',1)
                else:
                    continue

            # split line in elements
            elems = line.split(sep)

            # found bad header
            if not len(elems) == len(variables):
                #print('len elems!=len vars');
                #print(len(elems))
                #print(len(variables))
                variables = [ 'V'+repr(x) for x in range(0,len(elems)) ]
                hdrfound = True

            if hdrfound:
                break


    return variables

def convertTypes(variables,elems):
    data={}
    v=''
    ind=0
    for x in elems:
        try:
            v=int(x)
        except ValueError:
            try:
                v=float(x)
            except ValueError:
                v=x
        data[variables[ind]] = v
        ind = ind+1
    return data

def convertToVar(ref,variables):
    if ref is None:
        return None
    icol=0
    try:
        icol = int(ref)
    except ValueError:
        # it's not an int
        for v in variables:
            if ref==v:
                return ref
        print('error: unknown variable \'{}\''.format(ref))
        print('\nvalid variables are:')
        print(','.join(variables))
        return None
    
    # it's an int
    if icol<len(variables):
        return variables[icol]
    return 'V'+repr(icol)

def read_raw_data(filename,prefixtag,starttag,nlimit,variables,sep,hdrX,hdrY,hdrZ,colX,colY,colZ,colD,islblX,condition=None):
    '''
    Input:
     filename          column wise text file, lines with '#' are interpreted as comments and ignored
     gethdr            first commented line is assumed to contain a list of variable names
     sep               separator in the above list
     colX,colY,colD    columns containing the X and Y values (start at '0') as wells as the uncertainty (colD)
     islblX            if true, the X values are regarded as string labels
    '''
    # header list separator
    if sep=='': sep=None

    # return values
    valX=None
    valY=None
    valZ=None
    valD=None
    lblX=None

    if variables is None:
        variables = []

    hasprefix = prefixtag is not None
    foundstart = False
    if starttag is None:
        foundstart = True
    
    # if X is a label, disable logX
    if islblX:
        lblX=[]
        logX=False

    codeCond = None # byte compiled code as provided by condition
    codeVars = None
    nline=0
    ndata=0
    varsChecked = False
    varsOK = False;
    vardict=None
    nskipped=0
    nx=0
    with open(filename) as fh:
        for line in fh:
            nline=nline+1
            line = line.strip().rstrip('\n')
            if len(line)==0: continue

            # if start tag is not found (only the case if start tag is actually given)
            if not foundstart:
                if line.startswith(starttag):
                    foundstart = True;
                continue

            
            # skip commented or empty lines
            if len(line)==0 or line[0] in '#':
                continue

            if hasprefix:
                if line.startswith(prefixtag):
                    line = line.replace(prefixtag,'',1)
                else:
                    continue

            # split line in elements
            elems = line.split(sep)

            # bad avriable list
            if not varsChecked:
                varsOK = (len(elems) == len(variables))
                varsChecked = True

            data = convertTypes(variables,elems)

            # check condition
            if condition is not None:
                cond,vardict,codeVars,codeCond = evaluateCondition( condition, codeVars, codeCond, vardict, variables, data )
                if not cond:
                    nskipped=nskipped+1
                    continue

            # get X value
            if colX is not None:
                try:
                    lbl=None
                    if lblX is not None:
                        X=nx
                        nx = nx+1
                        lblX.append(data[colX])
                    else:
                        X=data[colX]
                
                    if valX is None:
                        valX=[]
                    valX.append(X)
                except:
                    print('error reading line {} - check column X {}'.format(nline,colX),file=sys.stderr)

            # get Y value(s)
            if colY is not None:
                try:
                    if valY is None:
                        valY=[]
                    valY.append([])
                    for y in colY:
                        valY[-1].append(data[y])
                except:
                    print('error reading line {} - check column Y {}'.format(nline,colY),file=sys.stderr)

            # get Z value
            if colZ is not None:
                try:
                    if valZ is None:
                        valZ=[]
                    valZ.append(data[colZ])
                except:
                    print('error reading line {} - check column Z {}'.format(nline,colZ),file=sys.stderr)

            # get dY value(s)
            if colD is not None:
                try:
                    if valD is None:
                        valD=[]
                    valD.append([])
                    for dy in colD:
                        valD[-1].append(data[dy])
                except:
                    print('error reading line {} - check column D {}'.format(nline,colD),file=sys.stderr)

            ndata = ndata+1
            if nlimit>0 and ndata>=nlimit:
                break;

        if hdrX is None:
            hdrX = 'X'
            if colX is not None:
                hdrX = colX

        if hdrY is None or len(hdrY)==0:
            hdrY = 'Y'
            if colY is not None:
                hdrY = []
                for y in colY:
                    hdrY.append(y)

        if hdrZ is None:
            hdrZ = 'Z'
            if colZ is not None:
                hdrZ = colZ

        
    if condition is not None and nskipped>0:
        print('info: skipped {} lines due to condition {}'.format(nskipped,condition))

    return (hdrX,hdrY,hdrZ,lblX,valX,valY,valZ,valD)

def transform(vmin,vmax,shiftx,dolog,values):
    # make np array
    tv = np.array(values)

    # shift values
    if shiftx is not None:
        tv = tv+shiftx

    # log values
    if dolog:
        try:
            tv = np.log10(tv)
        except:
            print('error: array contains non-positivce values in log mode')
            return None

    # check range
    if vmin is None:
        vmin=np.min(tv)
    if vmax is None:
        vmax=np.max(tv)
    selection = np.logical_and(tv>=vmin,tv<=vmax)
    
    return (vmin,vmax,selection,tv)

def process(plottype,nbins,normed,xmin,xmax,logx,logy,logz,shiftx,dograd,lblX,valX,valY,valZ,valD):
    '''
    post-process data
    if dohist:
       ignore valX, create histogram of valY's and generate X array on the union
    else:
       if no valX:
          generate valX as [0..1]
       else:
          keep valX

    valX = valX+shiftx
    if logx: valX = log10(valX)
    if logy: valY = log10(valY), valD = 1/valY/ln(10)
    filter on xmin,xmax

    valY = [ [v11,v12,...],[v21,v22,...],... ]
    Npoints = len(valY)
    Nvars = len(valY[0])
    '''

    dohist = (plottype=='HIST')
    doheat = (plottype=='HEAT')
    
    Nvariables = len(valY[0])
    Npoints = len(valY)

    # these two should be filled below
    X = None
    Y = None
    Z = None
    D = None
    H = None # returns the bins used
    Yorg = None # set to the Y values as given
    select = None
    # first transform X and Y
    # --- first case -> create a histogram
    if dohist:
        # set range and transform Y - this will be 
        result=transform(xmin,xmax,shiftx,logx,valY)
        if result is None:
            return None
        xmin=result[0]
        xmax=result[1]
        select=result[2]
        tY=result[3]
        # this will flatten the 2d array
        selY=tY[select]
        
        # for now ignore errors (D)
        # later add D = sqrt(H(bin))
        D=None
        
        # first get overall, sets the range
        # hoverall is count per bin
        # hbins : edges in X
        hoverall,hbins=np.histogram(selY,bins=nbins,density=normed)
        H=hbins
        
        # if logY
        if logy:
            # logy -> log(count)
            try:
                hoverall=np.log10(hoverall)
            except:
                print('error: overall histogram contains empty bins in logY mode')
                return None
        
        # create X array
        X = hbins[:-1]+0.5*np.diff(hbins)

        # this may change here
        Npoints=len(X)
        
        # get Y values
        Y=[]
        Yorg=[]
        if Nvariables>1:
            # then get each sub-graph
            for ind in range(Nvariables):
                y0 = tY[:,ind]
                histY,hbinsy = np.histogram(tY[:,ind],bins=hbins,range=(xmin,xmax))
                if logy:
                    histY=np.log10(histY)
                    y0 = np.log10(y0)
                
                if len(histY)==Npoints:
                    Y.append(histY)
                    Yorg.append(y0)
                else:
                    print('error: X and Y must contain the same nr of elements (nX = {}, nY = {})'.format(Npoints,len(histY)))
                    return None
        else:
            Y.append(hoverall)
            y0 = tY[:,0]
            if logy:
                y0 = np.log10(y0)
            Yorg.append(y0)

    # --- heatmap
    elif doheat:
        X = np.array(valX)
        Y = np.array(valY)
        Y = Y.reshape((len(valY[0]),len(valY)))
        Z = np.array(valZ)
    
    # --- regular 2d plot
    else:
        arrY = np.array(valY)
        
        if valX is None:
            #print('dbg: no valX, npoints = {}'.format(Npoints))
            # only Y is there
            # get array 0..1
            vX=np.linspace(0.0,1.0,Npoints)

            ###print('linspace of X {}'.format(len(vX)))

            # transform it
            result=transform(xmin,xmax,shiftx,logx,vX)
            if result is None:
                return None
            select=result[2]
            X=result[3]

            ###print('transformed linspace of X {}'.format(len(X)))
        
            Y=[]
            for ind in range(Nvariables):
                y=arrY[:,ind]
                y=y[select]
                if logy:
                    y = np.log10(y)
                Y.append(y)

            #print('dbg: X = {} Y = {}, Nvars = {}'.format(len(X),len(Y[0]),len(Y)))
            
        elif lblX is None:
            # both X and Y are there, no lblX
            result = transform(xmin,xmax,shiftx,logx,valX)
            if result is None:
                return None
            select=result[2]
            X = result[3]
            X=X[select]
            Y=[]
            for ind in range(Nvariables):
                y=arrY[:,ind]
                if select is not None:
                    y=y[select]
                if logy:
                    y = np.log10(y)
                Y.append(y)
        else:
            # label X - ignore xmin,xmax,shiftx
            X = np.array(valX)
            Y=[]
            select=None
            for ind in range(Nvariables):
                y=arrY[:,ind]
                if logy:
                    y = np.log10(y)
                Y.append(y)

        '''
        n, _ = np.histogram(x, bins=nbins)
        sy, _ = np.histogram(x, bins=nbins, weights=y)
        sy2, _ = np.histogram(x, bins=nbins, weights=y*y)
        mean = sy / n
        std = np.sqrt(sy2/n - mean*mean)
        
        plt.plot(x, y, 'bo')
        plt.errorbar((_[1:] + _[:-1])/2, mean, yerr=std, fmt='r-')
        plt.show()
        '''
    # get errors
    if valD is not None:
        r = 1.0/np.log(10.0)
        D=[]
        for ind in range(Nvariables):
            d=np.array(valD)
            d=d[:,ind]
            if select is not None:
                d=d[select]
            if logy:
                # correct error for log values
                Z = np.power(10.0,Y[ind])
                d = r*d/Z
                D.append(d)

    # do gradient
    if dograd:
        #y=Y.transpose()
        #print(Y.shape)
        #print(y.shape)
        Y = np.gradient(Y,axis=1)

    return X,Y,Z,D,Yorg,H
    

def genBackground(X,bkgstr):
    B = None
    if bkgstr is None:
        return None
    if X is None:
        return None
    # parse bkgstr
    fitfun=None
    params=[]
    for x in bkgstr.split(','):
        if fitfun is None:
            fitfun=x
            continue
        params.append(float(x))
    #
    if fitfun is None:
        print('error: no method in CSV string')
        return None
    localfn=globals()
    method_fit = localfn.get(fitfun)
    if method_fit is None:
        print('error: method not found %s'%fitfun)
        return None
    B = method_fit(params,X)
    
    return np.array(B)

def plotHeatMap(X,Y,Z,nbins):
    if nbins is None:
        nbins=30
    #print('dbg: create linspace x y')
    # create x-y points to be used in heatmap
    xi = np.linspace(X.min(),X.max(),nbins)
    yi = np.linspace(Y.min(),Y.max(),nbins)

    #print('dbg: create griddata z')
    # Z is a matrix of x-y values
    zi = griddata((X, Y), Z, (xi[None,:], yi[:,None]), method='cubic',rescale=False)
    #zi = griddata((X, Y), Z, (xi[None,:], yi[:,None]), method='linear',rescale=True)
    #zi = griddata((X, Y), Z, (xi, yi) ) #, method='cubic')

    zmax = Z.max()
    zmin = Z.min()
    dz = zmax-zmin
    zimax = zi.max()
    zimin = zi.min()
    dzi = zimax -zimin
    zcoef = dz/dzi

    zi = zcoef*(zi-zimin) +zmin
    # I control the range of my colorbar by removing data 
    # outside of my range of interest
    #zmin = 3
    #zmax = 12
    #zi[(zi<zmin) | (zi>zmax)] = None

    #print('dbg: create contour!')
    # Create the contour plot
    CS = plt.contourf(xi, yi, zi, 128, cmap=plt.cm.rainbow)

    #vmax=zmax, vmin=zmin)
    plt.colorbar()


def plotHist2D(X,Y,W,nbins):
    plt.hist2d(X,Y,bins=nbins,weights=W)
    #vmax=zmax, vmin=zmin)
    plt.colorbar()

    

parser = argparse.ArgumentParser(description='Simple plot tool')

#parser.add_argument('-f', '--file', action='append', dest='filename',
#                  help='read data from FILE', metavar='FILE', required=True, default=[])
parser.add_argument('-f', '--file', dest='filename',
                    help='read data from FILE', metavar='FILE', required=True, nargs='+')
parser.add_argument('-o', '--output', dest='output',
                    help='save plot to FILE', metavar='FILE', default=None)

parser.add_argument('--variables',     dest='dumpvars',    help='dump variable names - will skip all other arguments except -f',  action='store_true', default=False )
parser.add_argument('-s', '--separator', dest='separator',help='data separator', default=None)
parser.add_argument('-d', '--shiftx', dest='shiftx', help='shift X axis such that first point is at X=0',action='store_true', default=False )
parser.add_argument('-X', '--colX', dest='colX', type=str, default=None,
                    help='use X data in COLUMN (starting at 0)', metavar='COLUMN')
parser.add_argument('-Y', '--colY', dest='colY', type=str, required=False, nargs='+',default=[],
                    help='use Y data in COLUMN (starting at 0)', metavar='COLUMN')
parser.add_argument('-Z', '--colZ', dest='colZ', type=str, default=None,
                    help='use Z data in COLUMN (starting at 0)', metavar='COLUMN')
parser.add_argument('-D', '--colD', dest='colD', type=str, nargs='+',default=None,
                  help='use dY data in COLUMN', metavar='COLUMN')

parser.add_argument('-c', '--condition', dest='condition', default=None,
                    help='condition expression', metavar='string', required=False)
parser.add_argument('-L','--limit', dest='limit', default=0,
                    help='limit nr of lines read', metavar='int', type=int, required=False)
parser.add_argument('-T', '--tag', dest='starttag', default=None,
                    help='start data tag - if usehdr then the remaining string of the tag line is assumed to be the variable definitions', metavar='STRING')
parser.add_argument('-P', '--prefix', dest='prefixtag', default=None,
                    help='data line prefix tag - processes only lines starting with this prefix', metavar='STRING')
parser.add_argument('--usehdr',     dest='usehdr',    help='assume first line (or line defined by --tag option) is a header with variables',  action='store_true', default=False )

parser.add_argument('-l', '--label', dest='label', default='', nargs='+',
                  help='data label', metavar='STRING')
parser.add_argument('-x', '--xaxis', dest='xaxis', default=None,
                  help='x-axis label', metavar='STRING')
parser.add_argument('-y', '--yaxis', dest='yaxis', default=None,
                  help='y-axis label', metavar='STRING')

parser.add_argument('-S', '--style', dest='style', default='.',
                  help='line style as described in matplotlib', metavar='STRING')

parser.add_argument('-t', '--title', dest='title', default=None,
                    help='plot title', metavar='STRING')

parser.add_argument('-B', '--box', dest='box', help='box plot',action='store_true', default=False )
parser.add_argument('-b', '--nbins', dest='nbins', type=int, default=None,
                    help='nr of bins (with --box and --hist)', metavar='COLUMN')

parser.add_argument('--linspace', dest='linspace',    help='assuming colX is not given, this will generate an X = [0..1]',  action='store_true', default=False )
parser.add_argument('--hist',     dest='hist',    help='this will generate histogram(s) of Y',  action='store_true', default=False )
parser.add_argument('--normed',   dest='normed',  help='normalise histogram(s)',  action='store_true', default=False )
parser.add_argument('--heatmap',  dest='heatmap', help='generate heatmap - requires -X,Y,Z',  action='store_true', default=False )
parser.add_argument('--hist2d',   dest='hist2d',  help='plot 2D hist - requires -X,Y, Z is used as weight',  action='store_true', default=False )
parser.add_argument('--lblx',   dest='lblx',    help='X values are interpreted as labels',  action='store_true', default=False )
parser.add_argument('--logx',   dest='logx',    help='plot with log scale X',  action='store_true', default=False )
parser.add_argument('--logy',   dest='logy',    help='plot with log scale Y',  action='store_true', default=False )

parser.add_argument('--tlogx',   dest='tlogx',    help='transform x->log10(x)',  action='store_true', default=False )
parser.add_argument('--tlogy',   dest='tlogy',    help='transform y->log10(y)',  action='store_true', default=False )
parser.add_argument('--tlogz',   dest='tlogz',    help='transform z->log10(z)',  action='store_true', default=False )

parser.add_argument('--gradient', dest='gradient',  help='compute gradient',  action='store_true', default=False )

parser.add_argument('--linfit',   dest='linfit',    help='make linear fit',    action='store_true', default=False )
parser.add_argument('--linfit2',  dest='linfit2',   help='make x2 fit',        action='store_true', default=False )
parser.add_argument('--gaussfit', dest='gaussfit',  help='make gaussian fit',  action='store_true', default=False )

parser.add_argument('--fitres',   dest='fitres',    help='print fit result in a computer readable format (<slope> <intercept> <err>)',  action='store_true', default=False )

parser.add_argument('--backgr', dest='backgr',  help='background fit string: as obtained using --fitres',metavar='CSV fit string', default=None )


parser.add_argument("-M", "--xmax", dest="xmax", default=None,
                    help="X max value", metavar="float", type=float)
parser.add_argument("-m", "--xmin", dest="xmin", default=None,
                    help="X min value", metavar="float", type=float)

parser.add_argument('--noplot',   dest='noplot',    help='do not plot',  action='store_true', default=False )
options = parser.parse_args()

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

# get variables per file
variables=[]
reflist=None
fref=None
for f in options.filename:
    variables = getVariables(f,options.usehdr,options.separator,options.prefixtag,options.starttag)
    if reflist is None:
        reflist = variables
        fref = f
    
    if not reflist == variables:
        print('error: variable list is different in file {} - reference is {}'.format(f,fref))

if options.dumpvars:
    print('%-40s %s'%(f,repr(variables)))
    sys.exit(0)

if not (options.colD is None) and not (len(options.colD)==len(options.colY)):
    print( 'Must give the same number of columns to Y and D')
    sys.exit(-1)

if len(options.label)>0 and not (len(options.label)==len(options.colY)):
    print( 'Must give the same number of legend string as filenames')
    print( options.label)
    print( options.filename)
    sys.exit(-1)

if options.colY is None or len(options.colY)==0:
    print('error: --colY must be set to something')
    sys.exit(255)

# convert column refs to variable names
colX = None
if options.colX is not None:
    colX = convertToVar(options.colX,variables)
    if colX is None:
        print('error: invalid colX = {}'.format(options.colX))
        sys.exit(255)

colY = None
if options.colY is not None:
    colY=[]
    for ind in range(0,len(options.colY)):
        c = convertToVar(options.colY[ind],variables)
        if c is None:
            sys.exit(255)
        colY.append(c)

colD = None
if options.colD is not None:
    colD=[]
    for ind in range(0,len(options.colD)):
        c = convertToVar(options.colD[ind],variables)
        if c is None:
            sys.exit(255)
        colD.append(c)

colZ = None
if options.colZ is not None:
    colZ = convertToVar(options.colZ,variables)
    if colZ is None:
        print('error: invalid colZ = {}'.format(options.colZ))
        sys.exit(255)

if options.heatmap:
    colD = None
    if (colX is None) or (colY is None) or (colZ is None):
        print('error: X,Y and Z must all be defined for heatmap mode')
        sys.exit(255)
    if not len(colY)==1:
        print('error: heatmap only supports one Y column')
        sys.exit(255)
    if not len(options.filename)==1:
        print('error: heatmap only supports one input file')
        sys.exit(255)

if options.hist2d:
    colD = None
    if (colX is None) or (colY is None):
        print('error: X,Y must all be defined for heatmap mode')
        sys.exit(255)
    if not len(colY)==1:
        print('error: heatmap only supports one Y column')
        sys.exit(255)
    if not len(options.filename)==1:
        print('error: heatmap only supports one input file')
        sys.exit(255)

dolblx = options.lblx
dohist = options.hist
dolinsp= options.linspace
doheatmap = options.heatmap
dohist2d = options.hist2d


doflag=0
plottype=''
if dolblx:
    doflag=doflag+1
if dohist:
    doflag=doflag+1
    plottype='HIST'
if dolinsp:
    doflag=doflag+1
if doheatmap:
    doflag=doflag+1
    plottype='HEAT'
if dohist2d:
    doflag=doflag+1
    plottype='HEAT'

if doflag>1:
    print('error: only one of --heatmap, --linspace, --lblx and --hist may be set')
    sys.exit(255)
    
if dolblx and colX is None:
    print('error: option --lblx requires --colX to be set')
    sys.exit(255)
    
if dohist:
    # ignore colX and make sure colY is set
    colX = None
    if colY is None:
        print('error: option --hist requires --colY')
        sys.exit(255)

# limit of nr data lines read
nlimit = options.limit
if nlimit<0:
    nlimit=0

first=True
indf=0
doplot=not options.noplot

if doplot:
    fig=plt.figure(1)
    ncols=len(options.filename);
    #colormap = plt.cm.gist_ncar
    #plt.gca().set_color_cycle([colormap(i) for i in np.linspace(0, 0.9, ncols)])

hdrX = options.xaxis
hdrY = options.yaxis
hdrZ = None
anyplot=False

label = options.label

Hbins = options.nbins
if Hbins is None:
    Hbins='auto'

first=True
for f in options.filename:

    print('info: read file {}'.format(f))
    # read data
    (hdrX,hdrY,hdrZ,lblX,valX,valY,valZ,valD)=read_raw_data(f,
                                                            options.prefixtag,
                                                            options.starttag,
                                                            nlimit,
                                                            variables,
                                                            options.separator,
                                                            hdrX, hdrY, hdrZ,
                                                            colX,
                                                            colY,
                                                            colZ,
                                                            colD,
                                                            options.lblx,
                                                            options.condition)

    if valY is None or len(valY)==0:
        print('warning: no data found in file '+f)
        continue

    # process data
    result = process(plottype, Hbins, options.normed, options.xmin, options.xmax,
                     options.tlogx, options.tlogy,options.tlogz, options.shiftx,options.gradient,lblX,valX,valY,valZ,valD)

    if result is None:
        print('error: skipping file {}'.format(f))
        continue

    X=result[0]
    Y=result[1]
    Z=result[2]
    D=result[3]
    Yorg=result[4]  # only set if dohist is true
    Hbins=result[5] # idem

    print(len(Y))
    B = genBackground(X,options.backgr)

    # get label, default is filename
    labeltxt=f
    if label is not None:
        labeltxt = None
    #if label is not None and (len(label)==len(options.filename)):
    #    labeltxt=options.label[indf]

    # get X sample
    nline=options.nbins
    if nline is None:
        nline=100
    if nline<100:
        nline=100
    xmax=np.max(X)
    xmin=np.min(X)
    xline=np.linspace(xmin,xmax,nline)

    if doplot and first:
        first=False
        ax = fig.add_subplot(111)
        # possibly set log scales
        if options.logx:
            ax.set_xscale('log')
        if options.logy:
            ax.set_yscale('log')

    # loop over all Y's -> Y = np.array()
    indy=0
    D=None

    # check if any fitting
    dofit=False
    if options.linfit or options.linfit2 or options.gaussfit:
        dofit=True

    uselbl = labeltxt
    for vY in Y:

        if labeltxt is None:
            uselbl = label[indy]
        
        # correct Y if B is given
        Ycorr = vY
        if B is not None:
            Ycorr = vY-B

        # select error vector
        if D is not None:
            vD=D[indy]

        # check if fitting is set
        if options.linfit:
            fparams=[0.0,1.0]
            (fiterr,tfiterr,params)=fit( fit_pol1, fparams,X,Ycorr, options.tlogy);
            yline=fit_pol1(params,xline)
            if options.fitres:
                print(('%s,%e,%e,%e,%e')%('fit_pol1',params[1],params[0],fiterr,tfiterr))
            else:
                print(('%-20s : Y = %8.2gx + %8.2g (%8.2g)')%(uselbl, params[1],params[0],fiterr))

        if options.linfit2:
            fparams=[0.0,0.0,1.0]
            (fiterr,tfiterr,params)=fit( fit_pol2, fparams,X,Ycorr, options.tlogy);
            yline=fit_pol2(params,xline)
            if options.fitres:
                print(('%s,%e,%e,%e,%e')%('fit_pol2',params[2],params[1],params[0],fiterr))
            else:
                print(('%-20s : Y = %8.2gx2 + %8.2gx + %8.2g (%8.2g,%8.2g)')%(uselbl, params[2],params[1],params[0],fiterr,tfiterr))

        if options.gaussfit:
            fparams=[np.max(Ycorr),np.mean(X),np.mean(X)]
            (fiterr,tfiterr,params)=fit( fit_gauss, fparams,X,Ycorr, options.tlogy);
#            params[0] = 500
#            params[1] = 0.05
#            params[2] = 0.01

            yline=fit_gauss(params,xline)
            if options.fitres:
                print(('%s,%e,%e,%e,%e')%('fit_gauss',params[2],params[1],params[0],fiterr))
            else:
                print(('%-20s : Y = %8.2gx2 + %8.2gx + %8.2g (%8.2g,%8.2g)')%(uselbl, params[2],params[1],params[0],fiterr,tfiterr))

        # do the plotting
        if doplot:
            anyplot=True
            if D is None:
                if options.box:
                    mean,std,epos = bin_error(X,Ycorr,options.nbins)
                    plt.errorbar(epos, mean, yerr=std, fmt='r-')
                    #output,midp=bin_by(Ycorr,X,20)
                    #plt.boxplot(output,positions=midp,notch=True)
                elif dohist:
                    if indf==0:
                        plt.hist(Yorg,Hbins,label=uselbl,density=options.normed)
                    else:
                        plt.hist(Yorg,Hbins,label=uselbl,density=options.normed,histtype='step')
                elif doheatmap:
                    plotHeatMap(X,vY,Z,options.nbins)
                elif dohist2d:
                    plotHist2D(X,vY,Z,options.nbins)
                else:
                    plt.plot(X,vY,options.style,label=uselbl) #,color=colours[indf-1])
                    if options.lblx:
                        plt.xticks(X, lblX, rotation='vertical')

                    
            else:
                plt.errorbar(valX,Ycorr,D,label=uselbl)

            if dofit:
                plt.plot(xline,yline,'-')
        
        # loop over variables
        indy=indy+1
    # loop over files
    indf=indf+1
    
if doplot and anyplot:
    xaxis = hdrX
    yaxis = hdrY
    if type(hdrY) is list:
        yaxis = ','.join(hdrY)
    if options.tlogx:
        xaxis='log10('+xaxis+')'
    if options.tlogy:
        yaxis='log10('+yaxis+')'
    if dohist:
        yaxis = 'count'
    fsize='medium'
    plt.tick_params(axis='both', which='major', labelsize=12)
    plt.xlabel(xaxis,fontsize=fsize)
    plt.ylabel(yaxis,fontsize=fsize)
    if doheatmap:
        plt.title('Heatmap of {}'.format(colZ))
    if not doheatmap:
        if options.title is not None:
            print('setting title {}'.format(options.title))
            plt.title(options.title)
        plt.grid()
        plt.legend(loc='best',fontsize=fsize)

    if options.output is None:
        plt.show()
    else:
        plt.savefig(options.output,dpi=400)

