from __future__ import annotations

from nmreval.utils.text import convert

from ..Qt import QtCore, QtWidgets, QtGui
from .._py.fitmodelwidget import Ui_FitParameter
from .._py.save_fitmodel_dialog import Ui_SaveDialog
from ..lib.iconloading import get_icon
from ..lib.tables import TableWidget


class FitModelWidget(QtWidgets.QWidget, Ui_FitParameter):
    value_requested = QtCore.pyqtSignal(object)
    value_changed = QtCore.pyqtSignal(str)
    state_changed = QtCore.pyqtSignal()
    replace_single_value = QtCore.pyqtSignal(object)

    def __init__(self, label: str = 'Fitparameter', parent=None, fixed: bool = False):
        super().__init__(parent)
        self.setupUi(self)

        self.name = label

        self.parametername.setText(convert(label) + ' ')

        self.parameter_line.setText('1')
        self.parameter_line.setMaximumWidth(160)
        self.lineEdit.setMaximumWidth(100)
        self.lineEdit_2.setMaximumWidth(100)

        self.label_3.setText(f'< {convert(label)} <')

        self.checkBox.stateChanged.connect(self.enableBounds)
        self.global_checkbox.stateChanged.connect(lambda: self.state_changed.emit())
        self.parameter_line.editingFinished.connect(self.update_parameter)
        self.parameter_line.values_requested.connect(lambda: self.value_requested.emit(self))
        self.parameter_line.replace_single_values.connect(lambda: self.replace_single_value.emit(None))
        self.parameter_line.editingFinished.connect(lambda: self.value_changed.emit(self.parameter_line.text()))
        self.fixed_check.toggled.connect(self.set_fixed)

        if fixed:
            self.fixed_check.hide()

        self.parameter_pos = None
        self.func_idx = None

        self._linetext = '1'

        self.menu = QtWidgets.QMenu(self)

    def set_parameter_string(self, p: str):
        self.parameter_line.setText(p)
        self.parameter_line.setToolTip(p)

    def set_bounds(self, lb: float, ub: float, cbox: bool = True):
        self.checkBox.setCheckState(QtCore.Qt.Checked if cbox else QtCore.Qt.Unchecked)
        for val, bds_line in [(lb, self.lineEdit), (ub, self.lineEdit_2)]:
            if val is not None:
                bds_line.setText(str(val))
            else:
                bds_line.setText('')

    def enableBounds(self, value: int):
        self.lineEdit.setEnabled(value == 2)
        self.lineEdit_2.setEnabled(value == 2)

    def set_parameter(self, p: float | None, bds: tuple[float, float, bool] = None,
                      fixed: bool = None, glob: bool = None):
        ptext = f'{p:.4g}'

        self.set_parameter_string(ptext)

        if bds is not None:
            self.set_bounds(*bds)

        if fixed is not None:
            self.fixed_check.setCheckState(QtCore.Qt.Unchecked if fixed else QtCore.Qt.Checked)

        if glob is not None:
            self.global_checkbox.setCheckState(QtCore.Qt.Checked if glob else QtCore.Qt.Unchecked)

    def get_parameter(self):
        try:
            p = float(self.parameter_line.text().replace(',', '.'))
        except ValueError:
            p = self.parameter_line.text().replace(',', '.')

        if self.checkBox.isChecked():
            lb_text = self.lineEdit.text()
            lb = None
            if lb_text:
                try:
                    lb = float(lb_text.replace(',', '.'))
                except ValueError:
                    lb = lb_text

            ub_text = self.lineEdit_2.text()
            rb = None
            if ub_text:
                try:
                    rb = float(ub_text.replace(',', '.'))
                except ValueError:
                    rb = ub_text
        else:
            lb = rb = None

        bounds = (lb, rb)

        return p, bounds, not self.fixed_check.isChecked(), self.global_checkbox.isChecked()

    @QtCore.pyqtSlot(bool)
    def set_fixed(self, state: bool):
        # self.global_checkbox.setVisible(not state)
        self.frame.setVisible(not state)

    @QtCore.pyqtSlot()
    def update_parameter(self):
        new_value = self.parameter_line.text()
        if not new_value:
            self.parameter_line.setText('1')

        try:
            float(new_value)
            is_text = False
        except ValueError:
            is_text = True
            self.global_checkbox.setCheckState(False)

        self.set_fixed(is_text)


