"""Bounded experimental L20 dynamic diagnostic, NOT a training or hardware entry.

Requires inspected PhysX 110.1.13 + current Isaac Lab ProxyArray/xyzw API. Runtime
coupling is tested from passive follower response, not inferred from schema names.
Use an externally bounded timeout (default budget <=300s; longer runs need approval).
Do not launch alongside another Kit/GPU job.
"""

import argparse
import hashlib
import json
import math
import random
import traceback
from dataclasses import replace
from pathlib import Path


def validate_replay_request(steps, full_episode, has_hdf5, workspace_radius):
    """Stdlib-only launch bounds, evaluated before Kit or tracking imports."""
    maximum = 32000 if full_episode else 1200
    if not isinstance(steps, int) or not 100 <= steps <= maximum:
        raise ValueError(f"Require 100 <= steps <= {maximum}")
    if full_episode and not has_hdf5:
        raise ValueError("Full-episode mode requires --hdf5")
    if workspace_radius is not None and (not math.isfinite(workspace_radius) or workspace_radius <= 0):
        raise ValueError("Workspace radius must be finite positive")


def validate_replay_duration(steps, dt, duration, full_episode):
    """Reject extrapolation and partial/off-grid full-episode requests before sampling."""
    if not steps * dt <= duration:
        raise ValueError("Reference shorter than requested steps")
    if full_episode and abs(steps * dt - duration) > 1e-9:
        raise ValueError("Full episode requires exact on-grid end and steps")


def configure_presentation(args):
    """Resolve opt-in GUI without silently overriding explicit display conflicts."""
    if args.gui:
        if args.headless:
            raise ValueError("--gui conflicts with --headless")
        if (getattr(args, "visualizer_explicit", False) or args.visualizer is not None) and args.visualizer != ["kit"]:
            raise ValueError("--gui requires the Kit visualizer; conflicts with --visualizer/--viz")
        args.headless, args.visualizer = False, ["kit"]
    else:
        args.headless = True  # Preserve the diagnostic's legacy headless default.
        if args.visualizer is None:
            args.visualizer = "none"


