'''
@ 3' UTR location of each transcript can be extracted from the results file of dapars2
@ which return the 3UTR location in its output (the fourth column)
 run on python3 environment
 usage: python extract_3UTR_location.py --dapars_res /path/to/Dapars2_res.all_chromosomes.txt --output /path/to/3utr_location.txt
'''
import argparse

parser = argparse.ArgumentParser(description='')
parser.add_argument('--dapars_res',type=str,help="provide the output file generated by dapars2,the merged all chromosomes file is required if running dapars2 on each chromosome separately")
parser.add_argument('--output',type=str, default="3utr_location.txt",help="specify 3' UTR location file")

args = parser.parse_args()

gene2loc = {} # a dict to record the location
gene_order = [] # store every gene in this list to keep order for output

# -- load dapars result file and extract gene and location info
fh = open(args.dapars_res,'r')
header = fh.readline()
for line in fh.readlines():
    line = line.strip()
    w = line.split("\t")
    chrom,loc_code = w[3].split(":")
    start,end = loc_code.split("-")

    gene2loc[w[0]] = (chrom,start,end)
    gene_order.append(w[0])
fh.close()
print(len(gene_order),"genes have been processed!")

# -- output
fho = open(args.output,'w')
print("Gene\tChr\tStart\tEnd", file=fho)
for gene in gene_order:
    if gene in gene2loc:
        print("%s\t%s\t%s\t%s" % (gene,gene2loc[gene][0],gene2loc[gene][1],gene2loc[gene][2]), file=fho)
fho.close()

print("Done!")
