"""tiffslide.deepzoom

tiffslide.deepzoom is not intended to be used as a drop-in replacement for
for openslide's DeepZoomGenerator. It's aiming to provide the basis for a
minimal overhead tile server for viewing tiled whole slide images in the
browser without having to re-tile the existing layers in the slide.

"""
import math
from io import BytesIO
from os import PathLike
from warnings import warn
from xml.etree.ElementTree import Element
from xml.etree.ElementTree import ElementTree
from xml.etree.ElementTree import SubElement

from PIL import Image
from PIL import ImageFile
from tifffile import TiffPage
from tifffile import TIFF

from tiffslide import TiffSlide

# improve robustness when encountering corrupted tiles
ImageFile.LOAD_TRUNCATED_IMAGES = True


def __getattr__(name):
    if name == "DeepZoomGenerator":
        warn("tiffslide.deepzoom does not provide a drop-in replacement for openslide's deepzoom generator")
    raise AttributeError(name)


class _DeepZoomGenerator:

    def get_dzi(self):
        raise NotImplementedError('implement in subclass')

    def get_tile(self, level, x, y):
        raise NotImplementedError('implement in subclass')

    @property
    def level_size(self):
        raise NotImplementedError('implement in subclass')


class MinimalComputeAperioDZGenerator(_DeepZoomGenerator):
    """Minimal-compute tiffslide Deep Zoom tile generator

    Note:
      - this is providing tiles of existing layers in the svs directly from disk
      - it would be a zero-compute DZG if we'd create a openseadragon TileSource
        that supports arbitrary (not 2x) spaced tile layers

    """

    def __init__(self, svs_filename: PathLike):
        self._slide = TiffSlide(svs_filename)

        # Get information from baseline layers
        self._baseline = baseline = self._slide.ts_tifffile.series[0]
        assert baseline.keyframe.tilewidth == baseline.keyframe.tilelength
        self._tile_size = baseline.keyframe.tilewidth
        self._im_levels = tuple(lvl.shape[1::-1] for lvl in baseline.levels)

        # generate the levels for deep zoom
        dz_levels = dz_lvl, = [self._im_levels[0]]
        while dz_lvl[0] > 1 or dz_lvl[1] > 1:
            dz_lvl = tuple(max(1, int(math.ceil(z / 2))) for z in dz_lvl)
            dz_levels.append(dz_lvl)
        self._dz_levels = tuple(reversed(dz_levels))

        self._mapped_levels = {}
        for im_idx, im_lvl in enumerate(self._im_levels):
            for dz_idx, dz_lvl in enumerate(self._dz_levels):
                if abs(im_lvl[0] - dz_lvl[0]) <= 1 and abs(im_lvl[1] - dz_lvl[1]) <= 1:
                    self._mapped_levels[dz_idx] = im_idx

    @property
    def level_size(self):
        """return a dict mapping deep zoom level to tile index size"""
        return {
            idx: tuple(int(math.ceil(z / self._tile_size)) for z in lvl)
            for idx, lvl in enumerate(self._dz_levels)
        }

    def get_dzi(self):
        """return the dzi XML metadata"""
        image = Element(
            "Image",
            TileSize=str(self._tile_size),
            Overlap="0",     # tiles in AperioSVS files have 0 overlap
            Format="jpeg",   # tiles are stored in jpeg format
            xmlns="http://schemas.microsoft.com/deepzoom/2008"
        )
        width, height = self._im_levels[0]
        SubElement(image, "Size", Width=str(width), Height=str(height))
        tree = ElementTree(element=image)

        with BytesIO() as buffer:
            tree.write(buffer, encoding="UTF-8")
            return buffer.getvalue().decode("UTF-8")

    def _get_svs_tile(self, svs_level: int, x: int, y: int) -> bytes:
        """return a single tile from an svs as a jpeg into a buffer"""
        # note: this code relies on the image being an AperioSVS file.
        #   contributions to add support for other formats are welcome :)
        series = self._slide.ts_tifffile.series[0]
        assert series.is_pyramidal

        # try to get correct level, raise IndexError if not
        try:
            page_series = series.levels[svs_level]
        except IndexError:
            raise IndexError(f"requires 0 <= level < {len(series.levels)} got {svs_level}")

        # extract the current page from page_series
        page: TiffPage
        page, = page_series.pages

        # more assumptions to ensure programmer sanity
        assert page.compression == TIFF.COMPRESSION.JPEG
        assert page.is_tiled
        assert page.planarconfig == TIFF.PLANARCONFIG.CONTIG

        # calculate indices
        st_length, st_width = page.tilelength, page.tilewidth
        im_length, im_width = page.shaped[2:4]
        idx_width = (im_width + st_width - 1) // st_width
        idx_length = (im_length + st_length - 1) // st_length

        if not ((0 <= x < idx_width) and (0 <= y < idx_length)):
            raise IndexError(
                f"({x}, {y}) level={svs_level} out of bounds: max ({idx_width}, {idx_length})"
            )

        # index for reading tile
        tile_index = y * idx_width + x

        (data, _), = page.parent.filehandle.read_segments(
            [page.dataoffsets[tile_index]],
            [page.databytecounts[tile_index]],
        )

        jpeg_tables_tag, = (tag for tag in page.tags.values() if tag.name == "JPEGTables")

        with BytesIO() as buffer:
            buffer.write(jpeg_tables_tag.value[:-2])
            # to directly provide the stored tiles from disk, we need to fix that svs tiles
            # use a jpeg colorspace that doesn't show up correctly in the browser if the default
            # jpeg headers are used
            # See https://stackoverflow.com/questions/8747904/extract-jpeg-from-tiff-file/9658206#9658206
            buffer.write(b"\xFF\xEE\x00\x0E\x41\x64\x6F\x62\x65\x00\x64\x80\x00\x00\x00\x00")  # colorspace fix
            buffer.write(data[2:])
            tile_data = buffer.getvalue()

        # the outer edges need to be cropped to be interpreted correctly by default openseadragon
        out_width = ((im_width - 1) % st_width) + 1 if x == idx_width - 1 else st_width
        out_length = ((im_length - 1) % st_length) + 1 if y == idx_length - 1 else st_length
        if out_width < st_width or out_length < st_length:  # cut output if needed
            with BytesIO(tile_data) as buffer:
                im = Image.open(buffer).crop((0, 0, out_width, out_length))
            with BytesIO() as buffer:
                im.save(buffer, format="JPEG")
                return buffer.getvalue()

        return tile_data

    def get_tile(self, level: int, x: int, y: int) -> bytes:
        """return the jpeg tile as bytes"""

        if level in self._mapped_levels:
            # FAST PATH:
            # -> there's a direct mapping to a svs tiled level
            # note: if we optimize the frontend, this will be the only path that's hit!
            svs_lvl = self._mapped_levels[level]
            return self._get_svs_tile(svs_level=svs_lvl, x=x, y=y)

        elif 8 <= level <= max(self._mapped_levels):
            # SLOW PATH:
            # -> the frontend requests a dzi layer that's not available in the svs
            # we need to compute new tiles from lower levels
            dst = Image.new('RGB', (2 * self._tile_size, 2 * self._tile_size))

            out_width = out_height = 0
            for ix, iy in [(0, 0), (0, 1), (1, 0), (1, 1)]:
                nx, ny = 2 * x + ix, 2 * y + iy

                try:
                    # the recursive call is only really a problem for the small (very low magnification)
                    # deep zoom tile levels, the bigger ones (high magnification) hit after one recursion
                    # fixme: we could optimize this by doing oneshot crop resizing from the svs thumbnail layer
                    data = self.get_tile(level + 1, nx, ny)
                except IndexError:
                    continue

                with BytesIO(data) as buffer:
                    im = Image.open(buffer)
                    if ix == 0:
                        out_height += im.height
                    if iy == 0:
                        out_width += im.width
                    dst.paste(im, (ix * self._tile_size, iy * self._tile_size))

            if out_width == 0 or out_height == 0:
                raise IndexError(f"tile index ({x}, {y}) at INTERMEDIATE level={level} out of bounds")

            elif (out_width, out_height) != dst.size:
                dst = dst.crop((0, 0, out_width, out_height))
                thumb_size = (max(1, out_width // 2), max(1, out_height // 2))

            else:
                thumb_size = (self._tile_size, self._tile_size)

            dst.thumbnail(thumb_size, Image.ANTIALIAS)
            with BytesIO() as buffer:
                dst.save(buffer, format="JPEG")
                return buffer.getvalue()

        else:
            raise IndexError(f"requested level {level} invalid")


def _test_tifffile_timing():
    # build some load stats for a slide
    import random
    import sys
    import time
    from contextlib import contextmanager
    from operator import itemgetter

    if len(sys.argv) != 2:
        print("please provide a svs file as a cli argument")
        sys.exit(1)
    else:
        svs_fn = sys.argv[1]

    num_samples = 100

    @contextmanager
    def timer(label: str, samples: int = 1):
        t0 = time.time()
        yield
        avg = (time.time() - t0) / samples
        print(f"{label} took {avg} seconds")

    with timer("create dz"):
        dz = MinimalComputeAperioDZGenerator(svs_fn)
    print('tiffslide downsamples:', dz._slide.level_downsamples)
    print('mapped levels:', dz._mapped_levels)

    for test_lvl, lvl_size in sorted(dz.level_size.items(), key=itemgetter(0), reverse=True):
        if test_lvl < 8:
            print("skipping small levels")
            break

        test_idx = range(lvl_size[0] * lvl_size[1])
        if len(test_idx) > num_samples:
            test_idx = random.sample(test_idx, num_samples)

        with timer(f"accessing level {test_lvl}", samples=len(test_idx)):
            for flat_idx in test_idx:
                addr = flat_idx % lvl_size[0], flat_idx // lvl_size[0]
                dz.get_tile(test_lvl, addr[0], addr[1])

    print("done.")


if __name__ == "__main__":
    _test_tifffile_timing()
