import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
# Load models locally (no external text-encoder service) and prefer the HF cache.
os.environ.setdefault("TEXT_ENCODER_MODE", "local")

import base64
import gzip
import json
import random
import sys
import time
import xml.etree.ElementTree as ET
from pathlib import Path

import spaces  # must precede torch / CUDA-touching imports
import torch
import numpy as np
import trimesh
import gradio as gr

# The vendored ARDY package lives next to this file.
sys.path.insert(0, str(Path(__file__).resolve().parent))

from ardy.model import load_model  # noqa: E402
from ardy.model.load_model import load_text_encoder  # noqa: E402
from ardy.motion_rep.tools import length_to_mask  # noqa: E402
from ardy.tools import seed_everything, to_numpy  # noqa: E402

# Two rigs are offered in the playground:
#   "human" -> ARDY-Core-RP-20FPS-Horizon40, 27-joint skeleton @ 20 fps
#   "robot" -> ARDY-G1-RP-25FPS-Horizon52, 34-joint Unitree G1 robot @ 25 fps
# The robot rig matches the sibling Space hugging-apps/ardy-g1-motion-generation.
DEFAULT_RIG = "human"
MAX_SEED = 2**31 - 1

# -----------------------------------------------------------------------------
# Model loading (module scope).
#
# ZeroGPU has no live GPU at startup: it intercepts .to("cuda") *placement* but
# NOT arbitrary CUDA compute. LLM2Vec's PEFT load runs a LoRA `merge_and_unload`
# (real matmuls), so we build everything on CPU first, then .to("cuda") — the
# placement call is what the ZeroGPU hijack packs to disk and streams into VRAM
# on the first @spaces.GPU request.
# -----------------------------------------------------------------------------
# transformers / PEFT / safetensors infer their load device from
# torch.cuda.is_available(), which ZeroGPU reports True at startup even though no
# real GPU is attached — so a plain load tries to place weights on cuda and dies
# with "No CUDA GPUs are available". Force every loader onto CPU by masking
# is_available() during construction, then restore it and .to("cuda") (which the
# ZeroGPU hijack packs to disk + streams into VRAM on the first request).
#
# A single text encoder is built once and shared across both motion models
# (load_text_encoder is explicitly designed for this — see its docstring).
_real_cuda_available = torch.cuda.is_available
torch.cuda.is_available = lambda: False
try:
    print("Loading ARDY text encoder (LLM2Vec-Llama-3-8B) on CPU…", flush=True)
    _text_encoder = load_text_encoder(mode="local", device="cpu")

    print("Loading ARDY human (core) motion model on CPU…", flush=True)
    MODEL_HUMAN = load_model("core", device="cpu", text_encoder=_text_encoder)

    print("Loading ARDY robot (G1) motion model on CPU…", flush=True)
    MODEL_ROBOT = load_model("g1", device="cpu", text_encoder=_text_encoder)
finally:
    torch.cuda.is_available = _real_cuda_available

print("Moving models to CUDA (ZeroGPU-intercepted placement)…", flush=True)
_text_encoder = _text_encoder.to("cuda")
for _m in (MODEL_HUMAN, MODEL_ROBOT):
    _m.to("cuda")
    # self.device was captured at construction (device="cpu"); generation creates
    # many tensors on it, so retarget it to cuda to match the moved weights.
    _m.device = "cuda"
    _m.eval()


