from PyQt6.QtWidgets import QWidget, QVBoxLayout, QLabel, QScrollArea, QTableView
from PyQt6.QtCore import QAbstractTableModel, QModelIndex, Qt

import numpy as np
from cpdctrl.led_script import LedScript

class LedScriptTableModel(QAbstractTableModel):
    """
    A table model for the led script.

    It only allows editing values that are in the future
    """
    def __init__(self, led_script: LedScript, parent=None):
        super().__init__(parent)
        self.led_script = led_script
        self.indices = [i[0] for i in LedScript.ARRAY_DTYPE]
        self.dt = 0

    def rowCount(self, parent=None):
        return self.led_script.script[self.indices[0]].shape[0]

    def columnCount(self, parent=None):
        return len(self.indices)

    def data(self, index: QModelIndex, role: int):
        if role == Qt.ItemDataRole.DisplayRole:
            return str(self.led_script.script[self.indices[index.column()]][index.row()])

    def headerData(self, section: int, orientation: Qt.Orientation, role: int):
        if role == Qt.ItemDataRole.DisplayRole:
            if orientation == Qt.Orientation.Horizontal:
                return self.indices[section]
            else:
                return str(section)

    def setData(self, index: QModelIndex, value, role: int):
        if role == Qt.ItemDataRole.EditRole:
            self.led_script.script[self.indices[index.column()]][index.row()] = value
            return True
        return False

    def addRows(self, rowCount: int, parent: QModelIndex):
        self.beginInsertRows(parent, self.rowCount(), self.rowCount() + rowCount - 1)
        for i in self.indices:
            np.append(self.led_script.script[i], np.zeros((rowCount, self.led_script.script[i].shape[1])), axis=0)

        self.endInsertRows()
    def removeRows(self, row: int, count: int, parent: QModelIndex):
        self.beginRemoveRows(parent, row, row + count - 1)
        rows = [i for i in range(row, row+count)]
        for i in self.indices:
            np.delete(self.led_script.script[i], rows, axis=0)
        self.endRemoveRows()

    def flags(self, index: QModelIndex):
        flags = Qt.ItemFlag.ItemIsSelectable
        if index.row() >= self.led_script.get_current_index(self.dt):
            flags |= Qt.ItemFlag.ItemIsEnabled
            if index.column() in [0, 2]:
                flags |= Qt.ItemFlag.ItemIsEditable
        return flags




class LedScriptViewer(QWidget):
    def __init__(self, led_script: LedScript, parent=None):
        super().__init__(parent)
        self.led_script = None
        self.model = LedScriptTableModel(led_script, self)

        self.l_vbox = QVBoxLayout()
        self.setLayout(self.l_vbox)

        self.l_vbox.addWidget(QLabel("You may edit the dt and led values of future rows."))

        self.w_scroll = QScrollArea()
        self.l_vbox.addWidget(self.w_scroll)

        self.w_table = QTableView(self)
        self.w_table.setModel(self.model)
        self.w_table.show()
        self.w_scroll.setWidget(self.w_table)

        self.w_scroll.setWidgetResizable(True)
        self.w_scroll.setFixedHeight(200)
        self.l_vbox.addStretch(1)

    def set_script(self, led_script: LedScript):
        self.model = LedScriptTableModel(led_script, self)
        self.w_table.setModel(self.model)
        self.w_table.show()

    def set_relative_time(self, time: float):
        pass