#!/usr/bin/env python

import re
import sys
import time
from Bio import SeqIO
from StringIO import StringIO
from Bio import AlignIO
import os
import argparse

from __init__ import __version__

"""
Transcriptome Annotation by Modular Algorithms (TAMA)
TAMA Merge

Author: Richard I
. Kuo

This script merges transcriptome/genome annotations.

Last Updated: 2020/12/17

Added more informative error message for issues with strand information from input bed files. 

"""

tm_version = 'tm0.0'


start_time = time.time()
prev_time = start_time


##########################################################
###########
#####################################################################

ap = argparse.ArgumentParser(description='This script merges transcriptomes.')

ap.add_argument('-f', type=str, nargs=1, help='File list')
ap.add_argument('-p', type=str, nargs=1, help='Output prefix')

ap.add_argument('-e', type=str, nargs=1, help='Collapse exon ends flag: common_ends or longest_ends  (Default is common_ends)')

ap.add_argument('-a', type=str, nargs=1, help='5 prime threshold (Default is 10)')
ap.add_argument('-m', type=str, nargs=1, help='Exon ends threshold/ splice junction threshold (Default is 10)')
ap.add_argument('-z', type=str, nargs=1, help='3 prime threshold (Default is 10)')

ap.add_argument('-d', type=str, nargs=1, help='Flag for merging duplicate transcript groups (default no_merge quits when duplicates are found, merge_dup will merge duplicates)')

ap.add_argument('-s', type=str, nargs=1, help='Use gene and transcript ID from a merge source. Specify source name from filelist file here.')

ap.add_argument('-cds', type=str, nargs=1, help='Use CDS from a merge source. Specify source name from filelist file here.')

opts = ap.parse_args()

#check for missing args
missing_arg_flag = 0

if not opts.f:
    print("Filelist file missing")
    missing_arg_flag = 1
if not opts.p:
    print("Output prefix name missing")
    missing_arg_flag = 1
  

if not opts.e:
    print("Default collapse exon ends flag will be used: common_ends")
    collapse_flag = "common_ends"
else:
    collapse_flag = opts.e[0]
    
if not opts.a:
    print("Default 5 prime threshold: 20")
    fiveprime_threshold = 20
else:
    fiveprime_threshold = int(opts.a[0])
    
if not opts.m:
    print("Default exon end/splice junction threshold: 10")
    exon_diff_threshold = 10
else:
    exon_diff_threshold = int(opts.m[0])

if not opts.z:
    print("Default 3 prime threshold: 20")
    threeprime_threshold = 20
else:
    threeprime_threshold = int(opts.z[0])

if not opts.d:
    print("Default duplicate merge flag: no_merge")
    duplicate_flag = "no_merge"
else:
    duplicate_flag = str(opts.d[0])

if not opts.s:
    print("Default source ID merge flag: no_source_id")
    source_id_flag = "no_source_id"
else:
    source_id_flag = str(opts.s[0])

    source_id_flag_list = source_id_flag.split(",")

if not opts.cds:
    print("Default CDS merge flag: no_cds")
    source_cds_flag = "no_cds"
else:
    source_cds_flag = str(opts.cds[0])

    source_cds_flag_list = source_cds_flag.split(",")


if missing_arg_flag == 1:
    print("Please try again with complete arguments")


filelist_file = opts.f[0]
outfile_prefix = opts.p[0]

#####################################################################
#####################################################################


print("opening file list")
#filelist_file = sys.argv[1]
filelist_file_contents = open(filelist_file).read().rstrip("\n").split("\n")

#outfile_prefix = sys.argv[2]


bed_outfile_name = outfile_prefix + ".bed"
outfile_bed = open(bed_outfile_name,"w")

trans_report_outfile_name = outfile_prefix + "_trans_report.txt"
outfile_trans_report = open(trans_report_outfile_name,"w")
trans_report_line = "\t".join(["transcript_id","num_clusters","sources","start_wobble_list","end_wobble_list","exon_start_support","exon_end_support","all_source_trans"])
outfile_trans_report.write(trans_report_line)
outfile_trans_report.write("\n")

