import gzip
import os
import re
import subprocess
import sys

import numpy as np


def multitask_bed_generation(
    target_beds_file,
    feature_size=1000,
    merge_overlap=200,
    out_prefix="features",
    chrom_lengths_file=None,
    db_act_file=None,
    db_bed=None,
    ignore_auxiliary=False,
    no_db_activity=False,
    ignore_y=False,
):
    """Merge multiple bed files to select sample sequence regions with at least one
    peak.

    This function outputs a .bed file in the specified directory containing seven
    columns: chromosome, sequence start, sequence end, name, score, strand, and indexes
    of experiments that have a peak in this region.

    Parameters
    ----------
    target_beds_file: str
        Location of the sample file containing experiment label and their corresponding
        file locations. Should be a two column text file, first row contains label,
        second row contain directory for the .bed/.bed.gz file.
    feature_size: int, optional
        Length of the sequence region per sample in output. Default to 1000.
    merge_overlap: int, optional
        After adjusting peaks into feature_size, if two peak regions overlaps more than
        this amount, they will be re-centered and merged into a single sample. Defaults
        to 200.
    output_prefix: str, optional
        Location and naming of the output bed file. Default to 'features.bed'
    chrom_lenghts_file: str, optional
        Location of the chrom.sizes file. Default to None.
    db_act_file: str, optional
        Location of the existing database activity table. Defaults to None.
    db_bed: str, optional
        Location of the existing database .bed file. Defaults to None.
    ignore_auxiliary: bool, optional
        Ignore auxiliary chromosomes. Defaults to False.
    no_db_acticity: bool, optional
        Whether to pass along the activities of the database sequences. Defaults to
        False.
    ignor_y: bool, optional
        Ignore Y chromsosome features. Defaults to False.

    Returns
    -------
    None

    Examples
    --------
    >>> multitask_bed_generation(
        example_file,chrom_lengths_file='/data/hg38.chrom.size',
        feature_size=1000,merge_overlap=200,out_prefix='/data/multitask.bed')
    """

    if not target_beds_file:
        raise Exception(
            "Must provide file labeling the targets and providing BED file paths."
        )

    # determine whether we'll add to an existing DB
    db_targets = []
    db_add = False
    if db_bed is not None:
        db_add = True
        if not no_db_activity:
            if db_act_file is None:
                raise ValueError(
                    "Must provide both activity table or specify -n if you want to add"
                    " to an existing database"
                )
            else:
                # read db target names
                db_act_in = open(db_act_file)
                db_targets = db_act_in.readline().strip().split("\t")
                db_act_in.close()

    # read in targets and assign them indexes into the db
    target_beds = []
    target_dbi = []
    for line in open(target_beds_file):
        a = line.split()
        if len(a) != 2:
            print(a)
            print(
                "Each row of the target BEDS file must contain a label and BED file"
                " separated by whitespace",
                file=sys.stderr,
            )
            sys.exit(1)
        target_dbi.append(len(db_targets))
        db_targets.append(a[0])
        target_beds.append(a[1])

    # read in chromosome lengths
    chrom_lengths = {}
    if chrom_lengths_file is not None:
        chrom_lengths = {}
        for line in open(chrom_lengths_file):
            a = line.split()
            chrom_lengths[a[0]] = int(a[1])
    else:
        print(
            "Warning: chromosome lengths not provided, so regions near ends may be"
            " incorrect.",
            file=sys.stderr,
        )

    #################################################################
    # print peaks to chromosome-specific files
    #################################################################
    chrom_files = {}
    chrom_outs = {}

    peak_beds = target_beds
    if db_add:
        peak_beds.append(db_bed)

    for bi in range(len(peak_beds)):
        if peak_beds[bi][-3:] == ".gz":
            peak_bed_in = gzip.open(peak_beds[bi], "rt")
        else:
            peak_bed_in = open(peak_beds[bi])

        for line in peak_bed_in:
            if not line.startswith("#"):
                a = line.split("\t")
                a[-1] = a[-1].rstrip()

                # hash by chrom/strand
                chrom = a[0]
                strand = "+"
                if len(a) > 5 and a[5] in "+-":
                    strand = a[5]
                chrom_key = (chrom, strand)

                # adjust coordinates to midpoint
                start = int(a[1])
                end = int(a[2])
                mid = int(find_midpoint(start, end))
                a[1] = str(mid)
                a[2] = str(mid + 1)

                # open chromosome file
                if chrom_key not in chrom_outs:
                    chrom_files[chrom_key] = "%s_%s_%s.bed" % (
                        out_prefix,
                        chrom,
                        strand,
                    )
                    chrom_outs[chrom_key] = open(chrom_files[chrom_key], "w")

                # if it's the db bed
                if db_add and bi == len(peak_beds) - 1:
                    if no_db_activity:
                        # set activity to null
                        a[6] = "."
                        print("\t".join(a[:7]), file=chrom_outs[chrom_key])
                        # print >> chrom_outs[chrom_key], '\t'.join(a[:7])
                    else:
                        print(line, chrom_outs[chrom_key])
                        # print >> chrom_outs[chrom_key], line,

                # if it's a new bed
                else:
                    # specify the target index
                    while len(a) < 7:
                        a.append("")
                    a[5] = strand
                    a[6] = str(target_dbi[bi])
                    print("\t".join(a[:7]), file=chrom_outs[chrom_key])
                    # print >> chrom_outs[chrom_key], '\t'.join(a[:7])

        peak_bed_in.close()

    # close chromosome-specific files
    for chrom_key in chrom_outs:
        chrom_outs[chrom_key].close()

    # ignore Y
    if ignore_y:
        for orient in "+-":
            chrom_key = ("chrY", orient)
            if chrom_key in chrom_files:
                print("Ignoring chrY %s" % orient, file=sys.stderr)
                # print >> sys.stderr, 'Ignoring chrY %s' % orient
                os.remove(chrom_files[chrom_key])
                del chrom_files[chrom_key]

    # ignore auxiliary
    if ignore_auxiliary:
        # TODO: \d appears to be an invalid escape sequence. And re.compile will escape
        # \d anyway to \\d.
        primary_re = re.compile("chr\\d+$")
        for chrom_key in chrom_files.keys():
            chrom, strand = chrom_key
            primary_m = primary_re.match(chrom)
            if not primary_m and chrom != "chrX":
                print("Ignoring %s %s" % (chrom, strand), file=sys.stderr)
                # print >> sys.stderr, 'Ignoring %s %s' % (chrom,strand)
                os.remove(chrom_files[chrom_key])
                del chrom_files[chrom_key]

    #################################################################
    # sort chromosome-specific files
    #################################################################
    for chrom_key in chrom_files:
        chrom, strand = chrom_key
        chrom_sbed = "%s_%s_%s_sort.bed" % (out_prefix, chrom, strand)
        sort_cmd = "sortBed -i %s > %s" % (chrom_files[chrom_key], chrom_sbed)
        subprocess.call(sort_cmd, shell=True)
        os.remove(chrom_files[chrom_key])
        chrom_files[chrom_key] = chrom_sbed

    #################################################################
    # parse chromosome-specific files
    #################################################################
    final_bed_out = open("%s.bed" % out_prefix, "w")

    for chrom_key in chrom_files:
        chrom, strand = chrom_key

        open_peaks = []
        for line in open(chrom_files[chrom_key], "rt"):
            a = line.split("\t")
            a[-1] = a[-1].rstrip()

            # construct Peak
            peak_start = int(a[1])
            peak_end = int(a[2])
            peak_act = activity_set(a[6])
            peak = Peak(peak_start, peak_end, peak_act)
            peak.extend(feature_size, chrom_lengths.get(chrom, None))

            if len(open_peaks) == 0:
                # initialize open peak
                open_end = peak.end
                open_peaks = [peak]

            else:
                # operate on exiting open peak

                # if beyond existing open peak
                if open_end - merge_overlap <= peak.start:
                    # close open peak
                    mpeaks = merge_peaks(
                        open_peaks,
                        feature_size,
                        merge_overlap,
                        chrom_lengths.get(chrom, None),
                    )

                    # print to file
                    for mpeak in mpeaks:
                        print(mpeak.bed_str(chrom, strand), file=final_bed_out)
                        # print >> final_bed_out, mpeak.bed_str(chrom, strand)

                    # initialize open peak
                    open_end = peak.end
                    open_peaks = [peak]

                else:
                    # extend open peak
                    open_peaks.append(peak)
                    open_end = max(open_end, peak.end)

        if len(open_peaks) > 0:
            # close open peak
            mpeaks = merge_peaks(
                open_peaks, feature_size, merge_overlap, chrom_lengths.get(chrom, None)
            )

            # print to file
            for mpeak in mpeaks:
                print(mpeak.bed_str(chrom, strand), file=final_bed_out)
                # print >> final_bed_out, mpeak.bed_str(chrom, strand)

    final_bed_out.close()

    # clean
    for chrom_key in chrom_files:
        os.remove(chrom_files[chrom_key])

    #################################################################
    # construct/update activity table
    #################################################################
    final_act_out = open("%s_act.txt" % out_prefix, "w")

    # print header
    cols = [""] + db_targets
    print("\t".join(cols), file=final_act_out)
    # print >> final_act_out, '\t'.join(cols)

    # print sequences
    for line in open("%s.bed" % out_prefix):
        a = line.rstrip().split("\t")
        # index peak
        peak_id = "%s:%s-%s(%s)" % (a[0], a[1], a[2], a[5])

        # construct full activity vector
        peak_act = [0] * len(db_targets)
        for ai in a[6].split(","):
            if ai != ".":
                peak_act[int(ai)] = 1

        # print line
        cols = [peak_id] + peak_act
        print("\t".join([str(c) for c in cols]), file=final_act_out)
        # print >> final_act_out, '\t'.join([str(c) for c in cols])

    final_act_out.close()


