# -*- coding: utf-8 -*-

# Copyright © 2020, Institut Pasteur
#   Contributor: François Laurent

# This file is part of the TRamWAy software available at
# "https://github.com/DecBayComp/TRamWAy" and is distributed under
# the terms of the CeCILL license as circulated at the following URL
# "http://www.cecill.info/licenses.en.html".

# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL license and that you accept its terms.


from ..attribute import *
from ..artefact import *
from .abc import *
import os.path
import numpy as np


class ImagesInitializer(Initializer):
    """
    Initial value for the :class:`~tramway.analyzer.RWAnalyzer`
    :attr:`~tramway.analyzer.RWAnalyzer.images` attribute.

    *from_...* methods alters the parent attribute which specializes
    into an initialized :class:`Images` object.
    """

    __slots__ = ()

    def from_tiff_file(self, filepath):
        """
        Defines the raw single molecule imaging file.

        Loading is not performed while calling this method.
        """
        self.specialize(StandaloneTiffFile, filepath)

    def from_tiff_files(self, filepattern):
        """
        Defines the raw single molecule imaging files.

        Loading is not performed while calling this method.
        """
        self.specialize(TiffFiles, filepattern)


class ImageParameters(object):
    __slots__ = ()
    # default implementation for children classes that define `_pixel_size` and `_loc_offset`
    def __init__(self):
        self._pixel_size = self._loc_offset = None

    @property
    def pixel_size(self):
        r"""
        *float*: Pixel size in :math:`\mu m`
        """
        return self._pixel_size

    @pixel_size.setter
    def pixel_size(self, pxsize):
        self._pixel_size = pxsize

    @property
    def loc_offset(self):
        """
        *numpy.ndarray*: Offset between coordinates and the image, in pixels
        """
        return self._loc_offset

    @loc_offset.setter
    def loc_offset(self, offset):
        self._loc_offset = offset

    # access to shared parameters
    @property
    def frame_interval(self):
        """
        *float*: See :attr:`~tramway.analyzer.spt_data.SPTParameters.frame_interval`
        """
        return self._eldest_parent.spt_data.frame_interval

    @frame_interval.setter
    def frame_interval(self, dt):
        self._eldest_parent.spt_data.frame_interval = dt

    @property
    def dt(self):
        """
        *float*: See :attr:`~tramway.analyzer.spt_data.SPTParameters.dt`
        """
        return self.frame_interval

    @dt.setter
    def dt(self, dt):
        self.frame_interval = dt