class QSaveModelDialog(QtWidgets.QDialog, Ui_SaveDialog):
    def __init__(self, types=None, parent=None):
        super().__init__(parent=parent)
        self.setupUi(self)

        if types is None:
            types = []

        self.comboBox.blockSignals(True)
        self.comboBox.addItems(types)
        self.comboBox.addItem('New group...')
        self.comboBox.blockSignals(False)

        self.frame.hide()

    @QtCore.pyqtSlot(int, name='on_comboBox_currentIndexChanged')
    def new_group(self, idx: int):
        if idx == self.comboBox.count() - 1:
            self.frame.show()
        else:
            self.lineEdit_2.clear()
            self.frame.hide()

    @QtCore.pyqtSlot(name='on_toolButton_clicked')
    def accept_group(self):
        self.comboBox.insertItem(self.comboBox.count() - 1, self.lineEdit_2.text())
        self.comboBox.setCurrentIndex(self.comboBox.count() - 2)

    def accept(self):
        if self.lineEdit.text():
            self.close()


class FitModelTree(QtWidgets.QTreeWidget):
    icons = ['plus', 'mal_icon', 'minus_icon', 'geteilt_icon']

    treeChanged = QtCore.pyqtSignal()
    itemRemoved = QtCore.pyqtSignal(int)

    counterRole = QtCore.Qt.UserRole + 1
    operatorRole = QtCore.Qt.UserRole + 2

    def __init__(self, parent=None):
        super().__init__(parent=parent)
        self.setHeaderHidden(True)
        self.setDragEnabled(True)
        self.setDragDropMode(QtWidgets.QTreeWidget.InternalMove)
        self.setDefaultDropAction(QtCore.Qt.MoveAction)

        self.itemSelectionChanged.connect(lambda: self.treeChanged.emit())

    def keyPressEvent(self, evt):
        operators = [QtCore.Qt.Key_Plus, QtCore.Qt.Key_Asterisk,
                     QtCore.Qt.Key_Minus, QtCore.Qt.Key_Slash]

        if evt.key() == QtCore.Qt.Key_Delete:
            for item in self.selectedItems():
                self.remove_function(item)

        elif evt.key() == QtCore.Qt.Key_Space:
            for item in self.treeWidget.selectedItems():
                item.setCheckState(0, QtCore.Qt.Checked) if item.checkState(
                    0) == QtCore.Qt.Unchecked else item.setCheckState(0, QtCore.Qt.Unchecked)

        elif evt.key() in operators:
            idx = operators.index(evt.key())
            for item in self.selectedItems():
                item.setData(0, self.operatorRole, idx)
                item.setIcon(0, get_icon(self.icons[idx]))

        else:
            super().keyPressEvent(evt)

    def dropEvent(self, evt: QtGui.QDropEvent):
        super().dropEvent(evt)
        self.treeChanged.emit()

    def remove_function(self, item: QtWidgets.QTreeWidgetItem):
        """
        Remove function and children from tree and dictionary
        """
        while item.childCount():
            self.remove_function(item.child(0))

        if item.parent():
            item.parent().removeChild(item)
        else:
            self.invisibleRootItem().removeChild(item)

        idx = item.data(0, self.counterRole)
        self.itemRemoved.emit(idx)

    def add_function(self,
                     idx: int,
                     cnt: int,
                     op: int,
                     name: str,
                     color: QtGui.QColor | str | tuple,
                     parent: QtWidgets.QTreeWidgetItem = None,
                     children: list = None,
                     active: bool = True,
                     param_names: list[str] = None,
                     **kwargs):
        """
        Add function to tree and dictionary of functions.
        """
        if not isinstance(color, QtGui.QColor):
            if isinstance(color, tuple):
                color = QtGui.QColor.fromRgbF(*color)
            else:
                color = QtGui.QColor(color)

        it = QtWidgets.QTreeWidgetItem()
        it.setData(0, QtCore.Qt.UserRole, idx)
        it.setData(0, self.counterRole, cnt)
        it.setData(0, self.operatorRole, op)
        it.setText(0, name)
        if param_names is not None:
            it.setToolTip(0,
                          'Parameter names:\n' +
                          '\n'.join(f'{pn}({cnt})' for pn in param_names))
        it.setForeground(0, QtGui.QBrush(color))

        it.setIcon(0, get_icon(self.icons[op]))
        it.setCheckState(0, QtCore.Qt.Checked if active else QtCore.Qt.Unchecked)

        if parent is None:
            self.addTopLevelItem(it)
        else:
            parent.addChild(it)

        if children is not None:
            for c in children:
                self.add_function(**c, parent=it)

        self.setCurrentIndex(self.indexFromItem(it, 0))

    def sizeHint(self):
        w = super().sizeHint().width()
        return QtCore.QSize(w, 100)

    def get_selected(self):
        try:
            it = self.selectedItems()[0]
            function_nr = it.data(0, QtCore.Qt.UserRole)
            idx = it.data(0, self.counterRole)

        except IndexError:
            function_nr = None
            idx = None

        return function_nr, idx

    def get_functions(self, full: bool = True, pos: int = -1, return_pos: bool = False, parent=None):
        """
        Create nested list of functions in tree. Parameters saved are idx (Index of function in list of all functions),
        cnt (counter of number to associate with functione values), ops (+, -, *, /), and maybe children.
        """
        if parent is None:
            parent = self.invisibleRootItem()

        funcs = []
        for i in range(parent.childCount()):
            pos += 1
            it = parent.child(i)

            child = {
                'idx': it.data(0, QtCore.Qt.UserRole),
                'op': it.data(0, self.operatorRole),
                'pos': pos,
                'active': (it.checkState(0) == QtCore.Qt.Checked),
                'children': []
            }

            if full:
                child['name'] = it.text(0)
                child['cnt'] = it.data(0, self.counterRole)
                child['color'] = it.foreground(0).color().getRgbF()

            if it.childCount():
                child['children'], pos = self.get_functions(full=full, parent=it, pos=pos, return_pos=True)

            funcs.append(child)

        if return_pos:
            return funcs, pos
        else:
            return funcs