def activity_set(act_cs):
    """Return a set of ints from a comma-separated list of int strings.
    Attributes:
        act_cs (str) : comma-separated list of int strings
    Returns:
        set (int) : int's in the original string
    """
    ai_strs = [ai for ai in act_cs.split(",")]

    if ai_strs[-1] == "":
        ai_strs = ai_strs[:-1]

    if ai_strs[0] == ".":
        aset = set()
    else:
        aset = set([int(ai) for ai in ai_strs])

    return aset


def find_midpoint(start, end):
    """ Find the midpoint coordinate between start and end """
    mid = (start + end) / 2
    return int(mid)


def merge_peaks(peaks, peak_size, merge_overlap, chrom_len):
    """Merge and the list of Peaks.
    Repeatedly find the closest adjacent peaks and consider
    merging them together, until there are no more peaks
    we want to merge.
    Attributes:
        peaks (list[Peak]) : list of Peaks
        peak_size (int) : desired peak extension size
        chrom_len (int) : chromsome length
    Returns:
        Peak representing the merger
    """
    max_overlap = merge_overlap
    while len(peaks) > 1 and max_overlap >= merge_overlap:
        # find largest overlap
        max_i = 0
        max_overlap = peaks[0].end - peaks[1].start
        for i in range(1, len(peaks) - 1):
            peaks_overlap = peaks[i].end - peaks[i + 1].start
            if peaks_overlap > max_overlap:
                max_i = i
                max_overlap = peaks_overlap

        if max_overlap >= merge_overlap:
            # merge peaks
            peaks[max_i].merge(peaks[max_i + 1], peak_size, chrom_len)

            # remove merged peak
            peaks = peaks[: max_i + 1] + peaks[max_i + 2 :]

    return peaks


