# -*- coding: utf-8 -*-
# Copyright 2014-2016 The HyperSpyUI developers
#
# This file is part of HyperSpyUI.
#
# HyperSpyUI is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# HyperSpyUI 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with HyperSpyUI.  If not, see <http://www.gnu.org/licenses/>.
"""
Created on Wed Jan 21 14:49:08 2015

@author: Vidar Tonaas Fauske
"""


from python_qt_binding import QtGui, QtCore
from QtCore import *
from QtGui import *

from hyperspyui.plugins.plugin import Plugin
from hyperspyui.widgets.extendedqwidgets import ExToolWindow
from scipy.ndimage import rotate
import numpy as np
from hyperspyui.util import win2sig, dummy_context_manager
from hyperspyui.signalwrapper import SignalWrapper


def tr(text):
    return QCoreApplication.translate("ImageRotation", text)


class ImageRotation_Plugin(Plugin):
    name = 'Image Rotation'

    def create_actions(self):
        self.settings.set_enum_hint('new_or_replace', ['new', 'replace'])
        self.add_action('rotate', "Rotate", self.show_rotate_dialog,
                        icon='rotate.svg',
                        tip="Rotate an image",
                        selection_callback=self.selection_rules)

    def create_menu(self):
        self.add_menuitem('Image', self.ui.actions['rotate'])

    def create_toolbars(self):
        self.add_toolbar_button("Image", self.ui.actions['rotate'])

    def selection_rules(self, win, action):
        """
        Callback to determine if rotation is valid for the passed window.
        """
        s = win2sig(win, self.ui.signals)
        ok = False
        if s is not None:
            if win == s.navigator_plot and \
                    s.signal.axes_manager.navigation_dimension >= 2:
                ok = True
            elif win == s.signal_plot and \
                    s.signal.axes_manager.signal_dimension == 2:
                ok = True
        if ok:
            action.setEnabled(True)
        else:
            action.setEnabled(False)

    def rotate_signal(self, angle, signal=None, reshape=False, out=None,
                      record=True, *args, **kwargs):
        if signal is None:
            signal, axes, _ = self.ui.get_selected_plot()
            if isinstance(axes, str):
                axm = signal.signal.axes_manager
                if axes.startswith("nav"):
                    axes = (axm._axes.index(axm.navigation_axes[0]),
                            axm._axes.index(axm.navigation_axes[1]))
                elif axes.startswith("sig"):
                    axes = (axm._axes.index(axm.signal_axes[0]),
                            axm._axes.index(axm.signal_axes[1]))
            kwargs['axes'] = axes
        if isinstance(signal, SignalWrapper):
            signal = signal.signal
        if record:
            self.record_code(
                r"<p>.rotate_signal({0}, reshape={1}, {2}, {3})".format(
                    angle, reshape, args, kwargs))
        # Reframe into 0-360 deg
        angle %= 360
        const_mode = 'mode' not in kwargs or kwargs['mode'] == 'constant'
        # Check if multiple of 90 deg -> Fast rotation
        if round(angle % 90, 2) == 0 and (reshape or const_mode):
            angle = 90 * (angle // 90)
            axes = kwargs.pop('axes', (1, 0))
            k = angle // 90
            if k > 0:
                if k == 1:  # 90 deg
                    # Invert axes[0]
                    s = (slice(None),)* axes[0] + (slice(None, None, -1),) + \
                        (Ellipsis,)
                elif k == 2:  # 180 deg
                    # Invert both axes
                    axisA = min(axes[0], axes[1])
                    axisB = max(axes[0], axes[1])
                    adiff = (axisB-axisA-1)
                    s = (slice(None),)* axisA + (slice(None, None, -1),) + \
                        (slice(None),)* adiff + (slice(None, None, -1),) + \
                        (Ellipsis,)
                elif k == 3:  # 270 deg
                    # Invert axes[1]
                    s = (slice(None),)* axes[1] + (slice(None, None, -1),) + \
                        (Ellipsis,)
                data = signal.data[s]
            else:
                data = signal.data
            if k % 2 == 1:     # Rotating 90 or 270
                # TODO: Swap HyperSpy axes as well!
                data = np.swapaxes(data, axes[0], axes[1])
                if not reshape:
                    # By checking for constant mode before, padding is easy
                    cval = kwargs.pop('cval', 0.0)
                    bkgr = np.ones_like(signal.data) * cval
                    fa, la = min(axes), max(axes)
                    fd, ld = np.shape(data)[fa], np.shape(data)[la]
                    crop = np.abs(ld - fd) // 2
                    if fd < ld:  # Expand first, crop last
                        s = (slice(None),) * fa + (slice(0, fd),) + \
                            (slice(None),) * (la - fa - 1) + \
                            (slice(crop, crop + fd),) + (Ellipsis,)
                        t = (slice(None),) * fa + (slice(crop, crop + fd),) + \
                            (Ellipsis,)
                    else:       # Crop first, expand last
                        s = (slice(None),) * fa + (slice(crop, crop + ld),) + \
                            (slice(None),) * (la - fa - 1) + \
                            (slice(0, ld),) + (Ellipsis,)
                        t = (slice(None),) * la + (slice(crop, crop + ld),) + \
                            (Ellipsis,)
                    bkgr[t] = data[s]
                    data = bkgr
        else:
            data = rotate(signal.data, angle, reshape=reshape, *args, **kwargs)
        if out is None:
            sig = signal._deepcopy_with_new_data(None)
        else:
            sig = out
            old_shape = out.data.shape
        sig.data = data
        if out is None:
            if reshape:
                sig.get_dimensions_from_data()
                diff = np.array(sig.data.shape) - signal.data.shape
                for i, ax in enumerate(sig.axes_manager._axes):
                    dx = diff[i] * 0.5 * ax.scale
                    ax.offset -= dx
            return sig
        else:
            if sig.data.shape != old_shape:
                if hasattr(out.axes_manager, 'events') and hasattr(
                        out.axes_manager.events, 'any_axis_changed'):
                    cm = out.axes_manager.events.any_axis_changed.suppress
                else:
                    cm = dummy_context_manager
                with cm():
                    out.get_dimensions_from_data()
                    diff = np.array(out.data.shape) - old_shape
                    for i, ax in enumerate(out.axes_manager._axes):
                        dx = diff[i] * 0.5 * ax.scale
                        ax.offset -= dx
                # TODO: TAG: Functionality check
                if hasattr(out.axes_manager, 'events') and hasattr(
                        out.axes_manager.events, 'any_axis_changed'):
                    out.axes_manager.events.any_axis_changed.trigger(
                        out.axes_manager)
            # TODO: TAG: Functionality check
            if hasattr(out, 'events') and hasattr(out.events, 'data_changed'):
                out.events.data_changed.trigger(out)

    def on_dialog_accept(self):
        self.settings['angle'] = self.dialog.num_angle.value()
        if self.dialog.opt_new.isChecked():
            self.settings['new_or_replace'] = 'new'
        else:
            self.settings['new_or_replace'] = 'replace'
        self.settings['reshape'] = self.dialog.chk_reshape.isChecked()
        self.settings['preview'] = self.dialog.gbo_preview.isChecked()
        self.settings['grid'] = self.dialog.chk_grid.isChecked()
        self.settings['grid_spacing'] = self.dialog.num_grid.value()
        self.dialog.deleteLater()
        self.dialog = None

    def show_rotate_dialog(self):
        signal, space, _ = self.ui.get_selected_plot()
        if space not in ("navigation", "signal"):
            return
        self.dialog = ImageRotationDialog(signal, space, self.ui, self)
        if 'angle' in self.settings:
            self.dialog.num_angle.setValue(self.settings['angle', float])
        if 'new_or_replace' in self.settings:
            v = self.settings['new_or_replace']
            if v == 'new':
                self.dialog.opt_new.setChecked(True)
            elif v == 'replace':
                self.dialog.opt_replace.setChecked(True)
        if 'reshape' in self.settings:
            self.dialog.chk_reshape.setChecked(self.settings['reshape', bool])
        if 'preview' in self.settings:
            self.dialog.gbo_preview.setChecked(self.settings['preview', bool])
        if 'grid' in self.settings:
            self.dialog.chk_grid.setChecked(self.settings['grid', bool])
        if 'grid_spacing' in self.settings:
            self.dialog.num_grid.setValue(self.settings['grid_spacing', int])
        self.dialog.accepted.connect(self.on_dialog_accept)
        self.dialog.show()


class ImageRotationDialog(ExToolWindow):

    def __init__(self, signal, axes, parent, plugin):
        super(ImageRotationDialog, self).__init__(parent)
        self.ui = parent
        self.create_controls()
        self.accepted.connect(self.ok)
        self.rejected.connect(self.close_new)
        self.signal = signal
        self.plugin = plugin
        self.new_out = None
        self._connected_updates = False
        if isinstance(axes, str):
            axm = signal.signal.axes_manager
            if axes.startswith("nav"):
                axes = (axm._axes.index(axm.navigation_axes[0]),
                        axm._axes.index(axm.navigation_axes[1]))
            elif axes.startswith("sig"):
                axes = (axm._axes.index(axm.signal_axes[0]),
                        axm._axes.index(axm.signal_axes[1]))
        self.axes = axes
        self.setWindowTitle(tr("Rotate"))

        # TODO: TAG: Functionality check
        if not hasattr(signal.signal, 'events'):
            self.gbo_preview.setVisible(False)

        # TODO: Add dynamic rotation, e.g. one that rotates when source
        # signal's data_changed event triggers

    def connect(self):
        # TODO: Don't have to con/dis those in gbo
        self.opt_new.toggled.connect(self.close_new)
        self.num_angle.valueChanged.connect(self.update)
        self.chk_grid.toggled.connect(self.update)
        self.num_grid.valueChanged.connect(self.update)
        self.chk_reshape.toggled.connect(self.update)
        self.opt_new.toggled.connect(self.update)
        self.opt_replace.toggled.connect(self.update)

    def disconnect(self):
        self.num_angle.valueChanged.disconnect(self.update)
        self.chk_grid.toggled.disconnect(self.update)
        self.num_grid.valueChanged.disconnect(self.update)
        self.chk_reshape.toggled.disconnect(self.update)
        self.opt_new.toggled.disconnect(self.update)
        self.opt_replace.toggled.disconnect(self.update)

    def ok(self):
        # Draw figure if not already done
        # TODO: TAG: Functionality check
        if not hasattr(self.signal.signal, 'events') or \
                not self.gbo_preview.isChecked():
            self.update()
        angle = self.num_angle.value()
        reshape = self.chk_reshape.isChecked()
        self.plugin.record_code(
            r"<p>.rotate_signal({0}, reshape={1}, axes={2})".format(
                angle, reshape, self.axes))
        # Clean up event connections
        if self.new_out is not None:
            self.connect_update_plot(self.new_out.signal, disconnect=True)

    def close_new(self, value=False):
        if self.new_out is not None and not value:
            self.new_out.close()
            self.new_out = None
            self._connected_updates = False

    def set_preview(self, value):
        if not hasattr(self.signal.signal, 'events'):
            return
        if value:
            self.connect()
            self.update()
        else:
            self.disconnect()
            self.close_new()

    def _axes_in_nav(self):
        axm = self.signal.signal.axes_manager
        navidx = [axm._axes.index(ax) for ax in axm.navigation_axes]
        if self.axes[0] in navidx:
            return True
        return False

    def connect_update_plot(self, signal, disconnect=False):
        if self._connected_updates != disconnect:
            return  # Nothing to do, prevent double connections
        if self._axes_in_nav():
            f = signal._plot.navigator_plot.update
        else:
            f = signal._plot.signal_plot.update

        # TODO: TAG: Functionality check
        if hasattr(signal, 'events') and hasattr(
                signal.events, 'data_changed'):
            if disconnect:
                signal.events.data_changed.disconnect(f)
            else:
                signal.events.data_changed.connect(f, [])
        self._connected_updates = not disconnect

    def update(self):
        angle = self.num_angle.value()
        reshape = self.chk_reshape.isChecked()
        if self.opt_new.isChecked():
            if self.new_out is None:
                out = None
            else:
                out = self.new_out.signal
        elif self.opt_replace.isChecked():
            out = self.signal.signal
        else:
            return  # Indeterminate state, do nothing

        s = self.plugin.rotate_signal(angle, self.signal.signal, record=False,
                                      reshape=reshape, out=out, axes=self.axes)

        if out is None:
            s.metadata.General.title = self.signal.name + "[Rotated]"
            s.plot()
            self.connect_update_plot(s)
            if (self.gbo_preview.isChecked() and self.opt_new.isChecked() and
                                                         self.new_out is None):
                self.new_out = self.ui.lut_signalwrapper[s]
        else:
            s = out

        if self.chk_grid.isChecked() is True:
            pass    # TODO: Draw grid

    def create_controls(self):
        """
        Create UI controls.
        """
        vbox = QVBoxLayout()

        form = QFormLayout()
        self.num_angle = QDoubleSpinBox()
        self.num_angle.setValue(0.0)
        self.num_angle.setMinimum(-360)
        self.num_angle.setMaximum(360)
        form.addRow(tr("Angle:"), self.num_angle)
        vbox.addLayout(form)

        self.gbo_preview = QGroupBox(tr("Preview"))
        self.gbo_preview.setCheckable(True)
        self.gbo_preview.setChecked(False)
        gbo_vbox = QVBoxLayout()
        self.chk_grid = QCheckBox(tr("Grid"))
        self.chk_grid.setChecked(False)
        self.num_grid = QSpinBox()
        self.num_grid.setValue(4)
        self.num_grid.setMinimum(1)
        self.num_grid.setEnabled(False)
        self.chk_grid.toggled[bool].connect(self.num_grid.setEnabled)
        gbo_vbox.addWidget(self.chk_grid)
        gbo_vbox.addWidget(self.num_grid)
        self.gbo_preview.setLayout(gbo_vbox)
        vbox.addWidget(self.gbo_preview)

        self.gbo_preview.toggled[bool].connect(self.set_preview)

        self.gbo_output = QGroupBox(tr("Output"))
        self.opt_new = QRadioButton(tr("New signal"))
        self.opt_replace = QRadioButton(tr("In place"))
        self.opt_new.setChecked(True)
        gbo_vbox2 = QVBoxLayout()
        gbo_vbox2.addWidget(self.opt_new)
        gbo_vbox2.addWidget(self.opt_replace)
        self.gbo_output.setLayout(gbo_vbox2)
        vbox.addWidget(self.gbo_output)

        self.chk_reshape = QCheckBox(tr("Resize to fit"))
        self.chk_reshape.setChecked(False)
        vbox.addWidget(self.chk_reshape)

        self.btn_ok = QPushButton(tr("&OK"))
        self.btn_ok.setDefault(True)
        self.btn_ok.clicked.connect(self.accept)
        self.btn_cancel = QPushButton(tr("&Cancel"))
        self.btn_cancel.clicked.connect(self.reject)
        hbox = QHBoxLayout()
        hbox.addWidget(self.btn_ok)
        hbox.addWidget(self.btn_cancel)
        vbox.addLayout(hbox)

        vbox.addStretch(1)
        self.setLayout(vbox)
