import io
import gzip
import numpy as np
import os
import intervaltree
from future.utils import string_types
from collections import defaultdict
import logging
logger = logging.getLogger(__name__)


class GenomicRegion(object):
    """
    Class representing a genomic region.
    .. attribute:: chromosome
        Name of the chromosome this region is located on
    .. attribute:: start
        Start position of the region in base pairs
    .. attribute:: end
        End position of the region in base pairs
    .. attribute:: strand
        Strand this region is on (+1, -1)
    .. attribute:: ix
        Index of the region in the context of all genomic
        regions.
    """

    def __init__(self, start, end, chromosome=None, ix=None, strand=None):
        """
        Initialize this object.
        :param start: Start position of the region in base pairs
        :param end: End position of the region in base pairs
        :param chromosome: Name of the chromosome this region is located on
        :param ix: Index of the region in the context of all genomic
                   regions.
        """
        self.start = start
        self.end = end
        self.chromosome = chromosome
        self.ix = ix
        self.strand = strand

    @classmethod
    def from_string(cls, region_string):
        """
        Convert a string into a :class:`~GenomicRegion`.
        This is a very useful convenience function to quickly
        define a :class:`~GenomicRegion` object from a descriptor
        string.
        :param region_string: A string of the form
                              <chromosome>[:<start>-<end>[:<strand>]]
                              (with square brackets indicating optional
                              parts of the string). If any optional
                              part of the string is omitted, intuitive
                              defaults will be chosen.
        :return: :class:`~GenomicRegion`
        """
        chromosome = None
        start = None
        end = None

        # strip whitespace
        no_space_region_string = "".join(region_string.split())
        fields = no_space_region_string.split(':')

        if len(fields) > 3:
            raise ValueError(
                ("Genomic range string must be of the form"
                 " <chromosome>[:<start>-<end>:[<strand>]]"))

        # there is chromosome information
        if len(fields) > 0:
            chromosome = fields[0]

        # there is range information
        if len(fields) > 1 and fields[1] != '':
            start_end_bp = fields[1].split('-')
            if len(start_end_bp) > 0:
                try:
                    start = int(start_end_bp[0])
                except ValueError:
                    raise ValueError("Start of genomic range must be integer")

            if len(start_end_bp) > 1:
                try:
                    end = int(start_end_bp[1])
                except ValueError:
                    raise ValueError("End of genomic range must be integer")

                if not end > start:
                    raise ValueError(
                        "The end coordinate must be bigger than the start.")

        return cls(start=start, end=end, chromosome=chromosome)

    def overlaps(self, region):
        """
        Check if this region overlaps with the specified region.
        :param region: :class:`~GenomicRegion` object or string
        """
        if isinstance(region, str):
            region = GenomicRegion.from_string(region)

        if region.chromosome != self.chromosome:
            return False

        if region.start <= self.end and region.end >= self.start:
            return True
        return False

    def contains(self, region):
        """
        Check if the specified region is completely contained in this region.
        :param region: :class:`~GenomicRegion` object or string
        """
        if isinstance(region, str):
            region = GenomicRegion.from_string(region)

        if region.chromosome != self.chromosome:
            return False

        if region.start >= self.start and region.end <= self.end:
            return True
        return False

    def _equals(self, region):
        if region.chromosome != self.chromosome:
            return False
        if region.start != self.start:
            return False
        if region.end != self.end:
            return False
        return True

    def __eq__(self, other):
        return self._equals(other)

    def __ne__(self, other):
        return not self._equals(other)

    def __str__(self):
        return '{}:{}-{}:{}'.format(
            self.chromosome, self.start, self.end, self.strand)

    def copy(self):
        """
        Make a copy of this GenomicRegion object.
        """
        return GenomicRegion(
            chromosome=self.chromosome, start=self.start, end=self.end)


def is_gzipped(file_name):
    if file_name.endswith('.gz') or file_name.endswith('.gzip'):
        return True
    return False


def _open(file_name, mode='r'):
    if is_gzipped(file_name):
        return io.TextIOWrapper(gzip.open(file_name, mode=mode))
    return open(file_name, mode=mode)


