#!/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 argparse
import os
import sys
from collections import defaultdict, OrderedDict

from Bio import SeqIO
from intervaltree import IntervalTree


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.seq
    #     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 record in SeqIO.parse(fasta, "fasta"):
            # Retrieve the ID and sequence for each record
            _id = record.id
            _seq = record.seq
            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(_seq) + "\n")
                if not unmapped:  # Write to unmapped file
                    uf.write(">" + write_id + "_unmapped\n")
                    uf.write(str(_seq) + "\n")
                rcf.write(
                    write_id + "\t" + "Read unmapped" + "\t0\t" + str(len(_seq) - 1) + "\n")

            else:
                # Create an interval tree for each mapped record
                t = IntervalTree()
                t[1:len(str(_seq))] = "m"
                count = 0
                for qpair in seen_trip[_id]:
                    wf.write(">" + write_id + "_" + str(count) + "\n")
                    _lbound, _ubound = min(qpair), max(qpair)
                    t.chop(_lbound, _ubound)  # chop the interval seen
                    wf.write(str(_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 t:  # 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(_seq)):
                            _ubound += 1
                        wf.write(">" + write_id + "_unmapped_" + str(count) + "\n")
                        wf.write(str(_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 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()