def _rig_params(model):
    fps = float(model.motion_rep.fps)
    skeleton = model.skeleton
    parents = skeleton.joint_parents.cpu().numpy().astype(int).tolist()
    patch = model.num_frames_per_token
    gen_horizon = model.gen_horizon_len
    num_base_steps = int(model.diffusion.num_base_steps)
    # History carried between autoregressive windows. The reference streaming demo
    # keeps this SHORT so a new prompt takes effect within a window.
    hist_crop = max(patch, (4 // patch) * patch)
    root_idx = parents.index(-1) if -1 in parents else 0
    return {
        "model": model,
        "fps": fps,
        "skeleton": skeleton,
        "parents": parents,
        "patch": patch,
        "gen_horizon": gen_horizon,
        "num_base_steps": num_base_steps,
        "hist_crop": hist_crop,
        "root_idx": root_idx,
    }


RIGS = {
    "human": _rig_params(MODEL_HUMAN),
    "robot": _rig_params(MODEL_ROBOT),
}
# The largest base-step count across rigs bounds the diffusion-steps slider.
NUM_BASE_STEPS = max(r["num_base_steps"] for r in RIGS.values())
for _name, _r in RIGS.items():
    print(
        f"[{_name}] rig ready: {_r['skeleton'].nbjoints} joints, {_r['fps']} fps, "
        f"horizon {_r['gen_horizon']}, patch {_r['patch']}, "
        f"hist_crop {_r['hist_crop']}, base_steps {_r['num_base_steps']}",
        flush=True,
    )


def _normalize_rig(rig) -> str:
    """Map any UI/API rig value onto a valid RIGS key ('human' | 'robot')."""
    if rig is None:
        return DEFAULT_RIG
    key = str(rig).strip().lower()
    if key in RIGS:
        return key
    if key.startswith("hum") or "core" in key or "person" in key:
        return "human"
    if key.startswith("rob") or "g1" in key or "unitree" in key:
        return "robot"
    return DEFAULT_RIG


# -----------------------------------------------------------------------------
# HUMAN rig: skinned body mesh (ARDY "CoreSkin" linear-blend skinning).
#
# The reference viz (ardy/viz/viser_utils.py) renders a smooth humanoid body by
# skinning a bind mesh with the per-frame *global* joint transforms:
#     verts = CoreSkin.skin(global_rot_mats, posed_joints, rot_is_global=True)
# The browser holds the static skin data (bind vertices / faces / LBS
# indices+weights) and does the per-vertex blend, while the server sends only the
# tiny per-frame joint affine matrices
#     A[f,j] = fk[f,j] @ bind_rig_transform_inv[j]      (fk = [R_global | pos])
# so the payload stays small (~0.1 MB/clip).
# -----------------------------------------------------------------------------
_HUMAN_SKEL = RIGS["human"]["skeleton"]
_SKIN_PATH = Path(_HUMAN_SKEL.folder) / "skin_standard.npz"
_skin = np.load(_SKIN_PATH)
BIND_RIG_INV = np.linalg.inv(
    np.asarray(_skin["bind_rig_transform"], dtype=np.float64)
).astype(np.float32)  # [J, 4, 4]


def _build_skin_blob():
    """Pack the static skin data into one gzip+base64 blob (loaded once by the
    browser). Layout: bind_vertices f32[V,3] | faces u32[F,3] | lbs_idx u8[V,W]
    | lbs_wt f32[V,W]."""
    bind_v = np.asarray(_skin["bind_vertices"], dtype=np.float32)
    faces = np.asarray(_skin["faces"], dtype=np.uint32)
    idx = np.asarray(_skin["lbs_indices"], dtype=np.uint8)
    wt = np.asarray(_skin["lbs_weights"], dtype=np.float32)
    raw = (
        np.ascontiguousarray(bind_v).tobytes()
        + np.ascontiguousarray(faces).tobytes()
        + np.ascontiguousarray(idx).tobytes()
        + np.ascontiguousarray(wt).tobytes()
    )
    meta = {"V": int(bind_v.shape[0]), "F": int(faces.shape[0]), "W": int(idx.shape[1])}
    return base64.b64encode(gzip.compress(raw, 6)).decode("ascii"), meta


SKIN_B64, SKIN_META = _build_skin_blob()
print(f"CoreSkin ready: {SKIN_META['V']} verts / {SKIN_META['F']} faces, "
      f"blob {len(SKIN_B64) // 1024} KB", flush=True)


def _joint_affines_human(global_rot_mats: np.ndarray, posed_joints: np.ndarray) -> str:
    """Per-frame joint affine matrices A = fk @ bind_rig_inv, base64 f32 [T,J,12].

    global_rot_mats: [T, J, 3, 3]; posed_joints: [T, J, 3]."""
    T, J = posed_joints.shape[:2]
    fk = np.tile(np.eye(4, dtype=np.float32), (T, J, 1, 1))
    fk[..., :3, :3] = global_rot_mats.astype(np.float32)
    fk[..., :3, 3] = posed_joints.astype(np.float32)
    A = (fk @ BIND_RIG_INV)[..., :3, :]  # [T, J, 3, 4]
    A = np.ascontiguousarray(A.reshape(T, J, 12).astype(np.float32))
    return base64.b64encode(A.tobytes()).decode("ascii")


# -----------------------------------------------------------------------------
# ROBOT rig: G1 robot mesh rig (rigid per-joint STL meshes).
#
# Unlike the human skeleton (rendered with one skinned body mesh via LBS), the
# Unitree G1 robot is rendered by attaching a rigid STL mesh to each articulated
# joint. This mirrors ardy/viz/g1_rig.py (G1MeshRig): each mesh has a local
# transform (geom_pos, geom_rot) relative to its joint, read from the MuJoCo
# g1.xml, plus a coordinate change from MuJoCo to ARDY axes. We precompute — for
# each mesh — its geometry PRE-TRANSFORMED into the joint-local frame
# (v' = geom_rot @ v + geom_pos), so at render time the browser just applies the
# per-frame joint transform: world_v = joint_pos + joint_rot @ v'.
# -----------------------------------------------------------------------------
# G1 joint -> STL mesh mapping (mirrors ardy/viz/g1_rig.py G1_MESH_JOINT_MAP).
G1_MESH_JOINT_MAP = {
    "pelvis_skel": ["pelvis.STL", "pelvis_contour_link.STL"],
    "left_hip_pitch_skel": ["left_hip_pitch_link.STL"],
    "left_hip_roll_skel": ["left_hip_roll_link.STL"],
    "left_hip_yaw_skel": ["left_hip_yaw_link.STL"],
    "left_knee_skel": ["left_knee_link.STL"],
    "left_ankle_pitch_skel": ["left_ankle_pitch_link.STL"],
    "left_ankle_roll_skel": ["left_ankle_roll_link.STL"],
    "right_hip_pitch_skel": ["right_hip_pitch_link.STL"],
    "right_hip_roll_skel": ["right_hip_roll_link.STL"],
    "right_hip_yaw_skel": ["right_hip_yaw_link.STL"],
    "right_knee_skel": ["right_knee_link.STL"],
    "right_ankle_pitch_skel": ["right_ankle_pitch_link.STL"],
    "right_ankle_roll_skel": ["right_ankle_roll_link.STL"],
    "waist_yaw_skel": ["waist_yaw_link_rev_1_0.STL", "waist_yaw_link.STL"],
    "waist_roll_skel": ["waist_roll_link_rev_1_0.STL", "waist_roll_link.STL"],
    "waist_pitch_skel": [
        "torso_link_rev_1_0.STL",
        "torso_link.STL",
        "logo_link.STL",
        "head_link.STL",
    ],
    "left_shoulder_pitch_skel": ["left_shoulder_pitch_link.STL"],
    "left_shoulder_roll_skel": ["left_shoulder_roll_link.STL"],
    "left_shoulder_yaw_skel": ["left_shoulder_yaw_link.STL"],
    "left_elbow_skel": ["left_elbow_link.STL"],
    "left_wrist_roll_skel": ["left_wrist_roll_link.STL"],
    "left_wrist_pitch_skel": ["left_wrist_pitch_link.STL"],
    "left_wrist_yaw_skel": ["left_wrist_yaw_link.STL", "left_rubber_hand.STL"],
    "right_shoulder_pitch_skel": ["right_shoulder_pitch_link.STL"],
    "right_shoulder_roll_skel": ["right_shoulder_roll_link.STL"],
    "right_shoulder_yaw_skel": ["right_shoulder_yaw_link.STL"],
    "right_elbow_skel": ["right_elbow_link.STL"],
    "right_wrist_roll_skel": ["right_wrist_roll_link.STL"],
    "right_wrist_pitch_skel": ["right_wrist_pitch_link.STL"],
    "right_wrist_yaw_skel": ["right_wrist_yaw_link.STL", "right_rubber_hand.STL"],
}

_ROBOT_SKEL = RIGS["robot"]["skeleton"]
_MUJOCO_TO_ARDY = np.array(
    [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float64
)
_G1_SKEL_DIR = Path(_ROBOT_SKEL.folder)
_G1_MESH_DIR = _G1_SKEL_DIR / "meshes" / "g1"
_G1_XML = _G1_SKEL_DIR / "xml" / "g1.xml"


def _quat_wxyz_to_matrix(wxyz: np.ndarray) -> np.ndarray:
    w, x, y, z = wxyz
    n = np.sqrt(w * w + x * x + y * y + z * z)
    if n < 1e-12:
        return np.eye(3)
    w, x, y, z = w / n, x / n, y / n, z / n
    return np.array(
        [
            [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],
            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],
            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],
        ],
        dtype=np.float64,
    )


def _mesh_local_transforms() -> dict:
    """mesh_file -> (geom_pos[3], geom_rot[3x3]) in ARDY axes, parsed from g1.xml."""
    if not _G1_XML.exists():
        return {}
    root = ET.parse(_G1_XML).getroot()
    file_to_name = {}
    for mesh in root.findall(".//asset/mesh"):
        name, file = mesh.get("name"), mesh.get("file")
        if name and file:
            file_to_name[file] = name
    name_to_tf = {}
    for geom in root.findall(".//geom"):
        name = geom.get("mesh")
        if name is None:
            continue
        pos = geom.get("pos")
        quat = geom.get("quat")
        gp = np.zeros(3) if pos is None else np.array([float(v) for v in pos.split()])
        gr_ = np.eye(3) if quat is None else _quat_wxyz_to_matrix(
            np.array([float(v) for v in quat.split()])
        )
        name_to_tf[name] = (gp, gr_)
    out = {}
    for file, name in file_to_name.items():
        gp, gr_ = name_to_tf.get(name, (np.zeros(3), np.eye(3)))
        gp = _MUJOCO_TO_ARDY @ gp
        gr_ = _MUJOCO_TO_ARDY @ gr_ @ _MUJOCO_TO_ARDY.T
        out[file] = (gp, gr_)
    return out


def _build_g1_mesh_blob():
    """Pack all rigid G1 meshes, pre-transformed into joint-local frame, into one
    gzip+base64 blob. Returns (b64, meta). meta.parts lists per-mesh
    {joint, v_off, v_cnt}. Layout: all verts f32[Vtot,3] then all faces
    u32[Ftot,3] (face indices are GLOBAL into the concatenated vertex array)."""
    skeleton = _ROBOT_SKEL
    local_tf = _mesh_local_transforms()
    all_v = []
    all_f = []
    parts = []
    v_cursor = 0
    for joint_name, mesh_files in G1_MESH_JOINT_MAP.items():
        if joint_name not in skeleton.bone_index:
            continue
        joint_idx = int(skeleton.bone_index[joint_name])
        for mesh_file in mesh_files:
            mp = _G1_MESH_DIR / mesh_file
            if not mp.exists():
                continue
            mesh = trimesh.load_mesh(str(mp), process=True)
            if isinstance(mesh, trimesh.Scene):
                mesh = trimesh.util.concatenate(mesh.dump())
            verts = np.asarray(mesh.vertices, dtype=np.float64) @ _MUJOCO_TO_ARDY.T
            faces = np.asarray(mesh.faces, dtype=np.int64)
            gp, gr_ = local_tf.get(mesh_file, (np.zeros(3), np.eye(3)))
            # Pre-apply the mesh's joint-local transform: v' = geom_rot @ v + geom_pos.
            verts = (verts @ gr_.T) + gp
            vcnt = verts.shape[0]
            parts.append({"joint": joint_idx, "v_off": v_cursor, "v_cnt": vcnt})
            all_v.append(verts.astype(np.float32))
            all_f.append((faces + v_cursor).astype(np.uint32))  # global vertex indices
            v_cursor += vcnt
    V = np.concatenate(all_v, axis=0) if all_v else np.zeros((0, 3), np.float32)
    F = np.concatenate(all_f, axis=0) if all_f else np.zeros((0, 3), np.uint32)
    raw = np.ascontiguousarray(V).tobytes() + np.ascontiguousarray(F).tobytes()
    meta = {"Vtot": int(V.shape[0]), "Ftot": int(F.shape[0]), "parts": parts}
    b64 = base64.b64encode(gzip.compress(raw, 6)).decode("ascii")
    return b64, meta


G1_MESH_B64, G1_MESH_META = _build_g1_mesh_blob()
print(
    f"G1 rig ready: {len(G1_MESH_META['parts'])} meshes, "
    f"{G1_MESH_META['Vtot']} verts / {G1_MESH_META['Ftot']} faces, "
    f"blob {len(G1_MESH_B64) // 1024} KB",
    flush=True,
)


def _joint_transforms_robot(global_rot_mats: np.ndarray, posed_joints: np.ndarray) -> str:
    """Per-frame joint affine matrices [R_global | pos], base64 f32 [T,J,12].

    The browser applies world_v = pos + R_global @ v' per mesh (v' already in
    joint-local frame)."""
    T, J = posed_joints.shape[:2]
    A = np.zeros((T, J, 3, 4), dtype=np.float32)
    A[..., :3, :3] = global_rot_mats.astype(np.float32)
    A[..., :3, 3] = posed_joints.astype(np.float32)
    A = np.ascontiguousarray(A.reshape(T, J, 12).astype(np.float32))
    return base64.b64encode(A.tobytes()).decode("ascii")


# --- Autoregressive generation with a persistable latent state ---------------
# ARDY is autoregressive: it generates one `gen_horizon_len`-frame window at a
# time, conditioned on a history of previous frames. `autoregressive_step` is
# the streaming primitive — it returns the *normalized motion-feature tensor*
# for (history + new window), which can be fed straight back in as the next
# window's history. We thread that tensor to (a) fill a requested clip length
# and (b) CONTINUE a clip with a new prompt, exactly like the reference
# interactive demo. Persisting the tensor in a gr.State lets a second "Continue"
# click resume from where the first clip ended, on the same character.
def _generate_sequence(rig_key, prompt, num_new_frames, steps, cfg_weight, init_tensor):
    """Run the AR loop for one prompt on the selected rig. `init_tensor`:
    normalized feature tensor [1, Th, D] on cuda (prior motion to continue), or
    None to start fresh. Returns the full normalized feature tensor."""
    r = RIGS[rig_key]
    model = r["model"]
    gen_horizon = r["gen_horizon"]
    patch = r["patch"]
    hist_crop = r["hist_crop"]
    text_feat, text_pad_mask = model._encode_text([prompt])
    motion_tensor = init_tensor
    target_new = max(1, int(np.ceil(num_new_frames / gen_horizon))) * gen_horizon
    produced = 0
    while produced < target_new:
        if motion_tensor is None:
            history, hist_len = None, 0
        else:
            hist_len = (min(motion_tensor.shape[1], hist_crop) // patch) * patch
            history = motion_tensor[:, motion_tensor.shape[1] - hist_len:] if hist_len else None
            hist_len = history.shape[1] if history is not None else 0
        samples = model.autoregressive_step(
            num_frames=hist_len + gen_horizon,   # exactly history + one window (no future)
            num_denoising_steps=steps,
            motion_mask=None,
            observed_motion=None,
            cfg_weight=float(cfg_weight),
            texts=None,
            text_feat=text_feat,
            text_pad_mask=text_pad_mask,
            init_history_sequence=history,
            init_global_translation=None,   # first window -> defaults (origin / +Z heading)
            init_first_heading_angle=None,
        )
        new_tail = samples[:, hist_len:]     # the freshly generated window
        motion_tensor = new_tail if motion_tensor is None else torch.cat([motion_tensor, new_tail], dim=1)
        produced += new_tail.shape[1]
    return motion_tensor


def _pack_payload(rig_key, motion_tensor, prompt, seed):
    """Decode a normalized feature tensor to the browser payload (per-frame joint
    affines + root ground-track), tagged with the rig so the viewer loads the
    correct skeleton/model."""
    r = RIGS[rig_key]
    model = r["model"]
    with torch.no_grad():
        out = to_numpy(model.motion_rep.inverse(motion_tensor, is_normalized=True))
    posed = np.asarray(out["posed_joints"])[0]      # [T, J, 3] global joint positions
    grm = np.asarray(out["global_rot_mats"])[0]      # [T, J, 3, 3] global rotations
    if rig_key == "robot":
        affines = _joint_transforms_robot(grm, posed)
    else:
        affines = _joint_affines_human(grm, posed)
    return {
        "rig": rig_key,
        "fps": r["fps"],
        "num_frames": int(posed.shape[0]),
        "num_joints": int(posed.shape[1]),
        "affines": affines,                          # drives the browser rig
        "root": np.round(posed[:, r["root_idx"], :].astype(np.float32), 4).tolist(),
        "prompt": prompt,
        "seed": seed,
    }


def _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, init_np):
    rig_key = _normalize_rig(rig)
    r = RIGS[rig_key]
    prompt = (prompt or "").strip()
    if not prompt:
        raise gr.Error("Please enter a text prompt describing the motion.")
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)
    seed_everything(seed)
    steps = max(1, min(int(diffusion_steps), r["num_base_steps"]))
    num_new = max(r["patch"], int(round(float(duration) * r["fps"])))
    init_tensor = None if init_np is None else torch.from_numpy(init_np).to("cuda")

    t0 = time.perf_counter()
    with torch.no_grad():
        full = _generate_sequence(rig_key, prompt, num_new, steps, cfg_weight, init_tensor)
    payload = _pack_payload(rig_key, full, prompt, seed)
    tag = "continue" if init_np is not None else "generate"
    print(f"[{tag}:{rig_key}] '{prompt[:50]}' +{num_new}f -> {full.shape[1]}f total "
          f"steps={steps} seed={seed} {time.perf_counter() - t0:.1f}s", flush=True)
    return json.dumps(payload), seed, full.detach().cpu().numpy()


