from multiprocessing import Pool
import subprocess
import shlex
import os
import re
from .. import io

STRICTNESS = 0
arid2phase = {}

# TODO: Kill all threads gracefully if one encounters an error.
"""This whole program can hang with something like this:

Exception in thread Thread-3:
Traceback (most recent call last):
  File "/mnt/software/p/python/2.7.13-UCS4/centos-6/lib/python2.7/threading.py", line 801, in __bootstrap_inner
    self.run()
  File "/mnt/software/p/python/2.7.13-UCS4/centos-6/lib/python2.7/threading.py", line 754, in run
    self.__target(*self.__args, **self.__kwargs)
  File "/mnt/software/p/python/2.7.13-UCS4/centos-6/lib/python2.7/multiprocessing/pool.py", line 389, in _handle_results
    task = get()
TypeError: ('__init__() takes at least 3 arguments (1 given)', <class 'subprocess.CalledProcessError'>, ())

But the most typical cause of error is a missing input, so just check that early.
"""


def filter_stage1(input_):
    db_fn, fn, max_diff, max_ovlp, min_ovlp, min_len = input_
    try:
        ignore_rtn = []
        current_q_id = None
        ave_idt = 0.0
        all_over_len = 0.0
        overlap_data = {"5p": 0, "3p": 0}
        overlap_phase = {"5p": set(), "3p": set()}
        q_id = None

        for l in subprocess.check_output(shlex.split("LA4Falcon -mo %s %s" % (db_fn, fn)), encoding='ascii').splitlines():
            l = l.strip().split()
            q_id, t_id = l[:2]

            if q_id not in arid2phase:
                continue
            if t_id not in arid2phase:
                continue
            if arid2phase[t_id][0] != arid2phase[q_id][0]:
                continue

            if arid2phase[t_id][1] == arid2phase[q_id][1] and arid2phase[t_id][2] != arid2phase[q_id][2]:
                continue

            if q_id != None and q_id != current_q_id:

                left_count = overlap_data["5p"]
                right_count = overlap_data["3p"]

                if (left_count == 0 and min_ovlp <= right_count <= max_ovlp) or \
                    (min_ovlp <= left_count <= max_ovlp and right_count == 0):
                    # This is a good case. When a query has 0-count on one end
                    # it's actually the left-most or right-most read sequenced
                    # from a linear chromosome.
                    pass

                else:
                    if abs(left_count - right_count) > max_diff:
                        ignore_rtn.append(current_q_id)
                    elif left_count > max_ovlp or right_count > max_ovlp:
                        ignore_rtn.append(current_q_id)
                    elif left_count < min_ovlp or right_count < min_ovlp:
                        ignore_rtn.append(current_q_id)

                # remove unphased reads if they are sandwiched by the reads from the same phase
                if current_q_id not in arid2phase:
                    left_set = overlap_phase["5p"]
                    right_set = overlap_phase["3p"]
                    if len(left_set & right_set) > 0:
                        ignore_rtn.append(current_q_id)

                overlap_data = {"5p": 0, "3p": 0}
                overlap_phase = {"5p": set(), "3p": set()}
                current_q_id = q_id
                ave_idt = 0.0
                all_over_len = 0.0

            overlap_len = -int(l[2])
            idt = float(l[3])
            q_s, q_e, q_l = int(l[5]), int(l[6]), int(l[7])
            t_s, t_e, t_l = int(l[9]), int(l[10]), int(l[11])

            if idt < 90:
                continue

            if q_l < min_len or t_l < min_len:
                continue

            if l[-1] in ("contains", "overlap"):
                ave_idt += idt * overlap_len
                all_over_len += overlap_len
            if q_s == 0:
                overlap_data["5p"] += 1
                if t_id in arid2phase:
                    overlap_phase["5p"].add(arid2phase[t_id])
            if q_e == q_l:
                overlap_data["3p"] += 1
                if t_id in arid2phase:
                    overlap_phase["3p"].add(arid2phase[t_id])

        if q_id != None:
            left_count = overlap_data["5p"]
            right_count = overlap_data["3p"]

            if (left_count == 0 and min_ovlp <= right_count <= max_ovlp) or \
                (min_ovlp <= left_count <= max_ovlp and right_count == 0):
                # This is a good case. When a query has 0-count on one end
                # it's actually the left-most or right-most read sequenced
                # from a linear chromosome.
                pass

            else:

                if abs(left_count - right_count) > max_diff:
                    ignore_rtn.append(current_q_id)
                elif left_count > max_ovlp or right_count > max_ovlp:
                    ignore_rtn.append(current_q_id)
                elif left_count < min_ovlp or right_count < min_ovlp:
                    ignore_rtn.append(current_q_id)

            # remove unphased reads if they are sandwiched by the reads from the same phase
            if current_q_id not in arid2phase:
                left_set = overlap_phase["5p"]
                right_set = overlap_phase["3p"]
                if len(left_set & right_set) > 0:
                    ignore_rtn.append(current_q_id)

        return fn, ignore_rtn

    except (KeyboardInterrupt, SystemExit):
        return


