import os
import sys
import cv2
import numpy as np
import argparse

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 iou(box, clusters):
    x = np.minimum(clusters[:, 0], box[0])
    y = np.minimum(clusters[:, 1], box[1])
    intersection = x * y
    box_area = box[0] * box[1]
    cluster_area = clusters[:, 0] * clusters[:, 1]
    return intersection / (box_area + cluster_area - intersection + 1e-10)


def avg_iou(boxes, clusters):
    return np.mean([np.max(iou(box, clusters)) for box in boxes])


def iou_kmeans(boxes, k, max_iter=1000, tol=1e-6):
    np.random.seed(42)
    # Initialize clusters by randomly sampling boxes
    clusters = boxes[np.random.choice(len(boxes), k, replace=False)]
    # Start unassigned so the first median update is never skipped.
    last_clusters = np.full((len(boxes),), -1, dtype=int)
    for iteration in range(max_iter):
        distances = np.zeros((len(boxes), k))
        for i in range(k):
            cluster = clusters[i]
            inter = np.minimum(boxes[:, 0], cluster[0]) * np.minimum(boxes[:, 1], cluster[1])
            box_area = boxes[:, 0] * boxes[:, 1]
            cluster_area = cluster[0] * cluster[1]
            iou_vals = inter / (box_area + cluster_area - inter + 1e-10)
            distances[:, i] = 1 - iou_vals
        nearest_clusters = np.argmin(distances, axis=1)
        if (last_clusters == nearest_clusters).all():
            break
        for i in range(k):
            if np.any(nearest_clusters == i):
                clusters[i] = np.median(boxes[nearest_clusters == i], axis=0)
        last_clusters = nearest_clusters
    return clusters, iteration + 1


def get_all_wh_from_dataset(dataset_path, num_anchors, net_w, net_h):
    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)
                    wh_data.extend(wh)
    if not wh_data:
        raise ValueError("No non-empty label files found in the dataset.")
    wh_data_np = np.array(wh_data, dtype=np.float32)
    clusters, iterations = iou_kmeans(wh_data_np, num_anchors)
    scale_w = net_w / img_w
    scale_h = net_h / img_h
    anchors_scaled = np.round(clusters * [scale_w, scale_h]).astype(int)
    anchors_scaled = anchors_scaled[anchors_scaled[:, 0].argsort()]
    print("anchors=" + ",".join(f"{w},{h}" for w, h in anchors_scaled))
    mean_iou = avg_iou(wh_data_np, clusters)
    print(f"Average IoU: {mean_iou:.4f}")
    print(f"Iterations until convergence: {iterations}")
    return anchors_scaled


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Generate YOLO anchors using IoU-based k-means clustering."
    )
    parser.add_argument("dataset_path", type=str, help="Path to dataset directory (must contain a 'labels' folder).")
    parser.add_argument("--net_w", type=int, required=True, help="Network width (e.g., 416)")
    parser.add_argument("--net_h", type=int, required=True, help="Network height (e.g., 416)")
    parser.add_argument("--num_anchors", type=int, default=9, help="Number of anchors to generate (default: 9).")
    if len(sys.argv) == 1:
        parser.print_help(sys.stderr)
        sys.exit(1)
    args = parser.parse_args()
    get_all_wh_from_dataset(args.dataset_path, num_anchors=args.num_anchors, net_w=args.net_w, net_h=args.net_h)