# Copyright 2017-2018 The Tangram Developers. See the AUTHORS file at the
# top-level directory of this distribution and at
# https://github.com/renatoGarcia/tangram/blob/master/AUTHORS.
#
# This file is part of Tangram.
#
# Tangram is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Tangram is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with Tangram in the COPYING and COPYING.LESSER files.
# If not, see <http://www.gnu.org/licenses/>.

from typing import Set, Dict, Iterable, Tuple, List
from functools import reduce
from copy import copy
import random

from gi.repository import GObject
import cairo
import numpy as np
from scipy.ndimage import binary_dilation

from .._utility import ndarray_to_surface
from ..tangram_typing import Color, Mask
from .. import tsg


IntColor = Tuple[int, int, int, int]


def _to_int_palette(palette) -> List[IntColor]:
    return [(max(0, min(int(c.r*256), 255)),
             max(0, min(int(c.g*256), 255)),
             max(0, min(int(c.b*256), 255)),
             max(0, min(int(c.a*256), 255)))
            for c in palette]


class _Mosaic(tsg.Node):
    """
    Internal Mosaic class.

    To admit overlapping regions, in the internal array each region or set of regions will
    receive an unique code. The mapping from internal code to a set of regions will be
    done through the reg_map dictionary.
    """

    __gsignals__ = {
        'active-changed': (GObject.SignalFlags.RUN_FIRST, None, (GObject.TYPE_PYOBJECT,)),
    }

    def __init__(self, palettes_size: int, random_picker: bool = False) -> None:
        """
        Args:
            array: Array mapping Image positions to a Regions power set internal code.
        """
        super().__init__()

        self.__coded_regions: np.ndarray = np.zeros(shape=(), dtype=np.uint16)
        self.__code_to_indexes: Dict[int, Set[int]] = {0: frozenset()}
        self.__index_to_codes: Dict[int, Set[int]] = {}

        # Map from a region index to a color palette index.
        self.__colors: Dict[int, int] = {}
        self.__base_palette: List[IntColor] = [(0, 0, 0, 0) for _ in range(palettes_size)]
        self.__base_active_palette: List[IntColor] = [(0, 0, 0, 0) for _ in range(palettes_size)]
        self.__pinned_palette: List[IntColor] = [(0, 0, 0, 0) for _ in range(palettes_size)]
        self.__pinned_active_palette: List[IntColor] = [(0, 0, 0, 0) for _ in range(palettes_size)]

        self.__pinned_regions: Set[int] = set()
        self.__active_regions: Set[int] = set()

        self.__buffer: np.ndarray = np.zeros((), dtype=np.uint8)

        self.__random_picker = random_picker

        self.connect('enter-event', self.__on_hover_event)
        self.connect('motion-event', self.__on_hover_event)
        self.connect('leave-event', self.__on_leave_event)
        self.connect('notify::is-visible', self.__on_notify_is_visible)

        self.__is_frozen = False

    def set_envelope(self, box: tsg.Box) -> None:
        if box.is_degenerate():
            return

        if self.__coded_regions.ndim == 0:
            rows = int(box.get_height())
            cols = int(box.get_width())
            self.__coded_regions = np.zeros(shape=(rows, cols), dtype=np.uint16)
            self.__buffer = np.zeros((rows, cols, 4), dtype=np.uint8)
            self._set_surface_shape(box.get_size())
        elif box != self.envelope():
            raise RuntimeError("Trying to set a Mosaic to a invalid shaped Image. "
                               f"Mosaic shape: {self.envelope().get_size()}, "
                               f"Image shape: {box.get_size()}.")

    def contains(self, point: tsg.Point) -> bool:
        rows, cols = self.__coded_regions.shape
        r, c = int(point.y), int(point.x)
        if not 0 <= r < rows or not 0 <= c < cols:
            return False

        reg_code = self.__coded_regions[r, c]
        if reg_code != 0:
            return True
        else:
            return False

    def _draw(self, ctx: cairo.Context, _width: float, _height: float) -> None:
        surface = ndarray_to_surface(self.__buffer)
        pattern = cairo.SurfacePattern(surface)
        pattern.set_filter(cairo.Filter.NEAREST)
        ctx.set_source(pattern)
        ctx.paint()

    def __pick_color(self, mask: Mask) -> int:
        mask = binary_dilation(mask)
        codes = np.unique(self.__coded_regions[mask])
        regions = reduce(set.union, (self.__code_to_indexes[c] for c in codes), set())
        colors = {self.__colors[idx] for idx in regions}
        available_colors = [c for c in range(len(self.__base_palette)) if c not in colors]
        if not available_colors:
            raise ValueError("Coloring regions need more than available colors. "
                             f"({len(self.__base_palette)} available).")
        if self.__random_picker:
            return random.choice(available_colors)
        else:
            return available_colors[0]

    def __highlight_regions(self) -> None:
        """Update the highlighted regions."""
        if self.__is_frozen:
            return

        def paint(regions, palette):
            codes = reduce(set.union,
                           (self.__index_to_codes[idx] for idx in regions),
                           set())
            for code in codes:
                color = palette[self.__colors[code]]
                self.__buffer[self.__coded_regions == code] = color

        self.__buffer.fill(0)

        default_regions = self.get_regions() - self.__pinned_regions
        default_active_regions = default_regions & self.__active_regions
        default_only_regions = default_regions - self.__active_regions
        paint(default_active_regions, self.__base_active_palette)
        paint(default_only_regions, self.__base_palette)

        pinned_regions = self.__pinned_regions
        pinned_active_regions = pinned_regions & self.__active_regions
        pinned_only_regions = pinned_regions - self.__active_regions
        paint(pinned_active_regions, self.__pinned_active_palette)
        paint(pinned_only_regions, self.__pinned_palette)

        self._invalidate()

    def __on_notify_is_visible(self, node: tsg.Node, param: GObject.GParamSpec) -> None:
        is_visible = node.get_property(param.name)
        if not is_visible:
            if self.__active_regions:
                self.__active_regions.clear()
                self.emit('active-changed', set())
            self.__pinned_regions.clear()
        self.__highlight_regions()

    def __on_hover_event(self, _node: '_Mosaic', event: tsg.EventHover) -> None:
        point = event.local_point
        active_code = self.__coded_regions[int(point.y), int(point.x)]
        active_regions = copy(self.__code_to_indexes[active_code])

        if active_regions != self.__active_regions:
            self.__active_regions = copy(active_regions)
            self.__highlight_regions()
            self.emit('active-changed', active_regions)

        return False

    def __on_leave_event(self, _node: '_Mosaic', event: tsg.EventHover) -> None:
        self.__active_regions.clear()
        self.__highlight_regions()
        self.emit('active-changed', set())

        return False

    def shape(self) -> Tuple[int, int]:
        return self.__buffer.shape[:2]

    def set_base_palette(self, palette: List[Color]) -> None:
        assert len(palette) == len(self.__base_palette)
        self.__base_palette = _to_int_palette(palette)
        self.__highlight_regions()

    def set_base_active_palette(self, palette: List[Color]) -> None:
        assert len(palette) == len(self.__base_active_palette)
        self.__base_active_palette = _to_int_palette(palette)
        self.__highlight_regions()

    def set_pinned_palette(self, palette: List[Color]) -> None:
        assert len(palette) == len(self.__pinned_palette)
        self.__pinned_palette = _to_int_palette(palette)
        self.__highlight_regions()

    def set_pinned_active_palette(self, palette: List[Color]) -> None:
        assert len(palette) == len(self.__pinned_active_palette)
        self.__pinned_active_palette = _to_int_palette(palette)
        self.__highlight_regions()

    def regions_at_point(self, x: int, y: int) -> Set[int]:
        code = self.__coded_regions[y, x]
        return copy(self.__code_to_indexes[code])

    def regions_within(self, mask: Mask) -> Set[int]:
        if mask.dtype != np.bool:
            raise ValueError("The mask argument dtype must be bool")
        if mask.shape != self.__coded_regions.shape:
            raise ValueError("The mask argument must have the same shape as the mosaic")

        overlapping_region = self.__coded_regions[mask]
        overlapping_codes: Set[int] = set(np.unique(overlapping_region)) - {0}

        overlapping_indexes: Set[int] = set()
        for idxs in (self.__code_to_indexes[c] for c in overlapping_codes):
            overlapping_indexes |= idxs

        within_indexes: Set[int] = set()
        for idx in overlapping_indexes:
            codes = self.__index_to_codes[idx]
            idx_mask = reduce(np.logical_or, (self.__coded_regions == c for c in codes))
            # if idx_mask is wholly inside mask
            if not np.any(np.logical_xor(np.logical_or(idx_mask, mask), mask)):
                within_indexes.add(idx)

        return within_indexes

    def get_region_mask(self, region: int) -> Mask:
        codes = self.__index_to_codes[region]
        return reduce(np.logical_or, (self.__coded_regions == c for c in codes))

    def set_frozen(self, value: bool) -> None:
        # TODO: Do put this on Node class?
        if value == self.__is_frozen:
            return
        self.__is_frozen = value
        if not self.__is_frozen:
            self.__highlight_regions()

    def get_active(self) -> Set[int]:
        return copy(self.__active_regions)

    def set_active(self, value: Set[int]) -> None:
        active_regions = value & self.get_regions()
        if active_regions != self.__active_regions:
            self.__active_regions = copy(active_regions)
            self.__highlight_regions()
            # TODO: fix a infinite recusion on bind_mosaics when enabling this
            # self.emit('active-changed', active_regions)

    def get_pinned(self) -> Set[int]:
        return self.__pinned_regions.copy()

    def set_pinned(self, value: Iterable[int]) -> None:
        value = set(value)
        if value != self.__pinned_regions:
            self.__pinned_regions = value
            self.__highlight_regions()

    def get_regions(self) -> Set[int]:
        return set(self.__index_to_codes.keys())

    def remove_region(self, region: int) -> None:
        self.__pinned_regions.discard(region)
        self.__active_regions.discard(region)
        del self.__colors[region]

        codes: Set[int] = self.__index_to_codes.pop(region)
        for code in codes:
            regs: Set[int] = copy(self.__code_to_indexes[code])
            regs.remove(region)
            if not regs:
                self.__coded_regions[self.__coded_regions==code] = 0
                del self.__code_to_indexes[code]
            else:
                self.__code_to_indexes[code] = regs
                # TODO: merge any two codes translating to the same regions set

        self.__highlight_regions()

    def add_region(self, mask: Mask) -> int:
        """Adds a new region and rerturns its index.
        """
        assert mask.ndim == 2
        assert mask.dtype == np.bool

        if self.__coded_regions.ndim == 0:
            rows, cols = mask.shape
            self.__coded_regions = np.zeros(shape=(rows, cols), dtype=np.uint16)
            self.__buffer = np.zeros((rows, cols, 4), dtype=np.uint8)
            self._set_surface_shape(tsg.Size(width=cols, height=rows))
        elif mask.shape != self.__coded_regions.shape:
            raise RuntimeError("Trying to add a mask to a Mosaic with different shapes."
                               f"Mosaic shape: {self.__coded_regions.shape}, "
                               f"Mask shape: {mask.shape}.")

        if np.any(self.__coded_regions[mask]):
            # TODO: implement overlapping masks.
            raise ValueError("Overlapping masks not implemented yet.")

        code = max(self.__code_to_indexes.keys()) + 1
        if code > np.iinfo(self.__coded_regions.dtype).max:
            raise RuntimeError("The quantity of regions is greater than the supported.")
        index = max(self.__index_to_codes.keys(), default=0) + 1

        color = self.__pick_color(mask)
        self.__colors[index] = color

        self.__coded_regions[mask] = code
        self.__code_to_indexes[code] = {index}
        self.__index_to_codes[index] = {code}

        self.__highlight_regions()

        return index
