import os
import sys
import cv2
import numpy as np
import argparse
import matplotlib.pyplot as plt


def get_image_resolution(dataset_path):
    for root, dirs, files in os.walk(dataset_path):
        for file in files:
            if file.lower().endswith((".png", ".jpg", ".jpeg")):
                img_path = os.path.join(root, file)
                img = cv2.imread(img_path)
                if img is not None:
                    height, width = img.shape[:2]
                    return width, height
    raise ValueError("No valid image file found to determine resolution.")


def extract_wh_from_label(label_path, img_w, img_h):
    wh_list = []
    with open(label_path, "r") as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) != 5:
                continue
            _, _, _, w, h = map(float, parts)
            abs_w = w * img_w
            abs_h = h * img_h
            wh_list.append([abs_w, abs_h])
    return wh_list


def collect_boxes(dataset_path, use_scaled=False):
    wh_data = []
    img_w, img_h = get_image_resolution(dataset_path)
    for root, dirs, files in os.walk(dataset_path):
        if os.path.basename(root) == "labels":
            for file in files:
                if file.endswith(".txt"):
                    label_path = os.path.join(root, file)
                    if os.path.getsize(label_path) == 0:
                        continue
                    wh = extract_wh_from_label(label_path, img_w, img_h)
                    if use_scaled:
                        # Convert to fraction of image dimensions
                        wh = [[w / img_w, h / img_h] for w, h in wh]
                    wh_data.extend(wh)
    if not wh_data:
        raise ValueError("No non-empty label files found in the dataset.")
    return np.array(wh_data, dtype=np.float32), img_w, img_h


def visualize_boxes(boxes, img_w, img_h, use_scaled=False):
    widths = boxes[:, 0]
    heights = boxes[:, 1]
    areas = widths * heights
    plt.figure(figsize=(12, 5))
    # Scatter plot of width vs height
    plt.subplot(1, 2, 1)
    plt.scatter(widths, heights, alpha=0.6)
    if use_scaled:
        plt.xlim(0, 1)
        plt.ylim(0, 1)
        plt.xlabel("Width (fraction of image width)")
        plt.ylabel("Height (fraction of image height)")
        plt.title("Box Size Relative to Image Dimensions")
    else:
        plt.xlabel("Width (pixels)")
        plt.ylabel("Height (pixels)")
        plt.title("Box Width vs Height")
    plt.grid(True)
    # Histogram of box areas
    plt.subplot(1, 2, 2)
    plt.hist(areas, bins=20, color='orange', alpha=0.7)
    if use_scaled:
        plt.xlabel("Box area (fraction of image area)")
        plt.title("Histogram of Box Areas (Relative)")
    else:
        plt.xlabel("Box area (pixels^2)")
        plt.title("Histogram of Box Areas")
    plt.ylabel("Count")
    plt.grid(True)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Visualize the distinct box sizes in a YOLO dataset."
    )
    parser.add_argument(
        "dataset_path",
        type=str,
        help="Path to dataset directory (must contain a 'labels' folder)."
    )
    parser.add_argument(
        "--scaled",
        action="store_true",
        help="Use relative image dimensions (0–1 scale) instead of absolute pixel sizes"
    )
    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)
    args = parser.parse_args()
    boxes, img_w, img_h = collect_boxes(args.dataset_path, args.scaled)
    visualize_boxes(boxes, img_w, img_h, args.scaled)