import itertools
import logging

import h5py
import numpy as np
from pathlib import Path

from tqdm import tqdm
import cv2

from aliby.io.matlab import matObject
from aliby.utils import Cache, imread, get_store_path

logger = logging.getLogger(__name__)


def parse_local_fs(pos_dir, tp=None):
    """
    Local file structure:
    - pos_dir
        -- exptID_{timepointID}_{ChannelID}_{z_position_id}.png

    :param pos_dirs:
    :return: Image_mapper
    """
    pos_dir = Path(pos_dir)

    img_mapper = dict()

    def channel_idx(img_name):
        return img_name.stem.split("_")[-2]

    def tp_idx(img_name):
        return int(img_name.stem.split("_")[-3]) - 1

    def z_idx(img_name):
        return img_name.stem.split("_")[-1]

    if tp is not None:
        img_list = [img for img in pos_dir.iterdir() if tp_idx(img) in tp]
    else:
        img_list = [img for img in pos_dir.iterdir()]

    for tp, group in itertools.groupby(sorted(img_list, key=tp_idx), key=tp_idx):
        img_mapper[int(tp)] = {
            channel: {i: item for i, item in enumerate(sorted(grp, key=z_idx))}
            for channel, grp in itertools.groupby(
                sorted(group, key=channel_idx), key=channel_idx
            )
        }
    return img_mapper


class Timelapse:
    """
    Timelapse class contains the specifics of one position.
    """

    def __init__(self):
        self._id = None
        self._name = None
        self._channels = []
        self._size_c = 0
        self._size_t = 0
        self._size_x = 0
        self._size_y = 0
        self._size_z = 0
        self.image_cache = None
        self.annotation = None

    def __repr__(self):
        return self.name

    def full_mask(self):
        return np.full(self.shape, False)

    def __getitem__(self, item):
        cached = self.image_cache[item]
        # Check if there are missing values, if so reload
        # TODO only reload missing
        mask = np.isnan(cached)
        if np.any(mask):
            full = self.load_fn(item)
            shape = self.image_cache[
                item
            ].shape  # TODO speed this up by  recognising the shape from the item
            self.image_cache[item] = np.reshape(full, shape)
            return full
        return cached

    def get_hypercube(self):
        pass

    def load_fn(self, item):
        """
        The hypercube is ordered as: C, T, X, Y, Z
        :param item:
        :return:
        """

        def parse_slice(s):
            step = s.step if s.step is not None else 1
            if s.start is None and s.stop is None:
                return None
            elif s.start is None and s.stop is not None:
                return range(0, s.stop, step)
            elif s.start is not None and s.stop is None:
                return [s.start]
            else:  # both s.start and s.stop are not None
                return range(s.start, s.stop, step)

        def parse_subitem(subitem, kw):
            if isinstance(subitem, (int, float)):
                res = [int(subitem)]
            elif isinstance(subitem, list) or isinstance(subitem, tuple):
                res = list(subitem)
            elif isinstance(subitem, slice):
                res = parse_slice(subitem)
            else:
                res = subitem
                # raise ValueError(f"Cannot parse slice {kw}: {subitem}")

            if kw in ["x", "y"]:
                # Need exactly two values
                if res is not None:
                    if len(res) < 2:
                        # An int was passed, assume it was
                        res = [res[0], self.size_x]
                    elif len(res) > 2:
                        res = [res[0], res[-1] + 1]
            return res

        if isinstance(item, int):
            return self.get_hypercube(
                x=None, y=None, z_positions=None, channels=[item], timepoints=None
            )
        elif isinstance(item, slice):
            return self.get_hypercube(channels=parse_slice(item))
        keywords = ["channels", "timepoints", "x", "y", "z_positions"]
        kwargs = dict()
        for kw, subitem in zip(keywords, item):
            kwargs[kw] = parse_subitem(subitem, kw)
        return self.get_hypercube(**kwargs)

    @property
    def shape(self):
        return (self.size_c, self.size_t, self.size_x, self.size_y, self.size_z)

    @property
    def id(self):
        return self._id

    @property
    def name(self):
        return self._name

    @property
    def size_z(self):
        return self._size_z

    @property
    def size_c(self):
        return self._size_c

    @property
    def size_t(self):
        return self._size_t

    @property
    def size_x(self):
        return self._size_x

    @property
    def size_y(self):
        return self._size_y

    @property
    def channels(self):
        return self._channels

    def get_channel_index(self, channel):
        return self.channels.index(channel)


