#!/usr/bin/env python3
# This is one of script series in SimTool project. For more information,
# please see https://github.com/flipboards/simtool

import getopt
import sys
import pandas
import re 

def parse_titles(title_strings:list, header):
    ''' Parse title string (which can be $1 or string)
    '''
    index = []
    title = []

    for t in title_strings:
        assert t, 'Invalid string'
        if t[0] == '$':
            pos = int(t[1:]) - 1
            assert pos >= 0, 'Invalid column index'
            index.append(pos)
            title.append((header[pos] if header is not None else str(pos)))
        else:
            assert header is not None, 'No header, cannot use string as title'
            assert t in header, 'Cannot find header'
            title.append(t)
            index.append(header.get_loc(t))

    return index, title 


def main(argv):

    help_str = 'colstack [-s split] [-u (xcol1,xcol2,...:)ycol1,ycol2...] file ...'
    long_help_str = '''
colstack: Stack same title in different file.
Arguments:
-h, --help:     display this help message.
--notitle:      do not use title (and in output).
-s:             specify split string of input datasheet. Default is tab.
--sep:          specify output split string. Default is tab.
-t:             new title naming style. '%f': short filename; '%F': full filename, '%t': title. 
                Default is '%f' for single y column and '%f-%t' for multiple y columns.
-u:             which column to use. Similar with gnuplot grammar, column can be '$1' to specify index, or
                strings to specify title.
file:           files to read.
    '''

    short_args = 'hs:t:u:'
    long_args = ['notitle', 'help', 'sep=']

    try:
        opts, args = getopt.getopt(argv[1:], short_args, long_args)
    except getopt.GetoptError:
        print('Invalid argument. Use `-h` to get help', file=sys.stderr)
        return 1


    xcols = []
    ycols = []
    title = True
    title_style = '%f'
    title_style_multi = '%f-%t'
    sep_in = '\t'
    sep_out = '\t'

    for opt, arg in opts:
        if opt == '-h' or opt == '--help':
            print(long_help_str)
            return 1
        elif opt == '--notitle':
            title = False
        elif opt == '-s':
            sep_in = ('\t' if arg == '\\t' else arg )
        elif opt == '--sep':
            sep_out = ('\t' if arg == '\\t' else arg )
        elif opt == '-t':
            title_style = title_style_multi = arg 
        elif opt == '-u':
            if ':' in arg:
                x1, y1 = arg.split(':')
                xcols = x1.split(',')
                ycols = y1.split(',')
            else:
                ycols = arg.split(',')

    if not args:
        print('Error: No args', file=sys.stderr)
        return 1


    dataframes = [pandas.read_csv(name, sep=sep_in, index_col=False, header=(0 if title else None)) for name in args]

    if xcols:
        xindex, xtitle = parse_titles(xcols, dataframes[0].columns)
    
    if ycols:
        yindex, ytitle = parse_titles(ycols, dataframes[0].columns)
    else:
        yindex = range(len(dataframes[0].columns))
        ytitle = dataframes[0].columns

    if xcols:
        xoutput = dataframes[0][xtitle]  # select x columns
    else:
        xoutput = None

    youtput = []
    youtput_names = []
    replacer = re.compile(r'%f|%F|%t')

    cnt = 1
    for df, name in zip(dataframes, args):
        for yt, yidx in zip(ytitle, yindex):
            if title:
                replace_dict = {'%f':name.rsplit('\\', 1)[-1], '%F':name, '%t':yt}
                if len(ytitle) > 1:
                    ycolname = replacer.sub(lambda x:replace_dict[x.group()], title_style_multi)
                else:
                    ycolname = replacer.sub(lambda x:replace_dict[x.group()], title_style)
            else:
                ycolname = 'y%d' % cnt # won't show, but critical
                cnt += 1
            youtput.append(df.iloc[:, yidx])
            youtput_names.append(ycolname)

    yf = pandas.concat(youtput, axis=1, ignore_index=True)
    yf.columns = youtput_names
    if xoutput is not None:
        output = xoutput.join(yf)
    else:
        output = yf

    output.to_csv(sys.stdout, sep=sep_out, index=False, header=title)
    return 

if __name__ == '__main__':
    main(sys.argv)