class _RawImage(AnalyzerNode, ImageParameters):
    __slots__ = ("_stack", "_pixel_size", "_loc_offset")

    def __init__(self, stack, **kwargs):
        AnalyzerNode.__init__(self, **kwargs)
        ImageParameters.__init__(self)
        self._stack = stack

    @property
    def stack(self):
        if callable(self._stack):
            self._stack = self._stack()
        return self._stack

    @stack.setter
    def stack(self, stack):
        self._stack = stack

    @property
    def n_frames(self):
        return len(self.stack)

    @property
    def width(self):
        return first(self.stack).shape[1]

    @property
    def height(self):
        return first(self.stack).shape[0]

    def as_frames(self, index=None, return_time=False):
        """
        Generator function; iterates over the image frames and yields
        NumPy 2D arrays, or pairs of (*float*, NumPy 2D array).

        Arguments:

            index (*int*, *Set*, *Sequence* or *callable*):
                frame filter; see also :func:`~tramway.analyzer.attribute.indexer`.

            return_time (bool): return time along with image frames, as first item.

        """
        for f, frame in indexer(index, self.stack, return_index=True):
            if return_time:
                t = (f + 1) * self.dt
                yield t, frame
            else:
                yield frame

    def cropping_bounds(self, bounding_box, relative=True):
        """
        Discretize bounding box spatial coordinates to clip them
        onto image pixels.

        Pixels are assumed to be centered on their indices in the array,
        when plotted.

        .. note::

            `cropping_bounds` is supposed to work well in combination with
            `crop_frames`.
            However, a notable exception is out-of-bounds indices.
            `crop_frames` adjusts the size of the pixel matrix wrt the image,
            while `cropping_bounds` does not min-max the returned indices wrt
            the natural bounds of the image.

        Arguments:

            bounding_box (tuple): pair of ndarrays (lower bound, upper bound).

            relative (bool): if ``True``, ignores the coordinate offset if any.

        Returns:

            tuple: pair of ndarrays; pixel-clipped lower bounds and upper bounds.

        """
        loc_offset = self.loc_offset
        if loc_offset is None or relative:
            loc_offset = 0.0
        lb, ub = [b / self.pixel_size - loc_offset for b in bounding_box]
        lb = np.floor(lb) * self.pixel_size
        ub = np.ceil(ub) * self.pixel_size
        return lb, ub

    def crop_frames(self, bounding_box, index=None, return_time=False):
        """
        Generator function; iterates and crops the image frames, similarly to
        :meth:`as_frames`.

        .. note::

            Time bounds are not supported yet.

        Arguments:

            bounding_box (tuple): pair of NumPy arrays (lower bound, upper bound).

            index (*int*, *Set*, *Sequence* or *callable*):
                frame filter; see also :func:`~tramway.analyzer.attribute.indexer`.

            return_time (bool): return time along with cropped image frames, as first item.

        """
        lb, ub = bounding_box
        if 2 < lb.size:
            self.logger.warning("time supports are not supported yet")
            lb, ub = lb[:2], ub[:2]
            bounding_box = (lb, ub)
        loc_offset = self.loc_offset
        if loc_offset is None:
            loc_offset = 0.0
        lb, ub = [b / self.pixel_size - loc_offset for b in bounding_box]
        lb = np.floor(lb).astype(int)
        ub = np.ceil(ub).astype(int)
        if not np.all(lb <= ub):
            raise ValueError("image cropping failed: lower bound > upper bound")
        i_min = max(0, self.height - 1 - ub[1])  # range start (included)
        i_max = min(self.height - lb[1], self.height)  # range stop (excluded)
        j_min = max(0, lb[0])  # range start (included)
        j_max = min(ub[0] + 1, self.width)  # range stop (excluded)
        if return_time:
            for t, frame in self.as_frames(index, return_time):
                yield t, frame[i_min:i_max, j_min:j_max]
        else:
            for frame in self.as_frames(index, return_time):
                yield frame[i_min:i_max, j_min:j_max]

    def to_color_movie(
        self,
        output_file=None,
        fourcc="VP80",
        colormap="gray",
        locations=None,
        trajectories=None,
        frames=None,
        origin=None,
        markersize=2,
        linecolor="many",
        linewidth=None,
        magnification=None,
        playback_rate=None,
        light_intensity=1.0,
    ):
        """
        Generates a movie of the images with overlaid locations or trajectories.

        Arguments:

            output_file (str): path of the movie file.

            fourcc (str): 4-character code string.

            colormap (*str* or *matplotlib.colors.ListedColormap*): Matplotlib colormap;
                see also https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html.

            locations (pandas.DataFrame): particle locations with columns :const:`'x'`,
                :const:`'y'` and :const:`'t'`.

            trajectories (pandas.DataFrame): trajectories with columns :const:`'n'`,
                :const:`'x'`, :const:`'y'` and :const:`'t'`;
                note this differs from translocations in that displacements are not
                encoded along with locations and trajectory terminations are independent rows.

            frames (sequence): iterable of time (*float*, in seconds) and frame (2D pixel array)
                pairs; if `origin` defines time, `frames` can be a sequence of frames only;
                :meth:`as_frames` is called instead if undefined.

            origin (*numpy.ndarray* or *pandas.Series*): data lower bound if `frames` is defined;
                implicit columns are :const:`'x'` and :const:`'y'`;
                to define time, `origin` must be a `Series`
                with indices :const:`'x'`, :const:`'y'` and :const:`'t'`;
                `origin` is useful only for overlaying locations or trajectories.

            markersize (int): location marker size in pixels (square side).

            linecolor (*str* or 3-column float array): color for trajectories;
                value :const:`None` defaults to red.

            linewidth (*int* or *float*): trajectory line width.

            magnification (*int* or *str*): the original image pixels can be represented as
                square-patches of `magnification` video pixel side;
                if *str*:
                :const:`'1x'` = `round(pixel_size/localization_precision)`,
                :const:`'2x'` = `round(2*pixel_size/localization_precision)`;
                :const:`'2x'` is adequate for overlaid trajectories, even with the over-compressing
                :const:`'MJPG'` encoder.

            playback_rate (float): default playback rate;
                1 is real-time, 0.5 is half the normal speed.

            light_intensity (float): scale the pixel intensity values;
                works best with the 'gray' colormap.

        """
        import cv2

        # TODO: support PIL as an optional replacement for skimage
        from skimage.util import img_as_ubyte
        from . import draw
        from matplotlib import cm, colors

        if output_file is None:
            try:
                output_file = os.path.splitext(self.filepath)[0] + ".avi"
            except AttributeError:
                raise ValueError("output_file is not defined")

        dt = self.dt
        pxsize = self.pixel_size
        if pxsize is None:
            pxsize = 1.0

        if self.loc_offset is None:
            offset = np.zeros((1, 2), dtype=np.float)
        else:
            offset = np.asarray(self.loc_offset)
            if not offset.shape[1:]:
                offset = offset[np.newaxis, :]
        # x_offset, y_offset = offset[0,0], offset[0,1]

        if origin is None:
            xy0 = None
            t0 = None
        elif isinstance(origin, np.ndarray):
            xy0 = origin.reshape((1, 2)).astype(np.float)
            t0 = None
        else:  # if isinstance(origin, pd.Series):
            xy0 = origin[list("xy")].values[np.newaxis, :].astype(np.float)
            try:
                t0 = origin["t"].astype(np.float)
            except KeyError:
                t0 = None

        if magnification is None:
            magnification = 1
        elif isinstance(magnification, str) and magnification.endswith("x"):
            try:
                magnification = float(magnification[:-1])
            except ValueError:
                raise ValueError("cannot parse the magnification factor")
            try:
                magnification *= (
                    self.pixel_size / self._eldest_parent.localization_precision
                )
            except (AttributeError, TypeError):
                raise ValueError(
                    "failed to adjust magnification; pixel_size or localization_precision not defined"
                )
            magnification = int(np.round(magnification))

        if playback_rate is None:
            playback_rate = 1.0

        if light_intensity == 1:
            color_scale = None
        else:
            color_scale = int(np.round(1.0 / light_intensity))

        marker_color = line_color = np.array([[1, 0, 0]], dtype=np.float)  # red
        if locations is not None:
            if callable(locations):
                locations = locations()
            if markersize % 2:
                marker_size_delta = (markersize - 1) / 2

                def _floor(x):
                    return np.round(x) - marker_size_delta

                def _ceil(x):
                    return np.round(x) + marker_size_delta

            else:
                marker_size_delta = markersize / 2 - 1

                def _floor(x):
                    return np.floor(x) - marker_size_delta

                def _ceil(x):
                    return np.floor(x) + 1 + marker_size_delta

        if trajectories is not None:
            if callable(trajectories):
                trajectories = trajectories()
            if isinstance(linecolor, str):
                if linecolor == "many":
                    from tramway.plot.mesh import __colors__

                    line_color = np.stack(
                        [colors.to_rgb(c) for c in __colors__], axis=0
                    )
                else:
                    line_color = colors.to_rgb(linecolor)[np.newaxis, :]
            elif isinstance(linecolor, np.ndarray):
                line_color = linecolor
            lc = line_color

            def append_n(n, ijk):
                i, j, k = ijk
                return i, j, k, np.full(k.shape, n)

            trajectories_as_dict = {}
            # trajectories_as_dict = { n: trajectories[list('xyt')][trajectories['n']==n]
            #        for n in np.unique(trajectories['n']) }
        dt_max_err = (0.5 * dt) ** 2

        def isclose(ts, t):
            _dt = ts - t
            return _dt * _dt < dt_max_err

        vid = None  # will be initialized later
        vid_pxsize = pxsize / magnification
        mag_offset = 0.5 * (magnification - 1)
        vid_offset = offset / magnification - mag_offset

        if isinstance(colormap, str):
            colormap = cm.get_cmap(colormap)
        elif not callable(colormap):  # isinstance(colormap, colors.ListedColormap):
            colormap = cm.viridis

        if frames:
            if not (callable(frames) or t0 is None):

                def wrap(fs, t0, dt):
                    for i, f in enumerate(fs):
                        if isinstance(f, (tuple, list)):
                            t, f = f
                            if i == 0 and t != t0:
                                raise ValueError(
                                    "start time does not match with first frame time: {} != {}".format(
                                        t0, t
                                    )
                                )
                        else:
                            t = t0 + i * dt
                        yield t, f

                frames = wrap(frames, t0, dt)
        else:
            frames = self.as_frames

        t_prev = -1
        for t, frame in frames(return_time=True) if callable(frames) else frames:
            if t <= t_prev:
                raise ValueError("time does not strictly increase")
            t_prev = t

            if vid is None:
                # complete the initialization
                height, width = frame.shape
                ii_max = height * magnification - 1
                jj_max = width * magnification - 1

                vid = cv2.VideoWriter(
                    os.path.expanduser(output_file),
                    cv2.VideoWriter_fourcc(*fourcc),
                    playback_rate / dt,
                    (width * magnification, height * magnification),
                    True,
                )

            if color_scale:
                frame = frame // color_scale

            frame = colormap(frame)[:, :, :3]

            if magnification:
                frame = np.repeat(
                    np.repeat(frame, magnification, axis=0), magnification, axis=1
                )

            if trajectories is not None:
                # anti-aliasing
                lw = 0 if linewidth is None else linewidth
                aa = bool(linewidth)

                # get the ids of the active trajectories for frame f
                active_trajectories = trajectories["n"][isclose(trajectories["t"], t)]
                assert all_unique(active_trajectories)

                # extract the corresponding coordinate series truncated at time t
                trajs_f = []
                for n in active_trajectories:
                    try:
                        xyt_n = trajectories_as_dict[n]
                    except KeyError:
                        trajectories_as_dict[n] = xyt_n = trajectories[
                            trajectories["n"] == n
                        ]
                    xy_n = xyt_n[list("xy")][xyt_n["t"] < t + 0.5 * dt]
                    if len(xy_n) < 2:
                        continue
                    if xy0 is None:
                        traj = xy_n.values / vid_pxsize - vid_offset
                    else:
                        traj = (xy_n.values - xy0) / vid_pxsize + mag_offset
                    if np.any(traj < 0):  # this may occur with negative offsets
                        continue
                    if not aa:
                        traj = np.round(traj).astype(np.uint32)
                    trajs_f.append((n, traj))

                # overlay the truncated trajectories
                if trajs_f:
                    jj, ii, kk, nn = zip(
                        *[
                            append_n(
                                n,
                                draw.multiline_aa(traj[:, 0], traj[:, 1], thickness=lw),
                            )
                            for n, traj in trajs_f
                        ]
                    )
                    ii = np.concatenate(ii)
                    jj = np.concatenate(jj)
                    ok = (0 <= ii) & (ii <= ii_max) & (0 <= jj) & (jj <= jj_max)
                    ii, jj = ii[ok], jj[ok]
                    kk = np.concatenate(kk)[ok][:, np.newaxis]
                    if 1 < len(line_color):
                        nn = np.concatenate(nn)[ok]
                        lc = line_color[nn % len(line_color)]
                    frame[ii_max - ii, jj, :] = (1.0 - kk) * frame[
                        ii_max - ii, jj, :
                    ] + kk * lc

            if locations is not None:
                xy = locations[isclose(locations["t"], t)]
                if xy0 is None:
                    xy_f = xy[list("xy")].values / vid_pxsize - vid_offset
                else:
                    xy_f = (xy[list("xy")].values - xy0) / vid_pxsize + mag_offset
                for j, i in xy_f:
                    i = np.array([_floor(i), _ceil(i)], dtype=np.int)
                    j = np.array([_floor(j), _ceil(j)], dtype=np.int)
                    if np.any(i < 0) or np.any(j < 0):
                        continue
                    i, j = np.meshgrid(
                        np.arange(i[0], i[1] + 1),
                        np.arange(j[0], j[1] + 1),
                        indexing="ij",
                    )
                    try:
                        frame[ii_max - np.ravel(i), np.ravel(j), :] = marker_color
                    except IndexError:
                        pass

            vid.write(img_as_ubyte(frame)[:, :, ::-1])
        vid.release()

    @property
    def _mpl_impl(self):
        from .mpl import Mpl

        return Mpl

    @property
    def mpl(self):
        """tramway.analyzer.images.mpl.Mpl: Matplotlib utilities"""
        return self._mpl_impl(self)


