#!/usr/bin/env python3
import os
import random

import numpy as np
from ultralytics import YOLO
import cv2
from ultralytics.engine.results import Results


class Detector :

    def __init__(self, model_path):
        self.model = YOLO(model_path)
        self.used_height = 640
        self.used_width = 640

    def make_prediction(self, image : str | np.ndarray) -> list[Results]:
        return self.model.predict(source=image, conf=0.6)

if __name__ == "__main__":
    corner_model_path = "../../assets/models/edges.pt"
    pieces_model_path = "../../assets/models/unified-nano-refined.pt"

    corner_model = YOLO(corner_model_path)
    pieces_model = YOLO(pieces_model_path)

    img_folder = "../training/datasets/pieces/unified/test/images/"
    save_folder = "./results"
    os.makedirs(save_folder, exist_ok=True)

    test_images = os.listdir(img_folder)

    for i in range(0, 10):
        rnd = random.randint(0, len(test_images) - 1)
        img_path = os.path.join(img_folder, test_images[rnd])
        save_path = os.path.join(save_folder, test_images[rnd])

        image = cv2.imread(img_path)
        height, width = image.shape[:2]

        #fen = prediction_to_fen(results, height, width)
        #print("Predicted FEN:", fen)

        corner_result = corner_model.predict(source=image, conf=0.6)
        pieces_result = pieces_model.predict(source=image, conf=0.6)

        corner_annotated_image = corner_result[0].plot()
        pieces_annotated_image = pieces_result[0].plot(img=corner_annotated_image)

        cv2.imwrite(save_path, pieces_annotated_image)
        #cv2.namedWindow("YOLO Predictions", cv2.WINDOW_NORMAL)
        #cv2.imshow("YOLO Predictions", annotated_image)

        cv2.waitKey(0)
        cv2.destroyAllWindows()