""" 
Library for reading multiresolution micro-magellan
"""
import os
import numpy as np
import sys
import json
import dask.array as da
import warnings
import struct
import threading
from ndtiff.file_io import NDTiffFileIO, BUILTIN_FILE_IO

_POSITION_AXIS = "position"
_ROW_AXIS = "row"
_COLUMN_AXIS = "column"
_Z_AXIS = "z"
_TIME_AXIS = "time"
_CHANNEL_AXIS = "channel"

_AXIS_ORDER = {_ROW_AXIS: 7,
               _COLUMN_AXIS: 6,
               _POSITION_AXIS: 5, 
               _TIME_AXIS: 4, 
               _CHANNEL_AXIS:3, 
               _Z_AXIS:2}

def _get_axis_order_key(dict_item):
    axis_name = dict_item[0]
    if axis_name in _AXIS_ORDER.keys():
        return _AXIS_ORDER[axis_name]
    else:
        return 3  # stack next to channel axes

class _MultipageTiffReader:
    """
    Class corresponsing to a single multipage tiff file in a Micro-Magellan dataset.
    Pass the full path of the TIFF to instantiate and call close() when finished
    """

    # file format constants
    SUMMARY_MD_HEADER = 2355492
    EIGHT_BIT = 0
    SIXTEEN_BIT = 1
    EIGHT_BIT_RGB = 2
    UNCOMPRESSED = 0

    def __init__(self, tiff_path, file_io: NDTiffFileIO = BUILTIN_FILE_IO):
        """
        tiff_path: str
            The path to a .tiff file to load
        file_io: ndtiff.file_io.NDTiffFileIO
            A container containing various methods for interacting with files.
        """
        self.file_io = file_io
        self.tiff_path = tiff_path
        self.file = self.file_io.open(tiff_path, "rb")
        self.summary_md, self.first_ifd_offset = self._read_header()

    def close(self):
        """ """
        self.file.close()

    def _read_header(self):
        """
        Returns
        -------
        summary metadata : dict
        byte offsets : nested dict
            The byte offsets of TIFF Image File Directories with keys [channel_index][z_index][frame_index][position_index]
        first_image_byte_offset : int
            int byte offset of first image IFD
        """
        # read standard tiff header
        if self._read(0,2) == b"\x4d\x4d":
            # Big endian
            if sys.byteorder != "big":
                raise Exception("Potential issue with mismatched endian-ness")
        elif self._read(0,2) == b"\x49\x49":
            # little endian
            if sys.byteorder != "little":
                raise Exception("Potential issue with mismatched endian-ness")
        else:
            raise Exception("Endian type not specified correctly")
        if np.frombuffer(self._read(2,4), dtype=np.uint16)[0] != 42:
            raise Exception("Tiff magic 42 missing")
        first_ifd_offset = np.frombuffer(self._read(4,8), dtype=np.uint32)[0]

        # read custom stuff: header, summary md
        self._major_version = int.from_bytes(self._read(12,16), sys.byteorder)

        summary_md_header, summary_md_length = np.frombuffer(self._read(16,24), dtype=np.uint32)
        if summary_md_header != self.SUMMARY_MD_HEADER:
            raise Exception("Summary metadata header wrong")
        summary_md = json.loads(self._read(24, 24 + summary_md_length))
        return summary_md, first_ifd_offset

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
    
    def _read(self, start, end):
        """
        convert to python ints
        """
        self.file.seek(int(start), 0)
        return self.file.read(end - start)

    def read_metadata(self, index):
        return json.loads(
            self._read(
                index["metadata_offset"], index["metadata_offset"] + index["metadata_length"]
            )
        )

    def read_image(self, index):
        if index["pixel_type"] == self.EIGHT_BIT_RGB:
            bytes_per_pixel = 3
            dtype = np.uint8
        elif index["pixel_type"] == self.EIGHT_BIT:
            bytes_per_pixel = 1
            dtype = np.uint8
        elif index["pixel_type"] == self.SIXTEEN_BIT:
            bytes_per_pixel = 2
            dtype = np.uint16
        else:
            raise Exception("unrecognized pixel type")
        width = index["image_width"]
        height = index["image_height"]

        image = np.reshape(
            np.frombuffer(self._read(
            index["pixel_offset"], index["pixel_offset"] + width * height * bytes_per_pixel)
                , dtype=dtype),
            [height, width, 3] if bytes_per_pixel == 3 else [height, width],
        )
        return image