@spaces.GPU
def ui_generate(prompt, rig=DEFAULT_RIG, duration=5.0, diffusion_steps=NUM_BASE_STEPS,
                cfg_weight=2.0, seed=0, randomize_seed=True):
    """Start a fresh clip (resets the running sequence)."""
    return _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, None)


@spaces.GPU
def ui_continue(prompt, rig=DEFAULT_RIG, duration=5.0, diffusion_steps=NUM_BASE_STEPS,
                cfg_weight=2.0, seed=0, randomize_seed=True, state=None):
    """Append a new action, continuing from the previous clip's final pose."""
    return _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, state)


@spaces.GPU
def generate_motion(prompt: str, rig: str = DEFAULT_RIG, duration: float = 5.0,
                    diffusion_steps: int = NUM_BASE_STEPS, cfg_weight: float = 2.0,
                    seed: int = 0, randomize_seed: bool = True) -> tuple[str, int]:
    """Generate a 3D motion clip from a text prompt with ARDY.

    Args:
        prompt: Natural-language description of the motion (e.g. "a person walks in a circle").
        rig: Which character to animate — "human" (27-joint skeleton) or "robot" (Unitree G1).
        duration: Length of the generated motion in seconds.
        diffusion_steps: Number of denoising steps (1..num_base_steps).
        cfg_weight: Classifier-free-guidance weight for the text prompt.
        seed: Random seed for reproducibility.
        randomize_seed: If True, ignore `seed` and draw a fresh random one.

    Returns:
        A JSON string with the animated skeleton payload plus the seed used.
    """
    payload_json, seed, _ = _core_generate(
        rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, None
    )
    return payload_json, seed


