# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2016-2017 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/
"""Tool to display Spectra maps and estimate E0"""


__authors__ = ["H. Payno"]
__license__ = "MIT"
__date__ = "04/12/2019"


from silx.gui import qt
from silx.utils.enum import Enum
from est.core.types import XASObject
from silx.gui.plot.items.roi import RectangleROI
from silx.gui.plot.tools.roi import RegionOfInterestManager
import functools
import numpy

from est.gui.XasObjectViewer import MapViewer


class E0ComputationMethod(Enum):
    """Methods to compute E0"""

    MEAN = "mean"
    MEDIAN = "median"


class _ComputationScale(Enum):
    ALL = "all"
    CURRENT = "current"


class E0CalculatorDialog(qt.QDialog):
    """Dialog to compute E0 mean or median in a ROI"""

    def __init__(self, parent=None, xas_obj=None):
        qt.QDialog.__init__(self, parent=parent)
        if xas_obj is None:
            raise ValueError("xas_obj should not be None")

        self._methodToUse = None
        # set when the user quit the dialog and define to use mean or median

        # add the E0Calculator
        self._mainWidget = E0Calculator(parent=self, xas_obj=xas_obj)
        self.setLayout(qt.QVBoxLayout())
        self.layout().addWidget(self._mainWidget)

        # add the buttons
        style = qt.QApplication.instance().style()
        icon = style.standardIcon(qt.QStyle.SP_DialogApplyButton)

        _buttons = qt.QDialogButtonBox(parent=self)
        self._useMedian = qt.QPushButton(icon, "use median", self)
        _buttons.addButton(self._useMedian, qt.QDialogButtonBox.ActionRole)
        self._useMean = qt.QPushButton(icon, "use mean", self)
        _buttons.addButton(self._useMean, qt.QDialogButtonBox.ActionRole)
        _buttons.addButton(qt.QDialogButtonBox.Cancel)
        self.layout().addWidget(_buttons)

        # connect signal / slot
        _buttons.rejected.connect(self.reject)
        self._useMean.released.connect(
            functools.partial(self.validateMethodToUse, E0ComputationMethod.MEAN)
        )
        self._useMedian.released.connect(
            functools.partial(self.validateMethodToUse, E0ComputationMethod.MEDIAN)
        )

    def validateMethodToUse(self, method):
        """Define the method to use and close the dialog"""
        method = E0ComputationMethod.from_value(method)
        assert method in (None, E0ComputationMethod.MEDIAN, E0ComputationMethod.MEAN)
        self._methodToUse = method
        self.accept()

    def getE0(self):
        if self._methodToUse is None:
            return None
        else:
            return self._mainWidget.getE0(method=self._methodToUse)


