# coding: utf-8
# /*##########################################################################
#
# Copyright (c) 2016-2020 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.
#
# ###########################################################################*/

__authors__ = ["H. Payno"]
__license__ = "MIT"
__date__ = "12/02/2020"

import typing
import functools
from tomwer.utils import docstring
from silx.gui import qt
from tomwer.core.process.reconstruction.nabu import (_NabuStages,
                                                     _NabuReconstructionMethods,
                                                     _NabuFBPFilterType,
                                                     _NabuPaddingType,
                                                     NabuSliceMode)
from tomwer.gui.reconstruction.nabu.nabuconfig.base import _NabuStageConfigBase
from tomwer.gui.utils.inputwidget import SelectionLineEdit


import logging
_logger = logging.getLogger(__name__)


class SliceGroupBox(qt.QGroupBox):
    """GroupBox to define slice to be reconstructed"""
    def __init__(self, parent):
        qt.QGroupBox.__init__(self, 'slices', parent)
        self.setCheckable(True)

        self.setLayout(qt.QHBoxLayout())
        self._modeCB = qt.QComboBox()
        for mode in NabuSliceMode:
            self._modeCB.addItem(mode.value)
        self.layout().addWidget(self._modeCB)

        self._sliceQLE = SelectionLineEdit(parent=self)
        self.layout().addWidget(self._sliceQLE)

        spacer = qt.QWidget(parent=self)
        spacer.setSizePolicy(qt.QSizePolicy.Minimum, qt.QSizePolicy.Expanding)
        self.layout().addWidget(spacer)

        # set up
        self.setChecked(True)
        self.setMode(NabuSliceMode.MIDDLE)
        self._sliceQLE.setVisible(False)

        # connect signal / slot
        self._modeCB.currentIndexChanged.connect(self._updateSliceQLEVisibilty)

    def _updateSliceQLEVisibilty(self, *args, **kwargs):
        self._sliceQLE.setVisible(self.getMode() == NabuSliceMode.OTHER)

    def setMode(self, mode: NabuSliceMode):
        mode = NabuSliceMode.from_value(mode)
        item_index = self._modeCB.findText(mode.value)
        self._modeCB.setCurrentIndex(item_index)
        old = self.blockSignals(True)
        self._updateSliceQLEVisibilty()
        self.blockSignals(old)

    def getMode(self) -> NabuSliceMode:
        mode = NabuSliceMode.from_value(self._modeCB.currentText())
        return mode

    def _getSliceSelected(self):
        if self.getMode() is NabuSliceMode.MIDDLE:
            return NabuSliceMode.MIDDLE.value
        else:
            return self._sliceQLE.text()

    def getSlices(self):
        """Slice selected"""
        if self.isChecked():
            return self._getSliceSelected()
        else:
            return None

    def setSlices(self, slices):
        if slices is None:
            self.setChecked(False)
        else:
            self.setChecked(True)
            if slices != NabuSliceMode.MIDDLE.value:
                self._sliceQLE.setText(slices)
                self.setMode(NabuSliceMode.OTHER)
            else:
                self.setMode(NabuSliceMode.MIDDLE)


