#!/usr/bin/env python3
"""One-command demonstration of the whole cell: staging, conveyor, vision, pusher, plow.

    ./python.sh scripts/run_demo.py                       # 50 dispatches, lit, windowed
    ./python.sh scripts/run_demo.py --headless --repeats 2
    ./python.sh scripts/run_demo.py --preset harsh --no-vision
    ./python.sh scripts/run_demo.py --items bucket,barrel,box_300x200x200

This is the entry point for someone who has not built the scenario before: it stages the
cell, dispatches the item library, and writes one JSON with everything needed to judge the
result. Nothing has to be assembled by hand first.

It reports two things that are easy to conflate and must be kept apart:

* **classification** - did the vision stack name the class correctly? Measured against
  ground truth from the manifest, with a confusion matrix.
* **delivery** - did the item physically reach the tray its class routes to? A correct
  class that ends up on the floor is a delivery failure, not a vision failure, and the
  reverse happens too - a misread item can still land somewhere by luck.

Per dispatch it also records the plow's state *at the moment the item is level with it*:
the commanded angle, the angle the arm actually reached, and its angular rate. Those three
are different numbers because the plow is a compliant force drive, and the difference is
usually what explains a miss.
"""
from __future__ import annotations

import argparse
import json
import math
import sys
import time
from datetime import datetime
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))
USER_SITE = "/home/dasha/.local/lib/python3.12/site-packages"      # ultralytics lives here
if Path(USER_SITE).exists() and USER_SITE not in sys.path:
    sys.path.append(USER_SITE)


def parse_args(argv=None):
    p = argparse.ArgumentParser(description="Robozon sorting cell - full demonstration run")
    p.add_argument("--headless", action="store_true")
    p.add_argument("--preset", default="bright", choices=["bright", "dim", "harsh"],
                   help="lighting preset the run is staged under")
    p.add_argument("--no-floor", action="store_true", help="skip the catch floor")
    p.add_argument("--repeats", type=int, default=2,
                   help="passes over the library; 2 x 25 items = 50 dispatches")
    p.add_argument("--items", default=None, help="comma-separated subset, for a quick check")
    p.add_argument("--no-vision", action="store_true",
                   help="route on ground truth instead of running CRE-ROI v2b")
    p.add_argument("--pitch", type=float, default=2.5, help="metres between dispatches")
    p.add_argument("--speed", type=float, default=None, help="belt speed override, m/s")
    p.add_argument("--settle", type=float, default=6.0,
                   help="seconds to let the last item come to rest before scoring")
    p.add_argument("--out", default=None, help="where to write the run log")
    return p.parse_args(argv)


# --------------------------------------------------------------------------- scoring
def confusion(records):
    labels = ["B", "C", "D"]
    m = {g: {p: 0 for p in labels + ["?"]} for g in labels}
    for r in records:
        if r["gt"] in m:
            m[r["gt"]][r["pred"] if r["pred"] in m[r["gt"]] else "?"] += 1
    return m


