#!/usr/bin/env python3

import os
import argparse
from math import floor
from shutil import copyfile

DBG = True if os.environ.get('DBG') is not None else False


def parse_args():
    """ Set up the parsing of command-line arguments """

    parser = argparse.ArgumentParser(
        description="Benchmark: compute the union/intersection among coverage "
                    "intervals to compare performance between aligner or SNV "
                    "caller",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    requiredNamed = parser.add_argument_group('required named arguments')
    requiredNamed.add_argument(
        "-d", required=True, default=None, metavar='DIR1,DIR2,...',
        dest='directories',
        help="Sub-directories where results from vpipeBench are stored"
    )
    parser.add_argument(
        "-r", required=False, default=None, metavar='FASTA',
        dest='reference_file', type=str, help="Referene sequence"
    )
    parser.add_argument(
        "-wl", required=False, default=None, metavar='len1,len2,...',
        dest='window_len', type=str,
        help="Window length used by ShoRAH to construct overlapping windows"
    )
    parser.add_argument(
        "-ws", required=False, default=None, metavar='shift1,shift2,...',
        dest='window_shift', type=str,
        help="Length of the window shift used by ShoRAH to construct "
             "overlapping windows"
    )
    parser.add_argument(
        "--caller", required=False, default=None, action='store_true',
        dest="comparison_callers",
        help="Indicate whether to construct alignment intervals for comparing "
             "mutation callers"
    )
    parser.add_argument(
        "-o", required=False, default=None, metavar='DIR', dest='outdir',
        type=str, help="Output directory"
    )

    return parser.parse_args()


def read_fasta(fasta_file):
    "Parse fasta file"
    count = 0
    headers = []
    sequences = []
    aux = []
    with open(fasta_file, 'r') as infile:
        for line in infile:
            record = line.rstrip()
            if record and record[0] == '>':
                headers.append(record[1:])
                if count > 0:
                    sequences.append(''.join(aux))
                    aux = []
            else:
                aux.append(record)
            count += 1

    sequences.append(''.join(aux))
    return headers, sequences


def parse_region(region, shorah, reference_file, win_len, win_shift):

    def parse_region_aux(region, shorah, reference_file, win_len, win_shift):
        aux = region.split(":")
        reference_name = aux[0]
        aux = aux[1].split("-")
        start = int(aux[0])
        end = int(aux[1])

        if shorah:
            header, reference = read_fasta(reference_file)
            reference_len = len(reference[0])
            # ShoRAH was used for SNV calling
            # Assuming 3 windows were used for SNV calling, identify region
            # that is spanned by at least 2 windows (below, using
            # 0-based indexing and closed intervals)
            start_ = max(0, start - win_len - 1)
            end_ = min(reference_len, end + win_len)
            num_windows = floor((end_ - (start_ + win_len - 1)) / win_shift) + 1
            offset = 2 * win_shift

            start = max(0, start - offset - 1)
            # In order to identify the region which is covered by at least
            # two windows, add to the end of the first window the
            # increment multiply by the number of windows - 2 (i.e.,
            # discarding last window). In this case assuming half-open
            # interval [start, end)
            end = min(reference_len,
                      start_ + win_len + (num_windows - 2) * win_shift)

        return reference_name, start, end

    aux = region.split(",")
    reference_name = []
    interval = []
    for s in aux:
        ref_, start_, end_ = parse_region_aux(s, shorah, reference_file,
                                              win_len, win_shift)
        reference_name.append(ref_)
        interval.append((start_, end_))

    return reference_name, interval


def parse_intervals(input_file, union_intervals, intersection_intervals,
                    shorah, reference_file, window_len, window_shift):

    def merge_regions(intervals_new, intervals_old):
        """
        Merge intervals by:
        (1) sorting the intervals in decreasing order, and
        (2) iterate through intervals, comparing the current interval with the
            last one. If the interval overlap, merge both intervals into one
        """
        intervals = intervals_old + intervals_new
        intervals_sorted = sorted(intervals, key=lambda tup: tup[0])
        intervals_merged = []

        for interval in intervals_sorted:
            if not intervals_merged:
                intervals_merged.append(interval)
            else:
                last = intervals_merged[-1]
                # Test for overlap between current interval and the last one.
                # We know that last[0] <= interval[0]
                if interval[0] <= last[1]:
                    upper_bound = max(last[1], interval[1])
                    intervals_merged[-1] = (last[0], upper_bound)
                else:
                    intervals_merged.append(interval)

        return intervals_merged

    def find_intersect(intervals_new, intervals_old):

        intervals_intersected = []

        for i_new in intervals_new:
            for i_old in intervals_old:
                lower_bound = max(i_old[0], i_new[0])
                upper_bound = min(i_old[1], i_new[1])

                if lower_bound < upper_bound:
                    intervals_intersected.append((lower_bound, upper_bound))

        return intervals_intersected

    idx = 0
    with open(input_file, 'r') as infile:
        for line in infile:
            record = line.rstrip().split('\t')
            key = record[0]
            if len(record) > 1:
                win_len = int(window_len[idx]) if window_len else None
                win_shift = int(window_shift[idx]) if window_shift else None
                reference_name, intervals = parse_region(
                        record[1], shorah, reference_file, win_len, win_shift)
            else:
                reference_name, intervals = None, None
            idx += 1

            if key in union_intervals:
                # Update region if it was empty or if has to be enlarged
                if not union_intervals[key][1]:
                    if DBG:
                        print(f"Assigning {key}: {reference_name}, {intervals}")
                    union_intervals[key] = [reference_name, intervals]
                elif intervals:
                    intervals_new = merge_regions(intervals,
                                                  union_intervals[key][1])
                    # It assumes a unique reference sequence
                    union_intervals[key][0] = [
                        reference_name[0] for i in range(len(intervals_new))]
                    union_intervals[key][1] = intervals_new
                    if DBG:
                        print(f"Merging intervals for {key}: "
                              f"{union_intervals[key][0]}, "
                              f"{union_intervals[key][1]}")
            else:
                if DBG:
                    print(f"Assigning {key}: {reference_name}, {intervals}")
                union_intervals[key] = [reference_name, intervals]

            if key in intersection_intervals:
                if intersection_intervals[key][1] and intervals:
                    intervals_new = find_intersect(
                        intervals, intersection_intervals[key][1])
                    if intervals_new:
                        # It assumes a unique reference sequence
                        intersection_intervals[key][0] = [
                            reference_name[0] for i in range(len(intervals_new))]
                    intersection_intervals[key][1] = intervals_new
                elif not intervals:
                    # One empty interval is enough to set the intersect to the
                    # empty set
                    intersection_intervals[key] = None, None
            else:
                intersection_intervals[key] = [reference_name, intervals]

    return union_intervals, intersection_intervals


def read_coverage_file(input_file, reference_file, window_len, window_shift):

    out_dict = {}
    idx = 0
    with open(input_file, 'r') as infile:
        for line in infile:
            record = line.rstrip().split('\t')
            key = record[0]
            if len(record) > 1:
                win_len = int(window_len[idx]) if window_len else None
                win_shift = int(window_shift[idx]) if window_shift else None
                reference_name, intervals = parse_region(
                        record[1], True, reference_file, win_len, win_shift)
            idx += 1
            out_dict[key] = [reference_name, intervals]

    return out_dict


def write_output(output_file, intervals_dir):

    with open(output_file, 'w') as outfile:
        for key, value in intervals_dir.items():
            if value[0] and value[1]:
                str_ = []
                for idx, ref in enumerate(value[0]):
                    region = value[1][idx]
                    str_.append(f"{ref}:{region[0]}-{region[1]}")
                str_ = ','.join(str_)
                outfile.write(f"{key}\t{str_}\n")
            else:
                outfile.write(f"{key}\t\n")


def main():
    args = parse_args()

    union_dir_shorah = {}
    intersection_dir_shorah = {}

    union_dir = {}
    intersection_dir = {}

    if args.window_len:
        win_len = args.window_len.split(',')
    if args.window_shift:
        win_shift = args.window_shift.split(',')

    outdir = args.outdir if args.outdir is not None else os.getcwd()

    # Aggregarting results
    directories = args.directories.split(",")
    for indir in directories:
        if args.comparison_callers:
            # For the comparison of tools for mutation calling, the alignment
            # remains a constant. However, ShoRAH only reports SNVs that are
            # covered by at least two of the overlapping windows.
            # "union", in this context, results in using all positions above
            # certain coverage
            # "intersect", accounts for the overlapping windows.
            aligner = indir.split('-')[0]
            if 'shorah' in indir:
                input_file = os.path.join(indir, "variants",
                                          "coverage_intervals.tsv")
                intersect_caller = read_coverage_file(
                        input_file, args.reference_file, win_len, win_shift
                )
                output_file = os.path.join(
                        outdir, f"intersect_coverage_intervals_{aligner}.tsv"
                )
                write_output(output_file, intersect_caller)
            else:
                src = os.path.join(os.getcwd(), indir, "stats",
                                   "coverage_intervals.tsv")
                output_file = os.path.join(
                        outdir, f"union_coverage_intervals_{aligner}.tsv"
                )
                copyfile(src, output_file)
        else:
            if 'shorah' in indir:
                input_file = os.path.join(indir, "variants",
                                          "coverage_intervals.tsv")
                union_dir_shorah, intersection_dir_shorah = parse_intervals(
                    input_file, union_dir_shorah, intersection_dir_shorah,
                    True, args.reference_file, win_len, win_shift,
                )
            else:
                input_file = os.path.join(indir, "stats",
                                          "coverage_intervals.tsv")
                union_dir, intersection_dir = parse_intervals(
                    input_file, union_dir, intersection_dir, False, None,
                    None, None,
                )

    if not args.comparison_callers:
        output_file = os.path.join(outdir,
                                   "union_coverage_intervals_ShoRAH.tsv")
        write_output(output_file, union_dir_shorah)

        output_file = os.path.join(outdir,
                                   "intersect_coverage_intervals_ShoRAH.tsv")
        write_output(output_file, intersection_dir_shorah)

        output_file = os.path.join(outdir, "union_coverage_intervals.tsv")
        write_output(output_file, union_dir)

        output_file = os.path.join(outdir,
                                   "intersect_coverage_intervals.tsv")
        write_output(output_file, intersection_dir)


if __name__ == '__main__':
    main()
