#!/usr/bin/env python

'''
This script will find correspondence between qry and ref transcript through reading *.tmap derive from cuffcompare, and output coding *.txt file on each class code.
[Usage] ./tmap2coding tmap_file qry_coding.txt ref_coding.txt outpath
'''

'''
data:
ref_coding {id:coding potential}
qry_coding {id:coding potential}
tmap {class_code:{qry_id:ref_id}}
'''

import sys, os

def get_tras_name0(name,sep="|"):
    #refSeq and GENCODE
    name = name.split(sep)
    if len(name) == 1:
        return name[0]
    if len(name) == 5:
        #refSeq
        return name[3]
    #GENCODE
    return name[0]

def get_tras_name(name,seps=["|",None]):
    for sep in seps:
        name = get_tras_name0(name,sep)
    return name

def read_tmap(infl,sep="\t",has_header=True,col_ref_id=1,col_class_code=2,col_qry_id=4):

    data = {}
    if has_header:
        header = infl.readline()
    for line in infl:
        line_array = line.strip().split(sep)
        id_map = data.setdefault(line_array[col_class_code],{})
        if id_map.get(line_array[col_qry_id]):
            raise ValueError("%s has not only one ref_id" % line_array[col_qry_id])
        id_map[line_array[col_qry_id]] = line_array[col_ref_id]
    return data

def read_coding_txt(infl,sep="\t",has_header=True,col_id=0,col_coding=1,col_score=2,rmgene=False):
    #rmgene remove "gene=XXXX" when coding file contain these words in transcript ID field

    data = {}
    if has_header:
        header = infl.readline()
    for line in infl:
        line_array = line.strip().split(sep)
        key = line_array[col_id]
        if rmgene:
            key = key.split()[0]
        data[key] = (line_array[col_coding],line_array[col_score])
    return data

def read_ref_coding(infl,sep="\t",has_header=True,col_id=0,col_coding=1,rmgene=False):
    #rmgene get true trasncript id for ref_coding.txt file
    outdata = {}
    data = read_coding_txt(infl,sep,has_header,col_id,col_coding,rmgene=True)
    if rmgene:
        for key,value in data.items():
            outdata[get_tras_name(key)] = value
        return outdata
    else:
        return data
        
def write_coding(outfile,qry_data,ref_data,map_data,plus_rid,sep="\t"):
    # write one class code file.
    # plus_rid: Whether write ref transctipt id at the end of line.
    header=["name","label","score","actCls"]
    if plus_rid:
        header.append("ref_name")
    header = sep.join(header)+"\n"
    #score="-"
    outfile.write(header)
    for qid, value in qry_data.items():
        rid = map_data.get(qid)
        if not rid:
            sys.stderr.write(qid+" has no ref transcript\n")
            continue
        #else
        qscore = value[0]
        score = value[1]
        if not qscore:
            continue
        # rs[1] should be the score given by tool.
        rs = ref_data.get(rid)
        if not rs:
            sys.stderr.write("refID: %s can't find score"%rid+"\n")
            continue
        rscore = rs[0]
        #else
        outdata = [qid,qscore,score,rscore]
        if plus_rid:
            outdata.append(rid)
        outfile.write(sep.join(outdata)+"\n")

def write_codings(outpath,qry_data,ref_data,tmap_data,plus_rid,sep="\t"):
    #
    outpath = os.path.abspath(outpath)
    for ClsCode, map_data in tmap_data.items():
        outfile = open(outpath+"/coding_cls_%s.txt"%ClsCode,'w')
        write_coding(outfile,qry_data,ref_data,map_data,sep)
        outfile.close()

def main(argv):

    #if not len(args) == 4:
     #   sys.stdout.write("[Usage] ./tmap2coding tmap_file qry_coding.txt ref_coding.txt outpath\n")
      #  sys.exit(1)
   # tmap_f, q_f, r_f=map(open,args[:3])
   # qry_data = read_coding_txt(q_f)
   # ref_data = read_coding_txt(r_f)
   # tmap_data = read_tmap(tmap_f)
    import argparse, os

    parser = argparse.ArgumentParser(description="This script will read *.tmap from cuffcompare, and output coding *.txt file on each class code.")
    parser.add_argument('tmap_file',nargs='?',help="*.tmap file from cuffcompare",type=argparse.FileType('r'))
    parser.add_argument('q_f',nargs='?',help="query transcripts coding file",type=argparse.FileType('r'))
    parser.add_argument('r_f',nargs='?',help="refference transcripts coding file",type=argparse.FileType('r'))
    parser.add_argument('-o','--outpath',dest='outpath',nargs='?',help="output path")
    parser.add_argument('--rm_r',dest='rm_r',action='store_true',default=False,help="get transcirpt id of ref file")
    parser.add_argument('--rm_q',dest='rm_q',action='store_true',default=False,help="remove \"gene=XXX\" field of qry file")
    parser.add_argument('--real-code',dest='rc',action='store_true',default=False,help="Set actCls as source class rather predicted by last prediction on templates.")
    parser.add_argument('--plus_rid',dest='p_rid',action='store_true',default=False,help="write reference transcript id at the end of line.")
    args = parser.parse_args(argv[1:])

    if args.rc:
        col_coding = 3
    else:
        col_coding = 1
    qry_data = read_coding_txt(args.q_f,rmgene=args.rm_q)
    ref_data = read_ref_coding(args.r_f,rmgene=args.rm_r,col_coding=col_coding)
    tmap_data = read_tmap(args.tmap_file)
    if not os.path.exists(args.outpath):
        os.mkdir(args.outpath)
    elif not os.path.isdir(args.outpath):
        raise ValueError("%s exist and is not a directory"%args.outpath)
    write_codings(args.outpath,qry_data,ref_data,tmap_data,args.p_rid)

if __name__ == '__main__':

    import sys

    main(sys.argv)