import os
import cv2
import argparse

# YOLO head strides
YOLO_STRIDES = {
    1: 32,  # Head 1
    2: 16,  # Head 2
    3: 8,   # Head 3
}

# Colors per head (OpenCV uses BGR)
HEAD_COLORS = {
    1: (0, 0, 255),     # Red
    2: (255, 0, 0),     # Blue
    3: (0, 255, 255),   # Yellow
}


def draw_grid(img, grid_w, grid_h, color=(0, 255, 255)):
    h, w = img.shape[:2]
    step_x = w / grid_w
    step_y = h / grid_h
    for i in range(1, grid_w):
        x = int(i * step_x)
        cv2.line(img, (x, 0), (x, h), color, 1)
    for j in range(1, grid_h):
        y = int(j * step_y)
        cv2.line(img, (0, y), (w, y), color, 1)


def draw_boxes(img, label_path):
    h, w = img.shape[:2]
    if not os.path.exists(label_path):
        return
    with open(label_path, "r") as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) != 5:
                continue
            _, cx, cy, bw, bh = map(float, parts)
            x1 = int((cx - bw / 2) * w)
            y1 = int((cy - bh / 2) * h)
            x2 = int((cx + bw / 2) * w)
            y2 = int((cy + bh / 2) * h)
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)


def view_dataset(dataset_root, head, net_w, net_h):
    if head not in YOLO_STRIDES:
        raise ValueError(f"Invalid head: {head}. Must be 1, 2, or 3.")
    stride = YOLO_STRIDES[head]
    grid_w = net_w // stride
    grid_h = net_h // stride
    color = HEAD_COLORS[head]
    # Loop through all subfolders (each should be a dataset)
    for folder_name in sorted(os.listdir(dataset_root)):
        dataset_path = os.path.join(dataset_root, folder_name)
        if not os.path.isdir(dataset_path):
            continue
        images_dir = os.path.join(dataset_path, "images")
        labels_dir = os.path.join(dataset_path, "labels")
        if not os.path.isdir(images_dir) or not os.path.isdir(labels_dir):
            continue  # Skip folders that aren't YOLO datasets
        print(f"\nViewing dataset: {folder_name}")
        image_files = sorted(
            f for f in os.listdir(images_dir)
            if f.lower().endswith((".jpg", ".jpeg", ".png"))
        )
        for img_name in image_files:
            img_path = os.path.join(images_dir, img_name)
            label_path = os.path.join(
                labels_dir, os.path.splitext(img_name)[0] + ".txt"
            )
            img = cv2.imread(img_path)
            if img is None:
                continue
            img_resized = cv2.resize(img, (net_w, net_h))
            draw_boxes(img_resized, label_path)
            draw_grid(img_resized, grid_w, grid_h, color=color)
            window_label = f"{folder_name} | Head {head} Grid {grid_w} x{grid_h}"
            cv2.imshow(window_label, img_resized)
            key = cv2.waitKey(0)
            if key == 27:  # ESC
                cv2.destroyAllWindows()
                return
    cv2.destroyAllWindows()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Visualize YOLO grid for a given head and network resolution"
    )
    parser.add_argument("dataset_path", type=str, help="Path to YOLO dataset root")
    parser.add_argument(
        "--head", type=int, choices=[1, 2, 3], required=True,
        help="YOLO head (1, 2, or 3)"
    )
    parser.add_argument("--net_w", type=int, default=416, help="Network width")
    parser.add_argument("--net_h", type=int, default=416, help="Network height")
    args = parser.parse_args()
    view_dataset(args.dataset_path, args.head, args.net_w, args.net_h)