from PySide6.QtWidgets import (
    QApplication, QGraphicsView, QGraphicsScene, QGraphicsRectItem, QGraphicsSimpleTextItem
)
from PySide6.QtGui import QBrush, QColor, QPainter
from PySide6.QtCore import Qt, QTimer
import sys
import random


class KeyItem(QGraphicsRectItem):
    def __init__(self, label, x, y, width, height):
        super().__init__(x, y, width, height)
        self.label = label
        self.default_size = (width, height)
        self.original_pos = (x, y)
        self.setRect(0, 0, width, height)
        self.setPos(x, y)

        self.text = QGraphicsSimpleTextItem(label, self)
        self.text.setPos(10, 5)
        self.setBrush(QBrush(QColor("lightgray")))
        self.setFlag(QGraphicsRectItem.ItemIsSelectable)

    def set_scale_factor(self, scale):
        self.setZValue(scale)
        self.setTransformOriginPoint(self.default_size[0] / 2, self.default_size[1] / 2)
        self.setScale(scale)

    def mousePressEvent(self, QMouseEvent):
        print(self.label, "was clicked")


class KeyboardScene(QGraphicsScene):
    def __init__(self):
        super().__init__()
        self.keys = {}

        self.padding = 20  # <-- space around keys to prevent shifting
        self.total_width = 450  # estimate based on layout
        self.total_height = 200
        self.setSceneRect(-self.padding, -self.padding,
                          self.total_width + 2 * self.padding,
                          self.total_height + 2 * self.padding)

        self.layout_keys()

        self.timer = QTimer()
        self.timer.timeout.connect(self.simulate_prediction)
        self.timer.start(2000)

    def layout_keys(self):
        letters = "QWERTYUIOPASDFGHJKLZXCVBNM"
        x, y = 0, 0
        size = 40
        spacing = 5
        for i, char in enumerate(letters):
            if i == 10 or i == 19:
                y += size + spacing
                x = 0
            key = KeyItem(char, x, y, size, size)
            self.addItem(key)
            self.keys[char] = key
            x += size + spacing

    def simulate_prediction(self):
        most_likely = random.choice(list(self.keys.keys()))
        print(f"[Prediction] Most likely: {most_likely}")

        for char, key in self.keys.items():
            if char == most_likely:
                key.set_scale_factor(1.8)
                key.setBrush(QBrush(QColor("orange")))
            else:
                key.set_scale_factor(1.0)
                key.setBrush(QBrush(QColor("lightgray")))


class KeyboardView(QGraphicsView):
    def __init__(self):
        super().__init__()
        self.setScene(KeyboardScene())
        self.setRenderHint(QPainter.Antialiasing)
        self.setWindowTitle("Dynamic Keyboard")
        self.setAlignment(Qt.AlignLeft | Qt.AlignTop)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    view = KeyboardView()
    view.show()
    sys.exit(app.exec())