"""Full cell run: goods ride the line, vision classifies them, the pusher takes D and the
plow splits B and C into their trays. Sent into a live Isaac Sim:

    isaacsim_send.py --context sort --file scripts/run_plow_sorting.py \
        --args-json '{"vision": true, "pitch": 1.2}'

`scripts/run_plow_cell_vision.py` exercises the pusher only - it never constructs a
PlowSorter, so B and C simply ran off the end of the line. This one wires the plow in and,
more importantly, records *why* an item ended up where it did:

* **detection** - predicted vs ground-truth class, dims, roundness, view count, CRE time.
* **delivery** - the tray the item actually came to rest in, against the tray its class
  maps to.
* **kinematics** - for every item, a trace sampled while it crosses the plow: its pose and
  speed, the angle the plow was commanded to and the angle the arm actually reached. When
  an item does not arrive, that trace is what says whether the sensor missed it, the blade
  was still moving, or it was deflected and then stopped short.

The plow is a compliant force drive, so commanded and measured angle are different numbers
and both are logged; treating them as one is what hides a blade that never took up its
angle in time.
"""
import json
import sys
import time

REPO = "/home/dasha/robozon-sorter"
if REPO not in sys.path:
    sys.path.insert(0, REPO)

# The live Isaac process keeps every module it has ever imported, so an edited
# robozon_sorter/ on disk is invisible to a second run in the same session. Drop the
# package from sys.modules first or you spend the evening re-testing the old code.
for _m in [k for k in list(sys.modules) if k.startswith("robozon_sorter")]:
    del sys.modules[_m]
import importlib
importlib.invalidate_caches()   # a *new* module file is invisible until the finder is reset

import omni.timeline
import isaacsim.core.experimental.utils.app as app_utils

from robozon_sorter import config as C
from robozon_sorter.sim import plow_sort, plow_vision
from robozon_sorter.sim.mechanics import Cell
from robozon_sorter.sim.spawner import AutoFeeder

USE_VISION = bool(globals().get("vision", True))
PITCH = float(globals().get("pitch", 1.20))
SPEED = float(globals().get("speed", C.BELT_SPEED))
MAX_SECONDS = float(globals().get("max_seconds", 90.0))
OUT = globals().get("out", "/home/dasha/robozon-sorter/runs/plow_sorting.json")

# where each class is supposed to end up
EXPECT = {"D": "bin", "B": "container_B", "C": "container_C"}

# ---------------------------------------------------------------- scene
C.BELT_SPEED = SPEED
stage, info = plow_vision.load(belt_speed=SPEED, script_control=True)
items = {k: v["zone"] for k, v in info["items"].items()}

# plow_cell.prepare() deactivates /ConveyorTrack_01 as a stray duplicate. In this scene it
# is not a duplicate, it is the -Y sorting lane, so it has to come back on before the lanes
# are driven - otherwise everything the plow deflects toward -Y falls through the gap.
relit = plow_sort.keep_lanes_active(stage)
lanes = plow_sort.configure_lanes(stage, SPEED)
opened = plow_sort.open_junction(stage)
print(f"scene: {len(items)} items {sorted(set(items.values()))} | lane restored={relit} "
      f"| lanes driven={len(lanes)} | junction shells opened={len(opened)}")

vision = None
if USE_VISION:
    from robozon_sorter.cv.pipeline import CreRoiV2b
    vision = CreRoiV2b()
    vision.attach_cameras()
    print("CRE-ROI v2b ready")

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 = sorted(items)
route = {} if USE_VISION else dict(items)          # what the pusher acts on (class D)
classes = {} if USE_VISION else dict(items)        # what the plow acts on (B / C)

sorter = plow_sort.PlowSorter(stage, cell, classes, plow_sort.calibrate_mapping())
print(f"plow mapping {sorter.mapping} | sensor x={sorter.sense_x} | swing {sorter.swing} deg")

