<commit_before>"""
BGEN file reader.
"""

# This file is part of geneparse.
#
# The MIT License (MIT)
#
# Copyright (c) 2017 Pharmacogenomics Centre
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import zlib
import sqlite3
from os import path
from struct import unpack

import numpy as np

from .. import logging
from ..core import GenotypesReader, Genotypes, Variant

try:
    import zstd
    HAS_ZSTD = True
except ImportError:
    HAS_ZSTD = False


CHROM_STR_ENCODE = {"0{}".format(chrom): str(chrom) for chrom in range(1, 10)}
CHROM_STR_ENCODE["23"] = "X"
CHROM_STR_ENCODE["24"] = "Y"

CHROM_STR_DECODE = {v: k for k, v in CHROM_STR_ENCODE.items()}


class BGENReader(GenotypesReader):
    def __init__(self, filename, sample_filename=None, chromosome=None,
                 probability_threshold=0.9):
        """BGEN file reader.

        Args:
            filename (str): The name of the BGEN file.
            sample_filename (str): The name of the sample file (optional).
            probability_threshold (float): The probability threshold.

        """
        # The BGEN file
        self._bgen_file = open(filename, "rb")
        self._parse_header_block()

        # Getting the sample
        if not self._has_sample:
            # TODO: Parse sample file here
            if sample_filename is None:
                raise ValueError("No sample information in BGEN file, "
                                 "requires a 'sample_filename'")

        else:
            self._parse_sample_identifier_block()

        # If we have an index, we open it (sqlite3)
        if not path.isfile(filename + ".bgi"):
            raise ValueError("{}: no index file".format(filename))
        self._bgen_db = sqlite3.connect(filename + ".bgi")
        self._bgen_index = self._bgen_db.cursor()

        # The probability threshold
        self.prob_t = probability_threshold

        # The chromosome
        self.chrom = chromosome

    def close(self):
        if self._bgen_file:
            self._bgen_file.close()
        if self._bgen_db:
            self._bgen_db.close()

    def _parse_header_block(self):
        """Parses the BGEN file header."""
        # Getting the data offset (the start point of the data
        self._offset = unpack("I", self._bgen_file.read(4))[0]

        # Getting the header size
        self._header_size = unpack("I", self._bgen_file.read(4))[0]

        # Getting the number of samples and variants
        self.nb_variants = unpack("I", self._bgen_file.read(4))[0]
        self.nb_samples = unpack("I", self._bgen_file.read(4))[0]

        # Checking the magic number
        magic = self._bgen_file.read(4)
        if magic != b"bgen":
            # The magic number might be 0, then
            if unpack("I", magic)[0] != 0:
                raise ValueError(
                    "{}: invalid BGEN file.".format(self._bgen_file.name)
                )

        # Passing through the "free data area"
        self._bgen_file.read(self._header_size - 20)

        # Reading the flag
        flag = np.unpackbits(
            np.array([[_] for _ in self._bgen_file.read(4)], dtype=np.uint8),
            axis=1,
        )

        # Getting the compression type from the layout
        compression = self._bits_to_int(flag[0, -2:])
        self._is_compressed = False
        if compression == 0:
            # No decompression required
            self._decompress = self._no_decompress

        elif compression == 1:
            # ZLIB decompression
            self._decompress = zlib.decompress
            self._is_compressed = True

        elif compression == 2:
            if not HAS_ZSTD:
                raise ValueError("zstandard module is not installed")

            # ZSTANDARD decompression (needs to be check)
            self._decompress = zstd.ZstdDecompressor().decompress
            self._is_compressed = True

        # Getting the layout
        layout = self._bits_to_int(flag[0, -6:-2])
        if layout == 0:
            raise ValueError(
                "{}: invalid BGEN file".format(self._bgen_file.name)
            )
        elif layout == 1:
            self._layout = 1
        elif layout == 2:
            self._layout = 2
        else:
            raise ValueError(
                "{}: {} invalid layout type".format(self._bgen_file.name,
                                                    layout)
            )

        # Checking if samples are in the file
        self._has_sample = flag[-1, 0] == 1

    def _parse_sample_identifier_block(self):
        """Parses the sample identifier block."""
        # Getting the block size
        block_size = unpack("I", self._bgen_file.read(4))[0]
        if block_size + self._header_size > self._offset:
            raise ValueError(
                "{}: invalid BGEN file".format(self._bgen_file.name)
            )

        # Checking the number of samples
        n = unpack("I", self._bgen_file.read(4))[0]
        if n != self.nb_samples:
            raise ValueError(
                "{}: invalid BGEN file".format(self._bgen_file.name)
            )

        # Getting the sample information
        samples = []
        for i in range(self.nb_samples):
            size = unpack("H", self._bgen_file.read(2))[0]
            samples.append(self._bgen_file.read(size).decode())
        self.samples = tuple(samples)

    @staticmethod
    def _bits_to_int(bits):
        """Converts bits to int."""
        result = 0
        for bit in bits:
            result = (result << 1) | bit
        return result

    def _no_decompress(data):
        """No compression, so we return the data as is."""
        return data

    def _seek_to_first_variant(self):
        """Seeks to the first variant of the file."""
        self._bgen_index.execute(
            "SELECT file_start_position "
            "FROM Variant "
            "ORDER BY file_start_position "
            "LIMIT 1",<commit_msg>Now parses an external sample file if there is no sample block in the header<commit_after>"""
BGEN file reader.
"""

# This file is part of geneparse.
#
# The MIT License (MIT)
#
# Copyright (c) 2017 Pharmacogenomics Centre
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.


import zlib
import sqlite3
from os import path
from struct import unpack

import numpy as np

from .. import logging
from ..core import GenotypesReader, Genotypes, Variant

try:
    import zstd
    HAS_ZSTD = True
except ImportError:
    HAS_ZSTD = False


CHROM_STR_ENCODE = {"0{}".format(chrom): str(chrom) for chrom in range(1, 10)}
CHROM_STR_ENCODE["23"] = "X"
CHROM_STR_ENCODE["24"] = "Y"

CHROM_STR_DECODE = {v: k for k, v in CHROM_STR_ENCODE.items()}


class BGENReader(GenotypesReader):
    def __init__(self, filename, sample_filename=None, chromosome=None,
                 probability_threshold=0.9):
        """BGEN file reader.

        Args:
            filename (str): The name of the BGEN file.
            sample_filename (str): The name of the sample file (optional).
            probability_threshold (float): The probability threshold.

        """
        # The BGEN file
        self._bgen_file = open(filename, "rb")
        self._parse_header_block()

        # Getting the sample
        if not self._has_sample:
            # TODO: Parse sample file here
            if sample_filename is None:
                raise ValueError("No sample information in BGEN file, "
                                 "requires a 'sample_filename'")
            self._parse_sample_file(sample_filename)

        else:
            self._parse_sample_identifier_block()

        # If we have an index, we open it (sqlite3)
        if not path.isfile(filename + ".bgi"):
            raise ValueError("{}: no index file".format(filename))
        self._bgen_db = sqlite3.connect(filename + ".bgi")
        self._bgen_index = self._bgen_db.cursor()

        # The probability threshold
        self.prob_t = probability_threshold

        # The chromosome
        self.chrom = chromosome

    def close(self):
        if self._bgen_file:
            self._bgen_file.close()
        if self._bgen_db:
            self._bgen_db.close()

    def _parse_header_block(self):
        """Parses the BGEN file header."""
        # Getting the data offset (the start point of the data
        self._offset = unpack("I", self._bgen_file.read(4))[0]

        # Getting the header size
        self._header_size = unpack("I", self._bgen_file.read(4))[0]

        # Getting the number of samples and variants
        self.nb_variants = unpack("I", self._bgen_file.read(4))[0]
        self.nb_samples = unpack("I", self._bgen_file.read(4))[0]

        # Checking the magic number
        magic = self._bgen_file.read(4)
        if magic != b"bgen":
            # The magic number might be 0, then
            if unpack("I", magic)[0] != 0:
                raise ValueError(
                    "{}: invalid BGEN file.".format(self._bgen_file.name)
                )

        # Passing through the "free data area"
        self._bgen_file.read(self._header_size - 20)

        # Reading the flag
        flag = np.unpackbits(
            np.array([[_] for _ in self._bgen_file.read(4)], dtype=np.uint8),
            axis=1,
        )

        # Getting the compression type from the layout
        compression = self._bits_to_int(flag[0, -2:])
        self._is_compressed = False
        if compression == 0:
            # No decompression required
            self._decompress = self._no_decompress

        elif compression == 1:
            # ZLIB decompression
            self._decompress = zlib.decompress
            self._is_compressed = True

        elif compression == 2:
            if not HAS_ZSTD:
                raise ValueError("zstandard module is not installed")

            # ZSTANDARD decompression (needs to be check)
            self._decompress = zstd.ZstdDecompressor().decompress
            self._is_compressed = True

        # Getting the layout
        layout = self._bits_to_int(flag[0, -6:-2])
        if layout == 0:
            raise ValueError(
                "{}: invalid BGEN file".format(self._bgen_file.name)
            )
        elif layout == 1:
            self._layout = 1
        elif layout == 2:
            self._layout = 2
        else:
            raise ValueError(
                "{}: {} invalid layout type".format(self._bgen_file.name,
                                                    layout)
            )

        # Checking if samples are in the file
        self._has_sample = flag[-1, 0] == 1

    def _parse_sample_identifier_block(self):
        """Parses the sample identifier block."""
        # Getting the block size
        block_size = unpack("I", self._bgen_file.read(4))[0]
        if block_size + self._header_size > self._offset:
            raise ValueError(
                "{}: invalid BGEN file".format(self._bgen_file.name)
            )

        # Checking the number of samples
        n = unpack("I", self._bgen_file.read(4))[0]
        if n != self.nb_samples:
            raise ValueError(
                "{}: invalid BGEN file".format(self._bgen_file.name)
            )

        # Getting the sample information
        samples = []
        for i in range(self.nb_samples):
            size = unpack("H", self._bgen_file.read(2))[0]
            samples.append(self._bgen_file.read(size).decode())
        self.samples = tuple(samples)

        # Just a check with the header
        if len(self.samples) != self.get_number_samples():
            raise ValueError("{}: number of samples different between header "
                             "and sample block".format(self._bgen_file.name))

    def _parse_sample_file(self, fn):
        """Parses the sample file."""
        # Reading the samples
        with open(fn, "r") as f:
            sample_data = [_.split(" ") for _ in f.read().splitlines()]

        # Splitting the header and data
        header = {name: i for i, name in enumerate(sample_data[0])}
        sample_data = sample_data[2:]

        # Checking the header
        if ("ID_1" not in header) or ("ID_2" not in header):
            raise ValueError("{}: wrong header for sample file".format(fn))

        # Keeping only the IDs
        id1 = [_[header["ID_1"]] for _ in sample_data]
        id2 = [_[header["ID_2"]] for _ in sample_data]

        # We usually keep the ID_2 since it's the IID and not the FID
        self.samples = tuple(id2)
        if len(id2) != len(set(id2)):
            logging.info(
                "Setting the index as 'fid_iid' because the individual IDs "
                "are not unique."
            )
            self.samples = tuple(
                ["{}_{}".format(fid, iid) for fid, iid in zip(id1, id2)]
            )

            # Last check
            if len(self.samples) != len(set(self.samples)):
                raise ValueError("{}: duplicated fid_iid".format(fn))

        # Checking with the number of samples in the BGEN file
        if len(self.samples) != self.get_number_samples():
            raise ValueError("{}: invalid number of samples compared to BGEN "
                             "file".format(fn))

    @staticmethod
    def _bits_to_int(bits):
        """Converts bits to int."""
        result = 0
        for bit in bits:
            result = (result << 1) | bit
        return result

    def _no_decompress(data):
        """No compression, so we return the data as is."""
        return data

    def _seek_to_first_variant(self):
        """Seeks to the first variant of the file."""
        self._bgen_index.execute(
            "SELECT file_start_position "
            "FROM Variant "
            "ORDER BY file_start_position "
            "LIMIT 1",
{'fraction_non_alphanumeric': 0.09392861679156611, 'fraction_numerical': 0.008383081417502857, 'mean_word_length': 2.1206896551724137, 'pattern_counts': {'":': 2, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1327691', 'n_tokens_mistral': 4533, 'n_tokens_neox': 4315, 'n_words': 1428}