text: # MIT License
  #
  # Copyright (c) 2019 Tskit Developers
  #
  # 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.
  """
  Compression utilities for tskit tree sequences.
  """
  import contextlib
  import json
  import logging
  import os.path
  import tempfile
  import warnings
  import zipfile
  from typing import Mapping
  
  import humanize
  import numcodecs
  import numpy as np
  import tskit
  import zarr
  
  from . import exceptions
  from . import provenance
  
  logger = logging.getLogger(__name__)
  
  FORMAT_NAME = "tszip"
  FORMAT_VERSION = [1, 0]
  
  
  def minimal_dtype(array):
      """
      Returns the smallest dtype that can be used to represent the values in the
      specified array. If the array is not one of the integer types, the dtype of the
      array itself is returned directly.
      """
      dtype = array.dtype
      if array.shape[0] == 0:
          return dtype
      if dtype.kind == "u":
          maxval = np.max(array)
          dtypes = [np.uint8, np.uint16, np.uint32, np.uint64]
          for dtype in map(np.dtype, dtypes):
              if maxval <= np.iinfo(dtype).max:
                  break
      elif dtype.kind == "i":
          minval = np.min(array)
          maxval = np.max(array)
          dtypes = [np.int8, np.int16, np.int32, np.int64]
          for dtype in map(np.dtype, dtypes):
              if np.iinfo(dtype).min <= minval and maxval <= np.iinfo(dtype).max:
                  break
      return dtype
  
  
  def compress(ts, path, variants_only=False):
      """
      Compresses the specified tree sequence and writes it to the specified path.
      By default, fully lossless compression is used so that tree sequences are
      identical before and after compression. By specifying the ``variants_only``
      option, a lossy compression can be used, which discards any information
      that is not needed to represent the variants (which are stored losslessly).
  
      :param tskit.TreeSequence ts: The input tree sequence.
      :param str destination: The string or :class:`pathlib.Path` instance describing
          the location of the compressed file.
      :param bool variants_only: If True, discard all information not necessary
          to represent the variants in the input file.
      """
      destination = str(path)
      # Write the file into a temporary directory on the same file system so that
      # we can write the output atomically.
      destdir = os.path.dirname(os.path.abspath(destination))
      with tempfile.TemporaryDirectory(dir=destdir, prefix=".tszip_work_") as tmpdir:
          filename = os.path.join(tmpdir, "tmp.trees.tgz")
          logging.debug(f"Writing to temporary file {filename}")
          with zarr.ZipStore(filename, mode="w") as store:
              root = zarr.group(store=store)
              compress_zarr(ts, root, variants_only=variants_only)
          os.replace(filename, destination)
      logging.info(f"Wrote {destination}")
  
  
  def decompress(path):
      """
      Decompresses the tszip compressed file and returns a tskit tree sequence
      instance.
  
      :param str path: The location of the tszip compressed file to load.
      :rtype: tskit.TreeSequence
      :return: A :class:`tskit.TreeSequence` instance corresponding to the
          the specified file.
      """
      with load_zarr(path) as root:
          return decompress_zarr(root)
  
  
  class Column:
      """
      A single column that is stored in the compressed output.
      """
  
      def __init__(self, name, array, delta_filter=False):
          self.name = name
          self.array = array
          self.delta_filter = delta_filter
  
      def compress(self, root, compressor):
          shape = self.array.shape
          chunks = shape
          if shape[0] == 0:
              chunks = (1,)
          dtype = minimal_dtype(self.array)
          filters = None
          if self.delta_filter:
              filters = [numcodecs.Delta(dtype=dtype)]
          compressed_array = root.empty(
              self.name,
              chunks=chunks,
              shape=shape,
              dtype=dtype,
              filters=filters,
              compressor=compressor,
          )
          compressed_array[:] = self.array
          ratio = 0
          if compressed_array.nbytes > 0:
              ratio = compressed_array.nbytes / compressed_array.nbytes_stored
          logger.debug(
              "{}: output={} compression={:.1f}".format(
                  self.name,
                  humanize.naturalsize(compressed_array.nbytes_stored, binary=True),
                  ratio,
              )
          )
  
  
  def compress_zarr(ts, root, variants_only=False):
  
      provenance_dict = provenance.get_provenance_dict({"variants_only": variants_only})
  
      if variants_only:
          logging.info("Using lossy variants-only compression")
          # Reduce to site topology. Note that we will remove
          # any sites, individuals and populations here that have no references.
          ts = ts.simplify(reduce_to_site_topology=True)
  
      tables = ts.tables
  
      with warnings.catch_warnings():
          warnings.simplefilter("ignore")
          # When using a zipfile in Zarr we get some harmless warnings. See
          # https://zarr.readthedocs.io/en/stable/api/storage.html#zarr.storage.ZipStore
          root.attrs["format_name"] = FORMAT_NAME
          root.attrs["format_version"] = FORMAT_VERSION
          root.attrs["sequence_length"] = tables.sequence_length
          root.attrs["provenance"] = provenance_dict
  
      columns = {}
      for key, value in tables.asdict().items():
          if isinstance(value, dict):
              for sub_key, sub_value in value.items():
                  columns[f"{key}/{sub_key}"] = sub_value
          else:
              columns[key] = value
  
      if variants_only:
          time = np.unique(tables.nodes.time)
          columns["node/time"] = np.searchsorted(time, tables.nodes.time)
  
      # Encoding array is a tuple so must be converted
      columns["encoding_version"] = np.asarray(columns["encoding_version"])
  
      # Sequence length is stored as an attr for compatibility with older versions of tszip
      del columns["sequence_length"]
  
      # Schemas and metadata need to be converted to arrays
      for name in columns:
          if name.endswith("metadata_schema"):
              columns[name] = np.frombuffer(columns[name].encode("utf-8"), np.int8)
          if name.endswith("metadata"):
              columns[name] = np.frombuffer(columns[name], np.int8)
  
      # Some columns benefit from being quantised
      coordinates = np.unique(
          np.hstack(
              [
                  [0, ts.sequence_length],
                  tables.edges.left,
                  tables.edges.right,
                  tables.sites.position,
                  tables.migrations.left,
                  tables.migrations.right,
              ]
          )
      )
      columns["coordinates"] = coordinates
      for name in [
          "edges/left",
          "edges/right",
          "migrations/left",
          "migrations/right",
          "sites/position",
      ]:
          columns[name] = np.searchsorted(coordinates, columns[name])
  
      # Some columns benefit from additional options
      delta_filter_cols = ["edges/parent", "sites/position"]
  
      # Note: we're not providing any options to set this here because Blosc+Zstd seems to
      # have a clear advantage in compression performance and speed. There is very little
      # difference between compression level 6 and 9, and it's extremely fast in any case
      # so there's no point in adding complexity. The shuffle filter in particular makes
      # big difference.
      compressor = numcodecs.Blosc(
          cname="zstd", clevel=9, shuffle=numcodecs.Blosc.SHUFFLE
      )
      for name, data in columns.items():
          Column(
              name, data, delta_filter="_offset" in name or name in delta_filter_cols
          ).compress(root, compressor)
  
  
  def check_format(root):
      try:
          format_name = root.attrs["format_name"]
          format_version = root.attrs["format_version"]
      except KeyError as ke:
          raise exceptions.FileFormatError("Incorrect file format") from ke
      if format_name != FORMAT_NAME:
          raise exceptions.FileFormatError(
              "Incorrect file format: expected '{}' got '{}'".format(
                  FORMAT_NAME, format_name
              )
          )
      if format_version[0] < FORMAT_VERSION[0]:
          raise exceptions.FileFormatError(
              "Format version {} too old. Current version = {}".format(
                  format_version, FORMAT_VERSION
              )
          )
      if format_version[0] > FORMAT_VERSION[0]:
          raise exceptions.FileFormatError(
              "Format version {} too new. Current version = {}".format(
                  format_version, FORMAT_VERSION
              )
          )
  
  
  @contextlib.contextmanager
  def load_zarr(path):
      path = str(path)
      try:
          store = zarr.ZipStore(path, mode="r")
      except zipfile.BadZipFile as bzf:
          raise exceptions.FileFormatError("File is not in tgzip format") from bzf
      root = zarr.group(store=store)
      try:
          check_format(root)
          yield root
      finally:
          store.close()
  
  
  def decompress_zarr(root):
      coordinates = root["coordinates"][:]
      dict_repr = {"sequence_length": root.attrs["sequence_length"]}
  
      quantised_arrays = [
          "edges/left",
          "edges/right",
          "migrations/left",
          "migrations/right",
          "sites/position",
      ]
      for key, value in root.items():
          if isinstance(value, Mapping):
              for sub_key, sub_value in value.items():
                  if f"{key}/{sub_key}" in quantised_arrays:
                      dict_repr.setdefault(key, {})[sub_key] = coordinates[sub_value]
                  elif sub_key.endswith("metadata_schema"):
                      dict_repr.setdefault(key, {})[sub_key] = bytes(sub_value).decode(
                          "utf-8"
                      )
                  else:
                      dict_repr.setdefault(key, {})[sub_key] = sub_value
          elif key.endswith("metadata_schema"):
              dict_repr[key] = bytes(value).decode("utf-8")
          elif key.endswith("metadata"):
              dict_repr[key] = bytes(value)
          else:
              dict_repr[key] = value
  
      return tskit.TableCollection.fromdict(dict_repr).tree_sequence()
  
  
  def print_summary(path, verbosity=0):
      arrays = []
  
      def visitor(array):
          if isinstance(array, zarr.core.Array):
              arrays.append(array)
  
      with load_zarr(path) as root:
          root.visitvalues(visitor)
  
      arrays.sort(key=lambda x: x.nbytes_stored)
      max_name_len = max(len(array.name) for array in arrays)
      stored = [
          humanize.naturalsize(array.nbytes_stored, binary=True) for array in arrays
      ]
      max_stored_len = max(len(size) for size in stored)
      actual = [humanize.naturalsize(array.nbytes, binary=True) for array in arrays]
      max_actual_len = max(len(size) for size in actual)
  
      line = "File: {}\t{}".format(
          path, humanize.naturalsize(os.path.getsize(path), binary=True)
      )
      print(line)
      if verbosity > 0:
          print("format_version:", root.attrs["format_version"])
          prov = root.attrs["provenance"]
          print("provenance: ", end="")
          print(json.dumps(prov, indent=4, sort_keys=True))
      fmt = "{:<{}} {:<{}}\t{:<{}}\t{}"
      line = fmt.format(
          "name",
          max_name_len,
          "stored",
          max_stored_len,
          "actual",
          max_actual_len,
          "ratio",
      )
      print(line)
      for array, stored, actual in zip(arrays, stored, actual):
          ratio = 0
          if array.nbytes > 0:
              ratio = array.nbytes_stored / array.nbytes
          line = fmt.format(
              array.name,
              max_name_len,
              stored,
              max_stored_len,
              actual,
              max_actual_len,
              f"{ratio:.2f}",
          )
          print(line)
          if verbosity > 0:
              for line in str(array.info).splitlines():
                  print("\t", line)
  