import os
import cv2
from ultralytics import YOLO

# --------------------------
# Configuration
# --------------------------
model_path = "models/bck/best-3.pt"  # your trained YOLO model
images_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/universe/train/images"
labels_dir = "C:/Users/Laurent/Desktop/board-mate/rpi/training/datasets/universe/train/labels"
img_width = 640
img_height = 640

os.makedirs(labels_dir, exist_ok=True)

# --------------------------
# Load model
# --------------------------
model = YOLO(model_path)

# --------------------------
# Mapping YOLO class index -> piece name (optional)
# --------------------------
names = ['w_pawn','w_knight','w_bishop','w_rook','w_queen','w_king',
         'b_pawn','b_knight','b_bishop','b_rook','b_queen','b_king']

# --------------------------
# Process images
# --------------------------
for img_file in os.listdir(images_dir):
    if not img_file.lower().endswith((".png", ".jpg", ".jpeg")):
        continue

    img_path = os.path.join(images_dir, img_file)
    img = cv2.imread(img_path)
    if img is None:
        print(f"Failed to read {img_file}")
        continue

    height, width = img.shape[:2]

    # Run YOLO detection
    results = model(img)
    res = results[0]

    lines = []
    boxes = res.boxes.xyxy.cpu().numpy()  # [x1, y1, x2, y2]
    classes = res.boxes.cls.cpu().numpy()
    confs = res.boxes.conf.cpu().numpy()

    for box, cls, conf in zip(boxes, classes, confs):
        if conf < 0.5:  # skip low-confidence predictions
            continue

        x1, y1, x2, y2 = box
        x_center = (x1 + x2) / 2 / width
        y_center = (y1 + y2) / 2 / height
        w_norm = (x2 - x1) / width
        h_norm = (y2 - y1) / height

        lines.append(f"{int(cls)} {x_center:.6f} {y_center:.6f} {w_norm:.6f} {h_norm:.6f}")

    # Save YOLO .txt file with same basename as image
    txt_path = os.path.join(labels_dir, os.path.splitext(img_file)[0] + ".txt")
    with open(txt_path, "w") as f:
        f.write("\n".join(lines))

    print(f"Pre-labeled {img_file} -> {txt_path}")

print("All images have been pre-labeled!")