class _NabuReconstructionConfig(qt.QWidget, _NabuStageConfigBase):
    """
    Widget to define the configuration of nabu reconstruction processing
    """

    sigConfChanged = qt.Signal(str)
    """Signal emitted when the configuration change. Parameter is the option
    modified
    """
    def __init__(self, parent):
        qt.QWidget.__init__(self, parent, stage=_NabuStages.PRE)
        _NabuStageConfigBase.__init__(self, stage=_NabuStages.PRE)
        self.setLayout(qt.QGridLayout())

        self.__angle_files = ''
        self.__optimizationIteAlgo = 'chambolle-pock'

        # slices to be reconstructed online
        self._slicesWidget = SliceGroupBox(parent=self)
        self.layout().addWidget(self._slicesWidget, 0, 0, 1, 2)

        # method
        self._methodLabel = qt.QLabel('method', self)
        self.layout().addWidget(self._methodLabel, 1, 0, 1, 1)
        self._methodQCB = qt.QComboBox(parent=self)
        for method in _NabuReconstructionMethods:
            self._methodQCB.addItem(method.value)
        self.layout().addWidget(self._methodQCB, 1, 1, 1, 1)
        self.registerWidget(self._methodLabel, 'required')
        self.registerWidget(self._methodQCB, 'required')

        # angle_offset
        self._labelOffsetLabel = qt.QLabel('angle offset (in degree)', self)
        self.layout().addWidget(self._labelOffsetLabel, 2, 0, 1, 1)
        self._angleOffsetQDSB = qt.QDoubleSpinBox(self)
        self._angleOffsetQDSB.setMaximum(-180)
        self._angleOffsetQDSB.setMaximum(180)
        self.layout().addWidget(self._angleOffsetQDSB, 2, 1, 1, 1)
        self.registerWidget(self._labelOffsetLabel, 'advanced')
        self.registerWidget(self._angleOffsetQDSB, 'advanced')

        # fbp filter type
        self._fbpFilterCB = qt.QCheckBox('fbp filter', self)
        self.layout().addWidget(self._fbpFilterCB, 3, 0, 1, 1)
        self._fbpFilterType = qt.QComboBox(self)
        for filter_type in _NabuFBPFilterType:
            self._fbpFilterType.addItem(filter_type.value)
        self.layout().addWidget(self._fbpFilterType, 3, 1, 1, 1)
        self.registerWidget(self._fbpFilterCB, 'advanced')
        self.registerWidget(self._fbpFilterType, 'advanced')

        # padding type
        self._paddingTypeLabel = qt.QLabel('padding type', self)
        self.layout().addWidget(self._paddingTypeLabel, 4, 0, 1, 1)
        self._paddingType = qt.QComboBox(self)
        for fbp_padding_type in _NabuPaddingType:
            self._paddingType.addItem(fbp_padding_type.value)
        self.layout().addWidget(self._paddingType, 4, 1, 1, 1)
        self.registerWidget(self._paddingTypeLabel, 'optional')
        self.registerWidget(self._paddingType, 'optional')

        # half tomo
        self._halfTomoCB = qt.QCheckBox('half tomo')
        self.layout().addWidget(self._halfTomoCB, 5, 0, 1, 2)
        self.registerWidget(self._halfTomoCB, 'optional')

        # sub region
        self._subRegionSelector = _NabuReconstructionSubRegion(parent=self)
        self.layout().addWidget(self._subRegionSelector, 6, 0, 1, 2)

        # iterations
        self._iterationsLabel = qt.QLabel('iterations', self)
        self.layout().addWidget(self._iterationsLabel, 7, 0, 1, 1)
        self._iterationSB = qt.QSpinBox(parent=self)
        self.layout().addWidget(self._iterationSB, 7, 1, 1, 1)
        self._iterationSB.setMinimum(1)
        self._iterationSB.setMaximum(9999)
        # not supported for now so hidden
        self._iterationsLabel.hide()
        self._iterationSB.hide()

        # spacer for style
        spacer = qt.QWidget(self)
        spacer.setSizePolicy(qt.QSizePolicy.Minimum, qt.QSizePolicy.Expanding)
        self.layout().addWidget(spacer, 200, 1, 1, 1)

        # optimization algorithm:
        # set has default value for now, because has only one at the moment

        # weight total variation
        self._tvLabel = qt.QLabel('total variation weight', self)
        self.layout().addWidget(self._tvLabel, 8, 0, 1, 1)
        self._totalVariationWeight = qt.QDoubleSpinBox(self)
        self._totalVariationWeight.setMinimum(0.0)
        self._totalVariationWeight.setMaximum(1.0)
        self._totalVariationWeight.setDecimals(4)
        self._totalVariationWeight.setSingleStep(0.002)
        self.layout().addWidget(self._totalVariationWeight, 8, 1, 1, 1)
        # not supported for now so hidden
        self._tvLabel.hide()
        self._totalVariationWeight.hide()

        # preconditioning filter
        self._preconditioningFilter = qt.QCheckBox('preconditioning_filter', self)
        self._preconditioningFilter.setToolTip('Whether to enable "filter '
                                               'preconditioning" for iterative'
                                               ' methods')
        self.layout().addWidget(self._preconditioningFilter, 9, 0, 1, 2)
        # not supported for now so hidden
        self._preconditioningFilter.hide()

        # positivity constraint
        self._positivityConstraintCB = qt.QCheckBox('positivity constraint', self)
        self._positivityConstraintCB.setToolTip('Whether to enforce a '
                                                'positivity constraint in the '
                                                'reconstruction.')
        self.layout().addWidget(self._positivityConstraintCB, 10, 0, 1, 2)
        # not supported for now so hidden
        self._positivityConstraintCB.hide()

        # expose API
        self.getSlices = self._slicesWidget.getSlices
        self.setSlices = self._slicesWidget.setSlices

        # set up
        self._fbpFilterCB.setChecked(True)
        fbp_item = self._methodQCB.findText(_NabuReconstructionMethods.FBP.value)
        self._methodQCB.setCurrentIndex(fbp_item)
        self._angleOffsetQDSB.setValue(0)
        ramlak_item = self._fbpFilterType.findText(_NabuFBPFilterType.RAMLAK.value)
        self._fbpFilterType.setCurrentIndex(ramlak_item)
        padding_type_item = self._paddingType.findText(_NabuPaddingType.ZEROS.value)
        self._fbpFilterType.setCurrentIndex(padding_type_item)
        self._halfTomoCB.setChecked(False)
        self._iterationSB.setValue(200)
        self._totalVariationWeight.setValue(1.0e-2)
        self._preconditioningFilter.setChecked(True)
        self._positivityConstraintCB.setChecked(True)

        # connect signal / slot
        self._methodQCB.currentIndexChanged.connect(self._methodChanged)
        self._angleOffsetQDSB.editingFinished.connect(self._angleOffsetChanged)
        self._fbpFilterCB.toggled.connect(self._FBPFilterTypeChanged)
        self._fbpFilterType.currentIndexChanged.connect(self._FBPFilterTypeChanged)
        self._paddingType.currentTextChanged.connect(self._paddingTypeChanged)
        self._halfTomoCB.toggled.connect(self._halfTomoChanged)
        self._subRegionSelector.sigConfChanged.connect(self._signalConfChanged)
        self._iterationSB.valueChanged.connect(self._nbIterationChanged)
        self._totalVariationWeight.valueChanged.connect(self._weightTvChanged)
        self._preconditioningFilter.toggled.connect(self._preconditionningFilterChanged)
        self._positivityConstraintCB.toggled.connect(self._positivityConstraintChanged)

    def setConfigurationLevel(self, level):
        _NabuStageConfigBase.setConfigurationLevel(self, level)
        self._subRegionSelector.setConfigurationLevel(level=level)

    def _methodChanged(self, *args, **kwargs):
        self._signalConfChanged('method')

    def _angleOffsetChanged(self, *args, **kwargs):
        self._signalConfChanged('angle_offset')

    def _FBPFilterTypeChanged(self, *args, **kwargs):
        self._signalConfChanged('fbp_filter_type')

    def _paddingTypeChanged(self, *args, **kwargs):
        self._signalConfChanged('padding_type')

    def _halfTomoChanged(self, *args, **kwargs):
        self._signalConfChanged('enable_halftomo')

    def _nbIterationChanged(self, *args, **kwargs):
        self._signalConfChanged('iterations')

    def _weightTvChanged(self, *args, **kwargs):
        self._signalConfChanged('weight_tv')

    def _preconditionningFilterChanged(self, *args, **kwargs):
        self._signalConfChanged('preconditioning_filter')

    def _positivityConstraintChanged(self, *args, **kwargs):
        self._signalConfChanged('positivity_constraint')

    def setScan(self, scan):
        raise NotImplementedError()

    def getMethod(self) -> _NabuReconstructionMethods:
        return _NabuReconstructionMethods.from_value(
            self._methodQCB.currentText())

    def setMethod(self, method):
        method = _NabuReconstructionMethods.from_value(method)
        item_index = self._methodQCB.findText(method.value)
        self._methodQCB.setCurrentIndex(item_index)

    def getAngleOffset(self) -> float:
        return self._angleOffsetQDSB.value()

    def setAngleOffset(self, value: typing.Union[str, float]):
        self._angleOffsetQDSB.setValue(float(value))

    def getFBPFilterType(self) -> typing.Union[_NabuFBPFilterType, None]:
        if self._fbpFilterCB.isChecked():
            return _NabuFBPFilterType.from_value(self._fbpFilterType.currentText())
        else:
            return None

    def setFBPFilterType(self, filter):
        if type(filter) is str and filter.lower() == 'none':
            filter = None
        if filter is None:
            self._fbpFilterCB.setChecked(False)
        else:
            self._fbpFilterCB.setChecked(True)
            filter = _NabuFBPFilterType(filter)
            filter_index = self._fbpFilterType.findText(filter.value)
            self._fbpFilterType.setCurrentIndex(filter_index)

    def getFBPPaddingType(self) -> _NabuPaddingType:
        return _NabuPaddingType.from_value(self._paddingType.currentText())

    def setFBPPaddingType(self, padding):
        padding = _NabuPaddingType.from_value(padding)
        padding_index = self._paddingType.findText(padding.value)
        self._paddingType.setCurrentIndex(padding_index)

    def setHalfTomoEnabled(self, enable):
        self._halfTomoCB.setChecked(enable)

    def isHaflTomoEnabled(self):
        return self._halfTomoCB.isChecked()

    def getNIterations(self):
        return self._iterationSB.value()

    def setNIterations(self, n_iterations):
        self._iterationSB.setValue(n_iterations)

    def getTotalVariationWeight(self) -> float:
        return self._totalVariationWeight.value()

    def setTotalVariationWeight(self, weight: float):
        self._totalVariationWeight.setValue(weight)

    def isPreconditioningFilterEnable(self):
        return self._preconditioningFilter.isChecked()

    def setPreconditioningFilterEnable(self, enable: typing.Union[bool, int]):
        self._preconditioningFilter.setChecked(bool(enable))

    def isPositivityConstraintEnable(self):
        return self._positivityConstraintCB.isChecked()

    def setPositivityConstraintEnable(self, enable: typing.Union[bool, int]):
        self._positivityConstraintCB.setChecked(bool(enable))

    @docstring(_NabuStageConfigBase)
    def getConfiguration(self) -> dict:
        fbp_filter_type = self.getFBPFilterType()
        if fbp_filter_type is None:
            fbp_filter_type = 'none'
        else:
            fbp_filter_type = fbp_filter_type.value
        config = {
            'method': self.getMethod().value,
            'angles_file': self.__angle_files,
            # 'rotation_axis_position'  # if setted should be done upstream because managed by 'tomwer' axis
            'axis_correction_file': '',  # not managed for now
            'angle_offset': self.getAngleOffset(),
            'fbp_filter_type': fbp_filter_type,
            'padding_type': self.getFBPPaddingType().value,
            'enable_halftomo': int(self.isHaflTomoEnabled()),
            'iterations': self.getNIterations(),
            'optim_algorithm': self.__optimizationIteAlgo,
            'weight_tv': self.getTotalVariationWeight(),
            'preconditioning_filter': int(self.isPreconditioningFilterEnable()),
            'positivity_constraint': int(self.isPositivityConstraintEnable()),
            'rotation_axis_position': '',
        }
        config.update(self._subRegionSelector.getConfiguration())
        return config

    @docstring(_NabuStageConfigBase)
    def setConfiguration(self, config):
        if 'method' in config:
            self.setMethod(config['method'])
        if 'angles_file' in config:
            self.__angle_files = config['angles_file']
        if 'angle_offset' in config:
            self.setAngleOffset(value=config['angle_offset'])
        if 'fbp_filter_type' in config:
            self.setFBPFilterType(config['fbp_filter_type'])
        if 'padding_type' in config:
            self.setFBPPaddingType(config['padding_type'])
        if 'enable_halftomo' in config:
            self.setHalfTomoEnabled(enable=bool(config['enable_halftomo']))
        if 'iterations' in config:
            self.setNIterations(int(config['iterations']))
        if 'optim_algorithm' in config:
            self.__optimizationIteAlgo = config['optim_algorithm']
        if 'weight_tv' in config:
            self.setTotalVariationWeight(weight=config['weight_tv'])
        if 'preconditioning_filter' in config:
            self.setPreconditioningFilterEnable(int(config['preconditioning_filter']))
        if 'positivity_constraint' in config:
            self.setPositivityConstraintEnable(int(config['positivity_constraint']))
        self._subRegionSelector.setConfiguration(config=config)

    def _signalConfChanged(self, param):
        self.sigConfChanged.emit(param)


