#!/usr/bin/env python
# -*- coding: utf-8 -*-
from __future__ import division
import sys
import logging
import argparse
import json
from hicexplorer import HiCMatrix as hm
from hicexplorer.utilities import enlarge_bins
from scipy import sparse
import numpy as np
import multiprocessing
from hicexplorer._version import __version__

logging.basicConfig()
log = logging.getLogger("hicFindTADs")
log.setLevel(logging.INFO)


def parse_arguments(args=None):
    """
    get command line arguments
    """
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        conflict_handler='resolve',
        description="""
Uses a measure called TAD-separation score to identify the degree of separation between
the left and right regions at each Hi-C matrix bin. This is done for a
running window of different sizes. Then, TADs are called as those
positions having a local TAD-separation score minimum. The TAD-separation score is
measured using the z-score of the Hi-C matrix and is defined as the mean zscore of all
the matrix contacts between the left and right regions (diamond).

To find the TADs, the program  needs to compute first the
TAD scores at different window sizes. Then, the results of that computation
are used to call the TADs. This is convenient to test different filtering criteria quickly
as the demanding step is the computation of TAD-separation scores.

 An simple example usage is:

$ hicFindTads TAD_score -m hic_matrix.h5 -o TAD_score.txt

$ hicFindTads find_TADs -f TAD_score.txt --outPrefix TADs

For detailed help:

 hicFindTADs TAD_score -h

 or

 hicFindTADs find_TADs -h

""")

    parser.add_argument('--version', action='version',
                        version='%(prog)s {}'.format(__version__))

    subparsers = parser.add_subparsers(
        title="commands",
        dest='command',
        metavar='')

    tad_score_subparser = subparsers.add_parser('TAD_score',
                                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # define the arguments
    tad_score_subparser.add_argument('--matrix', '-m',
                                     help='Corrected Hi-C matrix to use for the computations',
                                     required=True)

    tad_score_subparser.add_argument('--outFileName', '-o',
                                     help='File name to store the computation of the TAD-separation score. The format'
                                          'of the output file is chrom start end TAD-sep1 TAD-sep2 TAD-sep3 .. etc. '
                                          'We call this format a bedgraph matrix and can be plotted using '
                                          '`hicPlotTADs`. Each of the TAD-separation scores in the file corresponds to '
                                          'a different window length starting from --minDepth to --maxDepth.',
                                     required=True)

    tad_score_subparser.add_argument('--minDepth',
                                     help='Minimum window length (in bp) to be considered to the left and to the right '
                                          'of each Hi-C bin. This number should be at least 3 times '
                                          'as large as the bin size of the Hi-C matrix.',
                                     metavar='INT bp',
                                     type=int)

    tad_score_subparser.add_argument('--maxDepth',
                                     help='Maximum window length to be considered to the left and to the right '
                                          'of the cut point in bp. This number should around 6-10 times '
                                          'as large as the bin size of the Hi-C matrix.',
                                     metavar='INT bp',
                                     type=int)

    tad_score_subparser.add_argument('--step',
                                     help='Step size when moving from --minDepth to --maxDepth. Note, the step size'
                                          'grows exponentially as '
                                          '`maxDeph + (step * int(x)**1.5) for x in [0, 1, ...]` until  it '
                                          'reaches `maxDepth`. For example, selecting  step=10,000, minDepth=20,000 '
                                          'and maxDepth=150,000 will compute TAD-scores for window sizes: '
                                          '20,000, 30,000, 40,000, 70,000 and 100,000',
                                     metavar='INT bp',
                                     type=int)

    tad_score_subparser.add_argument('--numberOfProcessors', '-p',
                                     help='Number of processors to use ',
                                     type=int,
                                     default=1)

    find_tads_subparser = subparsers.add_parser('find_TADs', formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    find_tads_subparser.add_argument('--tadScoreFile', '-f',
                                     help='File containing the TAD scores (generated by running hicFindTADs TAD_score)',
                                     required=True)

    find_tads_subparser.add_argument('--minBoundaryDistance',
                                     help='Minimum distance between boundaries (in bp). This parameter can be '
                                          'used to reduce spurious boundaries caused by noise. ',
                                     type=int)

    find_tads_subparser.add_argument('--pvalue',
                                     help='P-value threshold. The probability of a local minima to be a boundary '
                                          'is estimated by comparing the distribution (Wilcoxon ranksum) of '
                                          'the  zscores between the left and right '
                                          'regions (diamond) at the local minimum with the matrix zscores for a '
                                          'diamond at --minDepth to the left and a diamond --minDepth to the right. '
                                          'The reported pvalue is the Bonferroni correction all pvalues. Default is '
                                          '0.01',
                                     type=float,
                                     default=0.01
                                     )

    find_tads_subparser.add_argument('--delta',
                                     help='Minimum threshold of the difference between the TAD-separation score of a '
                                          'putative boundary and the mean of the TAD-sep. score of surrounding bins. '
                                          'The delta value reduces spurious boundaries that are shallow, which usually '
                                          'occur at the center of large TADs when the TAD-sep. score is flat. Higher '
                                          'delta threshold values produce more conservative boundary estimations. By '
                                          'default a value of 0.01 is used.',
                                     type=float,
                                     default=0.01
                                     )

    find_tads_subparser.add_argument('--outPrefix',
                                     help='File prefix to save the resulting files: 1. <prefix>_boundaries.bed, which '
                                          'contains the positions of boundaries. The genomic coordinates in this file '
                                          'correspond to the resolution used. Thus, for Hi-C bins of '
                                          '10.000bp the boundary position is 10.000bp long. For restriction fragment '
                                          'matrices the boundary position varies depending on the fragment length '
                                          'at the boundary. 2. <prefix>_domains.bed '
                                          'contains the TADs positions. This is a non-overlapping set of genomic '
                                          'positions. 3. <prefix>_boundaries.gff Similar to the boundaries bed file '
                                          'but with extra information (pvalue, delta). 4. <prefix>_score.bedgraph file '
                                          'contains the TAD-separation score '
                                          'measured at each Hi-C bin coordinate. Is useful to visualize in a genome '
                                          'browser. The delta and pvalue settings are saved as part of the name.',
                                     required=True)

    return parser


def get_cut_weight_by_bin_id(matrix, cut, depth, return_mean=False):
    """
    like get_cut_weight which is the 'diamond' representing the counts
    from a region -dept to +depth from the given bin position (cut):

    """
    if cut < 0 or cut > matrix.shape[0]:
        return None
    # the range [start:i] should have running window
    # length elements (i is excluded from the range)
    start = max(0, cut - depth)
    # same for range [i+1:end] (i is excluded from the range)
    end = min(matrix.shape[0], cut + depth)

    # the idea is to evaluate the interactions
    # between the upstream neighbors with the
    # down stream neighbors. In other words
    # the inter-domain interactions
    if return_mean is True:
        return matrix[start:cut, cut:end].mean()
    else:
        return matrix[start:cut, cut:end].todense().A1


def get_idx_of_bins_at_given_distance(hic_matrix, idx, window_len):
    """
    Returns the right and left indices that are at a given distance in bp
    from `idx`
    Parameters
    ----------
    idx reference index
    window_len distance in bp

    Returns
    -------
    tuple, with left and right bin indices
    """

    # find the bin that is closer to window_len
    # the range [start:i] should have running window
    # length elements (i is excluded from the range)
    chrom, cut_start, cut_end, _ = hic_matrix.getBinPos(idx)
    left_start = max(0, cut_start - window_len)
    left_idx = hic_matrix.getRegionBinRange(chrom, left_start, left_start + 1)[0]

    chr_end_pos = hic_matrix.get_chromosome_sizes()[chrom]
    right_end = min(chr_end_pos, cut_end + window_len) - 1
    right_idx = hic_matrix.getRegionBinRange(chrom, right_end, right_end)[0]

    return left_idx, right_idx


def get_cut_weight(hic_matrix, cut, window_len, return_mean=False):
    """
    like get_cut_weight which is the 'diamond' representing the counts
    from a region -window_len to +window_len from the given bin position (cut):

    """
    if cut < 0 or cut > hic_matrix.matrix.shape[0]:
        return None

    left_idx, right_idx = get_idx_of_bins_at_given_distance(hic_matrix, cut, window_len)

    if return_mean is True:
        return hic_matrix.matrix[left_idx:cut, cut:right_idx].mean()
    else:
        return hic_matrix.matrix[left_idx:cut, cut:right_idx].todense().A1


def get_triangle(hic_matrix, cut, window_len, return_mean=False):
    """
    like get_cut_weight which is the 'diamond' representing the counts
    from a region -window_len to +window_len from the given bin position (cut):

    """
    if cut < 0 or cut > hic_matrix.matrix.shape[0]:
        return None

    left_idx, right_idx = get_idx_of_bins_at_given_distance(hic_matrix, cut, window_len)

    def remove_lower_triangle(matrix):
        """
        remove all values in the lower triangle of a matrix
        """
        return matrix[np.triu_indices_from(matrix)].A1

    edges_left = remove_lower_triangle(hic_matrix.matrix[left_idx:cut, :][:, left_idx:cut].todense())
    edges_right = remove_lower_triangle(hic_matrix.matrix[cut:right_idx, :][:, cut:right_idx].todense())
#    if cut > 1000:
#        import ipdb;ipdb.set_trace()
    return np.concatenate([edges_left, edges_right])


def compute_matrix_wrapper(args):
    return compute_matrix(*args)


def get_incremental_step_size(min_win_size, max_win_size, start_step_len):
    """
    generates a list of incremental windows sizes (measured in bins)

    :param min_win_size: starting window size
    :param max_win_size: end window size
    :param start_step_len: start length
    :return: incremental_step list of bin lengths
    """
    incremental_step = []
    step = -1
    while 1:
        step += 1
        inc_step = min_win_size + int(start_step_len * (step**1.5))
        if step > 1 and inc_step == incremental_step[-1]:
            continue
        if inc_step > max_win_size:
            break
        incremental_step.append(inc_step)
    return incremental_step


def compute_matrix(bins_list, min_win_size=8, max_win_size=50, step_len=2):
    """
    Iterates over the Hi-C matrix computing at each bin
    interface the conductance at different window lengths
    :param hic_ma: Hi-C matrix object from HiCMatrix
    :param outfile: String, path of a file to save the conductance
                matrix in *bedgraph matrix* format
    :return: (chrom, start, end, matrix)
    """
    global hic_ma
    positions_array = []
    cond_matrix = []
    incremental_step = get_incremental_step_size(min_win_size, max_win_size, step_len)
    for cut in bins_list:

        chrom, chr_start, chr_end, _ = hic_ma.cut_intervals[cut]

        # get conductance
        # for multiple window lengths at a time
        mult_matrix = [get_cut_weight(hic_ma, cut, depth, return_mean=True) for depth in incremental_step]
        if np.any(np.isnan(mult_matrix)):
            # skip problematic cases
            continue

        cond_matrix.append(mult_matrix)

        positions_array.append((chrom, chr_start, chr_end))

    try:
        chrom, chr_start, chr_end = zip(*positions_array)
    except:
        import ipdb
        ipdb.set_trace()
    cond_matrix = np.vstack(cond_matrix)

    return chrom, chr_start, chr_end, cond_matrix


def peakdetect(y_axis, x_axis=None, lookahead=3, delta=0, chrom=None):
    """
    Based on the MATLAB script at:
    http://billauer.co.il/peakdet.html

    function for detecting local maximum and minimum in a signal.
    Discovers peaks by searching for values which are surrounded by lower
    or larger values for maximum and minimum respectively

    keyword arguments:
    :param: y_axis -- A list containig the signal over which to find peaks
    :param: x_axis -- (optional) A x-axis whose values correspond to the y_axis list
        and is used in the return to specify the position of the peaks. If
        omitted an index of the y_axis is used. (default: None)
    :param: lookahead -- (optional) distance to look ahead from a peak candidate to
        determine if it is the actual peak
    :param: delta -- (optional) this specifies a minimum difference between a peak and
        the following points, before a peak may be considered a peak. Useful
        to hinder the function from picking up false peaks towards to end of
        the signal. To work well delta should be set to delta >= RMSnoise * 5.


    :return: -- two lists [max_peaks, min_peaks] containing the positive and
        negative peaks respectively. Each cell of the lists contains a tuple
        of: (position, peak_value)
        to get the average peak value do: np.mean(max_peaks, 0)[1] on the
        results to unpack one of the lists into x, y coordinates do:
        x, y = zip(*tab)
    """
    max_peaks = []
    min_peaks = []
    dump = []   # Used to pop the first hit which almost always is false

    # check input data
    if x_axis is None:
        x_axis = np.arange(len(y_axis))

    if len(y_axis) != len(x_axis):
        raise (ValueError, 'Input vectors y_axis and x_axis must have same length')

    # store data length for later use

    if not (np.isscalar(delta) and delta >= 0):
        raise (ValueError, "delta must be a positive number")

    # maximum and minimum candidates are temporarily stored in
    # min_x and min_y respectively
    min_y, max_y = np.Inf, -np.Inf
    max_pos, min_pos = None, None
    search_for = None
    # Only detect peak if there is 'lookahead' amount of points after it
    prev_chrom = None
    for index, (x, y) in enumerate(zip(x_axis[:-lookahead], y_axis[:-lookahead])):
        assert -np.inf < y < np.inf, "Error, infinity value detected for value at position {}".format(index)
        if prev_chrom is None:
            prev_chrom = chrom[index]
        if prev_chrom != chrom[index]:
            # reset variables at the start of a chromosome
            min_y, max_y = np.Inf, -np.Inf
            max_pos, min_pos = None, None
            search_for = None

        prev_chrom = chrom[index]

        if y > max_y:
            max_y = y
            max_pos = x
        if y < min_y:
            min_y = y
            min_pos = x

        # look for max
        if y < max_y - delta and max_y != np.Inf and search_for != 'min':
            # Maximum peak candidate found
            # look ahead in signal to ensure that this is a peak and not jitter
            if y_axis[index:index + lookahead].max() < max_y:
                max_peaks.append([max_pos, max_y])
                dump.append(True)
                # set algorithm to only find minimum now
                max_y = y
                min_y = y
                min_pos = x
                search_for = 'min'
                continue

        # look for min
        if y > min_y + delta and min_y != -np.Inf and search_for != 'max':
            # Minimum peak candidate found
            # look ahead in signal to ensure that this is a peak and not jitter
            if y_axis[index:index + lookahead].min() > min_y:
                min_peaks.append([min_pos, min_y])
                dump.append(False)
                # set algorithm to only find maximum now
                min_y = y
                max_y = y
                max_pos = x
                search_for = 'max'

    # Remove the false hit on the first value of the y_axis
    try:
        if dump[0]:
            max_peaks.pop(0)
        else:
            min_peaks.pop(0)
        del dump
    except IndexError:
        # no peaks were found, should the function return empty lists?
        pass

    return [max_peaks, min_peaks]


def delta_wrt_window(min_idx_list, matrix_avg, chrom, window_len=10):
    """
    Computes the local 'delta' for each minima identified. The local
    delta is defined as the difference between the TAD-separation score at the minimum with respect
    to each of the TAD-separation scores of the `window_len` to the left, plus the `window_len` to
    the right of each minimum. The minimum itself is excluded.

    Parameters
    ----------
    matrix TAD score matrix for different window sizes
    chrom list of chrom names for each matrix row
    window_len the length of the window to look for the local TAD-score average

    Returns
    -------
    A dict of each minima delta to the mean of the TAD-separation score surrounding the minima. The keys
    are the min_idx

    """

    # compute the start and end points of the chromosomes
    # and store it as tuples in chrom_ranges list (e.g. chrom_ranges = [(0, 29254), (29254, 60006), ...]
    unique_chroms, chr_start_idx = np.unique(chrom, return_index=True)
    chr_start_idx = np.concatenate([chr_start_idx, [len(chrom) - 1]])
    chr_start_idx = np.sort(chr_start_idx)
    chrom_ranges = [(chr_start_idx[x], chr_start_idx[x + 1]) for x in range(len(chr_start_idx) - 1)]

    delta_to_mean = {}
    for min_idx in min_idx_list:
        # check that the min_idx is not to close to any of the chromosome boundaries
        close_to_chrom_border = True
        for start_range, end_range in chrom_ranges:
            if start_range < min_idx < end_range:
                if min_idx - window_len >= start_range and min_idx + window_len < end_range:
                    close_to_chrom_border = False
                    continue
        if close_to_chrom_border is True:
            delta_to_mean[min_idx] = np.nan

        else:
            local_tad_score = np.concatenate([matrix_avg[min_idx - window_len:min_idx + 3],
                                              matrix_avg[min_idx + 4: min_idx + window_len]])
            delta_to_mean[min_idx] = local_tad_score.mean() - matrix_avg[min_idx]

    return delta_to_mean


def find_consensus_minima(tad_score_matrix, lookahead=3, chrom=None):
    """
    Finds the minimum over the average values per column
    :param tad_score_matrix: TAD-separation score matrix
    :return:
    """

    tad_score_matrix_avg = tad_score_matrix.mean(axis=1)

    # compute local minima for the matrix average
    _max, _min = peakdetect(tad_score_matrix_avg, lookahead=lookahead, chrom=chrom)
    min_idx, value = zip(*_min)

    # get the delta for each boundary
    delta_to_mean = delta_wrt_window(min_idx, tad_score_matrix_avg, chrom)

    return min_idx, delta_to_mean


def hierarchical_clustering(boundary_list, clusters_cutoff=[]):
    """
    :param boundary_list: is a list of tuples each containing
    the location of a boundary. The order should be sorted
    and contain the following values:
        (chrom, start, value)
    :param clusters_cutoff: List of values to separate clusters. The
        clusters found at those value thresholds are returned.


    :return: z_value, clusters

    For z_value, the format used is similar as the scipy.cluster.hierarchy.linkage() function
    which is described as follows:

    A 4 by :math:`(n-1)` matrix ``z_value`` is returned. At the
    :math:`i`-th iteration, clusters with indices ``z_value[i, 0]`` and
    ``z_value[i, 1]`` are combined to form cluster :math:`n + i`. A
    cluster with an index less than :math:`n` corresponds to one of
    the :math:`n` original observations. The distance between
    clusters ``z_value[i, 0]`` and ``z_value[i, 1]`` is given by ``z_value[i, 2]``. The
    fourth value ``z_value[i, 3]`` represents the number of original
    observations in the newly formed cluster.

    The difference is that instead of a 4 times n-1 array, a
    6 times n-1 array is returned. Where positions 4, and 5
    correspond to the genomic coordinates of ``z_value[i, 0]`` and ``z_value[i, 1]``

    """
    # run the hierarchical clustering per chromosome
    if clusters_cutoff:
        # sort in reverse order
        clusters_cutoff = np.sort(np.unique(clusters_cutoff))[::-1]

    chrom, start, value = zip(*boundary_list)

    unique_chr, indices = np.unique(chrom, return_index=True)
    indices = indices[1:]  # the first element is not needed
    start_per_chr = np.split(start, indices)
    value_per_chr = np.split(value, indices)
    z_value = {}

    def get_domain_positions(boundary_position):
        """
        returns for each boundary a start,end position
        corresponding to each TAD
        :param boundary_position: list of boundary chromosomal positions
        :return: list of (start, end) tuples.
        """
        start_ = None
        domain_list = []
        for position in boundary_position:
            if start_ is None:
                start_ = position
                continue
            domain_list.append((start_, position))
            start_ = position

        return domain_list

    def find_in_clusters(clusters_, search_id):
        """
        Given a list of clusters (each cluster defined as as set,
        the function returns the position in which an id is found
        :param clusters_:
        :param search_id:
        :return:
        """
        for set_idx, set_of_ids in enumerate(clusters_):
            if search_id in set_of_ids:
                return set_idx

    def cluster_to_regions(clusters_, chrom_name):
        """
        Transforms a list of sets of ids from the hierarchical
        clustering to genomic positions
        :param clusters_: cluster ids
        :param chrom_name: chromosome name
        :return: list of tuples with (chrom_name, start, end)

        Example:

        clusters = [set(1,2,3), set(4,5,10)]

        """
        start_list = []
        end_list = []
        for set_ in clusters_:
            if len(set_) == 0:
                continue

            # the ids in the sets are created in such a
            # that the min id is the one with the smaller start position
            start_list.append(domains[min(set_)][0])
            end_list.append(domains[max(set_)][1])

        start_list = np.array(start_list)
        end_list = np.array(end_list)
        order = np.argsort(start_list)

        return zip([chrom_name] * len(order), start_list[order], end_list[order])

    return_clusters = {}  # collects the genomic positions of the clusters per chromosome
    # The values are a list, one for each cutoff.
    for chrom_idx, chrom_name in enumerate(unique_chr):
        z_value[chrom_name] = []
        return_clusters[chrom_name] = []
        clust_cutoff = clusters_cutoff[:]
        domains = get_domain_positions(start_per_chr[chrom_idx])
        clusters = [{x} for x in range(len(domains))]

        # initialize the cluster_x with the genomic position of domain centers
        cluster_x = [int(d_start + float(d_end - d_start) / 2) for d_start, d_end in domains]
        # number of domains should be equal to the number of values minus 1
        assert len(domains) == len(value_per_chr[chrom_idx]) - 1, "error"

        """
        domain:id
             0            1               2            3
         |---------|---------------|----------------|----|
        values:id
         0         1               3                3    4
        values id after removing flanks
                   0               1                2
         """
        values = value_per_chr[chrom_idx][1:-1]  # remove flanking values that do not join TADs

        # from highest to lowest merge neighboring domains
        order = np.argsort(values)[::-1]
        for idx, order_idx in enumerate(order):
            if len(clust_cutoff) and idx + 1 < len(order) and \
                    values[order_idx] >= clust_cutoff[0] > values[order[idx + 1]]:
                clust_cutoff = clust_cutoff[1:]  # remove first element
                return_clusters[chrom_name].append(cluster_to_regions(clusters, chrom_name))
            # merge domains order_idx - 1 and order_idx
            left = find_in_clusters(clusters, order_idx)
            right = find_in_clusters(clusters, order_idx + 1)
            z_value[chrom_name].append((left, right, values[order_idx],
                                        len(clusters[left]) + len(clusters[right]),
                                        cluster_x[left], cluster_x[right]))

            # set as new cluster position the center between the two merged
            # clusters
            gen_dist = int(float(abs(cluster_x[left] - cluster_x[right])) / 2)
            cluster_x.append(min(cluster_x[left], cluster_x[right]) + gen_dist)

            clusters.append(clusters[left].union(clusters[right]))
            clusters[left] = set()
            clusters[right] = set()

    # convert return_clusters from a per chromosome dictionary to
    # a per cut_off dictionary merging all chromosomes in to one list.
    ret_ = {}  # dictionary to hold the clusters per cutoff. The key of
    # each item is the str(cutoff)

    for idx, cutoff in enumerate(clusters_cutoff):
        cutoff = str(cutoff)
        ret_[cutoff] = []
        for chr_name in return_clusters:
            try:
                ret_[cutoff].extend(return_clusters[chr_name][idx])
            except IndexError:
                pass

    return z_value, ret_


def save_linkage(Z, file_name):
    """

    :param Z: Z has a format similar to the scipy.cluster.linkage matrix (see function
                hierarchical_clustering).
    :param file_name: File name to save the results
    :return: None
    """

    try:
        file_h = open(file_name, 'w')
    except IOError:
        log.error("Can't save linkage file:\n{}".format(file_name))
        return

    count = 0
    for chrom, values in Z.iteritems():
        for id_a, id_b, distance, num_clusters, pos_a, pos_b in values:
            count += 1
            file_h.write('{}\t{}\t{}\tclust_{}'
                         '\t{}\t.\t{}\t{}\t{}\n'.format(chrom,
                                                        int(pos_a),
                                                        int(pos_b),
                                                        count,
                                                        distance,
                                                        id_a, id_b,
                                                        num_clusters))


def get_domains(boundary_list):
    """
    returns for each boundary a chrom, start,end position
    corresponding to each TAD
    :param boundary_position: list of boundary chromosomal positions
    :return: list of (chrom, start, end, value) tuples.
    """
    prev_start = None
    prev_chrom = boundary_list[0][0]
    domain_list = []
    for chrom, start, value in boundary_list:
        if start is None:
            prev_start = start
            prev_chrom = chrom
            continue
        if prev_chrom != chrom:
            prev_chrom = chrom
            prev_start = None
            continue
        domain_list.append((chrom, prev_start, start, value))
        prev_start = start
        prev_chrom = chrom

    return domain_list


def toBytes(s):
    """
    Like toString, but for functions requiring bytes in python3
    """
    if sys.version_info[0] == 2:
        return s
    if isinstance(s, bytes):
        return s
    if isinstance(s, str):
        return bytes(s, 'ascii')
    if isinstance(s, list):
        return [toBytes(x) for x in s]
    return s


def save_bedgraph_matrix(outfile, chrom, chr_start, chr_end, score_matrix, args):
    """
    Save matrix as chrom, start, end ,row, values separated by tab
    I call this a bedgraph matrix (bm)

    :param outfile: string file name
    :param chrom: list of chrom names
    :param chr_start: list of start positions
    :param chr_end: list of end positions
    :param score_matrix: list of lists
    :return: None
    """
    params = dict(args._get_kwargs())
    params['zscore_matrix'] = args.outFileName + "_zscore_matrix.h5"
    params_str = json.dumps(params, separators=(',', ':'))

    with open(outfile, 'w') as f:
        f.write(toBytes("#" + params_str + "\n"))
        for idx in range(len(chrom)):
            matrix_values = "\t".join(
                np.char.mod('%f', score_matrix[idx, :]))
            f.write("{}\t{}\t{}\t{}\n".format(chrom[idx], chr_start[idx],
                                              chr_end[idx], matrix_values))


def save_clusters(clusters, file_prefix):
    """

    :param clusters: is a dictionary whose key is the cut of used to create it.
                     the value is a list of tuples, each representing
                      a genomec interval as ('chr', start, end).
    :param file_prefix: file prefix to save the resulting bed files
    :return: list of file names created
    """
    for cutoff, intervals in clusters.iteritems():
        fileh = open("{}_{}.bed".format(file_prefix, cutoff), 'w')
        for chrom, start, end in intervals:
            fileh.write("{}\t{}\t{}\t.\t0\t.\n".format(chrom, start, end))


def save_domains_and_boundaries(chrom, chr_start, chr_end, matrix, min_idx,
                                delta_of_min, pvalue_of_min, args):

    # a boundary is added to the start and end of each chromosome
    # np.unique return index is used to quickly get
    # the indices at which the name of the chromosome changes (chrom, start, end should be  sorted)
    unique_chroms, chr_start_idx = np.unique(chrom, return_index=True)

    # the trick to get the start positions using np.unique only works when
    # more than one chromosome is present
    if len(unique_chroms) == 1:
        chr_start_idx = [0]
        chr_end_idx = [len(chrom) - 1]
    else:
        chr_end_idx = chr_start_idx
        # the indices for the end of the chromosomes
        # are the the start indices - 1, with the exception
        # of the idx == 0 that is transformed into the length
        # of the chromosome to get the idx for the end of the
        # last chromosome
        chr_end_idx[chr_end_idx == 0] = len(chrom)
        chr_end_idx -= 1

    # put all indices together and sort
    min_idx = np.sort(np.concatenate([chr_start_idx, chr_end_idx, min_idx]))

    mean_mat_all = matrix.mean(axis=1)

    prefix = args.outPrefix
    filtered_min_idx = []
    for idx in min_idx:
        # filter by delta and pvalue_thresholds
        if idx not in delta_of_min:
            delta_of_min[idx] = np.nan
        if delta_of_min[idx] >= args.delta and pvalue_of_min[idx] <= args.pvalue:
            filtered_min_idx += [idx]
    log.info("Number of boundaries for delta {} and pval {}: {}".format(args.delta, args.pvalue,
                                                                        len(filtered_min_idx)))
    count = 1
    with open(prefix + '_boundaries.bed', 'w') as file_boundary_bin, open(prefix + '_domains.bed', 'w') as file_domains, open(prefix + '_boundaries.gff', 'w') as gff:
        for idx, min_bin_id in enumerate(filtered_min_idx):
            # skip if the start of the boundary
            # is the end of the chromosome
            if min_bin_id in chr_end_idx:
                continue

            if min_bin_id not in delta_of_min:
                delta_of_min[min_bin_id] = np.nan

            # 1. save boundaries at 1bp position
            right_bin_center = chr_start[min_bin_id] + int((chr_end[min_bin_id] - chr_start[min_bin_id]) / 2)

            left_bin_center = chr_start[min_bin_id - 1] + int((chr_end[min_bin_id - 1] - chr_start[min_bin_id - 1]) / 2)

            if chrom[min_bin_id] != chrom[min_bin_id - 1]:
                continue

            # 2. save the position of the boundary range
            file_boundary_bin.write("{}\t{}\t{}\tB{:05d}\t{}\t.\n".format(chrom[min_bin_id],
                                                                          left_bin_center,
                                                                          right_bin_center,
                                                                          min_bin_id,
                                                                          mean_mat_all[min_bin_id]))

            # safe gff file that can contain more information
            gff.write("{chrom}\tHiCExplorer\tboundary\t{start}\t{end}\t{score}"
                      "\t.\t.\tID=B{id:05d};delta={delta};pvalue={pvalue};tad_sep={score}\n".format(chrom=chrom[min_bin_id],
                                                                                                    start=left_bin_center,
                                                                                                    end=right_bin_center,
                                                                                                    delta=delta_of_min[min_bin_id],
                                                                                                    pvalue=pvalue_of_min[min_bin_id],
                                                                                                    score=mean_mat_all[min_bin_id],
                                                                                                    id=min_bin_id))

            start = chr_start[min_bin_id]
            # check that the next boundary exists and is in the same chromosome
            if idx + 1 == len(filtered_min_idx) or chrom[min_bin_id] != chrom[filtered_min_idx[idx + 1]]:
                continue

            end = chr_start[filtered_min_idx[idx + 1]]

            # 2. save domain intervals
            if count % 2 == 0:
                rgb = '51,160,44'
            else:
                rgb = '31,120,180'

            file_domains.write("{0}\t{1}\t{2}\tID_{6}_{3}\t{4}\t.\t{1}\t{2}\t{5}\n".format(chrom[min_bin_id],
                                                                                           start, end, count,
                                                                                           mean_mat_all[min_bin_id],
                                                                                           rgb, args.delta))

            count += 1

    # save track with mean values in bedgraph format
    with open(args.outPrefix + '_score.bedgraph', 'w') as tad_score:
        for idx in range(1, len(chrom)):
            right_bin_center = chr_start[idx] + int((chr_end[idx] - chr_start[idx]) / 2)
            left_bin_center = chr_start[idx - 1] + int((chr_end[idx - 1] - chr_start[idx - 1]) / 2)
            if right_bin_center < left_bin_center:
                # this condition happens at chromosome borders
                continue
            tad_score.write("{}\t{}\t{}\t{}\n".format(chrom[idx], left_bin_center, right_bin_center,
                                                      mean_mat_all[idx]))


def compute_spectra_matrix(args, matrix=None):

    if args.maxDepth is not None and args.minDepth is not None and args.maxDepth <= args.minDepth:
        exit("Please check that maxDepth is larger than minDepth.")

    global hic_ma
    if matrix is not None:
        hic_ma = matrix
    else:
        hic_ma = hm.hiCMatrix(args.matrix)

    # remove self counts
    log.info('removing diagonal values\n')
    hic_ma.diagflat(value=0)

    # mask bins without any information
    hic_ma.maskBins(hic_ma.nan_bins)
    orig_intervals = hic_ma.cut_intervals

    binsize = hic_ma.getBinSize()

    if args.maxDepth is None:
        if binsize < 1000:
            args.maxDepth = binsize * 60
        elif 1000 <= binsize < 20000:
            args.maxDepth = binsize * 20
        else:
            args.maxDepth = binsize * 10
    elif args.maxDepth < binsize * 5:
        sys.error("Please specify a --maxDepth that is at least 5 times larger than the matrix bin size")
        exit(1)

    if args.minDepth is None:
        if binsize < 1000:
            args.minDepth = binsize * 30
        elif 1000 <= binsize < 20000:
            args.minDepth = binsize * 10
        else:
            args.minDepth = binsize * 5
    elif args.minDepth < binsize * 3:
        log.error("Please specify a --minDepth that is at least 3 times larger than the matrix bin size")
        exit(1)

    if args.step is None:
        if binsize < 1000:
            args.step = binsize * 4
        else:
            args.step = binsize * 2

    elif args.step < binsize:
        log.error("Please specify a --step that is at least the size of the matrix bin size")
        exit(1)

    args.binsize = binsize
    print_args(args)

    # use zscore matrix
    log.info("Computing z-score matrix...\n")
    hic_ma.convert_to_zscore_matrix(maxdepth=args.maxDepth * 2.5, perchr=True)

    # extend remaining bins to remove gaps in
    # the matrix
    new_intervals = enlarge_bins(hic_ma.cut_intervals)

    # rebuilt bin positions if necessary

    if new_intervals != orig_intervals:
        hic_ma.interval_trees, hic_ma.chrBinBoundaries = \
            hic_ma.intervalListToIntervalTree(new_intervals)
        hic_ma.cut_intervals = new_intervals
        hic_ma.orig_bin_ids = None
        hic_ma.nan_bins = None

    hic_ma.save(args.outFileName + "_zscore_matrix.h5")

    if args.minDepth % hic_ma.getBinSize() != 0:
        log.warn('Warning. specified *depth* is not multiple of the '
                 'hi-c matrix bin size ({})\n'.format(hic_ma.getBinSize()))
    if args.step % hic_ma.getBinSize() != 0:
        log.warn('Warning. specified *step* is not multiple of the '
                 'hi-c matrix bin size ({})\n'.format(hic_ma.getBinSize()))

    binsize = hic_ma.getBinSize()

    log.info("Computing TAD-separation scores...\n")
    min_depth_in_bins = int(args.minDepth / binsize)
    max_depth_in_bins = int(args.maxDepth / binsize)
    step_in_bins = int(args.step / binsize)
    if step_in_bins == 0:
        exit("Please select a step size larger than {}".format(binsize))

    incremental_step = get_incremental_step_size(args.minDepth, args.maxDepth, args.step)

    log.info("computing spectrum for window sizes between {} ({} bp)"
             "and {} ({} bp) at the following window sizes {} {}\n".format(min_depth_in_bins,
                                                                           binsize * min_depth_in_bins,
                                                                           max_depth_in_bins,
                                                                           binsize * max_depth_in_bins,
                                                                           step_in_bins, incremental_step))
    if min_depth_in_bins <= 1:
        log.error('ERROR\nminDepth length too small. Use a value that is at least '
                  'twice as large as the bin size which is: {}\n'.format(binsize))
        exit(0)

    if max_depth_in_bins <= 1:
        log.error('ERROR\nmaxDepth length too small. Use a value that is larger '
                  'than the bin size which is: {}\n'.format(binsize))
        exit(0)

    # work only with the upper matrix
    # and remove all pixels that are beyond
    # 2 * max_depth_in_bis which are not required
    # (this is done by subtracting a second sparse matrix
    # that contains only the upper matrix that wants to be removed.
    limit = 2 * max_depth_in_bins
    hic_ma.matrix = sparse.triu(hic_ma.matrix, k=0, format='csr') - sparse.triu(hic_ma.matrix, k=limit, format='csr')
    hic_ma.matrix.eliminate_zeros()

    num_processors = args.numberOfProcessors

    func = compute_matrix_wrapper
    TASKS = []
    bins_to_consider = []
    for chrom in hic_ma.chrBinBoundaries.keys():
        bins_to_consider.extend(range(*hic_ma.chrBinBoundaries[chrom]))

    for idx_array in np.array_split(bins_to_consider, num_processors):
        TASKS.append((idx_array, args.minDepth, args.maxDepth, args.step))

    if num_processors > 1:
        pool = multiprocessing.Pool(num_processors)
        log.info("Using {} processors\n".format(num_processors))
        res = pool.map_async(func, TASKS).get(9999999)
    else:
        res = map(func, TASKS)

    chrom = []
    chr_start = []
    chr_end = []
    matrix = []
    for _chrom, _chr_start, _chr_end, _matrix in res:
        chrom.extend(_chrom)
        chr_start.extend(_chr_start)
        chr_end.extend(_chr_end)
        matrix.append(_matrix)

    matrix = np.vstack(matrix)
    return np.array(chrom), np.array(chr_start).astype(int), np.array(chr_end).astype(int), matrix


def load_spectrum_matrix(file):
    # load spectrum matrix:
    matrix = []
    chrom_list = []
    start_list = []
    end_list = []
    with open(file, 'r') as fh:
        for line in fh:
            if line.startswith("#"):
                # recover the parameters used to generate the spectrum_matrix
                parameters = json.loads(line[1:].strip())
                continue
            fields = line.strip().split('\t')
            chrom, start, end = fields[0:3]
            chrom_list.append(chrom)
            start_list.append(int(float(start)))
            end_list.append(int(float(end)))
            matrix.append(map(float, fields[3:]))

    matrix = np.vstack(matrix)
    chrom = np.array(chrom_list)
    start = np.array(start_list)
    end = np.array(end_list)
    return chrom, start, end, matrix, parameters


def print_args(args):
    """
    Print to stderr the parameters used
    Parameters
    ----------
    args

    Returns
    -------

    """
    for key, value in args._get_kwargs():
        sys.stderr.write("{}:\t{}\n".format(key, value))


def min_pvalue(matrix_file, min_idx, chrom, chr_start, chr_end, window_len=20000, delta=None):
    """
    For each putative local minima, find the -window_len diammond and the +window_len diamond
    and compare with the local minima using wilcoxon rank sum.

    Parameters
    ----------
    min_idx list of local minima
    window_len in base pairs

    Returns
    -------
    list of p-values per each local minima

    """
    hic_ma = hm.hiCMatrix(matrix_file)

    log.info("Computing p-values for window length: {}\n".format(window_len))
    pvalues = []

    from scipy.stats import ranksums

    new_min_idx = []
    for idx in min_idx:

        matrix_idx = hic_ma.getRegionBinRange(chrom[idx], chr_start[idx], chr_end[idx])
        if matrix_idx is None:
            continue
        else:
            matrix_idx = matrix_idx[0]

        new_min_idx += [idx]
        min_chr, min_start, min_end, _ = hic_ma.getBinPos(matrix_idx)
        assert chrom[idx] == min_chr and chr_start[idx] == min_start and chr_end[idx] == min_end

        left_idx, right_idx = get_idx_of_bins_at_given_distance(hic_ma, matrix_idx, window_len)

        left = get_cut_weight(hic_ma, left_idx, window_len)
        right = get_cut_weight(hic_ma, right_idx, window_len)
        boundary = get_cut_weight(hic_ma, matrix_idx, window_len)
        if left is None:
            left = []
        if right is None:
            right = []

        if len(left) == 0 and len(right) == 0:
            pval = np.nan

        elif boundary is None or len(boundary) == 0:
            pval = np.nan
        else:
            try:
                pval1 = ranksums(boundary, left)[1]
                pval2 = ranksums(boundary, right)[1]
                pval = min(pval1, pval2)
            except ValueError:
                # this condition happens when boundary, left and right are only composed of 'zeros'
                pval = np.nan

        pvalues += [pval]
        # if delta[idx] > 0.04 and pval > 0.001:
        #    import ipdb; ipdb.set_trace()

    assert len(pvalues) == len(new_min_idx)
    # bonferroni correction
    pvalues = np.array(pvalues) * len(pvalues)
    pvalues[pvalues > 1] = 1
    return dict(zip(new_min_idx, pvalues))


def min_pvalue_diamond_vs_triangle(matrix_file, min_idx, chrom, chr_start, chr_end, window_len=20000, delta=None):
    """
    For each putative local minima, find the -window_len triangle and the +window_len triangle
    and compare with the local minima dianmond using wilcoxon rank sum.

    Parameters
    ----------
    min_idx list of local minima
    window_len in base pairs

    Returns
    -------
    list of p-values per each local minima

    """
    hic_ma = hm.hiCMatrix(matrix_file)

    log.info("Computing p-values for window length: {}\n".format(window_len))
    pvalues = []

    from scipy.stats import ranksums

    new_min_idx = []
    for idx in min_idx:

        matrix_idx = hic_ma.getRegionBinRange(chrom[idx], chr_start[idx], chr_end[idx])
        if matrix_idx is None:
            continue
        else:
            matrix_idx = matrix_idx[0]

        new_min_idx += [idx]
        min_chr, min_start, min_end, _ = hic_ma.getBinPos(matrix_idx)
        assert chrom[idx] == min_chr and chr_start[idx] == min_start and chr_end[idx] == min_end

        triangle = get_triangle(hic_ma, matrix_idx, window_len)
        boundary = get_cut_weight(hic_ma, matrix_idx, window_len)
        if triangle is None or len(triangle) == 0:
            pval = np.nan
        elif boundary is None or len(boundary) == 0:
            pval = np.nan
        else:
            try:
                pval = ranksums(boundary[boundary.nonzero()], triangle[triangle.nonzero()])[1]
            except ValueError:
                # this condition happens when boundary, left and right are only composed of 'zeros'
                pval = np.nan

        pvalues += [pval]

    assert len(pvalues) == len(new_min_idx)
    # bonferroni correction
    pvalues = np.array(pvalues) * len(pvalues)
    pvalues[pvalues > 1] = 1

    return dict(zip(new_min_idx, pvalues))


def main(args=None):

    args = parse_arguments().parse_args(args)
    if args.command == 'TAD_score':
        chrom, chr_start, chr_end, matrix = compute_spectra_matrix(args)
        save_bedgraph_matrix(args.outFileName, chrom, chr_start, chr_end, matrix, args)
        return

    chrom, chr_start, chr_end, matrix, parameters = load_spectrum_matrix(args.tadScoreFile)

    # perform some checks
    avg_bin_size = np.median(chr_end - chr_start)

    # compute lookahead (in number of bins)
    if args.minBoundaryDistance is None:
        args.minBoundaryDistance = avg_bin_size * 4
    print_args(args)

    lookahead = int(args.minBoundaryDistance / avg_bin_size)
    if lookahead < 1:
        raise (ValueError, "minBoundaryDistance must be '1' or above in value")

    min_idx, delta = find_consensus_minima(matrix, lookahead=lookahead, chrom=chrom)

    pvalues = min_pvalue(parameters['zscore_matrix'], min_idx, chrom, chr_start, chr_end,
                         window_len=parameters['minDepth'], delta=delta)

    # pvalues = min_pvalue_diamond_vs_triangle(parameters['zscore_matrix'], min_idx, chrom, chr_start, chr_end,
    #                     window_len=parameters['minDepth'], delta=delta)
    if len(min_idx) <= 10:
        mat_mean = matrix.mean(axis=1)
        m_mean = mat_mean.mean()
        m_median = np.median(mat_mean)
        m_75 = np.percentile(mat_mean, 75)
        m_25 = np.percentile(mat_mean, 25)

        msg = ("Please check the parameters:\n"
               " delta: {}\n"
               " minBoundaryDistance: {}\n\n"
               "TAD Score values:\n"
               " mean: {:.3f}\n"
               " median: {:.3f}\n"
               " 1st quartile: {:.3f}\n"
               " 3rd quartile: {:.3f}\n".format(args.delta, args.minBoundaryDistance, m_mean, m_median, m_25, m_75))

        if len(min_idx) == 0:
            exit("\n*ERROR*\nNo boundaries were found. {}".format(msg))
        else:
            log.info("Only {} boundaries found. {}".format(len(min_idx), msg))

    save_domains_and_boundaries(chrom, chr_start, chr_end, matrix, min_idx, delta, pvalues, args)

    # turn of hierarchical clustering which is apparently not working.

    # if 2 == 1:
    #     boundary_list = [(hic_ma.cut_intervals[min_][0], hic_ma.cut_intervals[min_][2], mean_mat[min_]) for min_ in min_idx]

    #     Z, clusters = hierarchical_clustering(boundary_list, clusters_cutoff=[0.4, 0.3, 0.2])

    #     save_linkage(Z, args.outPrefix + '_linkage.bed')
    #     save_clusters(clusters, args.outPrefix)