class FitTableWidget(TableWidget):
    def __init__(self, parent=None):
        super().__init__(parent=parent)

        self.horizontalHeader().hide()
        self.verticalHeader().hide()
        self.setColumnCount(2)
        self.setSelectionBehavior(QtWidgets.QTableWidget.SelectRows)
        self.horizontalHeader().setStretchLastSection(True)
        self.hideColumn(1)

    def add_model(self, idx: str):
        model_count = 0
        for r in range(self.rowCount()):
            cb = self.cellWidget(r, 1)
            cb.addItem('Model ' + str(idx), userData=idx)
            model_count = cb.count()

        if model_count > 2:
            if self.isColumnHidden(1):
                self.showColumn(1)
                self.resizeColumnToContents(0)
                self.setColumnWidth(1, self.columnWidth(0) - self.columnWidth(1))

    def remove_model(self, idx: str):
        model_count = 0
        for r in range(self.rowCount()):
            cb = self.cellWidget(r, 1)
            if cb.currentData() == idx:
                cb.setCurrentIndex(0)
            cb.removeItem(cb.findData(idx))
            model_count = cb.count()

        if model_count == 2:
            self.hideColumn(1)
            self.resizeColumnToContents(0)

    def load(self, set_ids: list[str]):
        self.blockSignals(True)

        while self.rowCount():
            self.removeRow(0)

        self.setColumnCount(2)
        self.hideColumn(1)

        for (sid, name) in set_ids:
            item = QtWidgets.QTableWidgetItem(name)
            item.setCheckState(QtCore.Qt.Checked)
            item.setData(QtCore.Qt.UserRole+1, sid)
            row = self.rowCount()
            self.setRowCount(row+1)
            self.setItem(row, 0, item)

            item2 = QtWidgets.QTableWidgetItem('')
            self.setItem(row, 1, item2)
            cb = QtWidgets.QComboBox(parent=self)
            cb.addItem('Default')
            self.setCellWidget(row, 1, cb)

        self.blockSignals(False)

    def collect_data(self, default: str = None, include_name: bool = False) -> dict:
        data = {}

        for i in range(self.rowCount()):
            item = self.item(i, 0)
            if item.checkState() == QtCore.Qt.Checked:
                mod = self.cellWidget(i, 1).currentData()
                if mod is None:
                    mod = default

                if include_name:
                    arg = (item.data(QtCore.Qt.UserRole+1), item.text())
                else:
                    arg = item.data(QtCore.Qt.UserRole+1)

                if mod not in data:
                    data[mod] = []
                data[mod].append(arg)

        return data

    def data_list(self, include_name: bool = True) -> list:
        ret_val = []
        for i in range(self.rowCount()):
            item = self.item(i, 0)
            if include_name:
                ret_val.append((item.data(QtCore.Qt.UserRole+1), item.text()))
            else:
                ret_val.append(item.data(QtCore.Qt.UserRole+1))

        return ret_val