def filter_stage2(input_):
    db_fn, fn, max_diff, max_ovlp, min_ovlp, min_len, ignore_set, strictness = input_
    try:
        contained_id = set()
        for l in subprocess.check_output(shlex.split("LA4Falcon -mo %s %s" % (db_fn, fn)), encoding='ascii').splitlines():
            l = l.strip().split()
            q_id, t_id = l[:2]

            if q_id not in arid2phase:
                continue
            if t_id not in arid2phase:
                continue

            if arid2phase[t_id][0] != arid2phase[q_id][0]:
                continue
            if arid2phase[t_id][1] == arid2phase[q_id][1] and arid2phase[t_id][2] != arid2phase[q_id][2]:
                continue

            if strictness > 0:
                ###############################################
                # Added a couple more of strict filters:
                # The phase needs to match exactly to add a read to contained list:
                ###############################################
                if arid2phase[t_id][1] != arid2phase[q_id][1] or arid2phase[t_id][2] != arid2phase[q_id][2]:
                    continue
                # If one of the reads is unphased, do not consider this "containment" call:
                if arid2phase[t_id][1] == "-1" or arid2phase[q_id][1] == "-1":
                    continue
                ###############################################

            q_s, q_e, q_l = int(l[5]), int(l[6]), int(l[7])
            t_s, t_e, t_l = int(l[9]), int(l[10]), int(l[11])

            idt = float(l[3])
            if idt < 90:
                continue

            if q_l < min_len or t_l < min_len:
                continue

            if q_id in ignore_set:
                continue
            if t_id in ignore_set:
                continue
            if l[-1] == "contained":
                contained_id.add(q_id)
            if l[-1] == "contains":
                contained_id.add(t_id)
        return fn, contained_id

    except (KeyboardInterrupt, SystemExit):
        return


def filter_stage3(input_):

    db_fn, fn, max_diff, max_ovlp, min_ovlp, min_len, ignore_set, contained_set, bestn, strictness = input_

    overlap_data = {"5p": [], "3p": []}
    try:
        ovlp_output = []
        current_q_id = None
        for l in subprocess.check_output(shlex.split("LA4Falcon -mo %s %s" % (db_fn, fn)), encoding='ascii').splitlines():
            l = l.strip().split()
            q_id, t_id = l[:2]

            if q_id not in arid2phase:
                continue
            if t_id not in arid2phase:
                continue

            if arid2phase[t_id][0] != arid2phase[q_id][0]:
                continue
            if arid2phase[t_id][1] == arid2phase[q_id][1] and arid2phase[t_id][2] != arid2phase[q_id][2]:
                continue

            if strictness > 0:
                # if arid2phase[t_id][1] == "-1" or arid2phase[q_id][1] == "-1":
                #     continue
                ###############################################
                # Added a couple more of strict filters:
                # The phase needs to match exactly, skip this overlap.
                ###############################################
                if arid2phase[t_id][1] != arid2phase[q_id][1] or arid2phase[t_id][2] != arid2phase[q_id][2]:
                    continue
                # If one of the reads is unphased, do not consider this overlap.
                if arid2phase[t_id][1] == "-1" or arid2phase[q_id][1] == "-1":
                    continue
                ###############################################

            if current_q_id == None:
                current_q_id = q_id
                overlap_data = {"5p": [], "3p": []}

            elif q_id != current_q_id:
                left = overlap_data["5p"]
                right = overlap_data["3p"]
                left.sort()
                right.sort()

                for i in range(len(left)):
                    inphase, score, m_range, ovlp = left[i]
                    ovlp_output.append(ovlp)
                    # print " ".join(ovlp), read_end_data[current_q_id]
                    if i >= bestn and m_range > 1000:
                        break

                for i in range(len(right)):
                    inphase, score, m_range, ovlp = right[i]
                    ovlp_output.append(ovlp)
                    # print " ".join(ovlp), read_end_data[current_q_id]
                    if i >= bestn and m_range > 1000:
                        break

                overlap_data = {"5p": [], "3p": []}
                current_q_id = q_id

            if q_id in contained_set:
                continue
            if t_id in contained_set:
                continue
            if q_id in ignore_set:
                continue
            if t_id in ignore_set:
                continue

            overlap_len = -int(l[2])
            idt = float(l[3])
            q_s, q_e, q_l = int(l[5]), int(l[6]), int(l[7])
            t_s, t_e, t_l = int(l[9]), int(l[10]), int(l[11])

            if idt < 90:
                continue
            if q_l < min_len or t_l < min_len:
                continue
            if q_s == 0:
                l.extend([".".join(arid2phase.get(current_q_id, "NA")), ".".join(arid2phase.get(t_id, "NA"))])
                inphase = 1 if arid2phase.get(current_q_id, "NA") == arid2phase.get(t_id, "NA") else 0
                #nphase = 1 if arid2phase[current_q_id] == arid2phase[t_id] else 0
                overlap_data["5p"].append((-inphase, -overlap_len,  t_l - (t_e - t_s),  l))
            elif q_e == q_l:
                l.extend([".".join(arid2phase.get(current_q_id, "NA")), ".".join(arid2phase.get(t_id, "NA"))])
                inphase = 1 if arid2phase.get(current_q_id, "NA") == arid2phase.get(t_id, "NA") else 0
                #inphase = 1 if arid2phase[current_q_id] == arid2phase[t_id] else 0
                overlap_data["3p"].append((-inphase, -overlap_len, t_l - (t_e - t_s), l))

        left = overlap_data["5p"]
        right = overlap_data["3p"]
        left.sort()
        right.sort()

        for i in range(len(left)):
            inphase, score, m_range, ovlp = left[i]
            ovlp_output.append(ovlp)
            # print " ".join(ovlp), read_end_data[current_q_id]
            if i >= bestn and m_range > 1000:
                break

        for i in range(len(right)):
            inphase, score, m_range, ovlp = right[i]
            ovlp_output.append(ovlp)
            # print " ".join(ovlp), read_end_data[current_q_id]
            if i >= bestn and m_range > 1000:
                break

        return fn, ovlp_output
    except (KeyboardInterrupt, SystemExit):
        return


