import glob
import logging

import numpy as np

from gunpowder.array import Array
from gunpowder.array_spec import ArraySpec
from gunpowder.batch import Batch
from gunpowder.coordinate import Coordinate
from gunpowder.ext import pyklb
from gunpowder.profiling import Timing
from gunpowder.roi import Roi

from .batch_provider import BatchProvider

logger = logging.getLogger(__name__)


class KlbSource(BatchProvider):
    """A `KLB <https://bitbucket.org/fernandoamat/keller-lab-block-filetype>`_
    data source.

    Provides a single array from the given KLB dataset.

    Args:

        filename (``string``):

            The name of the KLB file. This string can be a glob expression
            (e.g., ``frame_*.klb``), in which case all files that match are
            sorted and stacked together to form an additional dimension (like
            time). The additional dimension will start at 0 and have a default
            voxel size of 1 (which can be overwritten using the ``array_spec``
            argument).

        array (:class:`ArrayKey`):

            ArrayKey that this source offers.

        array_spec (:class:`ArraySpec`, optional):

            An optional :class:`ArraySpec` to overwrite the array specs
            automatically determined from the KLB file. This is useful to set
            ``voxel_size``, for example. Only fields that are not ``None`` in
            the given :class:`ArraySpec` will be used.

        num_threads (``int``):

            An optional integer to pass to pyklb reader indicating the number
            of threads to use when reading klb files. Entering None causes
            uses the pyklb default, which now is based on the number of cores
            in the machine. This pyklb default is bad for jobs on the cluster that
            are limited to the number of cores requested, and 1 is recommended.

    """

    def __init__(self, filename, array, array_spec=None, num_threads=1):
        self.filename = filename
        self.array = array
        self.array_spec = array_spec
        self.num_threads = num_threads

        self.files = None
        self.ndims = None

    def setup(self):
        self.files = glob.glob(self.filename)
        self.files.sort()

        logger.info("Reading KLB headers of %d files...", len(self.files))
        headers = [pyklb.readheader(f) for f in self.files]
        spec = self.__read_spec(headers)

        self.provides(self.array, spec)

    def provide(self, request):
        timing = Timing(self)
        timing.start()

        batch = Batch()

        request_spec = request[self.array]

        logger.debug("Reading %s in %s...", self.array, request_spec.roi)

        voxel_size = self.spec[self.array].voxel_size

        # scale request roi to voxel units
        dataset_roi = request_spec.roi / voxel_size

        # shift request roi into dataset
        dataset_roi = dataset_roi - self.spec[self.array].roi.offset / voxel_size

        # create array spec
        array_spec = self.spec[self.array].copy()
        array_spec.roi = request_spec.roi

        # add array to batch
        batch.arrays[self.array] = Array(self.__read(dataset_roi), array_spec)

        logger.debug("done")

        timing.stop()
        batch.profiling_stats.add(timing)

        return batch

    def __read_spec(self, headers):
        num_files = len(headers)
        assert num_files > 0
        common_header = headers[0]
        for header in headers:
            for attr in ["imagesize_tczyx", "pixelspacing_tczyx"]:
                assert (common_header[attr] == header[attr]).all(), (
                    "Headers of provided KLB files differ in attribute %s" % attr
                )
            assert (
                common_header["datatype"] == header["datatype"]
            ), "Headers of provided KLB files differ in attribute datatype"

        size = Coordinate(common_header["imagesize_tczyx"])
        voxel_size = Coordinate(common_header["pixelspacing_tczyx"])
        dtype = common_header["datatype"]

        # strip leading 1 dimensions
        while size[0] == 1 and len(size) > 1:
            size = size[1:]
            voxel_size = voxel_size[1:]

        # append num_files dimension
        if num_files > 1:
            size = (num_files,) + size
            voxel_size = (1,) + voxel_size

        dims = Coordinate(size)
        self.ndims = len(dims)

        if self.array_spec is not None:
            spec = self.array_spec
        else:
            spec = ArraySpec()

        if spec.voxel_size is None:
            spec.voxel_size = Coordinate(voxel_size)

        if spec.roi is None:
            offset = Coordinate((0,) * self.ndims)
            spec.roi = Roi(offset, dims * spec.voxel_size)

        if spec.dtype is not None:
            assert spec.dtype == dtype, (
                "dtype %s provided in array_specs for %s, but differs from "
                "dataset dtype %s"
                % (self.array_specs[self.array].dtype, self.array, dtype)
            )
        else:
            spec.dtype = dtype

        if spec.interpolatable is None:
            spec.interpolatable = np.issubdtype(spec.dtype, np.floating) or (
                spec.dtype == np.uint8
            )
            logger.warning(
                "WARNING: You didn't set 'interpolatable' for %s. "
                "Based on the dtype %s, it has been set to %s. "
                "This might not be what you want.",
                self.array,
                spec.dtype,
                spec.interpolatable,
            )

        return spec

    def __read(self, roi):
        if len(self.files) == 1:
            return self.__read_file(self.files[0], roi)

        else:
            file_indices = range(roi.begin[0], roi.end[0])

            file_roi = Roi(roi.begin[1:], roi.shape[1:])

            return np.array(
                [self.__read_file(self.files[i], file_roi) for i in file_indices]
            )

    def __read_file(self, filename, roi):
        # pyklb reads max-inclusive, gunpowder rois are max exclusive ->
        # subtract (1, 1, ...) from max coordinate
        if self.num_threads:
            return pyklb.readroi(
                filename,
                roi.begin,
                roi.end - (1,) * roi.dims,
                numthreads=self.num_threads,
            )
        else:
            return pyklb.readroi(filename, roi.begin, roi.end - (1,) * roi.dims)

    def __repr__(self):
        return self.filename