def load_regions(file_name, sep=None, ignore_ix=False):
    """
    Load regions from regions bed file.
    :param file_name: Path to regions bed file.
    :param sep: Delimiter in regions bed, defaults to None (splits by tab)
    :returns: Tuple:
                (
                    List of :class:`~GenomicRegion` objects,
                    ix_converter dict: {region_id: position_in_list},
                    ix2region dict: {region_id: [chromosome, start, end]}
                (
    """
    regions = []
    ix2reg = {}
    ix_converter = None
    with _open(file_name, 'r') as f:
        for i, line in enumerate(f):
            line = line.rstrip()
            fields = line.split(sep)
            if len(fields) > 2:
                chromosome = fields[0]
                start = int(fields[1])
                end = int(fields[2])
                ix = i
                ix2reg[ix] = [chromosome, int(fields[1]), int(fields[2])]
                if not ignore_ix and len(fields) > 3 and fields[3] != '.':  # HicPro
                    if ix_converter is None:
                        ix_converter = dict()
                    if fields[3] in ix_converter:
                        raise ValueError(
                            "name column in region BED must "
                            "only contain unique values! ({})".format(
                                fields[3]))
                    ix_converter[fields[3]] = ix
                regions.append(
                    GenomicRegion(
                        chromosome=chromosome, start=start, end=end, ix=ix))

    if ix_converter is not None:
        ix_converter_rev = {str(v): int(k) for k, v in ix_converter.items()}
        ix2reg = {ix_converter_rev[str(k)]: v for k, v in ix2reg.items()}
    return regions, ix_converter, ix2reg


def region_interval_trees(regions):
    chromosome_intervals = defaultdict(list)
    for region in regions:
        interval = intervaltree.Interval(region.start - 1, region.end, data=region)
        chromosome_intervals[region.chromosome].append(interval)

    region_trees = dict()
    for chromosome, intervals in chromosome_intervals.items():
        region_trees[chromosome] = intervaltree.IntervalTree(intervals)

    return region_trees


def sub_regions(regions, region):
    """
    Get regions from a list the overlap with another region.
    :param regions: List of :class:`~GenomicRegion`
    :param region: :class:`~GenomicRegion` used for overlap calculation
    """
    if isinstance(region, string_types):
        region = GenomicRegion.from_string(region)

    # interval tree?
    try:
        rt = regions[region.chromosome]

        ixs = []
        sr = []
        for interval in rt[region.start - 1: region.end]:
            hit_region = interval.data
            ixs.append(hit_region.ix)
            sr.append(hit_region)
        start_ix = min(ixs)
        end_ix = max(ixs)

    # region list
    except TypeError:
        sr = []
        start_ix = None
        end_ix = None
        for i, r in enumerate(regions):
            if (r.chromosome == region.chromosome
                    and r.start <= region.end and r.end >= region.start):
                if start_ix is None:
                    start_ix = i
                end_ix = i
                sr.append(r.copy())
            else:
                if end_ix is not None:
                    break
    except IndexError:
        sr = []
        start_ix = None
        end_ix = None

    if len(sr) == 0:
        raise ValueError("Region not found in dataset! {}:{}-{}".format(
            region.chromosome, region.start, region.end))

    return sr, start_ix, end_ix


def sub_matrix_regions(hic_matrix, regions, region):
    """
    Get a square sub Hi-C matrix that overlaps a given region.
    :param hic_matrix: A square numpy array
    :param regions: List of :class:`~GenomicRegion`
    :param region: :class:`~GenomicRegion` used for overlap calculation
    """
    sr, start_ix, end_ix = sub_regions(regions, region)

    if start_ix is None:
        return np.empty((0, 0)), sr

    return np.copy(hic_matrix[start_ix:end_ix+1, start_ix:end_ix+1]), sr


def sub_matrix_from_edges_dict(edges, regions, region, default_weight=0.0):
    """
    Get a square sub Hi-C matrix that overlaps a given region.
    :param edges: A dict of the form { source_ix1: { sink_ix1: weight1, sink_ix2: weight2, ...}, ...}
    :param regions: List of :class:`~GenomicRegion` or dict of :class:`~intervaltree.IntervalTree`
    :param region: :class:`~GenomicRegion` used for overlap calculation
    :param default_weight: Default value with which to fill matrix
    """
    sr, start_ix, end_ix = sub_regions(regions, region)

    size = end_ix - start_ix + 1
    m = np.full((size, size), default_weight, dtype='float64')
    for i in range(start_ix, end_ix + 1):
        for j in range(i, end_ix + 1):
            try:
                m[i - start_ix, j - start_ix] = edges[i][j]
                m[j - start_ix, i - start_ix] = edges[i][j]
            except KeyError:
                pass

    return m, sr