# ---------------------------------------------------------------- logging
rec = {n: dict(item=n, gt=items[n], pred=None, dims=None, k=None, views=None,
               cre_ms=None, sensed=False, commanded=None, angle_at_plow=None,
               trace=[], max_speed=0.0, blowup=None,
               outcome=None, expected=EXPECT.get(items[n]), ok=None)
       for n in order}
events = []


def on_event(kind, name, payload):
    events.append((round(t_sim, 2), kind, name, payload))
    if kind in ("release", "gate", "divert", "error"):
        print(f"  {kind:8s} {name:18s} {payload if payload else ''}")


feeder = AutoFeeder(cell, order=order, pitch=PITCH, route=route, on_event=on_event)

# ---------------------------------------------------------------- physics hook
t_sim = 0.0


DECIMATE = 8          # 120 Hz / 8 = 15 samples/s: 60 of them span 4 s, the whole discharge
BLOWUP_MS = 5.0       # a belt runs at 1 m/s; anything past this is the solver, not the belt
_tick = 0


def _speed(name):
    try:
        v = cell._rp[name].get_velocities()[0].numpy()[0]
        return float((v[0] ** 2 + v[1] ** 2 + v[2] ** 2) ** 0.5)
    except Exception:
        return 0.0


def _step(dt):
    """the plow has to be serviced from the physics step, like the pusher: the sensor is a
    raycast and the blade target is ramped per-step.

    The trace is decimated: at full rate 60 samples cover 0.5 m and run out before the item
    even reaches the plow, which is how the first pass missed where goods were being thrown.
    """
    global t_sim, _tick
    t_sim += dt
    _tick += 1
    try:
        sorter.update(dt)
        for n in list(feeder.active):
            p = cell.pose(n)
            x, y, z = float(p[0]), float(p[1]), float(p[2])
            r = rec[n]
            if x < -5.6:                          # from the plow approach onward
                spd = _speed(n)
                if spd > r.get("max_speed", 0.0):
                    r["max_speed"] = round(spd, 2)
                if spd > BLOWUP_MS and r.get("blowup") is None:
                    r["blowup"] = dict(t=round(t_sim, 2), x=round(x, 3), y=round(y, 3),
                                       z=round(z, 3), speed=round(spd, 1),
                                       cmd=round(sorter.plow.commanded, 1),
                                       arm=round(sorter.plow.angle, 1))
                if _tick % DECIMATE == 0 and len(r["trace"]) < 60:
                    r["trace"].append(dict(t=round(t_sim, 2), x=round(x, 3), y=round(y, 3),
                                           z=round(z, 3), v=round(spd, 2),
                                           cmd=round(sorter.plow.commanded, 1),
                                           arm=round(sorter.plow.angle, 1)))
                if n in sorter.decided and not r["sensed"]:
                    r["sensed"] = True
                    r["commanded"] = round(sorter.decided[n], 1)
                if abs(x - C.PLOW_POS[0]) < 0.25 and r["angle_at_plow"] is None:
                    r["angle_at_plow"] = round(sorter.plow.angle, 1)
    except Exception as exc:
        events.append((round(t_sim, 2), "step-error", "", repr(exc)))


from omni.physx import get_physx_interface
sub = get_physx_interface().subscribe_physics_step_events(_step)
feeder.install()

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