def summarise(records, expect):
    graded = [r for r in records if r.get("pred") not in (None, "?")]
    hits = sum(1 for r in graded if r["pred"] == r["gt"])
    delivered = [r for r in records if r.get("delivered")]
    by_class = {}
    for cls in ("B", "C", "D"):
        same = [r for r in records if r["gt"] == cls]
        if same:
            by_class[cls] = dict(
                dispatched=len(same),
                delivered=sum(1 for r in same if r.get("delivered")),
                classified=sum(1 for r in same if r.get("pred") == cls),
                target=expect.get(cls))
    where = {}
    for r in records:
        where[r["outcome"]] = where.get(r["outcome"], 0) + 1
    return dict(
        dispatched=len(records),
        classification=dict(graded=len(graded), correct=hits,
                            accuracy=round(hits / len(graded), 3) if graded else None,
                            confusion=confusion(records)),
        delivery=dict(delivered=len(delivered),
                      rate=round(len(delivered) / len(records), 3) if records else None,
                      by_class=by_class, resting_places=where),
    )


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

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

    print("=" * 74)
    print(f"  Robozon sorting cell - {datetime.now():%Y-%m-%d %H:%M}")
    print(f"  lighting {args.preset}, belt {C.BELT_SPEED} m/s, pitch {args.pitch} m, "
          f"vision {'off' if args.no_vision else 'on'}")
    print("=" * 74)

    # ---- 1. scene ----------------------------------------------------------
    stage, info = PV.load(script_control=True)
    await app_utils.update_app_async(steps=50)
    staged = staging.stage_cell(stage, preset=args.preset, floor=not args.no_floor)
    PS.keep_lanes_active(stage)
    lanes = PS.configure_lanes(stage)
    opened = PS.open_junction(stage)
    print(f"[scene ] lighting={staged['lighting']['preset']} "
          f"floor={'yes' if 'floor' in staged else 'no'} lanes={len(lanes)} "
          f"junction shells opened={len(opened)}")

    # ---- 2. item library ---------------------------------------------------
    lib = ROOT / "assets" / "items"
    if not (lib / "manifest.json").exists():
        lib = C.MESHES
        print(f"[items ] assets/items missing - falling back to {lib.name} "
              "(run scripts/export_item_library.py for the full catalogue)")
    items_meta = PV.load_items(stage, meshes_dir=lib)
    classes = {k: v["zone"] for k, v in items_meta.items()}
    await app_utils.update_app_async(steps=30)
    from collections import Counter
    print(f"[items ] {len(classes)} loaded from {lib.name}: {dict(Counter(classes.values()))}")

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

    # ---- 3. vision ---------------------------------------------------------
    vision = None
    if not args.no_vision:
        from robozon_sorter.cv.pipeline import CreRoiV2b
        vision = CreRoiV2b()
        vision.attach_cameras()
        print(f"[vision] CRE-ROI v2b ready, gate pixels {vision.gate_px}")

    # ---- 4. dispatch order -------------------------------------------------
    if args.items:
        order = [n.strip() for n in args.items.split(",") if n.strip() in classes]
    else:
        order = [n for _ in range(args.repeats) for n in sorted(classes)]
    mapping = PS.calibrate_mapping()
    expect = {"D": "bin", "B": "container_B", "C": "container_C"}
    sorter = PS.PlowSorter(stage, cell, classes, mapping)
    print(f"[plan  ] {len(order)} dispatches, plow mapping {mapping}, expect {expect}")

    # ---- 5. run ------------------------------------------------------------
    route, records, seen = {}, {}, set()
    def rec(name):
        return records.setdefault(name + f"#{len([k for k in records if k.startswith(name)])}"
                                  if False else name, dict(item=name, gt=classes[name]))

    log_events = []
    def on_event(kind, name, payload):
        log_events.append(dict(kind=kind, item=name, **payload))
        if kind in ("release", "divert", "done"):
            print(f"   {kind:8s} {name:20s} {payload if payload else ''}")

    feeder = AutoFeeder(cell, order=order, pitch=args.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)

    dt_block, blocks = 15, 0
    prev_angle, prev_t = sorter.plow.angle, time.time()
    while blocks < 900 and len(feeder.finished) < len(order):
        await app_utils.update_app_async(steps=dt_block)
        blocks += 1
        sorter.update(dt_block / 60.0)

        for name in list(feeder.active):
            r = records.setdefault(name, dict(item=name, gt=classes[name], pred=None,
                                              dims=None, k=None, cre_ms=None, views=None,
                                              commanded=None, reached=None, rate=None,
                                              outcome=None, delivered=False))
            x = float(cell.pose(name)[0])

            # vision, once, while the item is under the portal
            if name not in seen and abs(x - C.CAM_X) < 0.08:
                if vision is not None:
                    playing = timeline.is_playing()
                    res = vision.measure()
                    if playing and not timeline.is_playing():
                        timeline.play()                      # Replicator stops the timeline
                        await app_utils.update_app_async(steps=2)
                    r.update(pred=res["cls"], dims=res["dims"], k=res["k"],
                             views=res["views"], cre_ms=res["cre_ms"])
                else:
                    r["pred"] = classes[name]
                route[name] = r["pred"]
                seen.add(name)
                mark = "ok  " if r["pred"] == r["gt"] else "MISS"
                print(f"   vision   {name:20s} pred={r['pred']} gt={r['gt']} {mark} "
                      f"dims={r['dims']} K={r['k']}")

            # plow state at the moment the item is level with the blade
            if r["commanded"] is None and abs(x - C.PLOW_POS[0]) < 0.25:
                now = time.time()
                ang = sorter.plow.angle
                r["commanded"] = sorter.decided.get(name)
                r["reached"] = round(ang, 2)
                r["rate"] = round((ang - prev_angle) / max(now - prev_t, 1e-6), 1)
                print(f"   plow     {name:20s} cmd={r['commanded']} reached={r['reached']} "
                      f"rate={r['rate']} deg/s")
            prev_angle, prev_t = sorter.plow.angle, time.time()

    # let the stragglers come to rest before scoring
    for _ in range(int(args.settle * 60 / dt_block)):
        await app_utils.update_app_async(steps=dt_block)
        sorter.update(dt_block / 60.0)

    for name in order:
        r = records.setdefault(name, dict(item=name, gt=classes[name], pred=None,
                                          outcome=None, delivered=False))
        p = cell.pose(name)
        r["outcome"] = sorter.lane_of(name)
        r["final"] = [round(float(v), 3) for v in p]
        r["expected"] = expect.get(r["gt"])
        r["delivered"] = (r["outcome"] == r["expected"])

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

    # ---- 6. report ---------------------------------------------------------
    recs = list(records.values())
    summary = summarise(recs, expect)
    print("\n" + "=" * 74)
    print(f"  dispatched {summary['dispatched']}")
    cls_s = summary["classification"]
    if cls_s["graded"]:
        print(f"  classification  {cls_s['correct']}/{cls_s['graded']} "
              f"= {cls_s['accuracy']:.0%}")
        print(f"      {'gt\\pred':>8} " + " ".join(f"{p:>5}" for p in ["B", "C", "D", "?"]))
        for g, row in cls_s["confusion"].items():
            print(f"      {g:>8} " + " ".join(f"{row[p]:>5}" for p in ["B", "C", "D", "?"]))
    d = summary["delivery"]
    print(f"  delivery        {d['delivered']}/{summary['dispatched']} = {d['rate']:.0%}")
    for cls, row in d["by_class"].items():
        print(f"      class {cls} -> {row['target']:<12} "
              f"delivered {row['delivered']}/{row['dispatched']}, "
              f"classified {row['classified']}/{row['dispatched']}")
    print(f"  came to rest    {d['resting_places']}")
    print("=" * 74)

    out = Path(args.out) if args.out else ROOT / "runs" / f"demo_{datetime.now():%Y%m%d_%H%M%S}.json"
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(dict(
        config=dict(preset=args.preset, floor=not args.no_floor, speed=C.BELT_SPEED,
                    pitch=args.pitch, vision=not args.no_vision, repeats=args.repeats,
                    mapping=mapping, expect=expect, library=lib.name),
        summary=summary, items=recs, events=log_events[-400:]), indent=2))
    print(f"  log -> {out}")
    return summary


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() else 1)