# -----------------------------------------------------------------------------
# Front-end: a self-contained Three.js playground, delivered as a Gradio-native
# custom HTML component (Gradio 6 `gr.HTML` templates + js_on_load).
#
# The motion JSON is carried as the component's own `value` prop. `js_on_load`
# imports three.js, builds the scene once, wires the controls, then registers a
# `watch('value', ...)` callback that Gradio fires whenever the component is
# updated as the output of a Python event (Generate button / Examples).
#
# Each payload is tagged with its `rig`. The viewer ships the static data for
# BOTH rigs (human skin blob + G1 rigid-mesh blob) and switches at load time:
#  - "human": one skinned body mesh (linear-blend skinning in the browser).
#  - "robot": rigid per-joint G1 STL meshes posed by the joint transforms.
# -----------------------------------------------------------------------------

PLAYER_TEMPLATE = """
<div class="ardy-playground">
  <div class="ardy-canvas-wrap">
    <div class="ardy-hint">Generate a motion to load it into the playground.</div>
  </div>
  <div class="ardy-controls">
    <button class="ardy-play ardy-btn" type="button">▶ Play</button>
    <input class="ardy-scrub" type="range" min="0" max="0" value="0" step="1" />
    <span class="ardy-frame">0 / 0</span>
    <label class="ardy-lbl">Speed
      <select class="ardy-speed">
        <option value="0.25">0.25×</option>
        <option value="0.5">0.5×</option>
        <option value="1" selected>1×</option>
        <option value="2">2×</option>
      </select>
    </label>
    <label class="ardy-lbl"><input type="checkbox" class="ardy-loop" checked/> Loop</label>
    <label class="ardy-lbl"><input type="checkbox" class="ardy-trail"/> Root trail</label>
  </div>
  <div class="ardy-caption"></div>
</div>
"""