class E0Calculator(qt.QMainWindow):
    """Interface to compute E0 mean and median values from a ROI"""

    sigE0Changed = qt.Signal(float)
    """Signal emitted when E0 is changing"""

    def __init__(self, parent=None, xas_obj=None):
        assert isinstance(xas_obj, (XASObject, type(None)))
        self._xas_obj = None

        qt.QMainWindow.__init__(self, parent=parent)
        # map view
        self._muView = MapViewer(parent=self, keys=("Mu",))
        self._muView.menuBar().hide()
        self._muView.keySelectionDocker.hide()
        # for now we don't wan't to manage the roi depending on the dimension.
        self._muView.setPerspectiveVisible(False)
        self.setCentralWidget(self._muView)
        self.setWindowFlags(qt.Qt.Widget)

        # add a ROI
        self._roiManager = RegionOfInterestManager(self._muView.getPlot())
        self._roiManager.setColor("red")  # Set the color of ROI

        self._roi = RectangleROI()
        self._updateRoi(xas_obj=xas_obj)
        self._roi.setName("ROI")
        self._roi.setEditable(True)
        self._roiManager.addRoi(self._roi)

        # dock widget to compute E0
        self._dockWidget = qt.QDockWidget(parent=self)
        self._e0CalculationWidget = _E0CalculationWidget(parent=self._dockWidget)
        self._dockWidget.setWidget(self._e0CalculationWidget)
        self.addDockWidget(qt.Qt.RightDockWidgetArea, self._dockWidget)
        self._dockWidget.setAllowedAreas(
            qt.Qt.RightDockWidgetArea | qt.Qt.LeftDockWidgetArea
        )
        self._dockWidget.setFeatures(qt.QDockWidget.NoDockWidgetFeatures)

        # expose scale
        self.getComputationScale = self._e0CalculationWidget.getComputationScale
        self.getE0 = self._e0CalculationWidget.getE0

        # connect Signal / Slot
        self._roi.sigRegionChanged.connect(functools.partial(self._updateE0, "ROI"))
        self._muView.sigFrameChanged.connect(functools.partial(self._updateE0, "frame"))
        self._e0CalculationWidget.sigComputationScaleChanged.connect(
            functools.partial(self._updateE0, "scale")
        )

        # set up
        if xas_obj is not None:
            self.setXasObject(xas_obj=xas_obj)

    def setXasObject(self, xas_obj, update_roi=True):
        """
        set XasObject to the E0 calculator

        :param xas_obj: dataset
        :param bool update_roi: if True, fir the ROI to the XAS map
        """
        self._xas_obj = xas_obj
        self._muView.setXasObject(xas_obj=self._xas_obj)
        if update_roi:
            self._updateRoi(self._xas_obj)
        # compute E0 for the current roi
        self._updateE0(origin="init")

    def getXasObject(self):
        return self._xas_obj

    def _updateRoi(self, xas_obj):
        if xas_obj is not None:
            dim_1, dim_2 = xas_obj.dim1, xas_obj.dim2
        else:
            dim_1, dim_2 = 100, 100
        self._roi.setGeometry(origin=(0, 0), size=(dim_2, dim_1))

    def _updateE0(self, origin):
        """

        :param origin: what is the reason of updating E0:
                       * `init`: construction of the widget, we want to run the
                         the first computation
                       * `frame`: the current frame change
                       * `scale`: the scale to compute E0 changed
                       * `ROI`: the region of interest changed
        :type: str
        :return:
        """
        assert origin in ("ROI", "frame", "scale", "init")
        if origin == "frame" and self.getComputationScale() == _ComputationScale.ALL:
            return
        if self._xas_obj is None:
            return

        if self.getComputationScale() == _ComputationScale.ALL:
            data = self._xas_obj.spectra.map_to("mu")
        else:
            data = self._muView.getActiveImage(just_legend=False).getData()

        e0 = self._compute0(roi=self._roi, data=data)
        self._e0CalculationWidget.setE0(e0)

    def _compute0(self, roi, data):
        if self._xas_obj is None:
            return
        roi_origin = roi.getOrigin()
        roi_size = roi.getSize()
        ymin, ymax = int(roi_origin[1]), int(roi_origin[1] + roi_size[1])
        xmin, xmax = int(roi_origin[0]), int(roi_origin[0] + roi_size[0])
        # clip roi
        ymin = max(ymin, 0)
        ymax = min(ymax, self._xas_obj.dim1)
        xmin = max(xmin, 0)
        xmax = min(xmax, self._xas_obj.dim2)

        if data.ndim == 3:
            # TODO: take into account the dimension displayed
            roi_data = data[:, ymin:ymax, xmin:xmax]
        elif data.ndim:
            roi_data = data[ymin:ymax, xmin:xmax]
        else:
            raise ValueError("given data dim %s is not managed" % data.ndim)

        return {
            E0ComputationMethod.MEAN: numpy.mean(roi_data),
            E0ComputationMethod.MEDIAN: numpy.median(roi_data),
        }