class _SubRegionEditor(qt.QObject):

    sigConfChanged = qt.Signal(str)
    """Signal emitted each type a parameter is edited"""

    def __init__(self, parent, layout, layout_row: int, name: str, min_param: str,
                 max_param: str):
        assert type(layout_row) is int
        qt.QObject.__init__(self)
        self._layout = layout
        self.__minParam = min_param
        self.__maxParam = max_param

        validator = qt.QIntValidator()
        validator.setBottom(0)

        self._subRegionLabel = qt.QLabel(name, parent)
        self.layout().addWidget(self._subRegionLabel, layout_row, 0, 1, 1)

        # min
        self._minCB = qt.QCheckBox('min', parent)
        self.layout().addWidget(self._minCB, layout_row, 1, 1, 1)
        self._minQLE = qt.QLineEdit('0', parent)
        self._minQLE.setValidator(validator)
        self.layout().addWidget(self._minQLE, layout_row, 2, 1, 1)

        # max
        self._maxCB = qt.QCheckBox('max', parent)
        self.layout().addWidget(self._maxCB, layout_row, 3, 1, 1)
        self._maxQLE = qt.QLineEdit('0', parent)
        self._maxQLE.setValidator(validator)
        self.layout().addWidget(self._maxQLE, layout_row, 4, 1, 1)

        # spacer for style
        self._spacer = qt.QWidget(parent)
        self._spacer.setSizePolicy(qt.QSizePolicy.Expanding, qt.QSizePolicy.Minimum)
        self.layout().addWidget(self._spacer, layout_row, 5, 1, 1)

        # set up
        self._minCB.setChecked(False)
        self._minQLE.setEnabled(False)
        self._maxCB.setChecked(False)
        self._maxQLE.setEnabled(False)

        # connect signal / slot
        self._minCB.toggled.connect(self._minQLE.setEnabled)
        self._maxCB.toggled.connect(self._maxQLE.setEnabled)
        self._minCB.toggled.connect(functools.partial(
            self._signalConfChanged, self.__minParam))
        self._maxCB.toggled.connect(functools.partial(
            self._signalConfChanged, self.__maxParam))
        self._minQLE.editingFinished.connect(functools.partial(
            self._signalConfChanged, self.__minParam))
        self._maxQLE.editingFinished.connect(functools.partial(
            self._signalConfChanged, self.__maxParam))

    def layout(self):
        return self._layout

    def getSubRegionMin(self) -> typing.Union[None, int]:
        """

        :return: None if region is unbounded else the value of the bound
        """
        if self._minCB.isChecked():
            return int(self._minQLE.text())
        else:
            return None

    def getSubRegionMax(self) -> typing.Union[None, int]:
        """

        :return: None if region is unbounded else the value of the bound
        """
        if self._maxCB.isChecked():
            return int(self._maxQLE.text())
        else:
            return None

    def setSubRegionMin(self, min: typing.Union[None, int]) -> None:
        """

        :param min: if min is None or -1 wr will expand it to - infinity
        """
        if min in (-1, 0):
            min = None
        if min is None:
            self._minCB.setChecked(False)
        else:
            self._minCB.setChecked(True)
            self._minQLE.setText(str(min))

    def setSubRegionMax(self, max) -> None:
        """

        :param max: if max is None or -1 wr will expand it to infinity
        """
        if type(max) is str:
            max = int(max)
        if max == -1:
            max = None
        if max is None:
            self._maxCB.setChecked(False)
        else:
            self._maxCB.setChecked(True)
            self._maxQLE.setText(str(max))

    def getSubRegion(self) -> tuple:
        """

        :return: min, max
        :rtype: tuple
        """
        return self.getSubRegionMin(), self.getSubRegionMax()

    def setSubRegion(self, min: typing.Union[None, int],
                     max: typing.Union[None, int]):
        """

        :param tuple min:
        :param tuple max:
        """
        self.setSubRegionMin(min)
        self.setSubRegionMax(max)

    def _signalConfChanged(self, param):
        self.sigConfChanged.emit(param)

    def setVisible(self, visible):
        for widget in (self._subRegionLabel, self._minCB, self._minQLE,
                       self._maxCB, self._maxQLE, self._spacer):
            widget.setVisible(visible)


