"""Teach an ROI once (human draws), persist it as a .hobj for headless reuse.

The rectangle case (default, the common "box a Mark" flow) uses an interactive
matplotlib viewer with **pan (right-drag), zoom (scroll wheel), and box-select
(left-drag)** — smooth even on multi-megapixel images. circle/region fall back to
HALCON's native draw_* window. Downstream procedures read_region + reduce_domain
the saved .hobj, so no human is needed after teaching.

Run (needs a visible desktop — this is the interactive seam):
    /mnt/c/Users/NAURA/.local/bin/uv.exe run python \
        .claude/skills/halcon/scripts/teach_roi.py \
        --image-file "C:/data/frames/a.bmp" \
        --out C:/workspace/agent-studio/halcon-001/roi.hobj

Controls (rectangle): left-drag = box · right-drag = pan · wheel = zoom ·
Enter = save · q/Esc = cancel.

Headless (no GUI, scripted): add --box r1,c1,r2,c2 to write a rectangle directly.

Requires matplotlib + numpy in the venv (one-time):
    /mnt/c/Users/NAURA/.local/bin/uv.exe pip install matplotlib numpy

Tip: draw TIGHTLY around the target features. A loose ROI captures extra
contours and breaks downstream count checks (e.g. "expected 16 circles").
"""
import argparse
import sys


def pick_reference_image(ha, image_dir, image_file):
    if image_file:
        return image_file
    files = ha.list_files(image_dir, "files")
    imgs = sorted(ha.tuple_regexp_select(files, r"\.(png|bmp|jpg|jpeg|tif|tiff)$"))
    if not imgs:
        raise SystemExit(f"No images found in {image_dir}")
    return imgs[0]


def load_gray(ha, image_file):
    # HALCON 默认 clip_region=true 会把区域裁到 128x128,全分辨率坐标建的 region
    # 会被裁空 -> 写盘回读 bbox=0,0,0,0。建/读区域前务必关掉。
    ha.set_system("clip_region", "false")
    img = ha.read_image(image_file)
    if ha.count_channels(img)[0] >= 3:
        img = ha.rgb1_to_gray(img)
    w, h = ha.get_image_size(img)
    return ha.himage_as_numpy_array(img), int(h[0]), int(w[0])


def pick_rectangle_interactive(arr, max_disp):
    """matplotlib pan/zoom/box-select; returns (r1,c1,r2,c2) in image coords or None."""
    import matplotlib.pyplot as plt
    from matplotlib.widgets import RectangleSelector

    plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei", "DejaVu Sans"]
    plt.rcParams["axes.unicode_minus"] = False

    # Downsample for display (fast redraws), map coords back to full res via extent.
    H, W = arr.shape
    step = max(1, round(max(H, W) / max_disp))
    disp = arr[::step, ::step]

    fig, ax = plt.subplots()
    ax.imshow(disp, cmap="gray", origin="upper", extent=(0, W, H, 0),
              interpolation="nearest", aspect="equal")
    ax.set_xlim(0, W); ax.set_ylim(H, 0)
    ax.set_title("left=box  right-drag=pan  wheel=zoom  Enter=save  q=cancel")

    box, pan = {}, {}

    def on_select(p, r):
        box.update(r1=p.ydata, c1=p.xdata, r2=r.ydata, c2=r.xdata)

    # keep the reference or RectangleSelector is GC'd and box-select stops working
    ax._roi_selector = RectangleSelector(
        ax, on_select, useblit=True, button=[1], minspanx=3, minspany=3,
        interactive=True)

    def on_scroll(ev):
        if ev.inaxes is not ax:
            return
        s = 0.8 if ev.button == "up" else 1.25
        x0, x1 = ax.get_xlim(); y0, y1 = ax.get_ylim()
        ax.set_xlim(ev.xdata - (ev.xdata - x0) * s, ev.xdata + (x1 - ev.xdata) * s)
        ax.set_ylim(ev.ydata - (ev.ydata - y0) * s, ev.ydata + (y1 - ev.ydata) * s)
        fig.canvas.draw_idle()

    def on_press(ev):
        if ev.button == 3 and ev.inaxes is ax:
            pan["p"] = (ev.xdata, ev.ydata)

    def on_motion(ev):
        if "p" in pan and ev.xdata is not None and ev.inaxes is ax:
            dx, dy = ev.xdata - pan["p"][0], ev.ydata - pan["p"][1]
            x0, x1 = ax.get_xlim(); y0, y1 = ax.get_ylim()
            ax.set_xlim(x0 - dx, x1 - dx); ax.set_ylim(y0 - dy, y1 - dy)
            fig.canvas.draw_idle()

    def on_release(ev):
        if ev.button == 3:
            pan.pop("p", None)

    def on_key(ev):
        if ev.key == "enter" and box:
            plt.close(fig)
        elif ev.key in ("q", "escape"):
            box.clear(); plt.close(fig)

    for name, fn in [("scroll_event", on_scroll), ("button_press_event", on_press),
                     ("motion_notify_event", on_motion),
                     ("button_release_event", on_release), ("key_press_event", on_key)]:
        fig.canvas.mpl_connect(name, fn)
    plt.show()
    if not box:
        return None
    return box["r1"], box["c1"], box["r2"], box["c2"]