# ---------------------------------------------------------------- run
print(f"\nrunning: pitch {PITCH} m at {SPEED} m/s, vision={USE_VISION}")
seen = set()
t0 = time.time()
settled = {}
while time.time() - t0 < MAX_SECONDS:
    await app_utils.update_app_async(steps=15)

    if vision is not None:
        for name in list(feeder.active):
            if name in seen:
                continue
            if abs(float(cell.pose(name)[0]) - C.CAM_X) < 0.10:
                was = timeline.is_playing()
                res = vision.measure()
                if was and not timeline.is_playing():   # Replicator's step stops the timeline
                    timeline.play()
                    await app_utils.update_app_async(steps=2)
                seen.add(name)
                r = rec[name]
                r.update(pred=res["cls"], dims=res["dims"], k=round(res.get("k", 0.0), 3),
                         views=res.get("views"), cre_ms=res.get("cre_ms"))
                route[name] = res["cls"]
                classes[name] = res["cls"]
                sorter.classes[name] = res["cls"]
                print(f"  vision   {name:18s} pred={res['cls']} gt={items[name]} "
                      f"{'ok' if res['cls'] == items[name] else 'MISS'} "
                      f"dims={res['dims']} K={res.get('k', 0):.2f}")

    for n in order:                                  # freeze the outcome once it stops
        if n in settled:
            continue
        p = cell.pose(n)
        where = sorter.lane_of(n)
        if where.startswith("container") or where == "floor":
            settled[n] = where
        elif where == "line" and float(p[0]) < C.MAIN_X0 + 0.35:
            settled[n] = "line-end"
    if len(settled) >= len(order):
        break

# outcomes: the D bin is the pusher's, read through mechanics; the trays are the plow's
for n in order:
    w = sorter.lane_of(n)
    if w == "line" and cell.where(n) == "bin":
        w = "bin"
    rec[n]["outcome"] = settled.get(n, w)
    rec[n]["ok"] = (rec[n]["outcome"] == rec[n]["expected"])
    p = cell.pose(n)
    rec[n]["final"] = [round(float(v), 3) for v in p[:3]]

app_utils.stop()
await app_utils.update_app_async(steps=10)
sub = None
feeder.remove()

# ---------------------------------------------------------------- report
print("\n===== DETECTION =====")
graded = [r for r in rec.values() if r["pred"] not in (None, "?")]
if graded:
    hit = sum(1 for r in graded if r["pred"] == r["gt"])
    cre = [r["cre_ms"] for r in graded if r["cre_ms"]]
    print(f"  class agreement {hit}/{len(graded)}"
          + (f" | CRE {sum(cre)/len(cre):.0f} ms/item" if cre else ""))
    for r in sorted(graded, key=lambda r: r["item"]):
        print(f"   {r['item']:18s} gt={r['gt']} pred={r['pred']} "
              f"{'ok' if r['pred'] == r['gt'] else 'MISS':4s} dims={r['dims']} K={r['k']}")
else:
    print("  (no vision this run)")

print("\n===== DELIVERY =====")
for r in sorted(rec.values(), key=lambda r: r["item"]):
    print(f"   {r['item']:18s} gt={r['gt']} -> {str(r['outcome']):12s} "
          f"want={str(r['expected']):12s} {'OK' if r['ok'] else 'FAIL'} "
          f"final={r['final']}")
good = [r for r in rec.values() if r["ok"]]
print(f"  delivered {len(good)}/{len(order)}")

bad = [r for r in rec.values() if not r["ok"]]
if bad:
    print("\n===== KINEMATICS ON FAILURES =====")
    for r in bad:
        print(f"  {r['item']} ({r['gt']}) -> {r['outcome']}")
        print(f"     sensed={r['sensed']} commanded={r['commanded']} "
              f"arm_at_plow={r['angle_at_plow']}")
        for s in r["trace"][:12]:
            print(f"     t={s['t']:6.2f} x={s['x']:+.2f} y={s['y']:+.2f} z={s['z']:+.2f} "
                  f"cmd={s['cmd']:+.1f} arm={s['arm']:+.1f}")

import os
os.makedirs(os.path.dirname(OUT), exist_ok=True)
json.dump(dict(config=dict(pitch=PITCH, speed=SPEED, vision=USE_VISION,
                           mapping=sorter.mapping, expect=EXPECT),
               items=list(rec.values()),
               events=[dict(t=t, kind=k, item=n, payload=str(p)) for t, k, n, p in events]),
          open(OUT, "w"), indent=2)
print(f"\nlog -> {OUT}")