#!/usr/bin/env python

'''

whole infile is parsed to data:
{key0:[array0_0,array0_1,array0_2,...],key1:[array1_0,array1_1,1array1_2,...],key2:[]}
key0=array0_0[key_col]=array0_1[key_col]=array0_*[key_col]
key1=array1_*[key_col]
each array*_* is a list return by split function working on one line.

'''

def read_tracking(infile,key_col=3,sep="\t",has_header=False):
    #key_col: which column should be regard as split key
    
    data={}
    if has_header:
        header = infile.readline()
    for line in infile:
        larray = line.strip().split(sep)
        key = larray[key_col]
        data.setdefault(key,[]).append(larray)
    return data

def write_lines(data,outfile,sep="\t",header=None):

    if header:
        outfile.write(header+"\n")
    for larray in data:
        outfile.write(sep.join(larray)+"\n")

def write_split_tracking(data,outfile,sep="\t",header=None):

    #outfile: prefix of outfiles
    for key, value in data.items():
        outf = open(outfile+"_"+str(key),'w')
        write_lines(value,outf,sep)
        outf.close()

def main():
    '''
    [Usage]: ./tracking_split.py infile prefix_outfile
    '''

    import sys
    if len(sys.argv)!=3:
        raise ValueError("parameters error!!!")
    infile = open(sys.argv[1])
    outfile = sys.argv[2]

    data = read_tracking(infile)
    write_split_tracking(data,outfile)

if __name__ == '__main__':

    main()