def load_annotation(filepath: Path):
    try:
        return matObject(filepath)
    except Exception as e:
        raise (
            "Could not load annotation file. \n"
            "Non MATLAB files currently unsupported"
        ) from e


class TimelapseOMERO(Timelapse):
    """
    Connected to an Image object which handles database I/O.
    """

    def __init__(self, image, annotation, cache, **kwargs):
        super(TimelapseOMERO, self).__init__()
        self.image = image
        # Pre-load pixels
        self.pixels = self.image.getPrimaryPixels()
        self._id = self.image.getId()
        self._name = self.image.getName()
        self._size_x = self.image.getSizeX()
        self._size_y = self.image.getSizeY()
        self._size_z = self.image.getSizeZ()
        self._size_c = self.image.getSizeC()
        self._size_t = self.image.getSizeT()
        self._channels = self.image.getChannelLabels()
        # Check whether there are file annotations for this position
        if annotation is not None:
            self.annotation = load_annotation(annotation)
        # Get an HDF5 dataset to use as a cache.
        compression = kwargs.get("compression", None)
        self.image_cache = cache.require_dataset(
            self.name,
            self.shape,
            dtype=np.float16,
            fillvalue=np.nan,
            compression=compression,
        )

    def get_hypercube(
        self, x=None, y=None, z_positions=None, channels=None, timepoints=None
    ):
        if x is None and y is None:
            tile = None  # Get full plane
        elif x is None:
            ymin, ymax = y
            tile = (None, ymin, None, ymax - ymin)
        elif y is None:
            xmin, xmax = x
            tile = (xmin, None, xmax - xmin, None)
        else:
            xmin, xmax = x
            ymin, ymax = y
            tile = (xmin, ymin, xmax - xmin, ymax - ymin)

        if z_positions is None:
            z_positions = range(self.size_z)
        if channels is None:
            channels = range(self.size_c)
        if timepoints is None:
            timepoints = range(self.size_t)

        z_positions = z_positions or [0]
        channels = channels or [0]
        timepoints = timepoints or [0]

        zcttile_list = [
            (z, c, t, tile)
            for z, c, t in itertools.product(z_positions, channels, timepoints)
        ]
        planes = list(self.pixels.getTiles(zcttile_list))
        order = (
            len(z_positions),
            len(channels),
            len(timepoints),
            planes[0].shape[-2],
            planes[0].shape[-1],
        )
        result = np.stack([x for x in planes]).reshape(order)
        # Set to C, T, X, Y, Z order
        result = np.moveaxis(result, -1, -2)
        return np.moveaxis(result, 0, -1)

    def cache_set(self, save_dir, timepoints, expt_name, quiet=True):
        # TODO deprecate when this is default
        pos_dir = save_dir / self.name
        if not pos_dir.exists():
            pos_dir.mkdir()
        for tp in tqdm(timepoints, desc=self.name):
            for channel in tqdm(self.channels, disable=quiet):
                for z_pos in tqdm(range(self.size_z), disable=quiet):
                    ch_id = self.get_channel_index(channel)
                    image = self.get_hypercube(
                        x=None,
                        y=None,
                        channels=[ch_id],
                        z_positions=[z_pos],
                        timepoints=[tp],
                    )
                    im_name = "{}_{:06d}_{}_{:03d}.png".format(
                        expt_name, tp + 1, channel, z_pos + 1
                    )
                    cv2.imwrite(str(pos_dir / im_name), np.squeeze(image))
        # TODO update positions table to get the number of timepoints?
        return list(itertools.product([self.name], timepoints))

    def run(self, keys, store, save_dir="./", **kwargs):
        """
        Parse file structure and get images for the timepoints in keys.
        """
        save_dir = Path(save_dir)
        if keys is None:
            # TODO save final metadata
            return None
        store = save_dir / store
        # A position specific store
        store = store.with_name(self.name + store.name)
        # Create store if it does not exist
        if not store.exists():
            # The first run, add metadata to the store
            with h5py.File(store, "w") as pos_store:
                # TODO Add metadata to the store.
                pass
        # TODO check how sensible the keys are with what is available
        #   if some of the keys don't make sense, log a warning and remove
        #   them so that the next steps of the pipeline make sense
        return keys

    def clear_cache(self):
        self.image_cache.clear()


