"""classes to fetch files based on overlaps against query regions.
"""
from bx.intervals import intersection


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
class RegionFileFetch(object):
    """Provide file paths based on their genomic coverage and a requested
    genomic region.

    Parameters
    ----------
    from_list : `NoneType` or `list` of `tuple`, optional, default: `NoneType`
        Initialise the `genomic_config.regions.RegionFileFetch` object with
        content from a `list`. The `list` must contain tuples of length 2.
        Where element ``[0]`` is a further tuple describing the genomic
        coverage of the file path at element ``[1]``. The coverage should have
        the structure (``chr_name`` `str`, ``start_pos`` `int`, ``end_pos``
        `int`). However, if a file has full genomic coverage then the tuple
        should have the structure (`NoneType`, `NoneType`, `NoneType`). If it
        covers a single chromosome then it should be (``chr_name`` `str`,
        `NoneType`, `NoneType`).

    Examples
    --------
    Initialise from a list.

    >>> from genomic_config import regions

    >>> files = [
    >>>     ((None, None, None), "/data/full_genome.vcf.gz"),
    >>>     (('1', None, None), "/data/chr1_file.vcf.gz"),
    >>>     (('2', 123455, 223455), "/data/chr2_123455_223455_file.vcf.gz")
    >>> ]

    >>> file_regions = regions.RegionFileFetch(from_list=files)
    """
    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def __init__(self, from_list=None):
        from_list = from_list or []
        self._multi_chr = []
        self._chr = {}
        self._regions = {}

        # Store all the regions that will be encoded by bx.Invervals as I can't
        # workout how to get them all back out if converting to a list, not
        # ideal but will work until I can research this some more.
        self._chunks = []

        for region, file_path in from_list:
            self.add_file(file_path, chr_name=region[0], start_pos=region[1],
                          end_pos=region[2])

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def add_file(self, file_path, chr_name=None, start_pos=None, end_pos=None):
        """Add a file to the object with specified coverage.

        Parameters
        ----------
        file_path : `str`
            The file path to add
        chr_name : `str` or ``NoneType``, optional, default: ``NoneType``
            The chromosome name associated with the file. If the file is
            multi-chromosome then this should be ``NoneType``.
        start_pos : `int` or ``NoneType``, optional, default: ``NoneType``
            The start position associated with the file. If defined, then the
            ``chr_name`` and ``end_pos`` should also be defined.
        end_pos : `int` or ``NoneType``, optional, default: ``NoneType``
            The end position associated with the file. If defined, then the
            ``chr_name`` and ``start_pos`` should also be defined.

        Examples
        --------
        Initialise from a list.

        >>> from genomic_config import regions

        >>> file_regions = regions.RegionFileFetch()

        >>> # Add a file with full genomic coverage
        >>> file_regions.add_file("/data/full_genome.vcf.gz")

        >>> # Add a single chromosome file
        >>> file_regions.add_file("/data/chr1_file.vcf.gz", chr_name="1")

        >>> # Add a file covering a region
        >>> file_regions.add_file(
        >>>     "/data/chr2_123455_223455_file.vcf.gz",
        >>>     chr_name="2",
        >>>     start_pos=123455,
        >>>     end_pos=223455
        >>> )
        """
        if chr_name is not None and end_pos is not None:
            chr_name = str(chr_name)
            start_pos = int(start_pos)
            end_pos = int(end_pos)

            if end_pos < start_pos:
                raise IndexError(
                    "end position should be >= start_position: {0} > "
                    "{1}".format(start_pos, end_pos)
                )

            # add the region
            try:
                self._regions[chr_name].insert(start_pos, end_pos, file_path)
            except KeyError:
                self._regions[chr_name] = intersection.IntervalTree()
                self._regions[chr_name].insert(start_pos, end_pos, file_path)

            # Store a non-bx copy of the region, just in case to_list is called
            self._chunks.append(((chr_name, start_pos, end_pos), file_path))
        elif chr_name is not None:
            chr_name = str(chr_name)
            try:
                self._chr[chr_name].append(file_path)
            except KeyError:
                self._chr[chr_name] = [file_path]
        else:
            self._multi_chr.append(file_path)

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def get_files(self, chr_name, start_pos, end_pos):
        """Retrieve files that overlap a supplied genomic region.

        Parameters
        ----------
        chr_name : `str`
            The chromosome name of the region.
        start_pos : `int`
            The start position of the region.
        end_pos : `int`
            The end position of the region.

        Returns
        -------
        matching_files : `list` of `str`
            All the file paths that overlap the supplied genomic region. If no
            files overlap the region then this will be an empty `list`.

        Notes
        -----
        Genome-wide files will match all regions and be returned. Whole
        chromosome files will match any region with that chromosome. Any
        specific chunks will only match if they overlap.

        Examples
        --------
        Create a set of regions.

        >>> from genomic_config import regions

        >>> files = [
        >>>     (('1', 1, 1000), "/data/chunk1.vcf.gz"),
        >>>     (('1', 1001, 2000), "/data/chunk2.vcf.gz"),
        >>>     (('1', 2001, 3000), "/data/chunk3.vcf.gz"),
        >>>     (('1', 3001, 4000), "/data/chunk4.vcf.gz"),
        >>>     (('1', 4001, 5000), "/data/chunk5.vcf.gz"),
        >>>     (('1', 5001, 6000), "/data/chunk6.vcf.gz"),
        >>>     (('1', 6001, 7000), "/data/chunk7.vcf.gz"),
        >>>     (('1', 7001, 8000), "/data/chunk8.vcf.gz"),
        >>>     (('1', 8001, 9000), "/data/chunk9.vcf.gz"),
        >>>     (('1', 9001, 10000), "/data/chunk10.vcf.gz"),
        >>> ]

        >>> file_regions = regions.RegionFileFetch(from_list=files)

        Now get the files that overlap the query regions:

        >>> file_regions.get_files('1' 500, 2500)
        ["/data/chunk1.vcf.gz", "/data/chunk2.vcf.gz", "/data/chunk3.vcf.gz"]

        >>> file_regions.get_files('2' 1, 1000000)
        []
        """
        chr_name = str(chr_name)
        end_pos += 1
        files = []
        files.extend(self._multi_chr)

        try:
            files.extend(self._chr[chr_name])
        except KeyError:
            pass

        try:
            files.extend(self._regions[chr_name].find(start_pos, end_pos))
        except KeyError:
            pass
        return files

    # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    def to_list(self):
        """Convert all the internal regions into a list format.

        Returns
        -------
        region_list : `NoneType` or `list` of `tuple`, optional, default: `NoneType`
            Initialise the `genomic_config.regions.RegionFileFetch` object with
            content from a `list`. The `list` must contain tuples of length 2.
            Where element ``[0]`` is a further tuple describing the genomic
            coverage of the file path at element ``[1]``. The coverage should
            have the structure (``chr_name`` `str`, ``start_pos`` `int`,
            ``end_pos`` `int`). However, if a file has full genomic coverage
            then the tuple should have the structure (`NoneType`, `NoneType`,
            `NoneType`). If it covers a single chromosome then it should be
            (``chr_name`` `str`, `NoneType`, `NoneType`).
        """
        # Regions that span the whole genome
        regions = [((None, None, None), i) for i in self._multi_chr.copy()]

        # Regions that span whole chromosomes
        for chr_name, region_files in self._chr.items():
            regions.extend([((chr_name, None, None), i) for i in region_files])

        # Add in any specific region chunks
        regions.extend(self._chunks)

        return regions
