from __future__ import annotations

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


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.ItemDataRole.UserRole + 1
    operatorRole = QtCore.Qt.ItemDataRole.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.DropAction.MoveAction)

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

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

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

        elif evt.key() == QtCore.Qt.Key.Key_Space:
            for item in self.selectedItems():
                cs = item.checkState(0)
                if cs == QtCore.Qt.CheckState.Unchecked:
                    item.setCheckState(0, QtCore.Qt.CheckState.Checked)
                else:
                    item.setCheckState(0, QtCore.Qt.CheckState.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.ItemDataRole.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.CheckState.Checked if active else QtCore.Qt.CheckState.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.ItemDataRole.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.ItemDataRole.UserRole),
                'op': it.data(0, self.operatorRole),
                'pos': pos,
                'active': (it.checkState(0) == QtCore.Qt.CheckState.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.CheckState.Checked)
            item.setData(QtCore.Qt.ItemDataRole.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.CheckState.Checked:
                mod = self.cellWidget(i, 1).currentData()
                if mod is None:
                    mod = default

                if include_name:
                    arg = (item.data(QtCore.Qt.ItemDataRole.UserRole+1), item.text())
                else:
                    arg = item.data(QtCore.Qt.ItemDataRole.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.ItemDataRole.UserRole+1), item.text()))
            else:
                ret_val.append(item.data(QtCore.Qt.ItemDataRole.UserRole+1))

        return ret_val
