"""Worker-only physical capability measurement on isolated metric lanes.

All lanes share the shipped chassis and tire materials; no AI or Gaussian
geometry. Results qualify this virtual profile only, never the real vehicles.
"""

import argparse
import hashlib
import json
import math
import time
from datetime import UTC, datetime
from pathlib import Path

from isaacsim import SimulationApp

parser = argparse.ArgumentParser()
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
app = SimulationApp({"headless": True, "hide_ui": True})
try:
    import omni.timeline
    import omni.usd
    from isaacsim.core.experimental.prims import Articulation
    from isaacsim.core.simulation_manager import SimulationManager
    from pxr import Gf, UsdGeom, UsdPhysics, UsdShade
    from rover_profile import PROFILE, DifferentialDrive, create_rover

    stage = omni.usd.get_context().get_stage()
    UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z)
    UsdGeom.SetStageMetersPerUnit(stage, 1)
    material = UsdShade.Material.Define(stage, "/World/Materials/Terrain")
    physics = UsdPhysics.MaterialAPI.Apply(material.GetPrim())
    physics.CreateStaticFrictionAttr(0.9)
    physics.CreateDynamicFrictionAttr(0.8)
    physics.CreateRestitutionAttr(0)

    def surface(prim):
        UsdPhysics.CollisionAPI.Apply(prim)
        UsdShade.MaterialBindingAPI.Apply(prim).Bind(
            material, UsdShade.Tokens.weakerThanDescendants, "physics"
        )

    def cube(path, center, scale):
        shape = UsdGeom.Cube.Define(stage, path)
        shape.CreateSizeAttr(1)
        shape.AddTranslateOp().Set(Gf.Vec3d(*center))
        shape.AddScaleOp().Set(Gf.Vec3f(*scale))
        surface(shape.GetPrim())

    cases = [("flat", 0)] + [("step", h) for h in (0.05, 0.1, 0.12, 0.14, 0.15, 0.18, 0.2, 0.25)]
    cases += [("slope", a) for a in (10, 15, 20, 25)]
    for i, (kind, value) in enumerate(cases):
        y = i * 5.0
        path = f"/World/Lane{i}"
        cube(path + "/Ground", (3, y, -0.1), (12, 3, 0.2))
        if kind == "step":
            cube(path + "/Step", (4.5, y, value / 2), (6, 3, value))
        elif kind == "slope":
            # A continuous supported ramp; x=1.5 is the toe, no hidden step.
            h = 6 * math.tan(math.radians(value))
            ramp = UsdGeom.Mesh.Define(stage, path + "/Ramp")
            ramp.CreatePointsAttr(
                [(1.5, y - 1.5, 0), (1.5, y + 1.5, 0), (7.5, y - 1.5, h), (7.5, y + 1.5, h)]
            )
            ramp.CreateFaceVertexCountsAttr([3, 3])
            ramp.CreateFaceVertexIndicesAttr([0, 2, 1, 1, 2, 3])
            ramp.CreateSubdivisionSchemeAttr(UsdGeom.Tokens.none)
            surface(ramp.GetPrim())
        create_rover(stage, [0, y, PROFILE["wheel_radius_m"] + 0.04], 0, root_path=path + "/Rover")
    SimulationManager.setup_simulation(dt=1 / 60, device="cpu")
    omni.timeline.get_timeline_interface().play()
    for _ in range(30):
        app.update()
    robots = [Articulation(f"/World/Lane{i}/Rover") for i in range(len(cases))]
    indices = [r.get_dof_indices(PROFILE["wheel_names"]) for r in robots]
    drives = [DifferentialDrive(r) for r in robots]
    traces = [[] for _ in cases]
    starts = [r.get_world_poses()[0].numpy()[0].tolist() for r in robots]
    started = time.monotonic()
    for drive in drives:
        drive.command(0.3)
    for tick in range(180):
        SimulationManager.step(steps=10)
        for i, robot in enumerate(robots):
            pos, quat = robot.get_world_poses()
            p, q = pos.numpy()[0], quat.numpy()[0]
            tilt = math.degrees(
                math.acos(max(-1, min(1, 1 - 2 * (float(q[1]) ** 2 + float(q[2]) ** 2))))
            )
            traces[i].append(dict(time_s=(tick + 1) / 6, xyz=p.tolist(), tilt_degrees=tilt))
            if p[0] >= 5.5 or tilt > 35:
                drives[i].command(0)
    stops = [r.get_world_poses()[0].numpy()[0].tolist() for r in robots]
    for drive in drives:
        drive.command(0)
    SimulationManager.step(steps=120)
    rows = []
    for i, (kind, value) in enumerate(cases):
        end = robots[i].get_world_poses()[0].numpy()[0].tolist()
        tilt = max(t["tilt_degrees"] for t in traces[i])
        rows.append(
            dict(
                kind=kind,
                value=value,
                start=starts[i],
                end=end,
                max_tilt_degrees=tilt,
                braking_drift_m=math.dist(stops[i], end),
                wheel_velocity_rps=robots[i]
                .get_dof_velocities(dof_indices=indices[i])
                .numpy()
                .tolist(),
                reached=end[0] >= 5.4,
                upright=tilt < 35,
                passed=end[0] >= 5.4 and tilt < 35 and math.dist(stops[i], end) < 0.1,
                trace=traces[i],
            )
        )
    report = dict(
        schema_version="missioncore.virtual-rover-capability/v1",
        utc=datetime.now(UTC).isoformat(),
        monotonic=time.monotonic(),
        elapsed_wall_s=time.monotonic() - started,
        profile=PROFILE,
        profile_sha256=hashlib.sha256(
            Path(__file__).with_name("rover_profile.py").read_bytes()
        ).hexdigest(),
        cases=rows,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps({"cases": [{k: v for k, v in row.items() if k != "trace"} for row in rows]}))
finally:
    app.close()