import cv2  # помщник, который умеет открывать видео, показывать его и сохранять
from ultralytics import YOLO  # умеет находить вещи


def detection_video(model: str, video_path: str) -> None:
    COUNT_OBJECTS: int = 0
    model = YOLO(model=model)  # загружается уже обученая нейросеть
    VEHICLE_CLASSES: dict[int, str] = {
        2: "car",
        3: "motorcycle",
        5: "bus",
        7: "truck",
    }  # какие обьекты надо искать
    cap = cv2.VideoCapture(video_path)  # создается обьект который умеет читать видео

    # Сохранение результата
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # выхватываем размеры из видео
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # выхватываем размеры из видео
    fps = cap.get(cv2.CAP_PROP_FPS)  # сколько кадров показывается за секунд

    writer = cv2.VideoWriter(
        "output2.mp4", cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height)
    )  # создаем новоё видео, на старте оно пустое, как альбом в который будем вклеивать кадры
    while cap.isOpened():
        ret, frame = cap.read()  # читаем кадр
        if not ret:
            break

        results = model(
            frame, verbose=False
        )  # передаем кадр в модель. то есть мы говорим нейросети найти что то
        for result in results:
            boxes = result.boxes

            for box in boxes:
                cls = int(box.cls[0])
                if cls not in VEHICLE_CLASSES:
                    continue
                COUNT_OBJECTS += 1
                conf = float(box.conf[0])  # Вероятность
                x1, y1, x2, y2 = map(int, box.xyxy[0])  # Кординаты
                label: str = f"{VEHICLE_CLASSES[cls]} {conf:.2f} | ID: {COUNT_OBJECTS}"  # Подпись
                cv2.rectangle(
                    frame, (x1, y1), (x2, y2), (0, 255, 0), 2
                )  # рисуем прямоугольник
                cv2.putText(
                    frame,
                    label,
                    (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6,
                    (0, 255, 0),
                    2,
                )  # пишем текст, типо что обьект
        writer.write(frame)  # сохраняем кадр, бёрем кадр и кладём его в новое видео
        cv2.imshow("DETECTION", frame)  # показыает оработанное видео
        if cv2.waitKey(1) & 0xFF == ord("q"):  # комманда для выхода из программы
            break
    cap.release()  # закрываем файл видео
    writer.release()  # закрывем файл, в который записывали результат.
    cv2.destroyAllWindows()  # закрываем все окна OpenCV

    print("===| TYT |===")


def detection_image(model: str, image_path: str) -> None:
    COUNT_OBJECTS: int = 0
    model = YOLO(model=model)
    VEHICLE_CLASSES: dict[int, str] = {
        2: "car",
        3: "motorcycle",
        5: "bus",
        7: "truck",
    }

    image = cv2.imread(filename=image_path)
    if image is None:
        raise FileNotFoundError("Файл не найден")
    results = model(image, verbose=False)
    for result in results:
        for box in result.boxes:
            cls: int = int(box.cls[0])
            if cls not in VEHICLE_CLASSES:
                continue
            conf: float = float(box.conf[0])
            # кординаты рамки
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            COUNT_OBJECTS += 1
            label = f"{VEHICLE_CLASSES[cls]} {conf:.2f} | ID: {COUNT_OBJECTS}"

            # рисуем рамку
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(
                img=image,
                text=label,
                org=(x1, y1 - 10),
                fontFace=cv2.FONT_HERSHEY_SIMPLEX,
                fontScale=0.6,
                color=(0, 255, 0),
                thickness=2,
            )

    cv2.imwrite("output.jpg", image)
    cv2.imshow("DETECTION", image)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def main() -> None:
    # detection_video(model="yolov8n.pt", video_path="video/1.mp4")
    detection_image(model="yolov8n.pt", image_path="images/1.jpg")


if __name__ == "__main__":
    print("===| RUN |===")
    main()