class _ResolutionLevel:
    def __init__(self, path=None, count=None, max_count=None, remote=False, summary_metadata=False, file_io: NDTiffFileIO = BUILTIN_FILE_IO):
        """
        Open all tiff files in directory, keep them in a list, and a tree based on image indices

        Parameters
        ----------
        path : str
        count : int
        max_count : int

        """
        self.file_io = file_io
        self.path_root = path + ("" if path[-1] == os.sep else os.sep)
        if remote:
            self.summary_metadata = summary_metadata
            self.index = {}
            self._readers_by_filename = {}
        else:
            self.index = self.read_index(path)
            tiff_names = [
                self.file_io.path_join(path, tiff) for tiff in self.file_io.listdir(path) if tiff.endswith(".tif")
            ]
            self._readers_by_filename = {}
            # populate list of readers and tree mapping indices to readers
            for tiff in tiff_names:
                print("\rOpening file {} of {}...".format(count + 1, max_count), end="")
                count += 1
                self._readers_by_filename[tiff.split(os.sep)[-1]] = _MultipageTiffReader(tiff, file_io=self.file_io)
            self.summary_metadata = list(self._readers_by_filename.values())[0].summary_md

    def has_image(self, axes):
        key = frozenset(axes.items())
        return key in self.index

    def add_index_entry(self, data):
        """
        Manually add a single index entry
        :param data: bytes object of a single index entry
        """
        _, axes, index_entry = self.read_single_index_entry(data, self.index)

        if index_entry["filename"] not in self._readers_by_filename:
            self._readers_by_filename[index_entry["filename"]] = _MultipageTiffReader(
                self.path_root + index_entry["filename"], file_io=self.file_io
            )
        return axes, index_entry

    def read_single_index_entry(self, data, entries, position=0):
        index_entry = {}
        (axes_length,) = struct.unpack("I", data[position : position + 4])
        if axes_length == 0:
            warnings.warn(
                "Index appears to not have been properly terminated (the dataset may still work)"
            )
            return None
        axes_str = data[position + 4 : position + 4 + axes_length].decode("utf-8")
        axes = json.loads(axes_str)
        position += axes_length + 4
        (filename_length,) = struct.unpack("I", data[position : position + 4])
        index_entry["filename"] = data[position + 4 : position + 4 + filename_length].decode(
            "utf-8"
        )
        position += 4 + filename_length
        (
            index_entry["pixel_offset"],
            index_entry["image_width"],
            index_entry["image_height"],
            index_entry["pixel_type"],
            index_entry["pixel_compression"],
            index_entry["metadata_offset"],
            index_entry["metadata_length"],
            index_entry["metadata_compression"],
        ) = struct.unpack("IIIIIIII", data[position : position + 32])
        position += 32
        entries[frozenset(axes.items())] = index_entry
        return position, axes, index_entry

    def read_index(self, path):
        print("\rReading index...          ", end="")
        with self.file_io.open(path + os.sep + "NDTiff.index", "rb") as index_file:
            data = index_file.read()
        entries = {}
        position = 0
        while position < len(data):
            print(
                "\rReading index... {:.1f}%       ".format(
                    100 * (1 - (len(data) - position) / len(data))
                ),
                end="",
            )
            entry = self.read_single_index_entry(data, entries, position)
            if entry is None:
                break
            position, axes, index_entry = entry
            if position is None:
                break

        print("\rFinshed reading index          ", end="")
        return entries

    def read_image(
        self,
        axes,
    ):
        """

        Parameters
        ----------
        axes : dict

        Returns
        -------
        image :
        """
        # determine which reader contains the image
        key = frozenset(axes.items())
        if key not in self.index:
            raise Exception("image with keys {} not present in data set".format(key))
        index = self.index[key]
        reader = self._readers_by_filename[index["filename"]]
        return reader.read_image(index)

    def read_metadata(self, axes):
        """

        Parameters
        ----------
        axes : dict

        Returns
        -------
        image_metadata
        """
        key = frozenset(axes.items())
        if key not in self.index:
            raise Exception("image with keys {} not present in data set".format(key))
        index = self.index[key]
        reader = self._readers_by_filename[index["filename"]]
        return reader.read_metadata(index)

    def close(self):
        for reader in self._readers_by_filename.values():
            reader.close()