# css_template rules are auto-scoped to this component by Gradio.
PLAYER_CSS_TEMPLATE = """
.ardy-playground { width: 100%; }
.ardy-canvas-wrap {
  position: relative; width: 100%; height: 480px;
  border-radius: 12px; overflow: hidden;
  background: #ffffff;
  border: 1px solid #e5e7eb;
}
.ardy-canvas-wrap canvas { display:block; width:100% !important; height:100% !important; }
.ardy-hint {
  position:absolute; top:50%; left:50%; transform:translate(-50%,-50%);
  color:#98a2b3; font-size:14px; text-align:center; pointer-events:none;
}
.ardy-controls {
  display:flex; align-items:center; gap:12px; flex-wrap:wrap;
  margin-top:10px; padding:8px 4px;
}
.ardy-controls .ardy-btn {
  background:#76B900; color:#fff;
  border:none; border-radius:8px; padding:6px 16px; cursor:pointer; font-weight:600;
}
.ardy-scrub { flex:1; min-width:160px; accent-color:#76B900; }
.ardy-frame { font-variant-numeric: tabular-nums; color: var(--body-text-color); min-width:70px; }
.ardy-lbl { font-size:13px; color: var(--body-text-color); display:flex; align-items:center; gap:4px; }
.ardy-caption { margin-top:6px; font-size:13px; color:#667085; }
"""

APP_CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