class TimelapseLocal(Timelapse):
    def __init__(
        self, position, root_dir, finished=True, annotation=None, cache=None, **kwargs
    ):
        """
        Linked to a local directory containing the images for one position
        in an experiment.
        Can be a still running experiment or a finished one.

        :param position: Name of the position
        :param root_dir: Root directory
        :param finished: Whether the experiment has finished running or the
        class will be used as part of a pipeline, mostly with calls to `run`
        """
        super(TimelapseLocal, self).__init__()
        self.pos_dir = Path(root_dir) / position
        assert self.pos_dir.exists()
        self._id = position
        self._name = position
        if finished:
            self.image_mapper = parse_local_fs(self.pos_dir)
            self._update_metadata()
        else:
            self.image_mapper = dict()
        self.annotation = None
        # Check whether there are file annotations for this position
        if annotation is not None:
            self.annotation = load_annotation(annotation)
        compression = kwargs.get("compression", None)
        self.image_cache = cache.require_dataset(
            self.name,
            self.shape,
            dtype=np.float16,
            fillvalue=np.nan,
            compression=compression,
        )

    def _update_metadata(self):
        self._size_t = len(self.image_mapper)
        # Todo: if cy5 is the first one it causes issues with getting x, y
        #   hence the sorted but it's not very robust
        self._channels = sorted(
            list(set.union(*[set(tp.keys()) for tp in self.image_mapper.values()]))
        )
        self._size_c = len(self._channels)
        # Todo: refactor so we don't rely on there being any images at all
        self._size_z = max([len(self.image_mapper[0][ch]) for ch in self._channels])
        single_img = self.get_hypercube(
            x=None, y=None, z_positions=None, channels=[0], timepoints=[0]
        )
        self._size_x = single_img.shape[2]
        self._size_y = single_img.shape[3]

    def get_hypercube(
        self, x=None, y=None, z_positions=None, channels=None, timepoints=None
    ):
        xmin, xmax = x if x is not None else (None, None)
        ymin, ymax = y if y is not None else (None, None)

        if z_positions is None:
            z_positions = range(self.size_z)
        if channels is None:
            channels = range(self.size_c)
        if timepoints is None:
            timepoints = range(self.size_t)

        def z_pos_getter(z_positions, ch_id, t):
            default = np.zeros((self.size_x, self.size_y))
            names = [
                self.image_mapper[t][self.channels[ch_id]].get(i, None)
                for i in z_positions
            ]
            res = [imread(name) if name is not None else default for name in names]
            return res

        # nested list of images in C, T, X, Y, Z order
        ctxyz = []
        for ch_id in channels:
            txyz = []
            for t in timepoints:
                xyz = z_pos_getter(z_positions, ch_id, t)
                txyz.append(np.dstack(list(xyz))[xmin:xmax, ymin:ymax])
            ctxyz.append(np.stack(txyz))
        return np.stack(ctxyz)

    def clear_cache(self):
        self.image_cache.clear()

    def run(self, keys, store, save_dir="./", **kwargs):
        """
        Parse file structure and get images for the time points in keys.
        """
        if keys is None:
            return None
        elif isinstance(keys, int):
            keys = [keys]
        self.image_mapper.update(parse_local_fs(self.pos_dir, tp=keys))
        self._update_metadata()
        # Create store if it does not exist
        store = get_store_path(save_dir, store, self.name)
        if not store.exists():
            # The first run, add metadata to the store
            with h5py.File(store, "w") as pos_store:
                # TODO Add metadata to the store.
                pass
        # TODO check how sensible the keys are with what is available
        #   if some of the keys don't make sense, log a warning and remove
        #   them so that the next steps of the pipeline make sense
        return keys