class _E0CalculationWidget(qt.QGroupBox):
    """Contains all the options to compute E0 and display the mean and the
    median values"""

    sigComputationScaleChanged = qt.Signal(str)
    """Signal emitted when the computation scale change"""

    def __init__(self, parent):
        qt.QGroupBox.__init__(self, parent=parent)
        self.setTitle("E0")
        self.setLayout(qt.QGridLayout())
        # median
        self.layout().addWidget(qt.QLabel("median", self), 0, 0)
        self._medianLE = qt.QLineEdit("0", parent=self)
        self._medianLE.setReadOnly(True)
        self.layout().addWidget(self._medianLE, 0, 1)
        # mean
        self.layout().addWidget(qt.QLabel("mean", self), 1, 0)
        self._meanLE = qt.QLineEdit("0", parent=self)
        self._meanLE.setReadOnly(True)
        self.layout().addWidget(self._meanLE, 1, 1)
        # calculation scale option
        self._computationScale = _E0CalculationScale(parent=self)
        self.layout().addWidget(self._computationScale, 2, 0, 1, 2)
        # spacer
        self._spacer = qt.QWidget(parent=self)
        self._spacer.setSizePolicy(qt.QSizePolicy.Minimum, qt.QSizePolicy.Expanding)
        self.layout().addWidget(self._spacer, 3, 0)

        # expose API
        self.getComputationScale = self._computationScale.getComputationScale

        # connect signal / slot
        self._computationScale.sigComputationScaleChanged.connect(
            self._computationScaleChanged
        )

    def getE0(self, method):
        """
        Return the current E0 value from the requested computation method

        :param method: method of compuation to get E0
        :type: Union[str,ComputationMethod]
        :return: E0
        :rtype: float
        """
        method = E0ComputationMethod.from_value(method)
        if method is E0ComputationMethod.MEAN:
            return float(self._meanLE.text())
        elif method is E0ComputationMethod.MEDIAN:
            return float(self._medianLE.text())
        else:
            raise ValueError("Given method %s is not managed" % method)

    def setE0(self, values):
        """

        :param values: dictionary with the different ComputationMethod as keys
                       and the value computed as values
        :type: dict
        :return:
        """
        assert E0ComputationMethod.MEDIAN in values
        assert E0ComputationMethod.MEAN in values
        self._meanLE.setText("%.3f" % values[E0ComputationMethod.MEAN])
        self._medianLE.setText("%.3f" % values[E0ComputationMethod.MEDIAN])

    def _computationScaleChanged(self, scale):
        self.sigComputationScaleChanged.emit(scale)


class _E0CalculationScale(qt.QGroupBox):
    """Interface to allow the user to select a the scale (from one frame or all
    frames) to compute E0"""

    sigComputationScaleChanged = qt.Signal(str)
    """Signal emitted when the computation scale change"""

    def __init__(self, parent):
        qt.QGroupBox.__init__(self, parent)
        self.__buttonGrp = qt.QButtonGroup(parent=self)
        self.setLayout(qt.QVBoxLayout())
        self.setTitle("compute E0 from:")
        self._frmCurrentFrameRB = qt.QRadioButton("from current frame", parent=self)
        self.layout().addWidget(self._frmCurrentFrameRB)
        self.__buttonGrp.addButton(self._frmCurrentFrameRB)
        self._frmAllMapsRB = qt.QRadioButton("from all frames", parent=self)
        self.layout().addWidget(self._frmAllMapsRB)
        self.__buttonGrp.addButton(self._frmAllMapsRB)

        # set up
        self.__buttonGrp.setExclusive(True)
        self._frmCurrentFrameRB.setChecked(True)

        # signal / slot connection
        self.__buttonGrp.buttonReleased.connect(self._statusChanged)

    def _statusChanged(self, *arg, **kwargs):
        self.sigComputationScaleChanged.emit(self.getComputationScale().value)

    def getComputationScale(self):
        if self._frmAllMapsRB.isChecked():
            return _ComputationScale.ALL
        elif self._frmCurrentFrameRB.isChecked():
            return _ComputationScale.CURRENT
        else:
            raise ValueError("At least all frame or current frame should be active")
