import time

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

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(1000)  # Обновление каждые 100 мс

    def update_image(self):
        (rgb, width, height) = self.get_pybullet_image()

        # logger.info(f"Image size: {width}x{height}, RGB data length: {len(rgb)}")

        # rgbim = Image.fromarray(rgb)
        # rgbim.save(f"test_img/rgbtest/{time.time() }.png")
        # Преобразуем RGB данные в массив NumPy
        rgb_array = np.array(rgb, dtype=np.uint8)

        # Изменяем форму массива на (height, width, 4) для QImage
        rgb_array = rgb_array.reshape(height, width, 4)
        # Преобразование RGB-данных в QImage
        image = QImage(rgb_array, width, height, QImage.Format_RGBA8888)
        image = image.convertToFormat(QImage.Format_RGB888)
        
        pixmap = QPixmap.fromImage(image)

        # Обновление изображения на QLabel
        self.label.setPixmap(pixmap)
        self.label.repaint()

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