gene_report_outfile_name = outfile_prefix + "_gene_report.txt"
outfile_gene_report = open(gene_report_outfile_name,"w")
gene_report_line = "\t".join(["gene_id","num_clusters","num_final_trans","sources","chrom", "start","end","source_genes","source_summary"])
outfile_gene_report.write(gene_report_line)
outfile_gene_report.write("\n")

merge_outfile_name = outfile_prefix + "_merge.txt"
outfile_merge = open(merge_outfile_name,"w")
#merge_line = "\t".join(["transcript_id","cluster_id","scaffold","strand","start","end","exon_starts","exon_ends"])
#outfile_merge.write(merge_line)
#outfile_merge.write("\n")



bed_dict = {} # bed_dict[scaffold][start][end][strand] = list of bed lines with added source id
source_dict = {} # source_dict[source id][filename, seq type, priority rank] = value

   

def track_time(start_time,prev_time):
    end_time = time.time()
    time_taken = int(end_time - prev_time)
    tt_hours = time_taken / 60
    tt_hours = tt_hours /60
    leftover_time = time_taken - (tt_hours * 60 * 60)
    tt_minutes = leftover_time / 60
    leftover_time = time_taken - (tt_minutes * 60)
    tt_seconds = leftover_time
    print("time taken since last check:\t" +  str(tt_hours) + ":" + str(tt_minutes) + ":" + str(tt_seconds) )
    
    time_total = int(end_time - start_time)
    tt_hours = time_total / 60
    tt_hours = tt_hours /60
    leftover_time = time_total - (tt_hours * 60 * 60)
    tt_minutes = leftover_time / 60
    leftover_time = time_total - (tt_minutes * 60)
    tt_seconds = leftover_time
    print("time taken since beginning:\t" +  str(tt_hours) + ":" + str(tt_minutes) + ":" + str(tt_seconds) )
    
    this_time = end_time
    
    return this_time


class Transcript:
    def __init__(self, trans_id):
        self.trans_id = trans_id
        self.uniq_trans_id = ""
        self.gene_id = ""
        self.merge_gene_id = ""
        self.scaffold = "none"
        self.trans_start = -1
        self.start_pos = self.trans_start
        self.trans_end = 0
        self.end_pos = self.trans_end 
        self.exon_start_list = [] 
        self.exon_end_list = []
        
        self.id_line = ""
        
        self.cds_start = 0
        self.cds_end = 0
        self.num_exons = 0
        self.strand = ""
        self.source_id = ""
        
        self.start_priority = 0
        self.end_priority = 0
        self.junction_priority = 0
    
    def add_bed_info(self,bed_line):
        line_split = bed_line.split("\t")
        self.scaffold = line_split[0]
        self.trans_start = int(line_split[1])
        self.start_pos = self.trans_start
        self.trans_end = int(line_split[2])
        self.end_pos = self.trans_end 
        self.id_line = line_split[3]
        id_line_split = self.id_line.split(";")
        self.gene_id = id_line_split[0]
        
        self.strand = line_split[5]
        self.num_exons = int(line_split[9])
        
        self.cds_start = int(line_split[6])
        self.cds_end = int(line_split[7])
        
        block_size_list = line_split[10].split(",")
        block_start_list = line_split[11].split(",")
        
        #add this for commas at the end of the exon block and start strings
        if block_size_list[-1] == "":
            block_size_list.pop(-1)
        
        if block_start_list[-1] == "":
            block_start_list.pop(-1)
        
        for i in xrange(len(block_size_list)):
            rel_exon_start = int(block_start_list[i])
            rel_exon_end = rel_exon_start + int(block_size_list[i])
            
            exon_start = self.trans_start + rel_exon_start
            exon_end = self.trans_start + rel_exon_end
            
            self.exon_end_list.append(exon_end)
            self.exon_start_list.append(exon_start)
            
            
    
    def add_source_id(self,source_id):
        self.source_id = source_id
        self.uniq_trans_id = "_".join([source_id,self.t