import time
import numpy as np

from PyQt6.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QSizePolicy
from PyQt6.QtGui import QImage, QPixmap
from PyQt6.QtCore import Qt, QTimer

from logger import logger


class Visualize(QWidget):
    def __init__(self, get_pybullet_image):
        super().__init__()
        self.get_pybullet_image = get_pybullet_image

        # Настройка компоновки
        self.layout = QVBoxLayout(self)
        self.layout.setContentsMargins(0, 0, 0, 0)  # Убираем отступы по краям

        self.h_layout = QHBoxLayout()
        self.h_layout.setContentsMargins(0, 0, 0, 0)  # Убираем отступы по краям

        self.label = QLabel(self)
        self.layout.addWidget(self.label)

        # Таймер для обновления изображения
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_image)
        self.timer.start(30) 

    def update_image(self):
        res = self.get_pybullet_image()

        if not res:
            return

        (rgb, width, height) = res
        bytesPerLine = 2 * width
        # print(type(rgb))
        image = QImage(np.asarray(rgb), width, height, bytesPerLine, QImage.Format.Format_RGBA8888).rgbSwapped()
        image = image.convertToFormat(QImage.Format.Format_RGB888)
        pixmap = QPixmap.fromImage(image)

        self.label.setPixmap(pixmap)
        self.label.setAlignment(Qt.AlignmentFlag.AlignHCenter)
        self.label.update()

    def paintEvent(self, event):
        p = self.palette()
        p.setColor(self.backgroundRole(), Qt.GlobalColor.magenta)
        self.setPalette(p)
        super().paintEvent(event)