content: # -*- 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 Fri Dec 12 23:44:01 2014
  
  @author: Vidar Tonaas Fauske
  """
  
  
  from hyperspyui.plugins.plugin import Plugin
  
  from qtpy import QtCore, QtWidgets
  from qtpy.QtWidgets import QDialog, QDialogButtonBox, QLineEdit, QLabel
  
  from hyperspy.learn.mva import LearningResults
  from hyperspyui.util import win2sig, fig2win, Namespace
  from hyperspyui.threaded import ProgressThreaded, ProcessCanceled
  from hyperspyui.widgets.extendedqwidgets import ExToolWindow
  
  
  def tr(text):
      return QtCore.QCoreApplication.translate("MVA", text)
  
  
  def align_yaxis(ax1, v1, ax2, v2):
      """adjust ax2 ylimit so that v2 in ax2 is aligned to v1 in ax1"""
      _, y1 = ax1.transData.transform((0, v1))
      _, y2 = ax2.transData.transform((0, v2))
      if y2 > y1:
          ratio = y1 / y2
      else:
          ratio = y2 / y1
      inv = ax2.transData.inverted()
      _, dy = inv.transform((0, 0)) - inv.transform((0, y1 - y2))
      miny2, maxy2 = ax2.get_ylim()
      ax2.set_ylim((miny2 + dy) / ratio, (maxy2 + dy) / ratio)
  
  
  def make_advanced_dialog(ui, algorithms=None):
      diag = ExToolWindow(ui)
  
      diag.setWindowTitle("Decomposition parameters")
  
      vbox = QtWidgets.QVBoxLayout()
      if algorithms:
          lbl_algo = QLabel(tr("Choose algorithm:"))
          cbo_algo = QLineEdit.QComboBox()
          cbo_algo.addItems(algorithms)
  
          vbox.addWidget(lbl_algo)
          vbox.addWidget(cbo_algo)
      else:
          lbl_comp = QLabel(tr(
              "Enter a comma-separated list of component numbers to use for "
              "the model:"))
          txt_comp = QLineEdit()
          vbox.addWidget(lbl_comp)
          vbox.addWidget(txt_comp)
  
      btns = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel,
                              QtCore.Qt.Horizontal)
      btns.accepted.connect(diag.accept)
      btns.rejected.connect(diag.reject)
      vbox.addWidget(btns)
  
      diag.setLayout(vbox)
  
      diag.algorithm = lambda: cbo_algo.currentText()
      diag.components = lambda: [int(s) for s in txt_comp.text().split(',')]
      return diag
  
  
  class MVA_Plugin(Plugin):
  
      """
      Implements MVA decomposition utilities.
      """
      name = 'MVA'    # Used for settings groups etc
  
      coc_values = {'convert': tr("Convert"),
                    'copy': tr("Copy")}
  
      # ----------- Plugin interface -----------
      def create_actions(self):
          self.settings.set_default('convert_or_copy', None)
          self.settings.set_enum_hint('convert_or_copy',
                                      self.coc_values.keys())
          self.add_action('plot_decomposition_results',
                          tr("Decompose"),
                          self.plot_decomposition_results,
                          icon='pca.svg',
                          tip=tr("Decompose signal using Principle Component "
                                 "analysis"),
                          selection_callback=self.selection_rules)
          self.add_action('pca', tr("Decomposition model"), self.pca,
                          icon='pca.svg',
                          tip=tr("Create a Principal Component Analysis "
                                 "decomposition model"),
                          selection_callback=self.selection_rules)
          self.add_action('bss', tr("BSS"), self.bss,
                          icon='bss.svg',
                          tip=tr("Run Blind Source Separation"),
                          selection_callback=self.selection_rules)
          self.add_action('bss_model', tr("BSS model"), self.bss_model,
                          icon='bss.svg',
                          tip=tr("Create a Blind Source Separation "
                                 "decomposition model"),
                          selection_callback=self.selection_rules)
          self.add_action('clear', tr("Clear"), self.clear,
                          tip=tr("Clear decomposition cache"),
                          selection_callback=self.selection_rules)
  
      def create_menu(self):
          self.add_menuitem('Decomposition',
                            self.ui.actions['plot_decomposition_results'])
          self.add_menuitem('Decomposition', self.ui.actions['pca'])
          self.add_menuitem('Decomposition', self.ui.actions['bss'])
          self.add_menuitem('Decomposition', self.ui.actions['bss_model'])
          self.add_menuitem('Decomposition', self.ui.actions['clear'])
  
      def create_toolbars(self):
          self.add_toolbar_button("Decomposition", self.ui.actions['pca'])
          self.add_toolbar_button("Decomposition", self.ui.actions['bss'])
  
      def selection_rules(self, win, action):
          """
          Callback to determine if action is valid for the passed window.
          """
          s = win2sig(win, self.ui.signals)
          if s is None or s.signal.data.ndim <= 1:
              action.setEnabled(False)
          else:
              action.setEnabled(True)
  
      # ------------ Action implementations --------------
  
      def _get_signal(self, signal):
          """
          Get a valid signal. If the signal is none, it ues the currently
          selected one. If the signal type is not float, it either converts it,
          or gets a copy of the correct type, depending on the 'convert_copy'
          setting.
          """
          if signal is None:
              signal = self.ui.get_selected_wrapper()
          s = signal.signal
  
          if s.data.dtype.char not in ['e', 'f', 'd']:  # If not float
              cc = self.settings.get_or_prompt(
                  'convert_or_copy',
                  [kv for kv in self.coc_values.items()],
                  title=tr("Convert or copy"),
                  descr=tr(
                      "Signal data has the wrong data type (float needed)." +
                      "Would you like to convert the current signal, or " +
                      "perform the decomposition on a copy?"))
              if cc is None:
                  # User canceled
                  raise ProcessCanceled()
              if cc == 'copy':
                  s = s.deepcopy()
                  s.metadata.General.title = signal.name + "[float]"
                  s.plot()
              s.change_dtype(float)
          return s, signal
  
      def _do_decomposition(self, s, force=False, algorithm=None):
          """
          Makes sure we have decomposition results. If results already are
          available, it will only recalculate if the `force` parameter is True.
          """
          if algorithm:
              s.decomposition(algorithm=algorithm)
          elif force or s.learning_results.explained_variance_ratio is None:
              s.decomposition()
          return s
  
      def _do_bss(self, s, n_components, algorithm=None):
          """
          Makes sure we have BSS results. If results already are available, it
          will only recalculate if the `force` parameter is True.
          """
          if algorithm:
              s.blind_source_separation(n_components, algorithm=algorithm)
          else:
              s.blind_source_separation(n_components)
  
      def get_bss_results(self, signal):
          factors = signal.get_bss_factors()
          loadings = signal.get_bss_loadings()
          factors.axes_manager._axes[0] = loadings.axes_manager._axes[0]
          return loadings, factors
  
      def _record(self, autosig, model, signal, n_components):
          if autosig:
              self.record_code(r"<p>.{0}(n_components={1})".format(
                               model, n_components))
          else:
              self.record_code(r"<p>.{0}({1}, n_components={2})".format(
                               model, signal, n_components))
  
      def _decompose_threaded(self, callback, label, signal=None,
                              algorithm=None, ns=None):
          if ns is None:
              ns = Namespace()
              ns.autosig = signal is None
              ns.s, signal = self._get_signal(signal)
  
          def do_threaded():
              ns.s = self._do_decomposition(ns.s, algorithm=algorithm)
  
          def on_error(message=None):
              em = QtWidgets.QErrorMessage(self.ui)
              msg = tr("An error occurred during decomposition")
              if message:
                  msg += ":\n" + message
              em.setWindowTitle(tr("Decomposition error"))
              em.showMessage(msg)
  
          t = ProgressThreaded(self.ui, do_threaded, lambda: callback(ns),
                               label=label)
          t.worker.error[str].connect(on_error)
          t.run()
  
      def _perform_model(self, ns, n_components):
          # Num comp. picked, get model, wrap new signal and plot
          if ns.model == 'pca':
              sc = ns.s.get_decomposition_model(n_components)
              sc.metadata.General.title = ns.signal.name + "[PCA-model]"
              sc.plot()
          elif ns.model == 'bss' or ns.model.startswith('bss.'):
              if ns.model.startswith('bss.'):
                  algorithm = ns.model[len('bss.'):]
                  self._do_bss(ns.s, n_components, algorithm=algorithm)
              else:
                  self._do_bss(ns.s, n_components)
              f, o = self.get_bss_results(ns.s)
              o.metadata.add_dictionary(ns.s.metadata.as_dictionary())
              f.metadata.General.title = ns.signal.name + "[BSS-Factors]"
              o.metadata.General.title = ns.signal.name + "[BSS-Loadings]"
              f.plot()
              o.plot()
          elif ns.model == 'bss_model':
              # Here we have to assume the user has actually performed the BSS
              # decomposition first!
              sc = ns.s.get_bss_model(n_components)
              sc.metadata.General.title = ns.signal.name + "[BSS-model]"
              sc.plot()
          if not ns.recorded:
              self._record(ns.autosig, ns.model, ns.signal, n_components)
  
      def _show_scree(self, ns, callback):
          ax = ns.s.plot_explained_variance_ratio()
  
          # Clean up plot and present, allow user to select components
          # by picker
          ax.set_title("")
          scree = ax.get_figure().canvas
          scree.draw()
          scree.setWindowTitle("Pick number of components")
  
          def clicked(event):
              n_components = int(round(event.xdata))
              # Close scree plot
              w = fig2win(scree.figure, self.ui.figures)
              w.close()
              callback(ns, n_components)
          scree.mpl_connect('button_press_event', clicked)
  
      def do_after_scree(self, model, signal=None, n_components=None):
          """
          Performs decomposition, then plots the scree for the user to select
          the number of components to use for a decomposition model. The
          selection is made by clicking on the scree, which closes the scree
          and creates the model.
          """
          ns = Namespace()
          ns.autosig = signal is None
          ns.model = model
          ns.s, ns.signal = self._get_signal(signal)
          if n_components is not None:
              self._record(ns.autosig, ns.model, ns.signal, n_components)
              ns.recorded = True
          else:
              ns.recorded = False
  
          def on_complete(ns):
              if n_components is None:
                  self._show_scree(ns, self._perform_model)
              else:
                  self._perform_model(ns, n_components)
  
          self._decompose_threaded(on_complete, "Performing %s" % model.upper(),
                                   n_components, ns=ns)
  
      def plot_decomposition_results(self, signal=None, advanced=False):
          """
          Performs decomposition if necessary, then plots the decomposition
          results according to the hyperspy implementation.
          """
          def on_complete(ns):
              ns.s.plot_decomposition_results()
              # Somewhat speculative workaround to HSPY not adding metadata
              sd = self.ui.hspy_signals[-1]
              sd.metadata.add_dictionary(ns.s.metadata.as_dictionary())
  
          if advanced:
              diag = make_advanced_dialog(
                  self.ui, ['svd', 'fast_svd', 'mlpca', 'fast_mlpca', 'nmf',
                            'sparse_pca', 'mini_batch_sparse_pca'])
              dr = diag.exec_()
              if dr == QDialog.Accepted:
                  self._decompose_threaded(
                      on_complete, "Decomposing signal",
                      algorithm=diag.algorithm())
          else:
              self._decompose_threaded(on_complete, "Decomposing signal")
  
      def pca(self, signal=None, n_components=None, advanced=False):
          """
          Performs decomposition, then plots the scree for the user to select
          the number of components to use for a decomposition model. The
          selection is made by clicking on the scree, which closes the scree
          and creates the model.
          """
          if advanced:
              diag = make_advanced_dialog(self.ui)
              dr = diag.exec_()
              if dr == QDialog.Accepted:
                  self.do_after_scree(
                      'pca', signal, n_components=diag.components())
          else:
              self.do_after_scree('pca', signal, n_components)
  
      def bss(self, signal=None, n_components=None, advanced=False):
          """
          Performs decomposition if neccessary, then plots the scree for the user
          to select the number of components to use for a blind source
          separation. The selection is made by clicking on the scree, which
          closes the scree and creates the model.
          """
          if advanced:
              diag = make_advanced_dialog(
                  self.ui, ['orthomax', 'sklearn_fastica', 'FastICA', 'JADE',
                            'CuBICA', 'TDSEP'])
              dr = diag.exec_()
              if dr == QDialog.Accepted:
                  model = 'bss.' + diag.algorithm()
                  self.do_after_scree(model, signal, n_components)
          else:
              self.do_after_scree('bss', signal, n_components)
  
      def bss_model(self, signal=None, n_components=None, advanced=False):
          """
          Performs decomposition if neccessary, then plots the scree for the user
          to select the number of components to use for a blind source
          separation model. The selection is made by clicking on the scree, which
          closes the scree and creates the model.
          """
          if advanced:
              diag = make_advanced_dialog(self.ui)
              dr = diag.exec_()
              if dr == QDialog.Accepted:
                  self.do_after_scree(
                      'bss_model', signal, n_components=diag.components())
          else:
              self.do_after_scree('bss_model', signal, n_components)
  
      def clear(self, signal=None):
          """
          Clears the learning results from the signal.
          """
          if signal is None:
              signal = self.ui.get_selected_signal()
          signal.learning_results = LearningResults()
  