#!/usr/bin/python3
# - input file(s)
# - X,Y,Z,dY,dZ columns
# - type: hist, scatter, line, box
import sys
import argparse
import io
import re
from collections import deque
import itertools
import math
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.linear_model import LinearRegression

def read_data_file( filename, label, options ):
    rmws = re.compile(rb'\s+')

    data = None
    hastag = options.starttag is not None
    haspref = options.prefixtag is not None
    foundtag = not hastag
    sep = options.separator.encode('utf-8')
    first=True
    colnames = None
    
    nlines=0
    datastream = io.BytesIO()
    with open(filename,'rb') as fh:
        for line in fh:
            if len(line)==0:
                continue

            # check for start tag
            if not foundtag:
                if line.startswith(options.starttag):
                    foundtag=True
                continue

            # check for commented lines
            if line[0] in b'#':
                continue

            # check prefix
            if haspref:
                if line.startswith(options.prefixtag):
                    line = line[len(options.prefixtag):].lstrip()
                else:
                    continue
            line = re.sub(rmws,b' ',line).rstrip()+b'\n'
            # data line
            if first:
                if not options.usehdr:
                    elems = line.split(sep)
                    colnames=[]
                    for i in range(0,len(elems)):
                        colnames+=['v{}'.format(i)]
                first=False

            datastream.write(line)
            nlines = nlines+1

        if nlines == 0:
            data = None
        else:
            if options.usehdr:
                header='infer'
                columns = None
            else:
                header = None

            print('read {} lines of data'.format(nlines))
            datastream.seek(0)
            data = pd.read_csv(datastream,encoding='utf8',sep=options.separator,header=header,names=colnames)
            print('converted to panda')
            
        data['__index__'] = list(range(len(data)))
        if options.normX is not None and options.normX>0:
            scale = options.normX/max(data['__index__'])
            data['__index__'] = data['__index__']*scale
            
        data[options.labelHdr] = [label]*len(data)
        
        
    return data

def eval_column( cstr,df ):
    if cstr in df.columns.array:
        return cstr
    try:
        v=int(cstr)
        return df.columns.array[v]
    except:
        print('error: invalid column label {}'.format(cstr))
    return None
    
def set_columns( col, df ):
    if col is None:
        return None
    r=[]
    if type(col)==str:
        v = eval_column(col,df)
        if v is None:
            return None
        else:
            r.append(v)
        return r

    for c in col:
        v = eval_column(c,df)
        if v is None:
            print('warning: skipping label {}'.format(c))
            continue
        else:
            r.append(v)

    if len(r)==0:
        return None
    return r

def fix_range( options, df ):
    if options.rngX is None:
        return df
    if colX is not None:
        xmin = min(df[colX]) if options.rngX[0] is None else options.rngX[0]
        xmax = max(df[colX]) if options.rngX[1] is None else options.rngX[1]
        selX = df[(colX>=xmin) and (colX<=xmax)]
    if colY is not None:
        ymin = min(df[colY]) if options.rngY[0] is None else options.rngY[0]
        ymax = max(df[colY]) if options.rngY[1] is None else options.rngY[1]
        selY = df[(colY>=ymin) and (colY<=ymax)]
    if colZ is not None:
        zmin = min(df[colZ]) if options.rngZ[0] is None else options.rngZ[0]
        zmax = max(df[colZ]) if options.rngZ[1] is None else options.rngZ[1]
        selZ = df[(colZ>=zmin) and (colZ<=zmax)]
        
def fix_columns( options, df ):
    colX = set_columns(options.colX,df)
    colX = None if colX is None else colX[0]
    colY = set_columns(options.colY,df)
    colZ = set_columns(options.colZ,df)
    colZ = None if colZ is None else colZ[0]
    if colZ is not None:
        colY+=[colZ]

    if len(options.scale)>0 and len(colY)!=len(options.scale):
        options.scale=[]
    
    return colX,colY,colZ

def mklabel(f):
    return f

def getRange(csvstr):
    if csvstr is None:
        return None
    elems = csvstr.split(',')
    if not len(elems)==2:
        return None
    
    hasLower = len(elems[0])>0
    hasUpper = len(elems[1])>0

    vmin = float(elems[0]) if hasLower else None
    vmax = float(elems[0]) if hasUpper else None
    return vmin,vmax
    
def read_data( options ):

    frames = [ read_data_file(f[0], f[1], options) for f in zip(options.filenames,options.labels) ]
    print('read {} frames'.format(len(frames)))
    print('concatenate frames')
    df = pd.concat(frames,keys=options.filenames)
    print('fix columns')
    colX,colY,colZ = fix_columns(options,df)
    return df,colX,colY,colZ