# Runs once, when the component first renders. `element`, `props`, and `watch`
# are injected by Gradio. We import three.js, decode the static skin data (human
# rig) and the static rigid-mesh data (robot rig), build the scene, and subscribe
# to value changes with `watch('value', ...)`. Each generated payload carries the
# per-frame joint affine matrices plus a `rig` tag; the viewer renders whichever
# rig the payload requests, swapping the on-screen mesh as needed.
PLAYER_JS_ON_LOAD = r"""
const root = element;
const q = (sel) => root.querySelector(sel);

const state = {
  ready:false,
  THREE:null, OrbitControls:null,
  renderer:null, scene:null, camera:null, controls:null,
  trailLine:null,
  data:null, verts:null, frame:0, playing:false, lastT:0,
  speed:1, loop:true, trail:false, pending:null,
  rig:null,                 // which rig mesh is currently mounted in the scene
  // human rig
  skin:null, humanMesh:null, humanGeom:null,
  // robot rig
  robotMesh:null, robotGeom:null,
  robotBaseVerts:null, robotFaces:null, robotVtot:0, robotParts:null,
};

// --- binary helpers ---------------------------------------------------------
function b64ToBytes(b64){
  const bin = atob(b64); const out = new Uint8Array(bin.length);
  for(let i=0;i<bin.length;i++) out[i]=bin.charCodeAt(i);
  return out;
}
async function gunzip(bytes){
  const ds = new DecompressionStream("gzip");
  const buf = await new Response(new Blob([bytes]).stream().pipeThrough(ds)).arrayBuffer();
  return buf;
}

// Decode the one-time static human skin blob into typed arrays.
async function decodeSkin(){
  const buf = await gunzip(b64ToBytes(ARDY_SKIN_B64));
  const V = ARDY_SKIN_META.V, F = ARDY_SKIN_META.F, W = ARDY_SKIN_META.W;
  let o = 0;
  const bindV = new Float32Array(buf.slice(o, o+V*3*4)); o += V*3*4;
  const faces = new Uint32Array(buf.slice(o, o+F*3*4));  o += F*3*4;
  const idx   = new Uint8Array(buf.slice(o, o+V*W));      o += V*W;
  const wt    = new Float32Array(buf.slice(o, o+V*W*4));  o += V*W*4;
  return {V, F, W, bindV, faces, idx, wt};
}

// Decode the one-time static robot rig blob: pre-transformed mesh vertices
// (joint-local frame) + global-indexed faces + per-mesh part table.
async function decodeRig(){
  const buf = await gunzip(b64ToBytes(G1_MESH_B64));
  const V = G1_MESH_META.Vtot, F = G1_MESH_META.Ftot;
  let o = 0;
  const baseVerts = new Float32Array(buf.slice(o, o+V*3*4)); o += V*3*4;
  const faces     = new Uint32Array(buf.slice(o, o+F*3*4));  o += F*3*4;
  return {V, F, baseVerts, faces, parts: G1_MESH_META.parts};
}

// HUMAN: per-vertex linear-blend skinning for every frame (once per clip).
// A[f,j] is a 3x4 affine (row-major, 12 floats); posed vertex =
// sum_k w_k * A[idx_k] @ [bind_x, bind_y, bind_z, 1].
function skinAllFrames(A, T, J){
  const s = state.skin, V = s.V, W = s.W, bindV = s.bindV, idx = s.idx, wt = s.wt;
  const frames = new Array(T);
  for(let f=0; f<T; f++){
    const out = new Float32Array(V*3);
    const Ab = f*J*12;
    for(let v=0; v<V; v++){
      const bx = bindV[v*3], by = bindV[v*3+1], bz = bindV[v*3+2];
      let x=0, y=0, z=0;
      for(let k=0; k<W; k++){
        const w = wt[v*W+k]; if(w===0) continue;
        const a = Ab + idx[v*W+k]*12;
        x += w*(A[a]*bx   + A[a+1]*by  + A[a+2]*bz  + A[a+3]);
        y += w*(A[a+4]*bx + A[a+5]*by  + A[a+6]*bz  + A[a+7]);
        z += w*(A[a+8]*bx + A[a+9]*by  + A[a+10]*bz + A[a+11]);
      }
      out[v*3]=x; out[v*3+1]=y; out[v*3+2]=z;
    }
    frames[f] = out;
  }
  return frames;
}

// ROBOT: per-frame rigid transform: for every mesh part, world_v = pos + R @ v'
// where (R,pos) is the driving joint's global transform this frame and v' is
// the vertex already baked into that joint's local frame.
function poseAllFrames(A, T, J){
  const V = state.robotVtot, base = state.robotBaseVerts, parts = state.robotParts;
  const frames = new Array(T);
  for(let f=0; f<T; f++){
    const out = new Float32Array(V*3);
    const Ab = f*J*12;
    for(let p=0; p<parts.length; p++){
      const jp = parts[p];
      const a = Ab + jp.joint*12;
      const r0=A[a],   r1=A[a+1], r2=A[a+2],  px=A[a+3];
      const r3=A[a+4], r4=A[a+5], r5=A[a+6],  py=A[a+7];
      const r6=A[a+8], r7=A[a+9], r8=A[a+10], pz=A[a+11];
      const vs = jp.v_off, ve = jp.v_off + jp.v_cnt;
      for(let v=vs; v<ve; v++){
        const bx=base[v*3], by=base[v*3+1], bz=base[v*3+2];
        out[v*3]   = px + r0*bx + r1*by + r2*bz;
        out[v*3+1] = py + r3*bx + r4*by + r5*bz;
        out[v*3+2] = pz + r6*bx + r7*by + r8*bz;
      }
    }
    frames[f] = out;
  }
  return frames;
}

// --- three.js scene ---------------------------------------------------------
function initScene(){
  const THREE = state.THREE, OrbitControls = state.OrbitControls;
  const wrap = q(".ardy-canvas-wrap");
  if(!wrap || state.renderer) return;
  const w = wrap.clientWidth || 800, h = wrap.clientHeight || 480;
  const scene = new THREE.Scene(); scene.background = new THREE.Color(0xffffff);
  const camera = new THREE.PerspectiveCamera(42, w/h, 0.05, 200);
  camera.position.set(3.8, 2.2, 4.7);
  const renderer = new THREE.WebGLRenderer({antialias:true});
  renderer.setSize(w, h); renderer.setPixelRatio(Math.min(window.devicePixelRatio,2));
  renderer.shadowMap.enabled = true; renderer.shadowMap.type = THREE.PCFSoftShadowMap;
  wrap.appendChild(renderer.domElement);
  const controls = new OrbitControls(camera, renderer.domElement);
  controls.target.set(0, 0.9, 0); controls.enableDamping = true;

  scene.add(new THREE.HemisphereLight(0xffffff, 0xdfe4ee, 1.4));
  const key = new THREE.DirectionalLight(0xffffff, 1.5);
  key.position.set(3, 6, 4); key.castShadow = true;
  key.shadow.mapSize.set(2048, 2048);
  const c = key.shadow.camera; c.near=0.5; c.far=25; c.left=-3; c.right=3; c.top=3; c.bottom=-3;
  key.shadow.bias = -0.0004;
  scene.add(key);
  scene.add(new THREE.DirectionalLight(0xeef2ff, 0.35).translateX(-4).translateZ(-2));

  const ground = new THREE.Mesh(
    new THREE.PlaneGeometry(40, 40),
    new THREE.ShadowMaterial({opacity:0.16})
  );
  ground.rotation.x = -Math.PI/2; ground.position.y = 0; ground.receiveShadow = true;
  scene.add(ground);

  const grid = new THREE.GridHelper(10, 20, 0xc4c9d4, 0xe4e7ee);
  grid.position.y = 0.0015; scene.add(grid);

  state.renderer=renderer; state.scene=scene; state.camera=camera; state.controls=controls;

  new ResizeObserver(()=>{
    const w2 = wrap.clientWidth, h2 = wrap.clientHeight;
    if(w2>0 && h2>0){ camera.aspect=w2/h2; camera.updateProjectionMatrix(); renderer.setSize(w2,h2); }
  }).observe(wrap);

  animate();
}

// Mount the mesh for the requested rig (lazily built, then shown/hidden). Only
// one rig mesh is visible at a time; both share the scene once created.
function mountRig(rig){
  const THREE = state.THREE;
  if(rig === "robot"){
    if(!state.robotMesh){
      const geom = new THREE.BufferGeometry();
      geom.setIndex(new THREE.BufferAttribute(state.robotFaces, 1));
      geom.setAttribute("position", new THREE.BufferAttribute(new Float32Array(state.robotVtot*3), 3));
      const mat = new THREE.MeshStandardMaterial({color:0xd7dde6, roughness:0.5, metalness:0.55});
      const mesh = new THREE.Mesh(geom, mat);
      mesh.castShadow = true; mesh.frustumCulled = false;
      state.scene.add(mesh); state.robotMesh = mesh; state.robotGeom = geom;
    }
    if(state.humanMesh) state.humanMesh.visible = false;
    state.robotMesh.visible = true;
    return state.robotGeom;
  } else {
    if(!state.humanMesh){
      const s = state.skin;
      const geom = new THREE.BufferGeometry();
      geom.setIndex(new THREE.BufferAttribute(s.faces, 1));
      geom.setAttribute("position", new THREE.BufferAttribute(new Float32Array(s.V*3), 3));
      const mat = new THREE.MeshStandardMaterial({color:0x98bdff, roughness:0.85, metalness:0.0});
      const mesh = new THREE.Mesh(geom, mat);
      mesh.castShadow = true; mesh.frustumCulled = false;
      state.scene.add(mesh); state.humanMesh = mesh; state.humanGeom = geom;
    }
    if(state.robotMesh) state.robotMesh.visible = false;
    state.humanMesh.visible = true;
    return state.humanGeom;
  }
}

function activeGeom(){
  return (state.rig === "robot") ? state.robotGeom : state.humanGeom;
}

function setFrame(f){
  if(!state.verts) return;
  const T = state.data.num_frames;
  f = Math.max(0, Math.min(T-1, f|0));
  state.frame = f;
  const geom = activeGeom();
  if(!geom) return;
  const pos = geom.getAttribute("position");
  pos.array.set(state.verts[f]);
  pos.needsUpdate = true;
  geom.computeVertexNormals();
  geom.computeBoundingSphere();
  updateTrail();
  const scrub = q(".ardy-scrub"); if(scrub) scrub.value = f;
  const lbl = q(".ardy-frame"); if(lbl) lbl.textContent = (f+1)+" / "+T;
}

function updateTrail(){
  const THREE = state.THREE;
  if(!state.data || !state.data.root){ if(state.trailLine) state.trailLine.visible=false; return; }
  if(!state.trail){ if(state.trailLine) state.trailLine.visible=false; return; }
  const T = state.data.num_frames, root = state.data.root;
  if(!state.trailLine){
    const g = new THREE.BufferGeometry();
    g.setAttribute("position", new THREE.BufferAttribute(new Float32Array(T*3),3));
    state.trailLine = new THREE.Line(g, new THREE.LineBasicMaterial({color:0xf59e0b}));
    state.scene.add(state.trailLine);
  }
  state.trailLine.visible = true;
  const attr = state.trailLine.geometry.getAttribute("position");
  for(let t=0;t<T;t++){ attr.setXYZ(t, root[t][0], 0.006, root[t][2]); }
  attr.needsUpdate = true;
  state.trailLine.geometry.setDrawRange(0, Math.max(1, state.frame+1));
}

function animate(){
  requestAnimationFrame(animate);
  if(!state.renderer) return;
  const now = performance.now();
  if(state.playing && state.verts){
    const dt = (now - state.lastT)/1000;
    const fps = state.data.fps * state.speed;
    if(dt >= 1/Math.max(1e-3,fps)){
      state.lastT = now;
      let nf = state.frame + 1;
      if(nf >= state.data.num_frames){
        if(state.loop){ nf = 0; } else { nf = state.data.num_frames-1; setPlaying(false); }
      }
      setFrame(nf);
    }
  }
  state.controls.update();
  state.renderer.render(state.scene, state.camera);
}

function setPlaying(p){
  state.playing = p;
  const b = q(".ardy-play");
  if(b) b.textContent = p ? "⏸ Pause" : "▶ Play";
  state.lastT = performance.now();
}

function loadData(data){
  initScene();
  const rig = (data.rig === "robot") ? "robot" : "human";
  state.rig = rig;
  mountRig(rig);
  state.data = data;
  // Decode per-frame affines and pre-pose every frame (one-time cost per clip).
  const T = data.num_frames, J = data.num_joints;
  const Abytes = b64ToBytes(data.affines);
  const A = new Float32Array(Abytes.buffer, Abytes.byteOffset, Abytes.byteLength/4);
  state.verts = (rig === "robot") ? poseAllFrames(A, T, J) : skinAllFrames(A, T, J);
  if(state.trailLine){ state.scene.remove(state.trailLine); state.trailLine.geometry.dispose(); state.trailLine=null; }
  const scrub = q(".ardy-scrub"); if(scrub){ scrub.max = T-1; scrub.value = 0; }
  const hint = q(".ardy-hint"); if(hint) hint.style.display = "none";
  const cap = q(".ardy-caption");
  if(cap) cap.textContent = '"' + data.prompt + '"  ·  ' + (rig==="robot"?"robot":"human") +
    '  ·  ' + T + ' frames @ ' + data.fps + ' fps  ·  seed ' + data.seed;
  root.dataset.ardyLoaded = "1";
  root.dataset.ardyRig = rig;
  root.dataset.ardyFrames = String(T);
  setFrame(0);
  setPlaying(true);
}

function applyPayload(payload){
  if(!payload) return;
  if(!state.ready){ state.pending = payload; return; }   // three.js / rigs still loading
  try { loadData(JSON.parse(payload)); }
  catch(e){ console.error("ARDY playground load error", e); }
}

function wireControls(){
  const bind = (sel, ev, fn) => {
    const el = q(sel); if(el && !el.dataset.wired){ el.dataset.wired="1"; el.addEventListener(ev, fn); }
  };
  bind(".ardy-play", "click", ()=> setPlaying(!state.playing));
  bind(".ardy-scrub", "input", (e)=>{ setPlaying(false); setFrame(parseInt(e.target.value)); });
  bind(".ardy-speed", "change", (e)=>{ state.speed = parseFloat(e.target.value); });
  bind(".ardy-loop", "change", (e)=>{ state.loop = e.target.checked; });
  bind(".ardy-trail", "change", (e)=>{ state.trail = e.target.checked; updateTrail(); });
}

// Gradio-native hand-off: render whenever the component's value prop updates as
// the output of a Python event (Generate / Continue / Examples).
if (typeof watch === "function") {
  watch("value", () => applyPayload(props.value));
}

// Import three.js (esm.sh, not jsDelivr: the OrbitControls addon has an internal
// bare `import ... from "three"` that a browser dynamic import() can't resolve
// without an import map; esm.sh rewrites it and dedupes three) and decode both
// rigs, then flush any value that already arrived.
Promise.all([
  import("https://esm.sh/three@0.160.0"),
  import("https://esm.sh/three@0.160.0/examples/jsm/controls/OrbitControls.js"),
  decodeSkin(),
  decodeRig(),
]).then(([THREE, oc, skin, rig]) => {
  state.THREE = THREE;
  state.OrbitControls = oc.OrbitControls;
  state.skin = skin;
  state.robotBaseVerts = rig.baseVerts;
  state.robotFaces = rig.faces;
  state.robotVtot = rig.V;
  state.robotParts = rig.parts;
  state.ready = true;
  wireControls();
  initScene();
  const start = state.pending || props.value;
  if (start) applyPayload(start);
}).catch((e)=> console.error("ARDY viewer init failed", e));
"""

