#!/usr/bin/env python3
"""Sweep-rate tuning as a standalone Isaac Sim process.

    /home/whatevenif/isaacsim/python.sh scripts/tune_sweep_standalone.py \
        --rates 120,200,300,450 --items 4 --headless

Same experiment as `tune_sweep_rate.py`, but it brings up its own `SimulationApp` instead
of being sent into a live Kit through the code-editor socket. That socket stopped returning
results on anything longer than a minute or two - the code kept running (one run wrote its
log in full) but the reply never arrived, so six runs in a row looked like failures. A
standalone process writes its log itself and can be read afterwards, which removes the
connection from the experiment entirely.

Each rate is judged on the lane-entry beams, not on delivery alone: an item left on the belt
and an item thrown to the floor both score zero and need opposite corrections.

    crossed  the push reached the lane   -> low means the sweep is too slow
    speed    how fast it crossed         -> high means it is throwing them
    homed    completed returns to centre -> the blade must park in the middle between items
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent


def parse_args(argv=None):
    p = argparse.ArgumentParser(description="tune the plow sweep rate")
    p.add_argument("--rates", default="120,200,300,450", help="deg/s, comma separated")
    p.add_argument("--items", type=int, default=4, help="how many B/C items per rate")
    p.add_argument("--speed", type=float, default=0.8, help="belt m/s")
    p.add_argument("--pitch", type=float, default=3.5, help="metres between items")
    p.add_argument("--seconds", type=float, default=60.0, help="budget per rate")
    p.add_argument("--headless", action="store_true")
    p.add_argument("--out", default=str(ROOT / "runs" / "sweep_tuning.json"))
    return p.parse_args(argv)


async def _run(app_utils, args):
    import omni.timeline
    from omni.physx import get_physx_interface

    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 = [float(r) for r in args.rates.split(",") if r.strip()]
    stage, info = plow_vision.load(belt_speed=args.speed, script_control=True,
                                   meshes_dir=str(ROOT / "assets" / "items"))
    staging.stage_cell(stage, preset="bright", floor=True)
    plow_sort.keep_lanes_active(stage)
    lanes = plow_sort.configure_lanes(stage, args.speed)
    plow_sort.open_junction(stage)
    items = {k: v["zone"] for k, v in info["items"].items()}
    print(f"scene ready: {len(items)} items, {len(lanes)} lanes/decks driven")

    await app_utils.update_app_async(steps=40)
    cell = Cell(stage, items.keys())
    order = [n for n in sorted(items) if items[n] in ("B", "C")][:args.items]
    want = {"B": "container_B", "C": "container_C"}
    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 = rate
        cell.park_all()
        await app_utils.update_app_async(steps=20)

        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):
            try:
                _s.update(dt)
                _b.tick(dt)
                _b.poll(rate=_r)
            except Exception:
                pass

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

        for _ in range(int(args.seconds * 4)):
            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=15)
        feeder.remove()
        sub = None

        cr = beams.report()
        sp = [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] == want[items[n]])
        row = dict(rate=rate, crossed=len(cr), of=len(order), delivered=delivered,
                   homed=sorter.homed,
                   mean_cross_speed=round(sum(sp) / len(sp), 2) if sp else None,
                   max_cross_speed=round(max(sp), 2) if sp else None,
                   where=where, crossings=cr)
        results.append(row)
        print(f"  rate {rate:6.0f} deg/s (tip {C.PLOW_ARM_LEN * rate * 3.14159 / 180:.2f} m/s)"
              f" -> crossed {len(cr)}/{len(order)}  delivered {delivered}  "
              f"homed {sorter.homed}  cross v mean {row['mean_cross_speed']} "
              f"max {row['max_cross_speed']}")

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

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(dict(rates=rates, belt=args.speed, pitch=args.pitch,
                                   items=order, results=results), indent=2))
    print(f"log -> {out}")
    return results


def main(argv=None):
    args = parse_args(argv)
    if str(ROOT) not in sys.path:
        sys.path.insert(0, str(ROOT))
    from isaacsim import SimulationApp
    app = SimulationApp({"headless": args.headless, "width": 1280, "height": 800})
    try:
        import asyncio
        import isaacsim.core.experimental.utils.app as app_utils
        return asyncio.get_event_loop().run_until_complete(_run(app_utils, args))
    finally:
        app.close()


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