def show_plot( plt, options ):
    plt.title(options.title)
    plt.xticks( rotation=60)
    if options.xlabel is not None:
        plt.xlabel(options.xlabel)
    if options.ylabel is not None:
        plt.ylabel(options.ylabel)
    plt.show()
    
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Simple plot tool')

    parser.add_argument('-f', '--file', dest='filenames',
                        help='read data from FILE', metavar='FILE', required=True, nargs='+')
    parser.add_argument('-l', '--labels', dest='labels',
                        help='label(s) for file(s)', metavar='string', required=False, nargs='+',default=None)

    parser.add_argument('--lheader', dest='labelHdr',
                        help='label box header', metavar='string', required=False, default='__labels__')

    parser.add_argument('--title', dest='title',
                        help='plot title', metavar='string', required=False, default=None)
    parser.add_argument('--xlabel', dest='xlabel',
                        help='x-axis label', metavar='string', required=False, default=None)
    parser.add_argument('--ylabel', dest='ylabel',
                        help='y-axis label', metavar='string', required=False, default=None)

    parser.add_argument('-s', '--separator', dest='separator',help='column separator', default=' ')
    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('-L', '--log', dest='log', type=str, default=None,
                        help='string consisting of X,Y or Z - these columns will be logged', metavar='COLUMN')
    parser.add_argument('-S', '--scale', dest='scale', type=float, required=False, nargs='+',default=[],
                        help='scale Y', metavar='COLUMN')
    parser.add_argument('-R', '--normX', dest='normX', type=float, required=False,default=None,
                        help='normalise max X to this value > 0', metavar='COLUMN')

    parser.add_argument('-n', '--nbins', dest='nbins', type=int, required=False,default=None,
                        help='nr of bins in histograms', metavar='int')
    parser.add_argument('--rngX', dest='rngX', type=str, default=None,
                        help='comma separated range for X', metavar='csv')
    parser.add_argument('--rngY', dest='rngY', type=str, default=None,
                        help='comma separated range for Y', metavar='csv')
    parser.add_argument('--rngZ', dest='rngZ', type=str, default=None,
                        help='comma separated range for Z', metavar='csv')
    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('--hist',     dest='dohist',    help='create a histogram',  action='store_true', default=False )
    parser.add_argument('--linreg',   dest='linreg',    help='linear regression',   action='store_true', default=False )

    options = parser.parse_args()

    if options.labels is None:
        options.labels = options.filenames
    else:
        if not len(options.labels) == len(options.filenames):
            print('warning: option -l does not have the same nr of elements as -f - resetting labels to file names')
            options.labels = options.filenames

    options.logX = False
    options.logY = False
    options.logZ = False

    if options.starttag is not None:
        options.starttag = options.starttag.encode('utf-8')
    if options.prefixtag is not None:
        options.prefixtag = options.prefixtag.encode('utf-8')
    if options.log is not None:
        options.log = options.log.upper()
        for x in options.log:
            if x not in ['X','Y','Z']:
                print('warning: --log ignore {}'.format(x))
                continue
            if   x=='X':
                options.logX = True
            elif x=='Y':
                options.logY = True
            elif x=='Z':
                options.logZ = True

    # set ranges
    options.rngX = getRange(options.rngX)
    options.rngY = getRange(options.rngY)
    options.rngZ = getRange(options.rngZ)


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

    def linreg(dfX,dfY):
        X=dfX.array.reshape(-1,1)
        Y=dfY.array.reshape(-1,1)
        linear_regressor = LinearRegression()
        lr = linear_regressor.fit(X, Y)
        Y_pred = linear_regressor.predict(X)
        return X,Y,Y_pred,lr.coef_,lr.intercept_

    def logval(v,mask):
        rval = np.log10(v)
        lm   = v<=0
        rval[lm]=0.0
        if mask is None:
            mask = lm
        else:
            mask = mask or lm
        return rval,mask

    data,colX,colY,colZ = read_data( options )

    mask = None
    if (colX is None) and (colY is None):
        print('--------------')
        print(' DATA SUMMARY ')
        print('--------------')
        print(data[data.columns].describe())
        sys.exit(0)

    if colX is not None:
        if options.logX:
            print('log column X')
            data[colX],mask = logval(data[colX],mask)

    if colZ is not None:
        if options.logZ:
            print('log column Z')
            data[colZ],mask = logval(data[colZ],mask)

    if colY is not None:
        if options.logY:
            for c in colY:
                print('log column Y {}'.format(c))
                data[c],mask = logval(data[c],mask)

        if len(options.scale)==len(colY):
            ind=0
            for c in colY:
                print('scale column {} with {}'.format(c,options.scale[ind]))
                data[c] = options.scale[ind]*data[c]
                ind = ind+1

    sns.set()
    if (colX is not None) and (colY is not None):
        print('scatter plot...')

        # create pivot table
        # will the columns: '<colX> variable value'
        # where value will be the various columns in colY
        df = data.melt(id_vars=colX,value_vars=colY)

        if colZ is None:
            X = data[colX]
            #sns.lineplot(x=colX,y='value',hue='variable',data=df)
            if type(colY)==list:
                for cY in colY:
                    sns.scatterplot(data=data,x=colX,y=cY,s=100,hue=options.labelHdr)
                    #sns.scatterplot(data=data,x=colX,y=cY)
                    if options.linreg:
                        aX,aY,pY,A,B = linreg(X,data[cY])
                        plt.plot(aX,pY,color='red',s=100)
                        print('{} = {}*{} + {}'.format(cY,A,colX,B))
            else:
                sns.scatterplot(data=data,x=colX,y=colY,s=100)
                if options.linreg:
                    aX,aY,pY,A,B = linreg(data[colX],data[cY])
                    plt.plot(aX,pY,color='red',s=100)
                    print('{} = {}*{} + {}'.format(cY,A,colX,B))
        else:
            sns.scatterplot(data=df,x=colX,y='value',s=100,hue=options.labelHdr)
        show_plot( plt, options )

    elif (colX is None) and (colY is not None):
        print('histogram on colY = {}'.format(colY))
        for cy in colY:
            if options.dohist:
                sns.histplot(data=data,x=cy,bins=options.nbins,hue=options.labelHdr)
            else:
                sns.scatterplot(data=data,x=data.__index__,y=options.colY[0],s=100, hue=options.labelHdr)
        show_plot( plt, options )

    elif (colX is not None) and (colY is None):
        print('histogram on colX = {}'.format(colX))
        sns.histplot(data=data,x=colX,bins=options.nbins,hue=options.labelHdr)
        show_plot( plt, options )