def build_parser(add_app_launcher_args):
    """Keep replay argument parsing testable without importing or launching Kit."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("asset", type=Path, help="Experimental prepared overlay, never original fixed preview")
    parser.add_argument("--manifest", type=Path, required=True, help="Original data/source manifest")
    parser.add_argument("--hdf5", type=Path, help="Omit for explicitly synthetic diagnostic reference")
    parser.add_argument("--episode", default="demo_000000")
    parser.add_argument("--steps", type=int, default=480, help="Per repetition at 240Hz; two repetitions")
    parser.add_argument(
        "--full-episode", action="store_true", help="Require complete HDF5 duration; allow up to 32000 steps"
    )
    parser.add_argument(
        "--workspace-radius", type=float, help="Metres from initial root; overrides only workspace_radius"
    )
    parser.add_argument("--execute-experimental", action="store_true", help="Acknowledge uncalibrated controller")
    parser.add_argument("--limits", type=Path, help="JSON fields of control.Limits, SI units; UNCALIBRATED")
    parser.add_argument("--gui", action="store_true", help="Opt-in bright-hand/dark-backdrop Kit preview")
    add_app_launcher_args(parser)
    # None distinguishes omitted flags from explicit --headless / --viz none.
    parser.set_defaults(headless=None, visualizer=None)
    return parser


def main():
    from isaaclab.app import AppLauncher

    parser = build_parser(AppLauncher.add_app_launcher_args)
    args = parser.parse_args()
    if not args.execute_experimental:
        parser.error("Require --execute-experimental")
    try:
        configure_presentation(args)
        validate_replay_request(args.steps, args.full_episode, args.hdf5 is not None, args.workspace_radius)
    except ValueError as error:
        parser.error(str(error))

    # Kit must choose/register its USD libraries before any pxr-dependent imports.
    launcher = AppLauncher(args)
    app = launcher.app
    exit_code = 0
    failure_context = {"phase": "setup", "completed_repetitions": 0}
    try:
        import numpy as np
        import torch
        from dex_workbench_tracking.cli import synthetic
        from dex_workbench_tracking.control import (
            Limits,
            reference_pose_to_xyzw,
            rotation_error,
            validate_reference,
            wrench,
            xyzw_pose_to_reference,
        )
        from dex_workbench_tracking.prepared import inspect_prepared, require_backend, validate_mimic
        from dex_workbench_tracking.trajectory import load, require, sample

        manifest = json.loads(args.manifest.read_text())
        prepared = inspect_prepared(args.asset, manifest)
        limits = Limits(**json.loads(args.limits.read_text())) if args.limits else Limits()
        if args.workspace_radius is not None:
            limits = replace(limits, workspace_radius=args.workspace_radius)
        data = load(args.hdf5, manifest) if args.hdf5 else synthetic(manifest)
        episode = data.episodes[args.episode]
        validate_reference(episode, limits)
        dt = 1 / 240
        validate_replay_duration(args.steps, dt, episode.time[-1], args.full_episode)
        reference = sample(episode, np.arange(args.steps + 1) * dt)
        # Hold-start reset uses zero velocities; require a genuinely moving diagnostic
        # for every mimic pair, so an ignored/disabled constraint cannot pass at zero.
        names = list(data.joint_names)
        for eq in prepared["mimic"]:
            require(
                np.ptp(reference.joint_position[:, names.index(eq["reference"])]) > 0.002,
                "Each mimic leader must move >0.002rad",
            )
        require(
            np.max(np.linalg.norm(reference.wrist_position - reference.wrist_position[0], axis=1)) > 0.001,
            "Wrist reference must translate >1mm",
        )
        require(
            max(np.linalg.norm(rotation_error(q, reference.wrist_quaternion[0])) for q in reference.wrist_quaternion)
            > 0.005,
            "Wrist reference must rotate >0.005rad",
        )
        from isaaclab_physx.physics import PhysxCfg

        import omni.kit.app
        from pxr import Usd, UsdPhysics

        import isaaclab.sim as sim_utils
        from isaaclab.actuators import ImplicitActuatorCfg
        from isaaclab.assets import Articulation, ArticulationCfg

        manager = omni.kit.app.get_app().get_extension_manager()
        extension = manager.get_enabled_extension_id("omni.physx")
        require(extension is not None, "PhysX extension not enabled")
        version = manager.get_extension_dict(extension)["package"]["version"]
        require_backend(version, bool(Usd.SchemaRegistry().FindAppliedAPIPrimDefinition("NewtonMimicAPI")))
        random.seed(42)
        np.random.seed(42)
        torch.manual_seed(42)
        sim = sim_utils.SimulationContext(
            sim_utils.SimulationCfg(dt=dt, gravity=(0, 0, -9.81), device=args.device, physics=PhysxCfg())
        )
        cfg = ArticulationCfg(
            prim_path="/World/Hand",
            spawn=sim_utils.UsdFileCfg(usd_path=str(args.asset.resolve())),
            init_state=ArticulationCfg.InitialStateCfg(pos=(0, 0, 0.4)),
            actuators={
                "model_independent": ImplicitActuatorCfg(
                    joint_names_expr=prepared["independent_joint_names"],
                    stiffness=limits.finger_stiffness,
                    damping=limits.finger_damping,
                    effort_limit_sim=limits.finger_effort,
                    velocity_limit_sim=limits.finger_velocity,
                )
            },
        )
        hand = Articulation(cfg)
        # Validate remapped USD relationships after reference spawning, before physics.
        remapped = dict(manifest)
        remapped["joints"] = [
            dict(j, path="/World/Hand" + j["path"][len(manifest["default_prim"]) :]) for j in manifest["joints"]
        ]
        validate_mimic(sim.stage, remapped, passive=True)
        anchor_path = "/World/Hand" + manifest["world_fixed_joints"][0]["path"][len(manifest["default_prim"]) :]
        anchor = sim.stage.GetPrimAtPath(anchor_path)
        require(anchor.IsValid() and not anchor.IsActive(), "Obsolete world anchor must be inactive")
        effective_roots = [p for p in sim.stage.Traverse() if p.HasAPI(UsdPhysics.ArticulationRootAPI)]
        require(len(effective_roots) == 1, "Expected exactly one effective articulation root")
        require(
            str(effective_roots[0].GetPath())
            == "/World/Hand" + manifest["root_body_path"][len(manifest["default_prim"]) :],
            "Effective articulation root is not the root body",
        )
        active_joints = [p for p in sim.stage.Traverse() if p.IsA(UsdPhysics.Joint)]
        require(len(active_joints) == len(names), "Unexpected active joint count")
        require(all(p.IsA(UsdPhysics.RevoluteJoint) for p in active_joints), "Unexpected active constraint")
        if args.gui:
            from dex_workbench_tracking.preview import create_preview

            eye, center = create_preview(sim.stage, episode.wrist_position)
            sim.set_camera_view(eye, center)
        sim.reset()
        if args.gui:
            sim.render()
            print("GUI_READY: cyan line is the wrist reference; bounded physical replay", flush=True)
        require(not hand.is_fixed_base and hand.num_instances == 1, "Expected one floating PhysX articulation")
        require(
            set(hand.joint_names) == set(names) and len(hand.joint_names) == len(names), "Runtime DOF mapping mismatch"
        )
        require(hand.num_bodies == len(manifest["bodies"]), "Runtime body count mismatch")
        require(hand.body_names[0] == manifest["root_link"], "Unexpected root body ordering")
        runtime_ids = [hand.joint_names.index(n) for n in names]
        master_ids = [hand.joint_names.index(n) for n in prepared["independent_joint_names"]]
        master_columns = [names.index(n) for n in prepared["independent_joint_names"]]
        follower_ids = [hand.joint_names.index(eq["joint"]) for eq in prepared["mimic"]]

        def array(proxy):
            return proxy.torch.detach().cpu().numpy().copy()

        def tensor(values):
            return torch.as_tensor(values, dtype=torch.float32, device=hand.device)

        def state():
            pose = array(hand.data.root_link_pose_w)[0]
            # Installed Lab uses xyzw; HDF5/controller use wxyz. Explicit boundary.
            pose = xyzw_pose_to_reference(pose)
            velocity = array(hand.data.root_link_vel_w)[0]
            q = array(hand.data.joint_pos)[0, runtime_ids]
            return pose, velocity, q

        for gains in (hand.data.joint_stiffness, hand.data.joint_damping):
            require((array(gains)[0, follower_ids] == 0).all(), "Runtime follower gains are nonzero")
        masses = array(hand.data.body_mass)[0]
        require(np.isfinite(masses).all() and (masses > 0).all(), "Invalid runtime mass")
        require(masses.sum() * 9.81 < limits.force * 0.8, "Insufficient bounded gravity support headroom")
        lower = np.array([j["lower_rad"] for j in manifest["joints"]])
        upper = np.array([j["upper_rad"] for j in manifest["joints"]])
        traces, reset_states, summaries = [], [], []
        for repetition in range(2):
            failure_context = {"phase": "reset", "completed_repetitions": repetition, "repetition": repetition}
            hand.reset()
            hand.permanent_wrench_composer.reset()
            hand.instantaneous_wrench_composer.reset()
            pose0 = reference_pose_to_xyzw(reference.wrist_position[0], reference.wrist_quaternion[0])
            hand.write_root_link_pose_to_sim_index(root_pose=tensor(pose0[None]))
            hand.write_root_link_velocity_to_sim_index(root_velocity=tensor(np.zeros((1, 6))))
            hand.write_joint_position_to_sim_index(
                position=tensor(reference.joint_position[0:1]), joint_ids=runtime_ids
            )
            hand.write_joint_velocity_to_sim_index(velocity=tensor(np.zeros((1, len(names)))), joint_ids=runtime_ids)
            hand.set_joint_position_target_index(
                target=tensor(reference.joint_position[0:1, master_columns]), joint_ids=master_ids
            )
            hand.set_joint_velocity_target_index(target=tensor(np.zeros((1, len(master_ids)))), joint_ids=master_ids)
            hand.update(dt)
            reset_pose, reset_vel, reset_q = state()
            np.testing.assert_allclose(reset_pose[:3], reference.wrist_position[0], atol=1e-6)
            require(
                np.linalg.norm(rotation_error(reference.wrist_quaternion[0], reset_pose[3:])) < 1e-6,
                "Reset orientation mismatch",
            )
            np.testing.assert_allclose(reset_q, reference.joint_position[0], atol=1e-6)
            np.testing.assert_allclose(reset_vel, 0, atol=1e-6)
            np.testing.assert_allclose(array(hand.data.joint_vel), 0, atol=1e-6)
            reset_states.append(np.r_[reset_pose, reset_vel, reset_q])
            trace, errors = [], []
            for step in range(args.steps):
                failure_context = {
                    "phase": "stepping",
                    "completed_repetitions": repetition,
                    "repetition": repetition,
                    "completed_steps_in_repetition": step,
                }
                require(app.is_running(), "Application stopped before finite test completed")
                pose, velocity, q = state()
                force, torque = wrench(
                    reference.wrist_position[step],
                    reference.wrist_quaternion[step],
                    pose,
                    velocity,
                    array(hand.data.root_com_pose_w)[0, :3],
                    array(hand.data.body_com_pose_w)[0, :, :3],
                    masses,
                    limits,
                )
                hand.permanent_wrench_composer.set_forces_and_torques_index(
                    forces=tensor(force[None, None]),
                    torques=tensor(torque[None, None]),
                    body_ids=torch.tensor([0], dtype=torch.int32, device=hand.device),
                    is_global=True,
                )
                hand.set_joint_position_target_index(
                    target=tensor(reference.joint_position[step : step + 1, master_columns]), joint_ids=master_ids
                )
                hand.write_data_to_sim()
                sim.step(render=args.gui and step % 8 == 0)
                hand.update(dt)
                pose, velocity, q = state()
                qdot = array(hand.data.joint_vel)[0]
                failure_context["completed_steps_in_repetition"] = step + 1
                require(all(np.isfinite(v).all() for v in (pose, velocity, q, qdot)), "Nonfinite dynamic state")
                failure_context["measured_speeds_m_s_rad_s_rad_s"] = [
                    float(np.linalg.norm(velocity[:3])),
                    float(np.linalg.norm(velocity[3:])),
                    float(np.max(np.abs(qdot))),
                ]
                require(
                    np.linalg.norm(velocity[:3]) < 0.5
                    and np.linalg.norm(velocity[3:]) < 3
                    and np.max(np.abs(qdot)) < 2,
                    "Measured velocity safety bound exceeded",
                )
                require((q >= lower - 0.01).all() and (q <= upper + 0.01).all(), "Joint limit violation >0.01rad")
                residual = max(
                    abs(
                        q[names.index(eq["joint"])]
                        - eq["multiplier"] * q[names.index(eq["reference"])]
                        - eq["offset_rad"]
                    )
                    for eq in prepared["mimic"]
                )
                position_error = np.linalg.norm(pose[:3] - reference.wrist_position[step + 1])
                angle_error = np.linalg.norm(rotation_error(reference.wrist_quaternion[step + 1], pose[3:]))
                joint_error = np.max(np.abs(q - reference.joint_position[step + 1]))
                failure_context["last_errors_m_rad_rad_rad"] = [
                    float(position_error),
                    float(angle_error),
                    float(joint_error),
                    float(residual),
                ]
                require(residual < 0.002, "Mimic runtime residual >0.002rad; parser/constraint not verified")
                require(
                    position_error < 0.05 and angle_error < 0.5 and joint_error < 0.2,
                    "Tracking safety envelope exceeded",
                )
                trace.append(np.r_[pose, q, velocity, qdot])
                errors.append([position_error, angle_error, joint_error, residual])
            trace, errors = np.array(trace), np.array(errors)
            for eq in prepared["mimic"]:
                for name in (eq["joint"], eq["reference"]):
                    require(np.ptp(trace[:, 7 + names.index(name)]) > 0.002, f"No nontrivial runtime motion: {name}")
            require(
                np.max(np.linalg.norm(trace[:, :3] - reset_pose[:3], axis=1)) > 0.0005, "No nontrivial root translation"
            )
            require(
                max(np.linalg.norm(rotation_error(p[3:7], reset_pose[3:])) for p in trace) > 0.002,
                "No nontrivial root rotation",
            )
            traces.append(trace)
            summaries.append(
                {
                    "repetition": repetition,
                    "max_errors_m_rad_rad_rad": errors.max(axis=0).tolist(),
                    "rms_errors_m_rad_rad_rad": np.sqrt((errors**2).mean(axis=0)).tolist(),
                }
            )
        failure_context = {"phase": "repeatability", "completed_repetitions": 2, "steps_per_repetition": args.steps}
        np.testing.assert_allclose(reset_states[0], reset_states[1], atol=1e-6, rtol=0)
        position_end = 7 + len(names)
        np.testing.assert_allclose(traces[0][:, :position_end], traces[1][:, :position_end], atol=1e-3, rtol=0)
        # Root linear/angular and joint velocities: 1e-3 m/s or rad/s absolute.
        np.testing.assert_allclose(traces[0][:, position_end:], traces[1][:, position_end:], atol=1e-3, rtol=0)
        hand.permanent_wrench_composer.reset()
        print(
            json.dumps(
                {
                    "status": "PASS",
                    "check": "bounded_experimental_dynamic_tracking",
                    "runtime_verified_for_this_run_only": True,
                    "provenance": data.metadata["provenance"],
                    "reference_source": "hdf5" if args.hdf5 else "analytic_in_memory",
                    "reference_hdf5_sha256": hashlib.sha256(args.hdf5.read_bytes()).hexdigest() if args.hdf5 else None,
                    "reference_description": data.metadata["source_description"],
                    "world_anchor_inactive": not anchor.IsActive(),
                    "hand_side": data.metadata["hand_side"],
                    "effective_articulation_root_count": len(effective_roots),
                    "active_state_joint_count": len(active_joints),
                    "runtime_is_fixed_base": hand.is_fixed_base,
                    "seed": 42,
                    "num_envs": 1,
                    "steps_per_repetition": args.steps,
                    "reference_duration_s": float(episode.time[-1]),
                    "replayed_duration_s": args.steps * dt,
                    "reference_coverage": "full" if abs(args.steps * dt - episode.time[-1]) <= 1e-9 else "partial",
                    "full_episode_requested": args.full_episode,
                    "repetitions": 2,
                    "physx_version": version,
                    "prepared": prepared,
                    "limits_uncalibrated": vars(limits),
                    "metrics": summaries,
                    "hardware_and_training_validated": False,
                }
            ),
            flush=True,
        )
    except BaseException as error:
        # Kit fast shutdown may not return; emit the failure before closing.
        exit_code = 1
        traceback.print_exc()
        print(json.dumps({"status": "FAIL", "error": str(error), "progress": failure_context}), flush=True)
        raise
    finally:
        app.close(exit_code=exit_code)


if __name__ == "__main__":
    main()