# -*- coding: utf-8 -*-
__author__ = "Konstantin Klementiev"
__date__ = "30 May 2023"
# !!! SEE CODERULES.TXT !!!

import copy
from functools import partial
import ast
import numpy as np

from ...core import singletons as csi
from ...core.logger import syslogger
from ...gui import gcommons as gco
from . import gbasefit as gbf

from silx.gui import qt


class FunctionFitModel(qt.QAbstractTableModel):
    """The Qt table model for a Function Fit in ParSeq. Each row in the model
    represents a fitting parameter in the formula.
    """
    HEADERS = ['var', 'value', 'step', 'tie', '[min, max]', 'error']

    def __init__(self, worker):
        super().__init__()
        self.worker = worker
        self.setParams()

    def rowCount(self, parent=qt.QModelIndex()):
        return len(self.params)

    def columnCount(self, parent=qt.QModelIndex()):
        return len(self.HEADERS)

    def flags(self, index):
        if not index.isValid():
            return qt.Qt.NoItemFlags
        if index.column() == 0:  # p name
            res = qt.Qt.ItemIsEnabled
        elif index.column() in (1, 2, 4):  # value, step, lim
            res = qt.Qt.ItemIsEnabled | qt.Qt.ItemIsEditable
        elif index.column() == 3:  # tie
            res = qt.Qt.ItemIsEnabled | qt.Qt.ItemIsEditable
        elif index.column() == 5:  # error
            res = qt.Qt.NoItemFlags
        return res

    def data(self, index, role=qt.Qt.DisplayRole):
        if len(self.params) == 0:
            return
        if not index.isValid():
            return
        column, row = index.column(), index.row()
        p = list(self.params.keys())[row]
        pv = self.params[p]  # dict(value, step, tie, lim, error)
        fmt = gco.getFormatStr(pv['step'])
        if role in (qt.Qt.DisplayRole, qt.Qt.EditRole):
            if column == 0:  # p name
                return p
            elif column == 1:  # value
                return fmt.format(pv['value'])
            elif column == 2:  # step
                return fmt.format(pv['step'])
            elif column == 3:  # tie
                if 'tie' in pv:
                    return pv['tie']
            elif column == 4:  # [min, max]
                if 'lim' in pv:
                    lim = list(pv['lim'])
                    lim[0] = '--' if lim[0] == -np.inf else fmt.format(lim[0])
                    lim[1] = '--' if lim[1] == -np.inf else fmt.format(lim[1])
                    return '[{0}, {1}]'.format(*lim)
            elif column == 5:  # error
                if 'error' in pv:
                    return fmt.format(pv['error'])
                else:
                    return '---'
        elif role == qt.Qt.TextAlignmentRole:
            return qt.Qt.AlignHCenter | qt.Qt.AlignVCenter
        elif role == qt.Qt.ToolTipRole:
            fmtLonger = gco.getFormatStr(pv['step']*0.001)
            if column == 1:
                return fmtLonger.format(pv['value'])
            elif column == 5 and 'error' in pv:
                return fmtLonger.format(pv['error'])
        elif role == qt.Qt.BackgroundRole:
            if column == 3:  # tie
                if 'tie' in pv:
                    if len(pv['tie']) < 2:
                        del pv['tie']
                    else:
                        if not self.worker.can_interpret_tie_str(
                                pv['tie'], self.params):
                            return qt.QColor(gbf.BAD_BKGND)
                        else:
                            return qt.QColor(gbf.GOOD_BKGND)
        elif role == gco.LIMITS_ROLE:  # return min, max, step
            if 'lim' in pv:
                return *pv['lim'], pv['step']
            else:
                return -np.inf, np.inf, pv['step']

    def setData(self, index, value, role=qt.Qt.EditRole):
        if len(self.params) == 0:
            return
        column, row = index.column(), index.row()
        p = list(self.params.keys())[row]
        pv = self.params[p]  # dict(value, step, tie, lim, error)
        if role == qt.Qt.EditRole:
            if column == 1:
                pv['value'] = float(value)
            elif column == 2:  # step
                pv['step'] = float(value)
            elif column == 3:  # tie
                if len(value) < 2 and 'tie' in pv:
                    del pv['tie']
                else:
                    pv['tie'] = str(value).strip()
            elif column == 4:  # [min, max]
                if len(value) == 0:
                    if 'lim' in pv:
                        del pv['lim']
                        self.dataChanged.emit(index, index)
                    return True
                try:
                    tmp = eval(value)
                except Exception:
                    return False
                if not isinstance(tmp, (list, tuple)):
                    return False
                pv['lim'] = [float(t) for t in tmp]
            self.dataChanged.emit(index, index)
            return True
        return False

    def headerData(self, section, orientation, role):
        if role == qt.Qt.DisplayRole:
            if orientation == qt.Qt.Horizontal:
                if section < len(self.HEADERS):
                    return self.HEADERS[section]
                else:
                    return section
            elif orientation == qt.Qt.Vertical:
                return str(section+1)
        elif role == qt.Qt.TextAlignmentRole:
            if orientation == qt.Qt.Horizontal:
                return qt.Qt.AlignCenter
        elif role == qt.Qt.ToolTipRole:
            if orientation == qt.Qt.Horizontal:
                if section == 0:
                    return 'parameter name'
                elif section == 1:
                    return 'parameter value'
                elif section == 2:
                    return 'step in the SpinBox value edit widget'
                elif section == 3:
                    return 'either "fixed" or str starting with = or > or < '\
                        '\nand followed by a Python expression of other vars'
                elif section == 4:
                    return 'parameter limits, applied to\nthe SpinBox value '\
                        'edit widget\nand the fitting method'
                elif section == 5:
                    return 'fitting error'

    def setParams(self, params=None, emit=True):
        self.beginResetModel()
        if params is None:
            self.params = {}
        else:
            self.params = params
        self.endResetModel()
        if emit:
            self.dataChanged.emit(qt.QModelIndex(), qt.QModelIndex())

    def moveItem(self, index, to):  # to = +1(up) or -1(down)
        if self.rowCount() == 0:
            return
        row = index.row()
        if row == 0 and to == 1:
            wro = self.rowCount()-1
        elif row == self.rowCount()-1 and to == -1:
            wro = 0
        else:
            wro = row - to
        self.beginResetModel()
        keys = list(self.params.keys())
        vals = list(self.params.values())
        self.beginResetModel()
        keys.insert(wro, keys.pop(row))
        vals.insert(wro, vals.pop(row))
        self.params = {k: v for k, v in zip(keys, vals)}
        self.endResetModel()
        self.dataChanged.emit(qt.QModelIndex(), qt.QModelIndex())
        return wro

    def removeItem(self, index):
        if self.rowCount() == 0:
            return
        self.beginResetModel()
        row = index.row()
        wro = row if row < self.rowCount()-1 else self.rowCount()-2
        key = list(self.params.keys())[row]
        del self.params[key]
        self.endResetModel()
        self.dataChanged.emit(qt.QModelIndex(), qt.QModelIndex())
        return wro

    def removeItems(self, indexes):
        keys = list(self.params.keys())
        vals = list(self.params.values())
        self.beginResetModel()
        for index in reversed(indexes):
            row = index.row()
            keys.pop(row)
            vals.pop(row)
        self.params = {k: v for k, v in zip(keys, vals)}
        self.endResetModel()
        self.dataChanged.emit(qt.QModelIndex(), qt.QModelIndex())