def merge_peaks_dist(peaks, peak_size, chrom_len):
    """Merge and grow the Peaks in the given list.
    Obsolete
    Attributes:
        peaks (list[Peak]) : list of Peaks
        peak_size (int) : desired peak extension size
        chrom_len (int) : chromsome length
    Returns:
        Peak representing the merger
    """
    # determine peak midpoints
    peak_mids = []
    peak_weights = []
    for p in peaks:
        mid = (p.start + p.end - 1) / 2.0
        peak_mids.append(mid)
        peak_weights.append(1 + len(p.act))

    # take the mean
    merge_mid = int(0.5 + np.average(peak_mids, weights=peak_weights))

    # extend to the full size
    merge_start = max(0, merge_mid - peak_size / 2)
    merge_end = merge_start + peak_size
    if chrom_len and merge_end > chrom_len:
        merge_end = chrom_len
        merge_start = merge_end - peak_size

    # merge activities
    merge_act = set()
    for p in peaks:
        merge_act |= p.act

    return Peak(merge_start, merge_end, merge_act)


class Peak:
    """Peak representation
    Attributes:
        start (int) : peak start
        end   (int) : peak end
        act   (set[int]) : set of target indexes where this peak is active.
    """

    def __init__(self, start, end, act):
        self.start = start
        self.end = end
        self.act = act

    def extend(self, ext_len, chrom_len):
        """Extend the peak to the given length
        Args:
            ext_len (int) : length to extend the peak to
            chrom_len (int) : chromosome length to cap the peak at
        """
        mid = find_midpoint(self.start, self.end)
        self.start = max(0, mid - ext_len / 2)
        self.end = self.start + ext_len
        if chrom_len and self.end > chrom_len:
            self.end = chrom_len
            self.start = self.end - ext_len

    def bed_str(self, chrom, strand):
        """Return a BED-style line
        Args:
            chrom (str)
            strand (str)
        """
        if len(self.act) == 0:
            act_str = "."
        else:
            act_str = ",".join([str(ai) for ai in sorted(list(self.act))])
        cols = (
            chrom,
            str(int(self.start)),
            str(int(self.end)),
            ".",
            "1",
            strand,
            act_str,
        )
        return "\t".join(cols)

    def merge(self, peak2, ext_len, chrom_len):
        """Merge the given peak2 into this peak
        Args:
            peak2 (Peak)
            ext_len (int) : length to extend the merged peak to
            chrom_len (int) : chromosome length to cap the peak at
        """
        # find peak midpoints
        peak_mids = [find_midpoint(self.start, self.end)]
        peak_mids.append(find_midpoint(peak2.start, peak2.end))

        # weight peaks
        peak_weights = [1 + len(self.act)]
        peak_weights.append(1 + len(peak2.act))

        # compute a weighted average
        merge_mid = int(0.5 + np.average(peak_mids, weights=peak_weights))

        # extend to the full size
        merge_start = max(0, merge_mid - ext_len / 2)
        merge_end = merge_start + ext_len
        if chrom_len and merge_end > chrom_len:
            merge_end = chrom_len
            merge_start = merge_end - ext_len

        # merge activities
        merge_act = self.act | peak2.act

        # set merge to this peak
        self.start = merge_start
        self.end = merge_end
        self.act = merge_act
