#!/usr/bin/env python3
"""
show_yolo_seg.py

Usage:
    python show_yolo_seg.py /path/to/image.jpg /path/to/labels.txt

Supports:
 - Segmentation polygons: "class x1 y1 x2 y2 ... xn yn"
 - YOLO bbox lines as fallback: "class x_center y_center width height"
Coordinates can be normalized [0..1] or absolute pixels (auto-detected).
"""
import sys
import cv2
import numpy as np
import matplotlib.pyplot as plt
import argparse
from pathlib import Path
import random


def parse_label_line(line):
    parts = line.strip().split()
    if not parts:
        return None
    cls = int(float(parts[0]))
    coords = [float(x) for x in parts[1:]]
    return cls, coords


def coords_are_normalized(coords):
    # If every coordinate is between 0 and 1 (inclusive-ish), assume normalized
    if not coords:
        return False
    return max(coords) <= 1.001


def yolo_bbox_to_xyxy(coords, img_w, img_h):
    # coords: [xc, yc, w, h] normalized or absolute
    xc, yc, w, h = coords[:4]
    if max(coords) <= 1.001:
        xc *= img_w
        yc *= img_h
        w *= img_w
        h *= img_h
    x1 = int(round(xc - w / 2))
    y1 = int(round(yc - h / 2))
    x2 = int(round(xc + w / 2))
    y2 = int(round(yc + h / 2))
    return x1, y1, x2, y2


def poly_to_pts(coords, img_w, img_h):
    # coords: [x1 y1 x2 y2 ...] either normalized or absolute
    if coords_are_normalized(coords):
        coords = [
            coords[i] * (img_w if i % 2 == 0 else img_h) for i in range(len(coords))
        ]
    pts = np.array(coords, dtype=np.int32).reshape(-1, 2)
    return pts


def random_color_for_class(cls):
    random.seed(cls)  # deterministic per class
    return tuple(int(x) for x in np.array([random.randint(0, 255) for _ in range(3)]))


def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
    # img: BGR numpy array
    overlay = img.copy()
    h, w = img.shape[:2]
    for cls, coords in labels:
        if not coords:
            continue
        # polygon case (>=6 coordinates)
        if len(coords) >= 6:
            pts = poly_to_pts(coords, w, h)
            color = random_color_for_class(cls)
            # fill on overlay
            cv2.fillPoly(overlay, [pts], color)
            # outline on base image
            cv2.polylines(img, [pts], isClosed=True, color=color, thickness=2)
            # put class text at first point
            x, y = int(pts[0, 0]), int(pts[0, 1]) - 6
            cv2.putText(
                img,
                str(cls),
                (x, max(6, y)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (255, 255, 255),
                2,
                cv2.LINE_AA,
            )
            if draw_bbox_for_poly:
                x, y, w_box, h_box = cv2.boundingRect(pts)
                cv2.rectangle(img, (x, y), (x + w_box, y + h_box), color, 1)
        # YOLO bbox case (4 coords)
        elif len(coords) == 4:
            x1, y1, x2, y2 = yolo_bbox_to_xyxy(coords, w, h)
            color = random_color_for_class(cls)
            cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
            cv2.putText(
                img,
                str(cls),
                (x1, max(6, y1 - 4)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (255, 255, 255),
                2,
                cv2.LINE_AA,
            )
        else:
            # Unknown / invalid format, skip
            continue

    # blend overlay for filled polygons
    cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0, img)
    return img


def load_labels_file(label_path):
    labels = []
    with open(label_path, "r") as f:
        for raw in f:
            line = raw.strip()
            if not line:
                continue
            parsed = parse_label_line(line)
            if parsed:
                labels.append(parsed)
    return labels


def main():
    parser = argparse.ArgumentParser(
        description="Show YOLO segmentation / polygon annotations"
    )
    parser.add_argument("image", type=str, help="Path to image file")
    parser.add_argument("labels", type=str, help="Path to YOLO label file (polygons)")
    parser.add_argument(
        "--alpha", type=float, default=0.4, help="Polygon fill alpha (0..1)"
    )
    parser.add_argument(
        "--no-bbox", action="store_true", help="Don't draw bounding boxes for polygons"
    )
    args = parser.parse_args()

    img_path = Path(args.image)
    lbl_path = Path(args.labels)

    if not img_path.exists():
        print("Image not found:", img_path)
        sys.exit(1)
    if not lbl_path.exists():
        print("Label file not found:", lbl_path)
        sys.exit(1)

    img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
    if img is None:
        print("Could not load image:", img_path)
        sys.exit(1)

    labels = load_labels_file(str(lbl_path))
    if not labels:
        print("No labels parsed from", lbl_path)
        # continue and just show image
    out = draw_annotations(
        img.copy(), labels, alpha=args.alpha, draw_bbox_for_poly=(not args.no_bbox)
    )

    # Convert BGR -> RGB for matplotlib display
    out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB)
    plt.figure(figsize=(10, 10 * out.shape[0] / out.shape[1]))
    plt.imshow(out_rgb)
    plt.axis("off")
    plt.title(f"{img_path.name} ({lbl_path.name})")
    plt.show()


if __name__ == "__main__":
    main()