def edges_from_sparse_matrix(file_name, ix_converter=None, sep="\t"):
    nan_values = 0
    total_values = 0

    with _open(file_name, 'r') as f:
        for line in f:
            if line.startswith("#"):
                continue
            total_values += 1

            line = line.rstrip()
            fields = line.split(sep)
            if ix_converter is None:
                source, sink, weight = int(fields[0]), int(fields[1]), float(fields[2])
            else:
                source = ix_converter[fields[0]]
                sink = ix_converter[fields[1]]
                weight = float(fields[2])

            if not np.isfinite(weight):
                nan_values += 1
                continue

            if source <= sink:
                yield source, sink, weight
            else:
                yield sink, source, weight

    if nan_values/total_values > 0.05:
        logger.warning("Could not import more than 5% of matrix entries ({}/{}), "
                       "as they were not finite".format(nan_values, total_values))


def edges_dict(edges):
    edges_d = defaultdict(dict)
    for source, sink, weight in edges:
        edges_d[source][sink] = weight
    return edges_d


def load_sparse_matrix(file_name, ix_converter=None, sep="\t"):
    return list(edges_from_sparse_matrix(file_name, ix_converter, sep))


def load_matrix(file_name, size=None, sep=None, ix_converter=None, is_oe=False):
    """Load Numpy matrix from Hi-C matrix file in sparse format.

    :param file_name: Path to Hi-C matrix file in sparse format.
    :param size: Size of the matrix (number of bins / edge length),
                 defaults to None
    :param sep: Delimiter in matrix file, defaults to None (splits by tab)
    :param ix_converter: ix_converter dict corresponding to the matrix file,
                         as produced by :func:`load_regions`, defaults to None
    :returns: Numpy matrix.
    """

    if size is None:
        raise ValueError("Must provide matrix size!")

    if is_oe:
        m = np.ones((size, size))
    else:
        m = np.zeros((size, size))

    for source, sink, weight in edges_from_sparse_matrix(
                                file_name, ix_converter=ix_converter, sep=sep):
        m[source, sink] = weight
        m[sink, source] = weight

    return m


def load_pairs(file_name, sep=None):
    """Load region pairs from bedpe file.

    Read input bedpe file and convert to dict mapping the comparison ID to the
    reference to query regions.
    :param file_name: Path to input bedpe file. Expected columns:
        chromosome_reference, start_reference, end_reference,
        chromosome_query, start_query, end_query,
        comparison_id, dummy column (not used), strand_reference, strand_query
    :param sep: Delimiter in bedpe file, defaults to None (splits by tab)
    :returns: Dict of form {ID: (reference_region, query_region)}, with
              *region being :class:`GenomicRegion` objects.
    """
    pairs = dict()
    for ID, cref, cqry in load_pairs_iter(file_name, sep=sep):
        pairs[ID] = (cref, cqry)
    return pairs


def load_pairs_iter(file_name, sep=None):
    """Load region pairs from bedpe file.

    Read input bedpe file and convert to dict mapping the comparison ID to the
    reference to query regions.
    :param file_name: Path to input bedpe file. Expected columns:
        chromosome_reference, start_reference, end_reference,
        chromosome_query, start_query, end_query,
        comparison_id, dummy column (not used), strand_reference, strand_query
    :param sep: Delimiter in bedpe file, defaults to None (splits by tab)
    :returns: Dict of form {ID: (reference_region, query_region)}, with
              *region being :class:`GenomicRegion` objects.
    """
    pairs = {}
    with _open(file_name, 'r') as f:
        for line in f:
            if line.startswith('#'):
                continue
            sline = line.rstrip()
            fields = sline.split(sep)
            if len(fields) != 10:
                raise ValueError((
                    "10 columns expected but {} found"
                    " in bedpe input"
                    ).format(len(fields)))
            c1, s1, e1, c2, s2, e2, ID, _, str1, str2 = line.split(sep)
            if ID in pairs:
                raise ValueError((
                    "Pair ID {} not unique in bedpe input!"))
            cref = GenomicRegion(
                chromosome=str(c1), start=int(s1) + 1,
                end=int(e1), strand=str(str1))
            cqry = GenomicRegion(
                chromosome=str(c2), start=int(s2) + 1,
                end=int(e2), strand=str(str2))
            yield ID, cref, cqry


def chunks(l, p):
    n = int(np.ceil(len(l)/p))
    """Yield successive len(l)/p-sized chunks from l."""
    for i in range(0, len(l), n):
        yield l[i:i + n]


def read_chromosome_sizes(file_name):
    chrom_sizes = {}
    with open(os.path.expanduser(file_name), 'r') as chrom_sizes_file:
        for line in chrom_sizes_file:
            line = line.rstrip()
            if line != '':
                chromosome, chromosome_length = line.split("\t")
                chrom_sizes[chromosome] = int(chromosome_length)
    return chrom_sizes
