import time

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

import numpy as np
from PIL import Image

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.label.setScaledContents(True)
        self.label.setSizePolicy(
            QSizePolicy.Expanding, QSizePolicy.Expanding
        )
        self.layout.addWidget(self.label)

        # Таймер для обновления изображения
        self.timer = QTimer()
        self.timer.timeout.connect(self.update_image)
        self.timer.start(int(1000/240))  # Обновление каждые 100 мс

    def update_image(self):
        res = self.get_pybullet_image()
        
        if not res: return
        
        (rgb, width, height) = res

        image = QImage(rgb, width, height, QImage.Format_RGBA8888)
        image = image.convertToFormat(QImage.Format_RGB888)
        pixmap = QPixmap.fromImage(image)

        self.label.setPixmap(pixmap)
        self.label.resize(image.size()) 
        self.label.repaint()

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