"""Find the plow sweep rate that actually lands goods on their lane.

    isaacsim_send.py --context tune --file scripts/tune_sweep_rate.py \
        --args-json '{"rates": [120, 200, 300, 450], "n": 4}'

Delivery alone cannot tune this: an item on the floor and an item still on the belt both
score zero and need opposite corrections. So each rate is judged on the lane-entry beams
(`sim/lane_beams`), which separate the two:

    crossed  the push reached the lane            <- too slow if this is low
    speed    how fast it was going when it did    <- throwing it if this is high

A usable rate crosses most items at a modest crossing speed. The sweep is a **push**, so
the blade returns to centre after each item and waits there - `PlowSorter` does that, and
the run reports how many times it completed a return, so a blade that stops homing shows up
as a number rather than as a mystery later.
"""
import json
import sys
import time

REPO = "/home/dasha/robozon-sorter"
if REPO not in sys.path:
    sys.path.insert(0, REPO)
import importlib
for _m in [k for k in list(sys.modules) if k.startswith("robozon_sorter")]:
    del sys.modules[_m]
importlib.invalidate_caches()

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

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

RATES = globals().get("rates", [120.0, 200.0, 300.0, 450.0])
N = int(globals().get("n", 4))
SPEED = float(globals().get("speed", 0.8))
PITCH = float(globals().get("pitch", 3.5))
BUDGET = float(globals().get("per_rate_seconds", 70.0))
OUT = globals().get("out", f"{REPO}/runs/sweep_tuning.json")

stage, info = plow_vision.load(belt_speed=SPEED, script_control=True,
                               meshes_dir=f"{REPO}/assets/items")
staging.stage_cell(stage, preset="bright", floor=True)
plow_sort.keep_lanes_active(stage)
plow_sort.configure_lanes(stage, SPEED)
plow_sort.open_junction(stage)
items = {k: v["zone"] for k, v in info["items"].items()}
await app_utils.update_app_async(steps=30)
cell = Cell(stage, items.keys())

order = [n for n in sorted(items) if items[n] in ("B", "C")][:N]
print(f"tuning on {len(order)} B/C items: {order}")
timeline = omni.timeline.get_timeline_interface()
results = []

for rate in RATES:
    C.PLOW_SWEEP_RATE = float(rate)
    cell.park_all()
    await app_utils.update_app_async(steps=15)
    sorter = plow_sort.PlowSorter(stage, cell, items, plow_sort.calibrate_mapping())
    beams = lane_beams.LaneBeams(stage, cell, plow=sorter.plow)

    def _step(dt, _s=sorter, _b=beams, _r=rate):
        _s.update(dt)
        _b.tick(dt)
        _b.poll(rate=_r)

    from omni.physx import get_physx_interface
    sub = get_physx_interface().subscribe_physics_step_events(_step)
    feeder = AutoFeeder(cell, order=order, pitch=PITCH, route={}).install()
    app_utils.play(commit=True)
    await app_utils.update_app_async(steps=20)

    t0 = time.time()
    while time.time() - t0 < BUDGET:
        await app_utils.update_app_async(steps=15)
        if len(beams.crossings) >= len(order):
            break
    app_utils.stop()
    await app_utils.update_app_async(steps=10)
    feeder.remove()
    sub = None

    cr = beams.report()
    speeds = [c["speed"] for c in cr]
    where = {n: sorter.lane_of(n) for n in order}
    delivered = sum(1 for n in order
                    if where[n] == {"B": "container_B", "C": "container_C"}[items[n]])
    row = dict(rate=rate, crossed=len(cr), of=len(order),
               mean_cross_speed=round(sum(speeds) / len(speeds), 2) if speeds else None,
               max_cross_speed=round(max(speeds), 2) if speeds else None,
               delivered=delivered, homed=sorter.homed, where=where,
               crossings=cr)
    results.append(row)
    print(f"  rate {rate:6.0f} deg/s -> crossed {len(cr)}/{len(order)}  "
          f"cross speed mean {row['mean_cross_speed']} max {row['max_cross_speed']}  "
          f"delivered {delivered}  homed {sorter.homed}")

print("\n===== SWEEP RATE TUNING =====")
print(f"{'rate':>7} {'crossed':>9} {'mean v':>8} {'max v':>7} {'delivered':>10} {'homed':>6}")
for r in results:
    print(f"{r['rate']:7.0f} {r['crossed']:4d}/{r['of']:<4d} "
          f"{str(r['mean_cross_speed']):>8} {str(r['max_cross_speed']):>7} "
          f"{r['delivered']:10d} {r['homed']:6d}")
best = max(results, key=lambda r: (r["delivered"], r["crossed"], -(r["max_cross_speed"] or 9)))
print(f"\nbest so far: {best['rate']:.0f} deg/s "
      f"(tip {C.PLOW_ARM_LEN * best['rate'] * 3.14159 / 180:.2f} m/s)")

import os
os.makedirs(os.path.dirname(OUT), exist_ok=True)
json.dump(dict(rates=RATES, belt=SPEED, pitch=PITCH, results=results),
          open(OUT, "w"), indent=2)
print(f"log -> {OUT}")