class RawImage(_RawImage):
    __slots__ = ()


Image.register(RawImage)


class _ImageFile(_RawImage):
    __slots__ = ("_filepath",)

    def __init__(self, filepath, **kwargs):
        self._filepath = filepath
        RawImage.__init__(self, self.read, **kwargs)

    @property
    def filepath(self):
        return self._filepath

    @filepath.setter
    def filepath(self, fp):
        self._filepath = fp

    def read(self):
        from skimage import io

        img = io.MultiImage(os.path.expanduser(self.filepath))
        if len(img) == 1 and 2 < len(img[0].shape):
            img = img[0]
        return img


class ImageFile(_ImageFile):
    __slots__ = ()


Image.register(ImageFile)


class _TiffFile(_ImageFile):
    __slots__ = ()


class TiffFile(ImageFile):
    __slots__ = ()


class StandaloneImage(object):
    __slots__ = ()

    def __len__(self):
        return 1

    def __iter__(self):
        yield self


class ImageIterator(AnalyzerNode, ImageParameters):
    """Partial :class:`Images` implementation for multiple SPT data items.

    Children classes must implement the :meth:`__iter__` method."""

    __slots__ = ()

    def __init__(self, **kwargs):
        AnalyzerNode.__init__(self, **kwargs)
        # ImageParameters.__init__(self)

    @property
    def pixel_size(self):
        it = iter(self)
        pxsize = next(it).pixel_size
        while True:
            try:
                _pxsize = next(it).pixel_size
            except StopIteration:
                break
            else:
                if _pxsize != pxsize:
                    raise AttributeError("not all the images share the same pixel size")
        return pxsize

    @pixel_size.setter
    def pixel_size(self, pxsize):
        for f in self:
            f.pixel_size = pxsize

    @property
    def loc_offset(self):
        it = iter(self)
        offset = next(it).loc_offset
        while True:
            try:
                _offset = next(it).loc_offset
            except StopIteration:
                break
            else:
                if not np.all(_offset == offset):
                    raise AttributeError(
                        "not all the images share the same localization offset"
                    )
        return offset

    @loc_offset.setter
    def loc_offset(self, offset):
        for f in self:
            f.loc_offset = offset


