#!/usr/bin/env python3

'''
Copyright 2021, Advait Balaji, All rights reserved. 
Free for research use. For commercial use, contact 
Dr. Todd Treangen (treangen@rice.edu)
'''


import os
import sys
import argparse
import numpy as np
from collections import defaultdict, Counter, OrderedDict
from intervaltree import Interval, IntervalTree
import pickle #debug
from Bio import SeqIO



def parseArguments(_args):
    parser = argparse.ArgumentParser(prog="split_reads", usage="%(prog)s [options]",
                                   description=f"SeqScreen-Nano: tool for pathogen identification from ONT data. Splits reads based on blastx.",
                                   epilog="Author: {__author__}\n")
    parser.add_argument("--fasta", type=str, required=True, help="Query file")
    parser.add_argument("--blastx", type=str, required=True, help="Blastx btab file")
    parseArgs = parser.parse_args(args=_args)
    return parseArgs
    

                    
                                

def main():
    
    parseArgs = parseArguments(sys.argv[1:])
    splitReads(parseArgs.fasta, parseArgs.blastx)

def splitReads(fasta, btab):
    
    sbtab = os.path.join(os.path.dirname(btab),os.path.basename(fasta)+".ur100.split.btab")
    
    int_trees = {}
    seqid_to_seqrecord = {}
    # Read in fasta file
    for record in SeqIO.parse(fasta, "fasta"):
        seqid_to_seqrecord[record.id] = record
        t=IntervalTree()
        t[1:len(str(record.seq))] = "m"
        int_trees[record.id] = t
                
    with open(btab, "r") as f, open(sbtab, "w+") as wf:
        seen_ids = OrderedDict() #set of ids
        seen_trip = defaultdict(OrderedDict) #Dict of id and qstart, qend OrderedDict()
        count = 0
        for line in f:
            tokens = line.split("\t")
            r_id, qstart, qend = tokens[0], int(tokens[6]), int(tokens[7])
            if not r_id in seen_ids:
                count = 0
                seen_ids[r_id] = None
                seen_trip[r_id][(qstart, qend)] = None
                wf.write(r_id+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
            else:
                if (qstart, qend) in seen_trip[r_id]:
                    wf.write(r_id+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
                else:
                    count+=1
                    wf.write(r_id+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
                    seen_trip[r_id][(qstart, qend)] = None
                    

    with open(os.path.join(os.path.dirname(btab),os.path.basename(fasta)+".split.fasta"), "w+") as wf:
        for _id in seqid_to_seqrecord:
            if not _id in seen_ids:
                wf.write(">"+_id+"_unmapped\n")
                wf.write(str(seqid_to_seqrecord[_id].seq)+"\n")
            else:
                count=0
                for qpair in seen_trip[_id]:
                    wf.write(">"+_id+"_"+str(count)+"\n")
                    _lbound, _ubound = min(qpair), max(qpair)
                    int_trees[_id].chop(_lbound, _ubound) #chop the interval seen
                    wf.write(str(seqid_to_seqrecord[_id].seq[_lbound-1:_ubound])+"\n")
                    count+=1
                count=0
                for _int in int_trees[_id]: #iterate over remaining intervals and check if they can be output
                    _lbound, _ubound = min((_int[0],_int[1])), max((_int[0], _int[1]))
                    if _ubound - _lbound >= 100: #at least 100bp
                        if _lbound == 1:
                            _lbound -= 1
                        if _ubound == len(str(seqid_to_seqrecord[_id].seq)):
                            _ubound += 1
                        wf.write(">"+_id+"_unmapped_"+str(count)+"\n")
                        wf.write(str(seqid_to_seqrecord[_id].seq[_lbound:_ubound-1])+"\n")
                        count+=1


if __name__ == '__main__':    
    main()
