#!/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")
    parser.add_argument("--unmapped", action="store_true", help="Running unmapped reads")
    parseArgs = parser.parse_args(args=_args)
    return parseArgs
    

                    
                                

def main():
    
    parseArgs = parseArguments(sys.argv[1:])
    
    if parseArgs.unmapped:
        splitReads(parseArgs.fasta, parseArgs.blastx, parseArgs.unmapped)
        mergeInfo(os.path.join(os.path.dirname(parseArgs.blastx),os.path.basename(parseArgs.fasta)))
    else:
        splitReads(parseArgs.fasta, parseArgs.blastx, parseArgs.unmapped)
                  

def splitReads(fasta, btab, unmapped):
    
    sbtab = os.path.join(os.path.dirname(btab),os.path.basename(fasta)+".ur100.split.first.btab")
    split_fasta = os.path.join(os.path.dirname(btab),os.path.basename(fasta)+".split.first.fasta")
    unmapped_fasta = os.path.join(os.path.dirname(btab),os.path.basename(fasta)+".unmapped.fasta")
    read_coordinate_file = os.path.join(os.path.dirname(btab), os.path.basename(fasta)+".rc.first.tsv")
    
    if unmapped:
        _fasta = os.path.basename(fasta).split(".unmapped.fasta")[0]
        sbtab = os.path.join(os.path.dirname(btab),_fasta+".ur100.split.second.btab")
        split_fasta = os.path.join(os.path.dirname(btab),_fasta+".split.second.fasta")
        read_coordinate_file = os.path.join(os.path.dirname(btab), _fasta+".rc.second.tsv")
                          

    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
                if unmapped:
                    wf.write(r_id.split("_unmapped")[0]+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
                else:
                    wf.write(r_id+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
            else:
                if (qstart, qend) in seen_trip[r_id]:
                    if unmapped:
                        wf.write(r_id.split("_unmapped")[0]+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
                    else:
                        wf.write(r_id+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
                else:
                    count+=1
                    if unmapped:
                        wf.write(r_id.split("_unmapped")[0]+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
                    else:
                        wf.write(r_id+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
                    seen_trip[r_id][(qstart, qend)] = None
                    
    

    with open(split_fasta, "w+") as wf, open(unmapped_fasta, "w+") as uf, open(read_coordinate_file, "w+") as rcf:
        rcf.write("RID\tMapStatus\tQstart\tQend\n")
        for _id in seqid_to_seqrecord:
            write_id = _id
            if unmapped:
                write_id = _id.split("_unmapped")[0]
            if not _id in seen_ids:                                                                              
                wf.write(">"+write_id+"_unmapped\n")
                wf.write(str(seqid_to_seqrecord[_id].seq)+"\n")
                if not unmapped: #Write to unmapped file
                    uf.write(">"+write_id+"_unmapped\n")
                    uf.write(str(seqid_to_seqrecord[_id].seq)+"\n")
                rcf.write(write_id+"\t"+"Read unmapped"+"\t0\t"+str(len(seqid_to_seqrecord[_id].seq)-1)+"\n")

            else:
                count=0
                for qpair in seen_trip[_id]:
                    wf.write(">"+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")
                    rcf.write(write_id+"\t"+"Region mapped [ORF]"+"\t"+str(_lbound-1)+"\t"+str(_ubound-1)+"\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(">"+write_id+"_unmapped_"+str(count)+"\n")
                        wf.write(str(seqid_to_seqrecord[_id].seq[_lbound:_ubound-1])+"\n")
                        rcf.write(write_id+"\t"+"Region unmapped"+"\t"+str(_lbound)+"\t"+str(_ubound-2)+"\n")
                        count+=1
               
                del seen_trip[_id]               
                del int_trees[_id]
                del seen_ids[_id]
                                                                                                            

def mergeInfo(base):
    
    base = base.split(".unmapped.fasta")[0]
    
    sbtab = base+".ur100.split.btab"
    split_file = base+".split.fasta"
    read_coordinate_file = base+".rc.tsv"
                                  
    with open(base+".ur100.split.first.btab", "r") as inpf1, open(base+".ur100.split.second.btab", "r") as inpf2, open(sbtab, "w+") as wf:
        for line in inpf1:
            wf.write(line)
        for line in inpf2:
            wf.write(line)
         
    with open(base+".split.first.fasta", "r") as inpf1, open(base+".split.second.fasta", "r") as inpf2, open(split_file, "w+") as wf:
        flag = False
        for line in inpf1:
            if ">" in line:
                if not line.strip().endswith("_unmapped"):
                    wf.write(line)
                    flag = True
            else:
                if flag:
                    wf.write(line)
                    flag = False
                    
        for line in inpf2:
            wf.write(line)
            
    with open(base+".rc.first.tsv", "r") as inpf1, open(base+".rc.second.tsv", "r") as inpf2, open(read_coordinate_file, "w+") as wf:
        for line in inpf1:
            tokens = line.split("\t")
            if tokens[1] != "Read unmapped":
                wf.write(line)

        for line in inpf2:
            wf.write(line)
            
                                                                                                                                              
if __name__ == '__main__':    
    main()