class StandaloneImageFile(_ImageFile, StandaloneImage):
    __slots__ = ()


Images.register(StandaloneImageFile)


class StandaloneTiffFile(_TiffFile, StandaloneImage):
    __slots__ = ()


Images.register(StandaloneTiffFile)


class RawImages(ImageIterator):
    __slots__ = ("_images",)

    def __init__(self, stacks, **kwargs):
        ImageIterator.__init__(self, **kwargs)
        self.images = stacks

    @property
    def images(self):
        return self._images

    @images.setter
    def images(self, stacks):
        self._images = tuple([self._bear_child(RawImage, stack) for stack in stacks])

    @property
    def reified(self):
        return True

    def __len__(self):
        return len(self.images)

    def __iter__(self):
        yield from self.images


# Images.register(RawImages)


class ImageFiles(ImageIterator):
    __slots__ = ("_files", "_filepattern")

    def __init__(self, filepattern, **kwargs):
        ImageIterator.__init__(self, **kwargs)
        self._files = []
        if isinstance(filepattern, str):
            self._filepattern = os.path.expanduser(filepattern)
        else:
            self._filepattern = [os.path.expanduser(pattern) for pattern in filepattern]

    @property
    def filepattern(self):
        return self._filepattern

    @filepattern.setter
    def filepattern(self, fp):
        if self._files:
            if fp != self._filepattern:
                raise AttributeError(
                    "the files have already been listed; cannot set the file pattern anymore"
                )
        else:
            self._filepattern = fp

    @property
    def files(self):
        if not self._files:
            self.list_files()
        return self._files

    @property
    def filepaths(self):
        return [f.filepath for f in self.files]

    @property
    def partially_reified(self):
        return self._files and any([f.reified for f in self._files])

    @property
    def fully_reified(self):
        return self._files and all([f.reified for f in self._files])

    @property
    def reified(self):
        return self.fully_reified

    def __len__(self):
        return len(self.files)

    def __iter__(self):
        yield from self.files

    def list_files(self, _element_cls=ImageFile):
        from glob import glob

        self._files = glob(self.filepattern)
        if not self._files:
            raise ValueError("no files found")
        self._files = [
            self._bear_child(_element_cls, filepath) for filepath in self._files
        ]


Images.register(ImageFiles)


class TiffFiles(ImageFiles):
    __slots__ = ()

    def list_files(self):
        ImageFiles.list_files(self, TiffFile)


def all_unique(values):
    return np.unique(values).size == values.size


__all__ = [
    "Images",
    "Image",
    "ImagesInitializer",
    "ImageParameters",
    "_RawImage",
    "RawImage",
    "ImageFile",
    "TiffFile",
    "StandaloneImageFile",
    "StandaloneTiffFile",
    "RawImages",
    "ImageFiles",
    "TiffFiles",
]