### This function outside class to prevent problems with pickling when running them in differnet process


class NDTiff_v2_0():
    """Class that opens a single NDTiffStorage dataset"""

    def __init__(self, dataset_path=None, full_res_only=True, remote_storage_monitor=None, file_io: NDTiffFileIO = BUILTIN_FILE_IO):
        """
        Creat a Object providing access to and NDTiffStorage dataset, either one currently being acquired or one on disk

        Parameters
        ----------
        dataset_path : str
            Abosolute path of top level folder of a dataset on disk
        full_res_only : bool
            One open the full resolution data, if it is multi-res
        remote_storage_monitor : JavaObjectShadow
            Object that allows callbacks from remote NDTiffStorage
        file_io: ndtiff.file_io.NDTiffFileIO
            A container containing various methods for interacting with files.
        """
        self.file_io = file_io
        self._tile_width = None
        self._tile_height = None
        self._lock = threading.Lock()
        if remote_storage_monitor is not None:
            # this dataset is a view of an active acquisiiton. The storage exists on the java side
            self.new_image_arrived = False # used by napari viewer to check for updates. Will be reset to false by them
            self._remote_storage_monitor = remote_storage_monitor
            self.summary_metadata = self._remote_storage_monitor.get_summary_metadata()
            if "GridPixelOverlapX" in self.summary_metadata.keys():
                self._tile_width = (
                    self.summary_metadata["Width"] - self.summary_metadata["GridPixelOverlapX"]
                )
                self._tile_height = (
                    self.summary_metadata["Height"] - self.summary_metadata["GridPixelOverlapY"]
                )

            dataset_path = remote_storage_monitor.get_disk_location()
            dataset_path += "" if dataset_path[-1] == os.sep else os.sep
            full_res_path = dataset_path + "Full resolution"
            with self._lock:
                self.res_levels = {
                    0: _ResolutionLevel(
                        remote=True, summary_metadata=self.summary_metadata, path=full_res_path,
                        file_io=self.file_io
                    )
                }
            self.axes = {}
            return
        else:
            self._remote_storage_monitor = None

        self.path = dataset_path
        res_dirs = [
            dI for dI in self.file_io.listdir(dataset_path) if self.file_io.isdir(self.file_io.path_join(dataset_path, dI))
        ]
        # map from downsample factor to datset
        with self._lock:
            self.res_levels = {}
        if "Full resolution" not in res_dirs:
            raise Exception(
                "Couldn't find full resolution directory. Is this the correct path to a dataset?"
            )
        num_tiffs = 0
        count = 0
        for res_dir in res_dirs:
            for file in self.file_io.listdir(self.file_io.path_join(dataset_path, res_dir)):
                if file.endswith(".tif"):
                    num_tiffs += 1
        for res_dir in res_dirs:
            if full_res_only and res_dir != "Full resolution":
                continue
            res_dir_path = self.file_io.path_join(dataset_path, res_dir)
            res_level = _ResolutionLevel(res_dir_path, count, num_tiffs, file_io=self.file_io)
            if res_dir == "Full resolution":
                with self._lock:
                    self.res_levels[0] = res_level
                # get summary metadata and index tree from full resolution image
                self.summary_metadata = res_level.summary_metadata

                self.overlap = (
                    np.array(
                        [
                            self.summary_metadata["GridPixelOverlapY"],
                            self.summary_metadata["GridPixelOverlapX"],
                        ]
                    )
                    if "GridPixelOverlapY" in self.summary_metadata
                    else None
                )

                self.axes = {}
                for axes_combo in res_level.index.keys():
                    for axis, position in axes_combo:
                        if axis not in self.axes.keys():
                            self.axes[axis] = set()
                        self.axes[axis].add(position)
                # Sort axes according to _AXIS_ORDER
                self.axes = dict(sorted(self.axes.items(), key=_get_axis_order_key, reverse=True))

                # figure out the mapping of channel name to position by reading image metadata
                print("\rReading channel names...", end="")
                self._read_channel_names()
                print("\rFinished reading channel names", end="")

                # remove axes with no variation
                # single_axes = [axis for axis in self.axes if len(self.axes[axis]) == 1]
                # for axis in single_axes:
                #     del self.axes[axis]

            else:
                with self._lock:
                    self.res_levels[int(np.log2(int(res_dir.split("x")[1])))] = res_level

        # get information about image width and height, assuming that they are consistent for whole dataset
        # (which isn't strictly neccesary)
        with self._lock:
            first_index = list(self.res_levels[0].index.values())[0]
        self._parse_first_index(first_index)

        print("\rDataset opened                ")

    def _read_channel_names(self):
        if _CHANNEL_AXIS in self.axes.keys():
            self._channel_names = {}
            for key in self.res_levels[0].index.keys():
                axes = {axis: position for axis, position in key}
                if (
                    _CHANNEL_AXIS in axes.keys()
                    and axes[_CHANNEL_AXIS] not in self._channel_names.values()
                ):
                    channel_name = self.res_levels[0].read_metadata(axes)["Channel"]
                    self._channel_names[channel_name] = axes[_CHANNEL_AXIS]
                if len(self._channel_names.values()) == len(self.axes[_CHANNEL_AXIS]):
                    break

    def _parse_first_index(self, first_index):
        """
        Read through first index to get some global data about images (assuming it is same for all)
        """
        if first_index["pixel_type"] == _MultipageTiffReader.EIGHT_BIT_RGB:
            self.bytes_per_pixel = 3
            self.dtype = np.uint8
        elif first_index["pixel_type"] == _MultipageTiffReader.EIGHT_BIT:
            self.bytes_per_pixel = 1
            self.dtype = np.uint8
        elif first_index["pixel_type"] == _MultipageTiffReader.SIXTEEN_BIT:
            self.bytes_per_pixel = 2
            self.dtype = np.uint16

        self.image_width = first_index["image_width"]
        self.image_height = first_index["image_height"]
        if "GridPixelOverlapX" in self.summary_metadata:
            self._tile_width = self.image_width - self.summary_metadata["GridPixelOverlapX"]
            self._tile_height = self.image_height - self.summary_metadata["GridPixelOverlapY"]

    def _add_index_entry(self, index_entry):
        """
        Add entry for a image that has been recieved and is now on disk
        """
        with self._lock:
            axes, index_entry = self.res_levels[0].add_index_entry(index_entry)

            # update the axes that have been seen
            for axis_name in axes.keys():
                if axis_name not in self.axes.keys():
                    self.axes[axis_name] = set()
                self.axes[axis_name].add(axes[axis_name])

            # update the map of channel names to channel indices
            self._read_channel_names()

        if not hasattr(self, 'image_width'):
            self._parse_first_index(index_entry)

        return axes

    def get_index_keys(self, res_level=0):
        """
        Return a list of every combination of axes that has a image in this dataset
        """
        frozen_set_list = list(self.res_levels[res_level].index.keys())
        # convert to dict
        return [{axis_name: position for axis_name, position in key} for key in frozen_set_list]

    def as_array(self, axes=None, stitched=False, **kwargs):
        """
        Read all data image data as one big Dask array with last two axes as y, x and preceeding axes depending on data.
        The dask array is made up of memory-mapped numpy arrays, so the dataset does not need to be able to fit into RAM.
        If the data doesn't fully fill out the array (e.g. not every z-slice collected at every time point), zeros will
        be added automatically.

        To convert data into a numpy array, call np.asarray() on the returned result. However, doing so will bring the
        data into RAM, so it may be better to do this on only a slice of the array at a time.

        Parameters
        ----------
        axes : list
            list of axes names over which to iterate and merge into a stacked array. The order of axes supplied in this 
            list will be the order of the axes of the returned dask array. If None, all axes will be used in PTCZYX order.
        stitched : bool
            If true and tiles were acquired in a grid, lay out adjacent tiles next to one another (Default value = False)
        **kwargs :
            names and integer positions of axes on which to slice data
        Returns
        -------
        dataset : dask array
        """
        if stitched and "GridPixelOverlapX" not in self.summary_metadata:
            raise Exception('This is not a stitchable dataset')
        w = self.image_width
        h = self.image_height
        self._empty_tile = (
            np.zeros((h, w), self.dtype)
            if self.bytes_per_pixel != 3
            else np.zeros((h, w, 3), self.dtype)
        )

        rgb = self.bytes_per_pixel == 3 and self.dtype == np.uint8

        if axes is None:
            axes = self.axes.keys()
        axes_to_slice = kwargs
        axes_to_stack = {key: self.axes[key] for key in axes if key not in kwargs.keys()}
        if stitched:
            if 'row' in axes_to_stack:
                del axes_to_stack['row']
            if 'column' in axes_to_stack:
                del axes_to_stack['column']
            if 'row' in axes_to_slice:
                del axes_to_slice['row']
            if 'column' in axes_to_slice:
                del axes_to_slice['column']

        def read_one_image(block_id, axes_to_stack=axes_to_stack, axes_to_slice=axes_to_slice):
            # a function that reads in one chunk of data
            axes = {key: block_id[i] for i, key in enumerate(axes_to_stack.keys())}
            if stitched:
                # Combine all rows and cols into one stitched image
                self.half_overlap = (self.overlap[0] // 2, self.overlap[1] // 2)
                # get spatial layout of position indices
                row_values = np.array(list(self.axes["row"]))
                column_values = np.array(list(self.axes["column"]))
                # make nested list of rows and columns
                blocks = []
                for row in row_values:
                    blocks.append([])
                    for column in column_values:
                        #remove overlap between tiles
                        if not self.has_image(**axes, **axes_to_slice, row=row, column=column):
                            tile = self._empty_tile
                        else:
                            tile = self.read_image(**axes, **axes_to_slice, row=row, column=column)
                        if self.half_overlap[0] != 0:
                            tile = tile[
                                self.half_overlap[0] : -self.half_overlap[0],
                                self.half_overlap[1] : -self.half_overlap[1],
                            ]
                        blocks[-1].append(tile)

                if rgb:
                    image = np.concatenate(
                        [
                            np.concatenate(row, axis=len(blocks[0][0].shape) - 2)
                            for row in blocks
                        ],
                        axis=0,
                    )
                else:
                    image = np.array(da.block(blocks))
            else:
                if not self.has_image(**axes, **axes_to_slice):
                    image = self._empty_tile
                else:
                    image = self.read_image(**axes, **axes_to_slice)
            for i in range(len(axes_to_stack.keys())):
                image = image[None]
            return image

        chunks = tuple([(1,) * len(axes_to_stack[axis]) for axis in axes_to_stack.keys()])
        if stitched:
            block_w = self._tile_width*len(self.axes["column"])
            block_h = self._tile_height*len(self.axes["row"])
        else:
            block_w = w
            block_h = h
        chunks += (block_h, block_w)
        if rgb:
            chunks += (3,)

        array = da.map_blocks(
            read_one_image,
            dtype=self.dtype,
            chunks=chunks,
            meta=self._empty_tile
        )

        return array

    def has_image(
        self,
        channel=0,
        z=None,
        time=None,
        position=None,
        channel_name=None,
        resolution_level=0,
        row=None,
        col=None,
        **kwargs
    ):
        """Check if this image is present in the dataset

        Parameters
        ----------
        channel : int
            index of the channel, if applicable (Default value = None)
        z : int
            index of z slice, if applicable (Default value = None)
        time : int
            index of the time point, if applicable (Default value = None)
        position : int
            index of the XY position, if applicable (Default value = None)
        channel_name : str
            Name of the channel. Overrides channel index if supplied (Default value = None)
        row : int
            index of tile row for XY tiled datasets (Default value = None)
        col : int
            index of tile col for XY tiled datasets (Default value = None)
        resolution_level :
            0 is full resolution, otherwise represents downampling of pixels
            at 2 ** (resolution_level) (Default value = 0)
        **kwargs :
            names and integer positions of any other axes

        Returns
        -------
        bool :
            indicating whether the dataset has an image matching the specifications
        """
        with self._lock:
            return self.res_levels[0].has_image(
                self._consolidate_axes(channel, channel_name, z, position, time, row, col, kwargs)
            )

    def read_image(
        self,
        channel=0,
        z=None,
        time=None,
        position=None,
        row=None,
        col=None,
        channel_name=None,
        resolution_level=0,
        **kwargs
    ):
        """
        Read image data as numpy array

        Parameters
        ----------
        channel : int
            index of the channel, if applicable (Default value = None)
        z : int
            index of z slice, if applicable (Default value = None)
        time : int
            index of the time point, if applicable (Default value = None)
        position : int
            index of the XY position, if applicable (Default value = None)
        channel_name :
            Name of the channel. Overrides channel index if supplied (Default value = None)
        row : int
            index of tile row for XY tiled datasets (Default value = None)
        col : int
            index of tile col for XY tiled datasets (Default value = None)
        resolution_level :
            0 is full resolution, otherwise represents downampling of pixels
            at 2 ** (resolution_level) (Default value = 0)
        **kwargs :
            names and integer positions of any other axes

        Returns
        -------
        image : numpy array or tuple
            image as a 2D numpy array, or tuple with image and image metadata as dict

        """
        with self._lock:
            axes = self._consolidate_axes(
                channel, channel_name, z, position, time, row, col, kwargs
            )

            res_level = self.res_levels[resolution_level]
            return res_level.read_image(axes)

    def read_metadata(
        self,
        channel=0,
        z=None,
        time=None,
        position=None,
        channel_name=None,
        row=None,
        col=None,
        resolution_level=0,
        **kwargs
    ):
        """
        Read metadata only. Faster than using read_image to retrieve metadata

        Parameters
        ----------
        channel : int
            index of the channel, if applicable (Default value = None)
        z : int
            index of z slice, if applicable (Default value = None)
        time : int
            index of the time point, if applicable (Default value = None)
        position : int
            index of the XY position, if applicable (Default value = None)
        channel_name :
            Name of the channel. Overrides channel index if supplied (Default value = None)
        row : int
            index of tile row for XY tiled datasets (Default value = None)
        col : int
            index of tile col for XY tiled datasets (Default value = None)
        resolution_level :
            0 is full resolution, otherwise represents downampling of pixels
            at 2 ** (resolution_level) (Default value = 0)
        **kwargs :
            names and integer positions of any other axes

        Returns
        -------
        metadata : dict

        """
        with self._lock:
            axes = self._consolidate_axes(
                channel, channel_name, z, position, time, row, col, kwargs
            )

            res_level = self.res_levels[resolution_level]
            return res_level.read_metadata(axes)

    def close(self):
        with self._lock:
            for res_level in self.res_levels.values():
                res_level.close()

    def get_channel_names(self):
        with self._lock:
            return self._channel_names.keys()

    def _consolidate_axes(self, channel, channel_name, z, position, time, row, col, kwargs):
        axes = {}
        if channel is not None:
            axes[_CHANNEL_AXIS] = channel
        if channel_name is not None:
            axes[_CHANNEL_AXIS] = self._channel_names[channel_name]
        if z is not None:
            axes[_Z_AXIS] = z
        if position is not None:
            axes[_POSITION_AXIS] = position
        if time is not None:
            axes[_TIME_AXIS] = time
        if row is not None:
            axes[_ROW_AXIS] = row
        if col is not None:
            axes[_COLUMN_AXIS] = col
        for other_axis_name in kwargs.keys():
            axes[other_axis_name] = kwargs[other_axis_name]
        return axes
