#!/usr/bin/env python3

import sys, os, re
import collections
    
ORIG_TRANSDECODER_FLAG = False

def main():


    import argparse

    parser = argparse.ArgumentParser(description="transdecoder orf selection algorithm", formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("--long_orfs_cds", dest="long_orfs_cds_filename", type=str, default="", required=True, help="longest_orfs.cds file")
    parser.add_argument("--long_orfs_scores", dest="long_orfs_scores_filename", type=str, default="", required=True, help="longest_orfs.cds.scores file")
    parser.add_argument("--single_best", dest="single_best_flag", action='store_true', default=False, help="select only single best orf")
    parser.add_argument("--all_good_orfs", dest="all_good_orfs", action='store_true', default=False, help="retain all good orfs allowing overlaps")
    parser.add_argument("--td_orig", dest="td_orig_flag", action='store_true', default=False, help='use the original TD algorithm')

    args = parser.parse_args()
    
    long_orfs_cds_file = args.long_orfs_cds_filename
    long_orfs_scored_file = args.long_orfs_scores_filename
        
    predicted_orf_coords = retrieve_orf_coords(long_orfs_cds_file)

    prediction_list = parse_predictions_and_scores(long_orfs_scored_file, predicted_orf_coords)


    def orig_transdecoder_algorithm(seq_length, frame_scores):
        return (fst_gt_zero(frame_scores) and fst_is_max_all(frame_scores))

    def custom_td_alg(orf_length, frame_scores, long_orf_size=1000, moderate_orf_size=500):

        pass_orf = True
        
        if orf_length >= long_orf_size:
            pass_orf = True # stays true
        elif orf_length >= moderate_orf_size:
            if not (fst_is_max3(frame_scores) and fst_gt_zero(frame_scores)):
                pass_orf = False
        else:
            if not (fst_is_max3(frame_scores) and fst_gt_zero(frame_scores)):
                pass_orf = False

        return pass_orf


    select_alg = custom_td_alg
    if args.td_orig_flag:
        select_alg = orig_transdecoder_algorithm
    

    # all good orfs
    trans_to_preds_list = select(prediction_list, predicted_orf_coords, select_alg)

    if not args.all_good_orfs:
        # by default, we prune overlaps
        trans_to_preds_list = select_best_non_overlapping_orfs(trans_to_preds_list)
        
    elif args.single_best_flag:
        trans_to_preds_list = select_single_orf_per_transcript(trans_to_preds_list)
    

    write_preds_to_file(trans_to_preds_list, sys.stdout)
    
    
    
    sys.exit(0)


def fst_is_max_all(frame_scores):
    return (frame_scores[0] > max(frame_scores[1:]))

def fst_gt_zero(frame_scores):
    return (frame_scores[0] > 0)

def fst_is_max3(frame_scores):
    return (frame_scores[0] > max(frame_scores[1:3]))



def retrieve_orf_coords(long_orfs_cds_file):

    orf_acc_to_coord_info = {}

    with open(long_orfs_cds_file) as fh:
        for line in fh:
            if not re.search("^>", line):
                continue
            match = re.search(" (\S+):(\d+)-(\d+)\(([+-])\)$", line)
            if not match:
                raise RuntimeError("Error, cannot extract orf coord info from line: {}".format(line))
            (transcript_id, lend, rend, orient) = (match.group(1), match.group(2), match.group(3), match.group(4))
            (lend, rend) = sorted((int(lend), int(rend)))

            match = re.search("^>(\S+)", line)
            orf_id = match.group(1)

            orf_acc_to_coord_info[orf_id] = {
                'transcript_id' : transcript_id,
                'lend' : lend,
                'rend' : rend,
                'orient' : orient }


    return orf_acc_to_coord_info


def parse_predictions_and_scores(long_orfs_scored_file, predicted_orf_coords):
        
    prediction_list = []
    
    fh = open(long_orfs_scored_file)

    for line in fh:
        if re.search("^#", line): continue
        line = line.rstrip()
        (orf_id, markov_order, orf_length, score_1, score_2, score_3, score_4, score_5, score_6)  = line.split("\t")

        orf_length = int(orf_length)
        score_1 = float(score_1)
        score_2 = float(score_2)
        score_3 = float(score_3)
        score_4 = float(score_4)
        score_5 = float(score_5)
        score_6 = float(score_6)


        orf_struct = predicted_orf_coords[orf_id]

        prediction = { 'orf_id' : orf_id,
                       'markov_order' : markov_order,
                       'orf_length' : orf_length,
                       'frame_scores' : (score_1, score_2, score_3, score_4, score_5, score_6),
                       'orf_struct' : orf_struct
                       }

        prediction_list.append(prediction)



    prediction_list.sort(key=lambda x: x['orf_length'])

    prediction_list.reverse()

    return prediction_list



def select(prediction_list, predicted_orf_coords, selection_alg):
    
    
    transcript_to_selected_orfs = collections.defaultdict(list)

    for prediction in prediction_list:

        orf_id = prediction['orf_id']
        markov_order = prediction['markov_order']
        orf_length = prediction['orf_length']
        frame_scores = prediction['frame_scores']
        (score_1, score_2, score_3, score_4, score_5, score_6) = frame_scores

        orf_struct = prediction['orf_struct']

        transcript_id = orf_struct['transcript_id']
        lend = orf_struct['lend']
        rend = orf_struct['rend']
        orient = orf_struct['orient']


        ###########################
        ## apply filtering criteria
        pass_orf = selection_alg(orf_length, frame_scores)

        if pass_orf:
            transcript_to_selected_orfs[transcript_id].append(prediction)


    
    return transcript_to_selected_orfs



def select_single_orf_per_transcript(transcript_to_selected_orfs):


    non_overlapping_orfs = select_best_non_overlapping_orfs(transcript_to_selected_orfs)

    ret_transcript_to_selected_orfs = collections.defaultdict(list)
    

    for transcript_id in non_overlapping_orfs:

        prediction_list = non_overlapping_orfs[transcript_id]
                
        prediction_list.sort(key=lambda x: x['orf_length'], reverse=True)
        #prediction_list.sort(key=lambda x: x['frame_scores'][0], reverse=True)

        top_selected_pred = prediction_list[0]
        
        ret_transcript_to_selected_orfs[transcript_id].append(top_selected_pred)
        
        
    return ret_transcript_to_selected_orfs


def select_best_non_overlapping_orfs(transcript_to_selected_orfs):


    ret_transcript_to_non_overlapping_orfs = collections.defaultdict(list)
    

    for transcript_id in transcript_to_selected_orfs:

        prediction_list = transcript_to_selected_orfs[transcript_id]

        
        if ORIG_TRANSDECODER_FLAG:
            prediction_list.sort(key=lambda x: x['orf_length'], reverse=True)
        else:
            prediction_list.sort(key=lambda x: x['frame_scores'][0], reverse=True)
        
        selected_preds = []
        
        for prediction in prediction_list:
        
            ## passed filters, report it

                                
            # ensure it doesn't overlap a selected one.
            found_overlap = False
            for pred in selected_preds:
                if prediction['orf_struct']['lend'] < pred['orf_struct']['rend'] and \
                   prediction['orf_struct']['rend'] > pred['orf_struct']['lend']:
                    found_overlap = True
                    break
            
            if not found_overlap:
                selected_preds.append(prediction)
        
        selected_preds.sort(key=lambda x: x['orf_length'], reverse=True)

        ret_transcript_to_non_overlapping_orfs[transcript_id].extend(selected_preds)


    return ret_transcript_to_non_overlapping_orfs



def write_preds_to_file(transcript_to_selected_orfs, ofh):


    for transcript_id in transcript_to_selected_orfs:

        pred_list = transcript_to_selected_orfs[transcript_id]

        for pred in pred_list:
            orf_struct = pred['orf_struct']
                        
            ofh.write("\t".join([transcript_id, "selected", "CDS",
                                 str(orf_struct['lend']), str(orf_struct['rend']), '.',
                                 orf_struct['orient'], str(pred['frame_scores'][0]), pred['orf_id']]) + "\n")
            
           



    return




if __name__ == '__main__':

    main()
    
