#!/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 Counter




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 splitReads(fasta, btab):
    
    sbtab = os.path.join(os.path.dirname(btab),os.path.basename(fasta)+".ur100.split.btab")
    r_ids = []
    
    with open(fasta, "r") as f:
        for line in f:
            if line[0] == ">":
                _id = line.split(" ")[0][1:]
                r_ids.append(_id)
                
    
    with open(btab, "r") as f, open(sbtab, "w+") as wf:
        seen_id = set() #set of ids
        seen_trip = set() #set of id, qstart, qend triplets
        count = 0
        for line in f:
            tokens = line.split("\t")
            r_id, qstart, qend = tokens[0], tokens[6], tokens[7]
            if not r_id in seen_id:
                count = 0
                seen_id.add(r_id)
                seen_trip.add((r_id, qstart, qend))
                wf.write(r_id+"_"+str(count)+"\t"+"\t".join(tokens[1:]))
            else:
                if (r_id, qstart, qend) in seen_trip:
                    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.add((r_id, qstart, qend))
                    
                                

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







if __name__ == '__main__':    
    main()