class FunctionFitTableView(qt.QTableView):
    columnWidths = [60, 80, 50, 80, 150, 55]

    def __init__(self, parent, model):
        super().__init__(parent)
        self.setModel(model)

        # horHeaders = self.horizontalHeader()  # QHeaderView instance
        horHeaders = gbf.UnderlinedHeaderView(qt.Qt.Horizontal, self)
        self.setHorizontalHeader(horHeaders)

        verHeaders = self.verticalHeader()  # QHeaderView instance
        verHeaders.setVisible(True)

        nC = model.columnCount()
        if 'pyqt4' in qt.BINDING.lower():
            horHeaders.setMovable(False)
            horHeaders.setResizeMode(0, qt.QHeaderView.Stretch)
            for i in range(1, nC):
                horHeaders.setResizeMode(i, qt.QHeaderView.Interactive)
            horHeaders.setClickable(True)
        else:
            horHeaders.setSectionsMovable(False)
            horHeaders.setSectionResizeMode(0, qt.QHeaderView.Stretch)
            for i in range(1, nC):
                horHeaders.setSectionResizeMode(i, qt.QHeaderView.Interactive)
            horHeaders.setSectionsClickable(True)
        horHeaders.setStretchLastSection(False)
        # horHeaders.setMinimumSectionSize(20)
        verHeaders.setDefaultSectionSize(20)

        for i, cw in enumerate(self.columnWidths[:nC]):
            self.setColumnWidth(i, int(cw*csi.screenFactor))

        kw = dict(alignment=qt.Qt.AlignHCenter | qt.Qt.AlignVCenter)
        self.setItemDelegateForColumn(1, gco.DoubleSpinBoxDelegate(self, **kw))

        self.setMinimumHeight(horHeaders.height() * max(2, model.rowCount()+1))
        # self.setMinimumWidth(
        #     int(sum(self.columnWidths[:nC])*csi.screenFactor) + 30)

        # self.setSelectionMode(qt.QAbstractItemView.NoSelection)
        self.setSelectionMode(qt.QAbstractItemView.SingleSelection)
        # self.setSelectionBehavior(qt.QAbstractItemView.SelectItems)
        # self.setFocusPolicy(qt.Qt.NoFocus)

        self.setContextMenuPolicy(qt.Qt.CustomContextMenu)
        self.customContextMenuRequested.connect(self.onCustomContextMenu)
        self.makeActions()

    def makeActions(self):
        if csi.mainWindow is not None:
            self.actionCopyFitParams = self._addAction(
                "Copy fit params to picked data",
                partial(self.startPickData, 1), "Ctrl+P")
        self.actionMoveUp = self._addAction(
            "Move up", partial(self.moveItem, +1), "Ctrl+Up")
        self.actionMoveDown = self._addAction(
            "Move down", partial(self.moveItem, -1), "Ctrl+Down")
        self.actionRemove = self._addAction("Remove", self.removeItem, "Del")

    def _addAction(self, text, slot, shortcut=None):
        action = qt.QAction(text, self)
        action.triggered.connect(slot)
        if shortcut:
            action.setShortcut(qt.QKeySequence(shortcut))
        action.setShortcutContext(qt.Qt.WidgetWithChildrenShortcut)
        self.addAction(action)
        return action

    def onCustomContextMenu(self, point):
        menu = qt.QMenu()
        if csi.mainWindow is not None:
            menu.addAction(self.actionCopyFitParams)
            menu.addSeparator()
        menu.addAction(self.actionMoveUp)
        menu.addAction(self.actionMoveDown)
        menu.addSeparator()
        menu.addAction(self.actionRemove)
        menu.exec_(self.viewport().mapToGlobal(point))

    def startPickData(self, pickReason):
        self.pickReason = pickReason
        self.model().worker.node.widget.preparePickData(self)

    def applyPendingProps(self):
        if self.pickReason == 1:  # copy fit params
            srcParams = self.parent().spectrum.fitParams
            for item in csi.selectedItems:
                dfparams = item.fitParams
                dfparams['ffit_formula'] = str(srcParams['ffit_formula'])
                for prop in self.model().worker.defaultParams:
                    if prop in srcParams:
                        dfparams[prop] = copy.deepcopy(srcParams[prop])

        csi.selectionModel.clear()
        csi.model.selectItems(self.parent().spectrum)

    def moveItem(self, to):
        newRow = self.model().moveItem(self.currentIndex(), to)
        if newRow is None:
            return
        newInd = self.model().index(newRow, 0)
        if newInd.isValid():
            self.setCurrentIndex(newInd)

    def removeItem(self):
        newRow = self.model().removeItem(self.currentIndex())
        if newRow is None:
            return
        newInd = self.model().index(newRow, 0)
        if newInd.isValid():
            self.setCurrentIndex(newInd)


