"""LabelMe JSON에서 특정 레이블 폴리곤만 원본 이미지에 표시.

사용:
    python tools/render_label_polygons.py \
        --image data/역사구간/.../DJI_20260306100900_0034.JPG \
        --json output/detect/DJI_20260306100900_0034/tiles1to24_railway_zone_001.json \
        --label catenary_pole \
        --output output/detect/DJI_20260306100900_0034/catenary_pole_only.jpg

    # 여러 레이블:
    python tools/render_label_polygons.py ... --label catenary_pole bracket

    # 레이블 생략 시 전체 표시:
    python tools/render_label_polygons.py ... --json foo.json --image foo.jpg
"""
import argparse
import json
import cv2
import numpy as np
from pathlib import Path


COLORS = [
    (0, 0, 255),    # red
    (0, 200, 0),    # green
    (255, 100, 0),  # blue-orange
    (0, 200, 255),  # yellow
    (200, 0, 200),  # magenta
    (0, 165, 255),  # orange
    (255, 0, 100),  # pink-blue
    (100, 255, 100),# light green
]


def render(image_path: Path, json_path: Path, output_path: Path, labels: list[str] | None):
    buf = np.fromfile(str(image_path), dtype=np.uint8)
    img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(f"이미지 읽기 실패: {image_path}")

    with open(json_path, encoding="utf-8") as f:
        data = json.load(f)

    shapes = data.get("shapes", [])
    if labels:
        shapes = [s for s in shapes if s["label"] in labels]

    # 레이블별 색상 매핑
    label_set = sorted({s["label"] for s in shapes})
    color_map = {lbl: COLORS[i % len(COLORS)] for i, lbl in enumerate(label_set)}

    for s in shapes:
        pts = np.array(s["points"], dtype=np.int32).reshape((-1, 1, 2))
        color = color_map[s["label"]]
        cv2.polylines(img, [pts], isClosed=True, color=color, thickness=8)
        x, y = int(s["points"][0][0]), int(s["points"][0][1])
        score = s.get("score")
        text = f"{s['label']} {score:.2f}" if score is not None else s["label"]
        cv2.putText(img, text, (x, max(y - 10, 20)),
                    cv2.FONT_HERSHEY_SIMPLEX, 2.5, color, 5)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    ret, buf_out = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 85])
    if not ret:
        raise RuntimeError("JPEG 인코딩 실패")
    output_path.write_bytes(buf_out.tobytes())
    print(f"저장: {output_path}  ({len(shapes)}개 폴리곤)")


def main():
    p = argparse.ArgumentParser(description="LabelMe JSON 폴리곤 시각화")
    p.add_argument("--image", required=True, type=Path)
    p.add_argument("--json",  required=True, type=Path)
    p.add_argument("--output", required=True, type=Path)
    p.add_argument("--label", nargs="*", default=None,
                   help="표시할 레이블 (생략 시 전체)")
    args = p.parse_args()

    render(args.image, args.json, args.output, args.label)


if __name__ == "__main__":
    main()