EXAMPLES = [
    ["A person walks forward confidently.", "human", 5.0],
    ["A person walks in a circle.", "human", 6.0],
    ["A person jumps up and down.", "human", 4.0],
    ["The robot walks forward confidently.", "robot", 5.0],
    ["The robot waves with the right hand.", "robot", 4.0],
    ["The robot crouches down and then stands back up.", "robot", 5.0],
]

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🕺 ARDY Motion Playground
            Interactive text-to-motion with **[ARDY](https://research.nvidia.com/labs/sil/projects/ardy/)**
            (Autoregressive Diffusion with Hybrid Representation) by NVIDIA.
            Pick a **rig** (human or robot), type a prompt and **Generate** a 3D motion clip,
            then **orbit, scrub, and play** it below.
            Chain actions with **Continue ▸** — the same character keeps going from where it stopped.
            """
        )

        # Running latent state (normalized feature tensor, CPU) — lets "Continue"
        # resume the same character from the end of the previous clip.
        seq_state = gr.State(None)

        prompt = gr.Textbox(
            label="Motion prompt",
            placeholder="e.g. a person walks in a circle then waves",
            lines=2,
        )
        rig = gr.Radio(
            choices=[("🕺 Human", "human"), ("🤖 Robot (Unitree G1)", "robot")],
            value=DEFAULT_RIG,
            label="Rig",
        )
        with gr.Row():
            run = gr.Button("Generate", variant="primary", scale=2)
            cont = gr.Button("Continue ▸", variant="secondary", scale=1)

        # The interactive 3D playground — a Gradio-native custom HTML component.
        # Its `value` (the motion JSON) is set directly by the Generate handler;
        # `watch('value', ...)` in js_on_load renders it. Ship the static data for
        # both rigs (human skin blob + G1 rigid-mesh blob) once, as a header
        # prepended to js_on_load; the per-frame affines ride in each payload.
        _rig_header = (
            f'const ARDY_SKIN_B64="{SKIN_B64}";\n'
            f"const ARDY_SKIN_META={json.dumps(SKIN_META)};\n"
            f'const G1_MESH_B64="{G1_MESH_B64}";\n'
            f"const G1_MESH_META={json.dumps(G1_MESH_META)};\n"
        )
        player = gr.HTML(
            value="",
            html_template=PLAYER_TEMPLATE,
            css_template=PLAYER_CSS_TEMPLATE,
            js_on_load=_rig_header + PLAYER_JS_ON_LOAD,
            elem_id="ardy-player",
        )

        with gr.Accordion("Advanced settings", open=False):
            duration = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Duration (seconds)")
            diffusion_steps = gr.Slider(
                1, NUM_BASE_STEPS, value=NUM_BASE_STEPS, step=1, label="Diffusion steps"
            )
            cfg_weight = gr.Slider(1.0, 6.0, value=2.0, step=0.5, label="Text guidance (CFG)")
            with gr.Row():
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                seed = gr.Number(label="Seed", value=0, precision=0)

        gr.Examples(
            examples=EXAMPLES,
            inputs=[prompt, rig, duration],
            outputs=[player, seed, seq_state],
            fn=ui_generate,
            cache_examples=False,
            run_on_click=True,
        )

        gr.Markdown(
            """
            <small>Rigs: **ARDY-Core-RP-20FPS-Horizon40** (human, 27-joint skeleton @ 20 fps)
            and **ARDY-G1-RP-25FPS-Horizon52** (Unitree G1 robot, 34-joint skeleton @ 25 fps).
            Text encoder: LLM2Vec-Llama-3-8B. Post-processing (foot-skate cleanup) is disabled in this demo.
            Motion is generated autoregressively; longer clips take longer.
            **Generate** starts a new clip; **Continue ▸** keeps the same character going,
            transitioning it into the new prompt (like the reference demo's prompt timeline).</small>
            """
        )

    _gen_inputs = [prompt, rig, duration, diffusion_steps, cfg_weight, seed, randomize_seed]
    # Generate starts fresh; Continue resumes from the running latent state.
    # The payload is written straight into the player's `value`; its js_on_load
    # `watch('value', ...)` renders it (loading the correct rig from the payload).
    run.click(fn=ui_generate, inputs=_gen_inputs,
              outputs=[player, seed, seq_state], api_name=False)
    cont.click(fn=ui_continue, inputs=_gen_inputs + [seq_state],
               outputs=[player, seed, seq_state], api_name=False)

    # Clean single-shot endpoint for the HTTP API / MCP tool (no session state).
    gr.api(generate_motion, api_name="generate")

demo.queue()

if __name__ == "__main__":
    # Gradio 6 moved theme/css from the Blocks constructor to launch(). The
    # player's JS/CSS now live on the gr.HTML component itself (js_on_load /
    # css_template), so no global `head=` script is needed.
    demo.launch(
        theme=gr.themes.Citrus(),
        css=APP_CSS,
        mcp_server=True,
        ssr_mode=False,
    )