"""
Quantiphyse - Widgets allowing the user to build simple ROIs

A new ROI is based on the grid of an existing data set - chosen by the user.

Copyright (c) 2013-2020 University of Oxford

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

from __future__ import division, unicode_literals, absolute_import, print_function

import collections

import numpy as np

try:
    from PySide import QtGui, QtCore, QtGui as QtWidgets
except ImportError:
    from PySide2 import QtGui, QtCore, QtWidgets

from quantiphyse.data import NumpyData
from quantiphyse.gui.options import OptionBox, DataOption, NumericOption, TextOption
from quantiphyse.gui.widgets import QpWidget, TitleWidget
from quantiphyse.gui.viewer.pickers import PickMode
from quantiphyse.utils import get_icon, QpException

from .tools import CrosshairsTool, PenTool, WalkerTool, PainterTool, EraserTool, RectTool, EllipseTool, PolygonTool, PickTool, BucketTool

DESC = """
Widget for creating test ROIs and basic manual segmentation
"""

TOOLS = [CrosshairsTool(), PenTool(), PainterTool(), EraserTool(), RectTool(), EllipseTool(), PolygonTool(), PickTool(), BucketTool()]

class RoiBuilderWidget(QpWidget):
    """
    Widget for building ROIs
    """
    ADD = 1
    ERASE = 2
    MASK = 3

    def __init__(self, **kwargs):
        super(RoiBuilderWidget, self).__init__(name="ROI Builder", icon="roi_builder", desc=DESC,  
                                               group="ROIs", **kwargs)
        self._history = collections.deque(maxlen=10)
        self._tool = None
        self.grid = None
        self.roi = None
        self.roiname = None

    def init_ui(self):
        layout = QtGui.QVBoxLayout()
        
        title = TitleWidget(self, help="roibuilder", batch_btn=False)
        layout.addWidget(title)

        self.options = OptionBox("Options")
        btn = QtGui.QPushButton("New")
        btn.clicked.connect(self._new_roi)
        self.options.add("ROI", DataOption(self.ivm, rois=True, data=False), btn, key="roi")
        self.options.add("Current label", NumericOption(minval=1, slider=False, intonly=True), key="label")
        self.options.add("Label description", TextOption(), key="label_text")
        self.options.option("roi").sig_changed.connect(self._roi_changed)
        self.options.option("label").sig_changed.connect(self._label_changed)
        self.options.option("label_text").sig_changed.connect(self._label_text_changed)
        layout.addWidget(self.options)
        
        # Add toolbox buttons in a grid
        hbox = QtGui.QHBoxLayout()
        self._toolbox = QtGui.QGroupBox()
        self._toolbox.setTitle("Toolbox")
        self.tools_grid = QtGui.QGridLayout()
        self._toolbox.setLayout(self.tools_grid)

        x, y, cols = 0, 0, 4
        for tool in TOOLS:
            self._add_tool(tool, y, x)
            x += 1
            if x == cols:
                y += 1
                x = 0

        self._undo_btn = QtGui.QPushButton()
        self._undo_btn.clicked.connect(self.undo)
        self._undo_btn.setEnabled(False)
        undo_icon = QtGui.QIcon(get_icon("undo"))
        self._undo_btn.setIcon(undo_icon)
        self._undo_btn.setToolTip("Undo last action")
        self._undo_btn.setFixedSize(32, 32)
        self.tools_grid.addWidget(self._undo_btn, y, x)

        hbox.addWidget(self._toolbox)
        self._toolbox.setEnabled(False)
        hbox.addStretch(1)
        layout.addLayout(hbox)

        # Tool options box - initially invisible
        hbox = QtGui.QHBoxLayout()
        self._tool_options = QtGui.QGroupBox()
        self._tool_options.setVisible(False)

        hbox.addWidget(self._tool_options)
        hbox.addStretch(1)
        layout.addLayout(hbox)

        layout.addStretch(1)
        self.setLayout(layout)

    def activate(self):
        self._roi_changed()
        self.ivl.set_picker(PickMode.SINGLE)
        if self._tool is not None:
            self._tool.selected()

    def deactivate(self):
        self.ivl.set_picker(PickMode.SINGLE)
        if self._tool is not None:
            self._tool.deselected()

    def modify(self, vol=None, slice2d=None, points=None, mode=None):
        """
        Make a change to the ROI we are building

        :param roi_new: Numpy array containing the updated ROI data on the current base grid
        :param vol: 3D 3D selection specified as binary Numpy array same shape as ROI
        :param slice2d: 2D slice selection specified as tuple of (2D binary Numpy array, axis index, position)
        :param points: Selection specified as list of 3D co-ordinates
        :param mode: ``ADD`` to add selection to the ROI (using current label), ``ERASE`` to erase
                     selection from the ROI (set to 0), ``MASK`` to preserve ROI labels in
                     selection but zero everything outside selection
        """
        label = self.options.option("label").value
        self.debug("label=%i", label)

        # For undo functionality: selection is an object specifying which
        # points or ROI region were selected, data_orig is a corresponding
        # object which contains the data before the operation occurred.
        selection, data_orig = None, None

        if points is not None:
            selection = points
            data_orig = [self.roidata[point[0], point[1], point[2]] for point in points]
            for point in points:
                if mode == self.ADD:
                    self.roidata[point[0], point[1], point[2]] = label
                elif mode == self.ERASE:
                    self.roidata[point[0], point[1], point[2]] = 0
                else:
                    raise ValueError("Invalid mode: %i" % mode)
        else:
            # change_subset is a Numpy selection of the points to be
            # affected, data_new is a corresponding binary array identifying
            # the selected points in this subset to modify.
            if vol is not None:
                # Selection left as None to indicate whole volume
                selected_points = vol
                change_subset = self.roidata
            elif slice2d is not None:
                selected_points, axis, pos = slice2d
                slices = [slice(None)] * 3
                slices[axis] = pos
                change_subset = self.roidata[slices]
                # Selection is the axis index and position
                selection = (axis, pos)
            else:
                raise ValueError("Neither volume nor slice nor points provided")

            data_orig = np.copy(change_subset)
            if mode == self.ADD:
                self.debug("Adding: %i", np.count_nonzero(selected_points))
                change_subset[selected_points > 0] = label
            elif mode == self.ERASE:
                self.debug("Erasing: %i", np.count_nonzero(selected_points))
                change_subset[selected_points > 0] = 0
            elif mode == self.MASK:
                self.debug("Masking: %i", np.count_nonzero(selected_points))
                change_subset[selected_points == 0] = 0
            else:
                raise ValueError("Invalid mode: %i" % mode)
        
        # Save the previous state of the data in the history list
        self._history.append((selection, data_orig))
        self._undo_btn.setEnabled(True)
        
        # Update the ROI - note that the regions may have been affected so make
        # sure they are regenerated
        self._update_regions()
        self.ivl.redraw()
        self.debug("Now have %i nonzero", np.count_nonzero(self.roidata))

    def _update_regions(self):
        """
        ROI regions may have been created or added so regenerate them but put 
        back existing label names
        """
        current_regions = self.ivm.data[self.roiname].metadata.pop("roi_regions", {})
        new_regions = self.ivm.data[self.roiname].regions
        for label, desc in current_regions.items():
            if label in new_regions:
                new_regions[label] = desc
        
    def undo(self):
        """
        Undo the last change
        """
        self.debug("ROI undo: %i", len(self._history))
        if not self._history: 
            return

        selection, data_orig = self._history.pop()

        # For selection, None indicates whole volume, tuple indicates
        # an (axis, pos) slice, otherwise we have a sequence of points
        if selection is None:
            self.roidata[:] = data_orig
        elif isinstance(selection, tuple):
            axis, pos = selection
            slices = [slice(None)] * 3
            slices[axis] = pos
            self.roidata[slices] = data_orig
        else:
            for point, orig_value in zip(selection, data_orig):
                self.roidata[point[0], point[1], point[2]] = orig_value

        self._update_regions()
        self.ivl.redraw()
        self.debug("Now have %i nonzero", np.count_nonzero(self.roidata))
        self._undo_btn.setEnabled(len(self._history) > 0)
      
    def _label_changed(self):
        self.debug("Label changed")
        roi = self.ivm.data.get(self.options.option("roi").value, None)
        if roi is not None:
            label = self.options.option("label").value
            self.debug(label)
            regions = roi.regions
            if label in regions:
                self.options.option("label_text").value = regions[label]
            else:
                self.options.option("label_text").value = "Region %i" % label
                self.debug("New ROI label - Force redraw")
                regions[label] = ""
                roi.view.cmap_range = roi.suggest_cmap_range()

    def _label_text_changed(self):
        self.debug("Label text changed")
        roi = self.ivm.data.get(self.options.option("roi").value, None)
        if roi is not None:
            label = self.options.option("label").value
            label_text = self.options.option("label_text").value
            self.debug(label)
            self.debug(label_text)
            regions = roi.regions
            regions[label] = label_text

    def _roi_changed(self):
        roi = self.ivm.data.get(self.options.option("roi").value, None)
        self._toolbox.setEnabled(roi is not None)
        if roi is not None:
            # FIXME this will only work if ROI is NumpyData. Otherwise we are
            # manipulating a numpy array which may just be a proxy for the file
            # storage.
            regions = roi.regions
            current_label = self.options.option("label").value
            if self.roiname != roi.name or current_label not in regions.keys():
                self.options.option("label").value = min(list(regions.keys()) + [1, ])
            self.roiname = roi.name
            self.grid = roi.grid
            self.roidata = roi.raw()

    def _new_roi(self):
        dialog = QtGui.QDialog(self)
        dialog.setWindowTitle("New ROI")
        vbox = QtGui.QVBoxLayout()
        dialog.setLayout(vbox)

        optbox = OptionBox()
        optbox.add("ROI name", TextOption(), key="name")
        optbox.add("Data space from", DataOption(self.ivm), key="grid")
        vbox.addWidget(optbox)

        buttons = QtGui.QDialogButtonBox(QtGui.QDialogButtonBox.Ok | QtGui.QDialogButtonBox.Cancel)
        buttons.accepted.connect(dialog.accept)
        buttons.rejected.connect(dialog.reject)
        vbox.addWidget(buttons)
        
        ok = dialog.exec_()
        if ok:
            roiname = optbox.option("name").value
            gridfrom = optbox.option("grid").value
            if not roiname or not gridfrom:
                raise QpException("Must provide a ROI name and a dataset to base it on")
            grid = self.ivm.data[gridfrom].grid
            roidata = np.zeros(grid.shape, dtype=np.int)
            self.ivm.add(NumpyData(roidata, grid=grid, roi=True, name=roiname), make_current=True)

            # Throw away old history. FIXME is this right, should we keep existing data and history?
            # Also should we cache old history in case we go back to this ROI?
            self._history = []
            self._undo_btn.setEnabled(False)
            self.options.option("roi").value = roiname

    def _tool_selected(self, tool):
        def _select():
            if self._tool is not None:
                self._tool.btn.setStyleSheet("")
                self._tool.deselected()
            
            self._tool = tool
            self._tool.btn.setStyleSheet("border: 2px solid QLinearGradient( x1: 0, y1: 0, x2: 0, y2: 1, stop: 0 #ffa02f, stop: 1 #d7801a);")
            # Replace the old tool options with the new one. Need to reparent the
            # existing layout to a temporary widget which will then get deleted 
            QtGui.QWidget().setLayout(self._tool_options.layout())
            self._tool_options.setLayout(self._tool.interface())
            self._tool_options.setTitle(tool.name)
            self._tool_options.setVisible(True)
            self._tool.selected()

        return _select
    
    def _add_tool(self, tool, x, y):
        tool.ivm = self.ivm
        tool.ivl = self.ivl
        tool.builder = self
        btn = QtGui.QPushButton()
        btn.setIcon(QtGui.QIcon(get_icon(tool.name.lower())))
        btn.setToolTip(tool.tooltip)
        btn.setFixedSize(32, 32)
        btn.clicked.connect(self._tool_selected(tool))
        tool.btn = btn
        self.tools_grid.addWidget(btn, x, y)