class _NabuReconstructionSubRegion(qt.QGroupBox, _NabuStageConfigBase):
    """Widget to select a sub region to reconstruct"""

    sigConfChanged = qt.Signal(str)
    """Signal emitted each type a parameter is edited"""

    def __init__(self, parent):
        qt.QGroupBox.__init__(self, parent, stage=_NabuStages.PROC)
        _NabuStageConfigBase.__init__(self, stage=_NabuStages.PROC)
        self.setTitle('sub region')

        self.setLayout(qt.QGridLayout())

        self._xSubRegion = _SubRegionEditor(parent=self,
                                            layout=self.layout(),
                                            layout_row=0,
                                            name='x',
                                            min_param='start_x',
                                            max_param='end_x')
        self.registerWidget(self._xSubRegion, 'optional')

        self._ySubRegion = _SubRegionEditor(parent=self,
                                            layout=self.layout(),
                                            layout_row=1,
                                            name='y',
                                            min_param='start_y',
                                            max_param='end_y')
        self.registerWidget(self._ySubRegion, 'optional')

        self._zSubRegion = _SubRegionEditor(parent=self,
                                            layout=self.layout(),
                                            layout_row=2,
                                            name='z',
                                            min_param='start_z',
                                            max_param='end_z')
        self.registerWidget(self._zSubRegion, 'optional')

        # set up

        # connect signal / slot
        self._xSubRegion.sigConfChanged.connect(self._signalConfChanged)
        self._ySubRegion.sigConfChanged.connect(self._signalConfChanged)
        self._zSubRegion.sigConfChanged.connect(self._signalConfChanged)

    def getConfiguration(self) -> dict:
        return {
            'start_x': self._xSubRegion.getSubRegionMin() or 0,
            'end_x': self._xSubRegion.getSubRegionMax() or -1,
            'start_y': self._ySubRegion.getSubRegionMin() or 0,
            'end_y': self._ySubRegion.getSubRegionMax() or -1,
            'start_z': self._zSubRegion.getSubRegionMin() or 0,
            'end_z': self._zSubRegion.getSubRegionMax() or -1,
        }

    def setConfiguration(self, config):
        if 'start_x' in config:
            self._xSubRegion.setSubRegionMin(int(config['start_x']))
        if 'end_x' in config:
            self._xSubRegion.setSubRegionMax(int(config['end_x']))
        if 'start_y' in config:
            self._ySubRegion.setSubRegionMin(int(config['start_y']))
        if 'end_y' in config:
            self._ySubRegion.setSubRegionMax(int(config['end_y']))
        if 'start_z' in config:
            self._zSubRegion.setSubRegionMin(int(config['start_z']))
        if 'end_z' in config:
            self._zSubRegion.setSubRegionMax(int(config['end_z']))

    def _signalConfChanged(self, param):
        self.sigConfChanged.emit(param)


if __name__ == '__main__':
    app = qt.QApplication([])
    widget = _NabuReconstructionConfig(None)
    widget.show()
    print(widget.getConfiguration())
    app.exec_()