from __future__ import annotations

from typing import Any

from numpy import ndarray, iscomplexobj, asarray

from ..Qt import QtGui, QtCore, QtWidgets
from .._py.valueeditor import Ui_MaskDialog
from ..lib.pg_objects import PlotItem


class ValueEditWidget(QtWidgets.QWidget, Ui_MaskDialog):
    requestData = QtCore.pyqtSignal(str)
    maskSignal = QtCore.pyqtSignal(str, list)
    itemChanged = QtCore.pyqtSignal(str, tuple, object)
    itemDeleted = QtCore.pyqtSignal(str, list)
    itemAdded = QtCore.pyqtSignal(str)
    values_selected = QtCore.pyqtSignal(str, str)
    split_signal = QtCore.pyqtSignal(str, int)

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

        self.setupUi(self)

        self.connected_figure = None
        self.shown_set = None
        self.items = {}

        self.model = ValueModel()
        self.model.itemChanged.connect(self.update_items)

        self.selection_model = QtCore.QItemSelectionModel(self.model)
        self.selection_model.selectionChanged.connect(self.show_position)

        self.tableView.setModel(self.model)
        self.tableView.setSelectionModel(self.selection_model)
        self.tableView.setContextMenuPolicy(QtCore.Qt.ContextMenuPolicy.CustomContextMenu)
        self.tableView.customContextMenuRequested.connect(self.ctx)

        self.selection_real = PlotItem(x=[], y=[], symbolSize=25, symbol='x',
                                       pen=None, symbolPen='#c9308e', symbolBrush='#c9308e')
        self.selection_imag = PlotItem(x=[], y=[], symbolSize=25, symbol='+',
                                       pen=None, symbolPen='#dcdcdc', symbolBrush='#dcdcdc')

    def __call__(self, items: dict):
        self.items = items

        self.graph_combobox.blockSignals(True)
        self.graph_combobox.clear()
        for k, v in items.items():
            self.graph_combobox.addItem(k[1], userData=k[0])
        self.graph_combobox.blockSignals(False)

        idx = self.graph_combobox.findData(self.connected_figure)
        if idx == -1:
            idx = 0

        self.graph_combobox.setCurrentIndex(idx)
        self.graph_combobox.currentIndexChanged.emit(idx)

        return self

    @QtCore.pyqtSlot(int, name='on_graph_combobox_currentIndexChanged')
    def _populate_sets(self, idx: int):
        if idx == -1:
            self.graph_combobox.setCurrentIndex(0)
            return

        self.set_combobox.blockSignals(True)
        self.set_combobox.clear()
        old_figure = self.connected_figure
        self.connected_figure = self.graph_combobox.currentData()

        if self.items:
            for sid, name in self.items[(self.graph_combobox.currentData(), self.graph_combobox.currentText())]:
                self.set_combobox.addItem(name, userData=sid)
        self.set_combobox.blockSignals(False)

        sidx = self.set_combobox.findData(self.shown_set)
        if sidx == -1:
            sidx = 0

        self.set_combobox.setCurrentIndex(sidx)
        self.set_combobox.currentIndexChanged.emit(sidx)

        self.values_selected.emit(old_figure, self.connected_figure)

    @QtCore.pyqtSlot(int, name='on_set_combobox_currentIndexChanged')
    def _populate_table(self, idx):
        self.selection_model.clearSelection()
        self.shown_set = self.set_combobox.itemData(idx)
        self.requestData.emit(self.set_combobox.itemData(idx))

    def set_data(self, data: list, mask: ndarray):
        self.selection_model.clearSelection()
        self.model.loadData(data, mask)
        self.spinBox.setMaximum(self.model.rowCount())

    def remove_graph(self):
        # remove everything
        self.connected_figure = None
        self.items = {}
        self.selection_model.clear()
        while self.model.rowCount():
            self.model.removeRow(0)
        self.graph_combobox.clear()
        self.set_combobox.clear()

    def ctx(self, pos: QtCore.QPoint):
        idx = self.tableView.indexAt(pos)
        if not idx.isValid():
            return

        menu = QtWidgets.QMenu()
        menu.addSeparator()

        hide_action = menu.addAction('Hide/Show')
        del_action = menu.addAction('Delete')
        split_action = menu.addAction('Split after selection')
        cpc_action = menu.addAction('Copy to clipboard')
        action = menu.exec(self.tableView.viewport().mapToGlobal(pos))

        if action == hide_action:
            self.mask_row()
        elif action == del_action:
            self.delete_item()
        elif action == cpc_action:
            self.copy_selection()
        elif action == split_action:
            self.split()

    def keyPressEvent(self, evt):
        if evt.matches(QtGui.QKeySequence.Copy):
            self.copy_selection()
        elif evt.key() == QtCore.Qt.Key.Key_Delete:
            self.delete_item()
        else:
            super().keyPressEvent(evt)

    def copy_selection(self) -> str:
        table = ''
        for r in self.selection_model.selectedRows():
            table += self.model.data(r) + '\t'
            table += '\t'.join([self.model.data(r.sibling(r.row(), i)) for i in [1, 2]]) + '\n'

        QtWidgets.QApplication.clipboard().setText(table)

        return table

    @QtCore.pyqtSlot(name='on_split_button_clicked')
    def split(self):
        self.split_signal.emit(self.set_combobox.currentData(),
                               self.selection_model.selectedRows()[-1].row()+1)

    @QtCore.pyqtSlot(name='on_mask_button_clicked')
    def mask_row(self):
        for r in self.selection_model.selectedRows():
            self.model.setData(r, not self.model.data(r, ValueModel.maskRole), ValueModel.maskRole)
        self.maskSignal.emit(self.set_combobox.currentData(), self.model.mask)

    @QtCore.pyqtSlot(name='on_unmaskbutton_clicked')
    def unmask(self):
        self.model.unmask()
        self.maskSignal.emit(self.set_combobox.currentData(), self.model.mask)

    @QtCore.pyqtSlot(name='on_delete_button_clicked')
    def delete_item(self):
        idx = [r.row() for r in self.selection_model.selectedRows()]
        success = False
        for i in sorted(idx, reverse=True):
            success = self.model.removeRow(i)

        if success:
            self.itemDeleted.emit(self.set_combobox.currentData(), idx)
            self.spinBox.setMaximum(self.spinBox.maximum()-len(idx))

    @QtCore.pyqtSlot(name='on_add_button_clicked')
    def new_value(self):
        success = self.model.addRows()
        if success:
            self.spinBox.setMaximum(self.spinBox.maximum()+1)
            self.itemAdded.emit(self.set_combobox.currentData())

    @QtCore.pyqtSlot(int, int, str)
    def update_items(self, col, row, val):
        sid = self.set_combobox.currentData()
        new_value = complex(val)
        new_value = new_value.real if new_value.imag == 0 else new_value

        # table view loses focus when itemChanged is emitted
        # if edit of item is cause of change resume editing at next item
        prev_state = self.tableView.state()
        idx = self.tableView.currentIndex()
        idx = idx.sibling((col+1)//3+row, (col+1) % 3)
        self.itemChanged.emit(sid, (col, row), new_value)
        if prev_state == self.tableView.State.EditingState:
            self.tableView.setCurrentIndex(idx)
            self.tableView.edit(idx)

    @QtCore.pyqtSlot(QtCore.QItemSelection, QtCore.QItemSelection)
    def show_position(self, *_):
        xvals = []
        yvals = []
        for idx in self.selection_model.selectedRows():
            xvals.append(float(self.model.data(idx)))
            try:
                yvals.append(float(self.model.data(idx.sibling(idx.row(), 1))))
            except ValueError:
                yvals.append(complex(self.model.data(idx.sibling(idx.row(), 1))))

        yvals = asarray(yvals)

        if iscomplexobj(yvals):
            self.selection_real.setData(x=xvals, y=yvals.real)
            self.selection_imag.setData(x=xvals, y=yvals.imag)
        else:
            self.selection_real.setData(x=xvals, y=yvals)
            self.selection_imag.setData(x=[], y=[])

    @QtCore.pyqtSlot(name='on_toolButton_clicked')
    def goto(self):
        self.tableView.scrollTo(self.model.index(self.spinBox.value()-1, 0))


class ValueModel(QtCore.QAbstractTableModel):
    """
    TableModel with lazy loading
    """
    itemChanged = QtCore.pyqtSignal(int, int, str)
    load_number = 20
    maskRole = QtCore.Qt.ItemDataRole.UserRole+321

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

        self._data = None
        self.total_rows = 0
        self.rows_loaded = 0
        self.mask = None
        self.headers = ['x', 'y', '\u0394y']
        for i, hd in enumerate(self.headers):
            self.setHeaderData(i, QtCore.Qt.Orientation.Horizontal, hd)

    def rowCount(self, *args, **kwargs) -> int:
        return self.total_rows

    def columnCount(self, *args, **kwargs) -> int:
        return len(self.headers)

    def loadData(self, data: list[ndarray], mask: ndarray):
        self.beginResetModel()
        self._data = []
        for x, y, y_err in zip(*data):
            self._data.append([x, y, y_err])
        self.total_rows = len(self._data)

        self.mask = mask.tolist()

        self.endResetModel()
        self.dataChanged.emit(
            self.index(0, 0),
            self.index(0, 1), [QtCore.Qt.ItemDataRole.DisplayRole]
        )

    def data(self, idx: QtCore.QModelIndex, role=QtCore.Qt.ItemDataRole.DisplayRole) -> Any:
        if not idx.isValid():
            return

        row = idx.row()
        if role in [QtCore.Qt.ItemDataRole.DisplayRole, QtCore.Qt.ItemDataRole.EditRole]:
            val = self._data[row][idx.column()]
            return self.as_string(val)

        elif role == QtCore.Qt.ItemDataRole.BackgroundRole:
            pal = QtGui.QGuiApplication.palette()
            if not self.mask[row]:
                return pal.color(QtGui.QPalette.Disabled, QtGui.QPalette.Base)
            else:
                return pal.color(QtGui.QPalette.Base)

        elif role == QtCore.Qt.ItemDataRole.ForegroundRole:
            pal = QtGui.QGuiApplication.palette()
            if not self.mask[row]:
                return pal.color(QtGui.QPalette.Disabled, QtGui.QPalette.Text)
            else:
                return pal.color(QtGui.QPalette.Text)

        elif role == ValueModel.maskRole:
            return self.mask[row]

        else:
            return

    def setData(self, idx: QtCore.QModelIndex, value: str | bool, role=QtCore.Qt.ItemDataRole.DisplayRole) -> Any:
        col, row = idx.column(), idx.row()

        if role == ValueModel.maskRole:
            self.mask[row] = bool(value)
            self.dataChanged.emit(self.index(0, 0), self.index(0, 1), [role])

            return True

        if value:
            if role == QtCore.Qt.ItemDataRole.EditRole:
                if value == self.as_string(self._data[row][col]):
                    return True

                try:
                    value = complex(value)
                except ValueError:
                    # not a number
                    return False

                value = value.real if value.imag == 0 else value
                self._data[row][col] = value.real if value.imag == 0 else value
                self.itemChanged.emit(col, row, str(value))
                self.dataChanged.emit(self.index(0, 0), self.index(0, 1), [role])

            else:
                return super().setData(idx, value, role=role)

            return True

        else:
            return False

    def headerData(self, section: int, orientation, role=QtCore.Qt.ItemDataRole.DisplayRole) -> Any:
        if role == QtCore.Qt.ItemDataRole.DisplayRole:
            if orientation == QtCore.Qt.Orientation.Horizontal:
                return self.headers[section]
            else:
                return str(section+1)

        return

    def canFetchMore(self, idx: QtCore.QModelIndex) -> bool:
        if not idx.isValid():
            return False

        return self.total_rows > self.rows_loaded

    def fetchMore(self, idx: QtCore.QModelIndex):
        remaining = self.total_rows - self.rows_loaded
        to_be_loaded = min(remaining, ValueModel.load_number)

        self.beginInsertRows(QtCore.QModelIndex(), self.rows_loaded, self.rows_loaded + to_be_loaded - 1)
        self.rows_loaded += to_be_loaded
        self.endInsertRows()

    def flags(self, idx: QtCore.QModelIndex) -> QtCore.Qt.ItemFlag:
        return QtCore.QAbstractTableModel.flags(self, idx) | QtCore.Qt.ItemFlag.ItemIsEditable

    def removeRows(self, pos: int, rows: int, parent=None, *args, **kwargs) -> bool:
        self.beginRemoveRows(parent, pos, pos+rows-1)

        for _ in range(rows):
            self._data.pop(pos)
            self.mask.pop(pos)

        self.endRemoveRows()
        self.total_rows -= rows

        return True

    def addRows(self, num=1):
        return self.insertRows(self.rowCount(), num)

    def insertRows(self, pos: int, rows: int, parent=QtCore.QModelIndex(), *args, **kwargs):
        self.beginInsertRows(parent, pos, pos+rows-1)

        for _ in range(rows):
            self._data.insert(pos, [0.0] * self.columnCount())
            self.mask.insert(pos, True)
        self.total_rows += rows

        self.endInsertRows()

        return True

    def unmask(self):
        self.mask = [True] * self.total_rows
        self.dataChanged.emit(self.index(0, 0), self.index(0, 1), [ValueModel.maskRole])

    @staticmethod
    def as_string(value) -> str:
        if isinstance(value, complex):
            return f'{value.real:.8g}{value.imag:+.8g}j'
        else:
            return f'{value:.8g}'