#!/usr/bin/env python3
"""Run the plow cell with the full vision stack.

    ./python.sh scripts/run_plow_cell_vision.py                # windowed
    ./python.sh scripts/run_plow_cell_vision.py --headless
    ./python.sh scripts/run_plow_cell_vision.py --no-vision    # route on ground truth

Items are released on the added infeed belt, measured by CRE-ROI v2b under the camera
portal, and the ones that come back class D are diverted by the Y-split pusher when they
break the laser beam. The plow and the authored kinematics are not touched.
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


def parse_args(argv=None):
    p = argparse.ArgumentParser(description="plow cell + CRE-ROI v2b")
    p.add_argument("--headless", action="store_true")
    p.add_argument("--no-vision", action="store_true",
                   help="skip inference and route on ground truth")
    p.add_argument("--speed", type=float, default=None, help="belt speed, m/s")
    p.add_argument("--pitch", type=float, default=None, help="metres between items")
    p.add_argument("--items", default=None, help="comma-separated release order")
    p.add_argument("--log", default=None)
    return p.parse_args(argv)


async def _run(app_utils, args):
    from robozon_sorter import config as C
    from robozon_sorter.sim import plow_vision
    from robozon_sorter.sim.mechanics import Cell
    from robozon_sorter.sim.spawner import AutoFeeder

    if args.speed:
        C.BELT_SPEED = args.speed

    stage, info = plow_vision.load(belt_speed=args.speed, script_control=True)
    items = {k: v["zone"] for k, v in info["items"].items()}
    print(f"plow cell ready: {len(items)} items {sorted(set(items.values()))}, "
          f"infeed release at x={info['spawn_x']}")

    vision = None
    if not args.no_vision:
        from robozon_sorter.cv.pipeline import CreRoiV2b
        vision = CreRoiV2b()
        vision.attach_cameras()
        print("CRE-ROI v2b ready; gate pixels:", vision.gate_px)

    await app_utils.update_app_async(steps=40)
    cell = Cell(stage, items.keys())
    cell.park_all()
    await app_utils.update_app_async(steps=15)

    order = [n.strip() for n in args.items.split(",")] if args.items else sorted(items)
    order = [n for n in order if n in items]
    pitch = args.pitch if args.pitch is not None else C.RELEASE_GAP

    log, seen = [], set()
    route = dict(items) if args.no_vision else {}

    def on_event(kind, name, payload):
        print(f"  {kind:8s} {name:18s} {payload if payload else ''}")
        if kind == "done":
            for rec in log:
                if rec["item"] == name and "outcome" not in rec:
                    rec["outcome"] = payload.get("where")

    feeder = AutoFeeder(cell, order=order, pitch=pitch, route=route,
                        on_event=on_event).install()

    import omni.timeline
    timeline = omni.timeline.get_timeline_interface()
    app_utils.play(commit=True)
    await app_utils.update_app_async(steps=20)

    print(f"\npitch {pitch} m at {C.BELT_SPEED} m/s\n{'kind':>10}  detail")
    for _ in range(400):
        await app_utils.update_app_async(steps=15)

        # classify each item once, while it sits under the portal
        if vision is not None:
            for name in list(feeder.active):
                if name in seen:
                    continue
                x = float(cell.pose(name)[0])
                if abs(x - C.CAM_X) < 0.08:
                    was_playing = timeline.is_playing()
                    res = vision.measure()
                    # Replicator's step stops the timeline; resume or the line freezes
                    if was_playing and not timeline.is_playing():
                        timeline.play()
                        await app_utils.update_app_async(steps=2)
                    gt = items[name]
                    route[name] = res["cls"]
                    seen.add(name)
                    print(f"  vision   {name:18s} pred={res['cls']} gt={gt} "
                          f"{'ok' if res['cls'] == gt else 'MISS'} dims={res['dims']} "
                          f"K={res['k']:.2f} views={res['views']} cre={res['cre_ms']}ms")
                    log.append(dict(item=name, gt=gt, **res))

        if len(feeder.finished) >= len(order):
            break

    app_utils.stop()
    await app_utils.update_app_async(steps=15)
    feeder.remove()
    cell.blade_to(C.BLADE_HOME_Y)

    print(f"\n  outcomes: {feeder.finished}")
    expected = {n: ("bin" if items[n] == "D" else "line-end") for n in order}
    wrong = [n for n in expected if feeder.finished.get(n) != expected[n]]
    print(f"  expected: {expected}")
    print("  routing matches ground truth" if not wrong else f"  differs on: {wrong}")
    graded = [r for r in log if r.get("cls") not in (None, "?")]
    if graded:
        hits = sum(1 for r in graded if r["cls"] == r["gt"])
        cre = [r["cre_ms"] for r in graded if r.get("cre_ms")]
        print(f"  vision agreed with ground truth on {hits}/{len(graded)}"
              + (f", CRE {sum(cre)/len(cre):.0f} ms/item" if cre else ""))
    if args.log:
        Path(args.log).write_text(json.dumps(log, indent=2))
        print(f"  log -> {args.log}")
    return log


def main(argv=None):
    args = parse_args(argv)
    try:
        import omni.usd
        inside = omni.usd.get_context().get_stage() is not None
    except Exception:
        inside = False

    app = None
    if not inside:
        from isaacsim import SimulationApp
        app = SimulationApp({"headless": args.headless, "width": 1600, "height": 900})

    import asyncio
    import isaacsim.core.experimental.utils.app as app_utils
    loop = asyncio.get_event_loop()
    try:
        return loop.run_until_complete(_run(app_utils, args))
    finally:
        if app is not None:
            app.close()


if __name__ == "__main__":
    sys.exit(0 if main() is not None else 1)