#!/usr/bin/python3
'''
This tool applies a key map to a column wise data set.
Pipeline usage: MERGEINPAR (louvain)
'''
from __future__ import print_function

import sys
import argparse

def parseColumnStr(colstr):
    colflags=[]
    skip=False
    for k in colstr:
        if k=='M':
            flag=0 # label to be mapped
        elif k=='K':
            flag=1 # keep column
        elif k=='+':
            flag=2 # keep this column and the rest
            skip=True
        elif k=='R':
            flag=3 # remove column
        elif k==' ':
            continue
        else:
            print('error: got bad column flag '+k,file=sys.stderr)
            return None
        colflags.append(flag)
        if skip:
            break

    return colflags
    
def readMap( fname, swap ):
    labelMap = {}
    nbad=0
    nok=0
    idnext=0
    if fname is None:
        return labelMap,idnext
    
    with open(fname,mode='r') as fsh:

        # read map
        for rawline in fsh.readlines():
            line = rawline.strip().rstrip('\n')
            if len(line)==0: continue
            if line[0]=='#': continue
            # split line in:
            #   <string id> <numerical id>
            elems=line.split()
            if len(elems)<2: continue
            idnum=0
            try:
                idnum=int(elems[1])
            except:
                print('error: unexpected format of map. second column should be integer - bad id found = %s'%(elems[1]),file=sys.stderr)
                nbad=nbad+1
                continue

            # does no checking of duplicates...
            nok=nok+1
            if swap:
                labelMap[str(idnum)] = elems[0]
#                print('map[{}] = {}'.format(idnum,elems[0]))
            else:
                labelMap[str(elems[0])]=idnum
#                print('map[{}] = {}'.format(elems[0],idnum))
            if idnum>=idnext:
                idnext=idnum+1

    if nbad>0:
        if nok==0:
            print('error: map file contained no good lines',file=sys.stderr)
        else:
            print('warning: map file contained both good (%d) and bad (%d) lines'%(nok,nbad),file=sys.stderr)

    return labelMap,idnext


def process(datafile,colstr,keep,labelMap):
    colflags = parseColumnStr(colstr)
    if colflags is None:
        return

    ncol = len(colflags)
    if ncol==0:
        return
    
    colind=range(0,ncol)
    nbad=0
    lnum=0
    # the data file and parse according to column string
    with open(datafile,mode='r') as fsh:
        for rawline in fsh.readlines():
            lnum=lnum+1
            line = rawline.strip().rstrip('\n')
            if len(line)==0: continue
            if line[0]=='#':
                if keep:
                    print(line)
                continue
            elems=line.split()
            nelems=len(elems)
            if nelems<ncol:
                nbad=nbad+1
                print('error at line (%d): bad line format'%(lnum,line),file=sys.stderr)
                continue

            outstr=''
            keepall=False
            for ind in colind:
                val = str(elems[ind])
                if keepall:
                    # keep all the rest
                    outstr = outstr +' '+'{}'.format(val)
                    continue

                if colflags[ind]==0:
                    # map value
                    if val in labelMap:
                        mapval=labelMap[val]
                    else:
                        outstr='###NOMAP '+outstr
                        mapval='BAD<{}>'.format(val)
                    outstr = outstr +' '+'{}'.format(mapval)
                elif colflags[ind]==1:
                    # keep column
                    outstr = outstr +' '+'{}'.format(val)
                elif colflags[ind]==2:
                    # keep this and the rest
                    outstr = outstr +' '+' '.join(str(x) for x in elems[ind:])
                    break

            # output the line
            if len(outstr)>0:
                print(outstr)
    
    return

#-----------------------------------------------------------------
# MAIN


parser = argparse.ArgumentParser(description='Apply map on arbitrary columns')

parser.add_argument('-f', '--file',    dest='fname',   help='white space delimited data file',    metavar='FILE', required=True)
parser.add_argument('-c', '--colstr',  dest='colstr',  help='string with one chr per column; M : map, R : remove, K : keep, + : keep the rest', metavar='string', required=True)
parser.add_argument('-i', '--inpmap',  dest='inpmap',  help='optional input map, format <str id> <num id>', metavar='FILE', required=True, default=None)
parser.add_argument('-s', '--swap',    dest='swap',    help='swap mapping order in input map, will always map key to value -> newval = map[oldval]', action='store_true', default=False)
parser.add_argument('-k', '--keep',    dest='keep',    help='keep comment lines (starting with \'#\')', action='store_true', default=False)

options = parser.parse_args()

# read input label, map if it exists
labelmap = {}
if options.inpmap is not None:
    labelmap,idnext = readMap(options.inpmap,options.swap)

    
# open data file
process( options.fname, options.colstr, options.keep, labelmap )