def save_rectangle(ha, out, r1, c1, r2, c2, H, W):
    r1, r2 = sorted((max(0, round(r1)), min(H - 1, round(r2))))
    c1, c2 = sorted((max(0, round(c1)), min(W - 1, round(c2))))
    ha.write_region(ha.gen_rectangle1(r1, c1, r2, c2), out)
    print(f"ROI saved: {out}")
    print(f"Rectangle Row1={r1} Col1={c1} Row2={r2} Col2={c2}")


def draw_native(ha, ref, shape, out):
    """circle / region via HALCON's own draw window (no pan/zoom)."""
    img = ha.read_image(ref)
    width, height = ha.get_image_size(img)
    width, height = width[0], height[0]
    win = ha.open_window(0, 0, 1024, 768, 0, "visible", "")
    ha.set_part(win, 0, 0, height - 1, width - 1)
    ha.disp_obj(img, win)
    ha.set_color(win, "green")
    print(f"Draw a {shape} with the left mouse button, release to finish ...", flush=True)
    if shape == "circle":
        row, col, radius = ha.draw_circle(win)
        region = ha.gen_circle(row, col, radius)
        print(f"Circle Row={row:.1f} Col={col:.1f} Radius={radius:.1f}")
    else:  # freehand region
        region = ha.draw_region(win)
    ha.write_region(region, out)
    ha.close_window(win)
    print(f"ROI saved: {out}")


def main() -> int:
    ap = argparse.ArgumentParser(description="Teach an ROI and write it to disk")
    src = ap.add_mutually_exclusive_group(required=True)
    src.add_argument("--image-dir", help="Directory; first sorted image is shown")
    src.add_argument("--image-file", help="Explicit reference image to display")
    ap.add_argument("--out", required=True, help="Output region file, e.g. .../roi.hobj")
    ap.add_argument("--shape", choices=["rectangle", "circle", "region"],
                    default="rectangle",
                    help="rectangle = interactive pan/zoom picker; circle/region = HALCON draw")
    ap.add_argument("--disp", type=int, default=1000,
                    help="rectangle picker: display long-side px cap (smaller=faster/blurrier)")
    ap.add_argument("--box", help="rectangle headless (no GUI): r1,c1,r2,c2 image coords")
    args = ap.parse_args()

    try:
        import halcon as ha
    except Exception as e:
        print("ERROR: cannot import 'halcon'. Run via the project venv using uv.",
              e, file=sys.stderr)
        return 1

    ref = pick_reference_image(ha, args.image_dir, args.image_file)
    print(f"Reference image: {ref}")

    if args.shape != "rectangle":
        draw_native(ha, ref, args.shape, args.out)
        return 0

    arr, H, W = load_gray(ha, ref)
    if args.box:
        r1, c1, r2, c2 = (float(v) for v in args.box.split(","))
        save_rectangle(ha, args.out, r1, c1, r2, c2, H, W)
        return 0

    rect = pick_rectangle_interactive(arr, args.disp)
    if rect is None:
        print("cancelled, nothing saved")
        return 0
    save_rectangle(ha, args.out, *rect, H, W)
    return 0


if __name__ == "__main__":
    sys.exit(main())