def assert_exists(fn):
    if not os.path.exists(fn):
        msg = 'File does not exist: {!r}'.format(fn)
        raise Exception(msg)


def run(args):
    max_diff = args.max_diff
    max_cov = args.max_cov
    min_cov = args.min_cov
    min_len = args.min_len
    bestn = args.bestn
    db_fn = args.db
    strictness = args.strictness

    assert_exists(db_fn)

    with open(args.rid_phase_map) as f:
        for row in f:
            row = row.strip().split()
            arid2phase[row[0]] = (row[1], row[2], row[3])  # ctg_id, phase_blk_id, phase_id
    assert arid2phase, 'Empty rid_phase_map: {!r}'.format(args.rid_phase_map)
    exe_pool = Pool(args.n_core)

    file_list = list(io.yield_abspath_from_fofn(args.fofn))
    inputs = []
    for fn in file_list:
        assert_exists(fn)
        inputs.append((db_fn, fn, max_diff, max_cov, min_cov, min_len))

    ignore_all = []
    for res in exe_pool.imap(filter_stage1, inputs):
        ignore_all.extend(res[1])

    inputs = []
    ignore_all = set(ignore_all)
    for fn in file_list:
        if len(fn) != 0:
            inputs.append((db_fn, fn, max_diff, max_cov, min_cov, min_len, ignore_all, strictness))
    contained = set()
    for res in exe_pool.imap(filter_stage2, inputs):
        contained.update(res[1])
        # print res[0], len(res[1]), len(contained)

    # print "all", len(contained)
    inputs = []
    ignore_all = set(ignore_all)
    for fn in file_list:
        if len(fn) != 0:
            inputs.append((db_fn, fn, max_diff, max_cov, min_cov, min_len, ignore_all, contained, bestn, strictness))
    for res in exe_pool.imap(filter_stage3, inputs):
        for l in res[1]:
            print(" ".join(l))


######
import argparse
import sys


def parse_args(argv):
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description='a simple multi-processes LAS ovelap data filter')
    parser.add_argument(
        '--n-core', type=int, default=4,
        help='number of processes used for generating consensus')
    parser.add_argument(
        '--fofn', type=str, help='file contains the path of all LAS file to be processed in parallel')
    parser.add_argument(
        '--db', type=str, help='read db file path')
    parser.add_argument(
        '--max-diff', type=int, help="max difference of 5' and 3' coverage")
    parser.add_argument(
        '--max-cov', type=int, help="max coverage of 5' or 3' coverage")
    parser.add_argument(
        '--min-cov', type=int, help="min coverage of 5' or 3' coverage")
    parser.add_argument(
        '--min-len', type=int, default=2500, help="min length of the reads")
    parser.add_argument(
        '--bestn', type=int, default=10,
        help="output at least best n overlaps on 5' or 3' ends if possible")
    parser.add_argument(
        '--rid-phase-map', type=str,
        help="the file that encode the relationship of the read id to phase blocks", required=True)
    parser.add_argument(
        '--strictness', type=int, default=STRICTNESS,
        help='If >0, keep *only* the edges which have both nodes of the same phase. Unphased edges are considered dangereous here and removed.')
    args = parser.parse_args(argv[1:])
    return args


def main(argv=sys.argv):
    args = parse_args(argv)
    run(args)


if __name__ == '__main__':  # pragma: no cover
    main()