class FunctionFitWidget(gbf.FitWidget):
    """
    .. imagezoom:: _images/func-fit.png
       :align: center
       :alt: &ensp;A Function Fit example.

    The formula operates numpy functions as attributes of `np`, e.g.
    `np.exp(x)`. Additionally, these functions are defined:

    gau(x, m, s): normalized Gaussian function, m=center, s=sigma

    lor(x, m, s): normalized Lorentzian function, m=center, s=HWHM

    """

    def __init__(self, parent, worker, plot):
        super().__init__(parent, worker, plot)
        self.spectrum = None

        layout = qt.QVBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)

        layoutF = qt.QHBoxLayout()
        self.formulaLabel = qt.QLabel('f(x) =')
        layoutF.addWidget(self.formulaLabel)
        self.formulaEdit = qt.QLineEdit()
        toolTip = 'Use numpy functions as attributes of np, e.g. `np.exp(x)`'
        extraToolTip = worker.getToolTip()
        if extraToolTip:
            toolTip += '\nExtra functions defined in this context:'
            toolTip += extraToolTip
        self.formulaEdit.setToolTip(toolTip)

        self.formulaEdit.returnPressed.connect(self.parseFormula)
        layoutF.addWidget(self.formulaEdit)
        layout.addLayout(layoutF)

        self.fitModel = FunctionFitModel(worker)
        table = FunctionFitTableView(self, self.fitModel)
        layout.addWidget(table)
        self.fitModel.dataChanged.connect(self.makeFit)

        self.addRangeAndStartWidgets(layout)
        # layout.addStretch()
        self.setLayout(layout)

    def setSpectrum(self, spectrum):
        dfparams = spectrum.fitParams
        fitFormula = dfparams['ffit_formula']
        self.formulaEdit.setText(fitFormula)
        super().setSpectrum(spectrum)

    def parseFormula(self):
        cs = self.spectrum
        dfparams = cs.fitParams
        fitVars = dfparams['ffit_params']
        try:
            st = ast.parse(self.formulaEdit.text())
            self.formulaEdit.setStyleSheet('')
            fvars = []
            for node in ast.walk(st):
                if type(node) is ast.Name:
                    if node.id in (
                            ['np', 'x'] + list(self.worker.customFunctions)):
                        continue
                    if node.id not in fvars:
                        fvars.append(node.id)
            for fv in sorted(fvars):
                if fv not in fitVars:
                    defPar = list(self.worker.defaultEntryDict.values())[0]
                    fitVars[fv] = dict(defPar)
            dfparams['ffit_formula'] = self.formulaEdit.text()
            delKeys = [k for k in fitVars.keys() if k not in fvars]
            for k in delKeys:
                del fitVars[k]
            self.fitModel.setParams(fitVars)
        except Exception as err:
            syslogger.error(
                'Error in FunctionFitWidget.parseFormula:\n'+str(err))
            self.formulaEdit.setStyleSheet(
                "QLineEdit {background-color: "+gbf.BAD_BKGND+";}")

    def makeFit(self):
        """Here, all tie formulas and off-bounds are ignored."""
        if not self.fitModel.params:
            return
        cs = self.spectrum
        self.worker.make_model_curve(cs)

        dfparams = cs.fitParams
        lcfRes = dfparams['ffit_result']
        self.fitR.setText('R={0:.5g}'.format(lcfRes['R']))
        self.updateFitResults()
        self.fitReady.emit()
