#!/usr/bin/env python3
"""Headless demo-video recorder for microduck_rl (Cloof tooling).

Runs the official ONNX policies through the SAME loop as
scripts/infer_policy.py (BAM actuators, 50 Hz control, decimation 4),
renders offscreen with MuJoCo EGL, and pipes frames into ffmpeg.

Demos:
  walk      velstand protocol: stand -> forward -> strafe -> turn -> forward
  recovery  velstand: get pushed 1 m/s (curriculum cap) and recover
  kick      scene_ball.xml: kick left, then kick right
"""
import argparse
import json
import math
import os
import subprocess
import sys

os.environ.setdefault("MUJOCO_GL", "egl")

import numpy as np  # noqa: E402

REPO = "/home/cloofsiu/microduck_rl"
POLICIES = "/home/cloofsiu/microduck_policies"
sys.path.insert(0, os.path.join(REPO, "scripts"))
import infer_policy as ip  # noqa: E402
import mujoco  # noqa: E402

DEMO_DATA = {
    "walk": {
        "scene": ip.MICRODUCK_XML,
        "sessions": {"walking_onnx_path": f"{POLICIES}/../microduck_sprint/policy.onnx"
                     if os.path.exists(f"{POLICIES}/../microduck_sprint/policy.onnx")
                     else f"{POLICIES}/velstand.onnx"},
        "plan": [
            (0.5, "vel", (0.35, 0.0, 0.0)),
            (9.8, "vel", (0.0, 0.0, 0.0)),
        ],
        "seconds": 11.5,
    },
    "local": {
        "scene": f"{REPO}/src/mjlab_microduck/robot/microduck/scene_walk.xml",
        "sessions": {"walking_onnx_path": f"{REPO}/output.onnx"},
        "plan": [
            (0.5, "vel", (0.30, 0.0, 0.0)),
            (9.8, "vel", (0.0, 0.0, 0.0)),
        ],
        "seconds": 11.5,
    },
    "recovery": {
        "scene": ip.MICRODUCK_XML,
        "sessions": {"walking_onnx_path": f"{POLICIES}/velstand.onnx"},
        "plan": [
            (1.5, "push", (1.15, math.radians(105))),
            (7.5, "vel", (0.35, 0.0, 0.0)),
            (9.8, "vel", (0.0, 0.0, 0.0)),
        ],
        "seconds": 10.5,
        "push": 1.15,
    },
    "kick": {
        "scene": ip.MICRODUCK_BALL_XML,
        "sessions": {
            "walking_onnx_path": f"{POLICIES}/velstand.onnx",
            "kick_left_onnx_path": f"{POLICIES}/ball_kick_left.onnx",
            "kick_right_onnx_path": f"{POLICIES}/ball_kick_right.onnx",
            "kick_duration": 3.0,
        },
        "plan": [
            (1.2, "kick", "kick_left"),
            (4.5, "kick", "kick_right"),
        ],
        "seconds": 7.5,
    },
}


def trunk_addresses(model):
    jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "trunk_base_freejoint")
    return int(model.jnt_qposadr[jid]), int(model.jnt_dofadr[jid])


def build(spec, action_scale=1.0, width=1152, height=864):
    bam_model = ip.load_bam_model(ip.BAM_KP_FW, 7.4, ip.BAM_MAX_CURRENT)
    model, data, bam_ctrl, _names = ip.load_mujoco_with_bam(
        spec["scene"], bam_model, 0.005, None, ip.BAM_VIN_MIN)
    model.vis.global_.offwidth = max(width, 640)
    model.vis.global_.offheight = max(height, 480)
    policy = ip.PolicyInference(
        model, data, bam_ctrl=bam_ctrl, action_scale=action_scale,
        use_projected_gravity=True, new_cmd_obs=True,
        **spec["sessions"])
    qpos_adr, qvel_adr = trunk_addresses(model)
    data.qpos[qpos_adr:qpos_adr + 2] = 0.0
    data.qpos[qpos_adr + 2] = 0.125
    data.qpos[qpos_adr + 3:qpos_adr + 7] = [1, 0, 0, 0]
    for i, idx in enumerate(policy.joint_qpos_indices):
        data.qpos[idx] = policy.default_pose[i]
    if spec["scene"] == ip.MICRODUCK_BALL_XML:
        bj = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "ball_free")
        bqa = int(model.jnt_qposadr[bj])
        data.qpos[bqa:bqa + 3] = [-0.5, 0.35, 0.035]
        data.qpos[bqa + 3:bqa + 7] = [1, 0, 0, 0]
    bam_ctrl.reset(data.qpos)
    policy.set_position_targets(policy.default_pose)
    mujoco.mj_forward(model, data)
    policy.vel_max_x, policy.vel_min_x = 0.3, -0.3
    policy.vel_max_y, policy.vel_min_y = 0.2, -0.2
    policy.vel_max_ang = 1.5
    return model, data, bam_ctrl, policy, qpos_adr, qvel_adr


def plan_op(plan, t):
    for (pt, kind, args) in plan:
        if abs(t - pt) < 0.01:
            return kind, args
    return None, None


def run(demo, out_dir, width, height, fps_div, push_mag=None):
    spec = dict(DEMO_DATA[demo])
    if push_mag is not None:
        spec["plan"] = [(t, k, (push_mag, a) if k == "push" else v)
                        for (t, k, v) in spec["plan"]]
    model, data, bam_ctrl, policy, qpos_adr, qvel_adr = build(spec, width=width, height=height)
    os.makedirs(out_dir, exist_ok=True)
    mp4 = os.path.join(out_dir, f"duck_{demo}.mp4")

    from mujoco import Renderer
    metrics = {"demo": demo, "steps": 0, "control_dt": 0.02,
               "vel_avg": [], "frames": 0}
    renderer = None
    ffmpeg = None
    try:
        renderer = Renderer(model, height=height, width=width)
        cam = mujoco.MjvCamera()
        cam.azimuth = 130
        cam.elevation = -14
        cam.distance = 0.78

        ffmpeg = subprocess.Popen([
            "ffmpeg", "-y", "-loglevel", "error",
            "-f", "rawvideo", "-pix_fmt", "rgb24",
            "-s", f"{width}x{height}", "-r", "25", "-i", "pipe:",
            "-c:v", "libx264", "-preset", "medium", "-crf", "21",
            "-pix_fmt", "yuv420p", "-movflags", "+faststart", mp4,
        ], stdin=subprocess.PIPE)

        decim = 4
        control_dt = decim * model.opt.timestep
        n_steps = int(spec["seconds"] / control_dt)
        metrics["steps"] = n_steps
        metrics["control_dt"] = control_dt
        frame_every = max(1, fps_div)
        vel_hist = []
        start_pos = None
        for step in range(n_steps):
            t = step * control_dt
            kind, args = plan_op(spec["plan"], t)
            if kind == "vel":
                policy.set_vel_cmd(*args)
                print(f"[{t:5.1f}s] cmd -> {args}", flush=True)
            elif kind == "push":
                mag, ang = args
                data.qvel[qvel_adr + 0] = mag * math.cos(ang)
                data.qvel[qvel_adr + 1] = mag * math.sin(ang)
                print(f"[{t:5.1f}s] PUSH {mag:.2f} m/s @ {math.degrees(ang):.0f} deg", flush=True)
            elif kind == "kick":
                policy.trigger_behavior(args)
                print(f"[{t:5.1f}s] trigger {args}", flush=True)

            policy.update_ground_pick_phase(control_dt)
            policy.update_behavior(control_dt)
            action = policy.infer()
            policy.apply_action(action)
            for _ in range(decim):
                bam_ctrl.update()
                mujoco.mj_step(model, data)

            quat = data.qpos[qpos_adr + 3:qpos_adr + 7].astype("float32")
            v_w = data.qvel[qvel_adr:qvel_adr + 3].astype("float32")
            v_b = policy.quat_rotate_inverse(quat, v_w)
            vel_hist.append((float(v_b[0]), float(v_b[1]), float(t)))
            if len(vel_hist) > 50:
                vel_hist.pop(0)
            if step % 50 == 0:
                avg = tuple(float(sum(v[i] for v in vel_hist) / len(vel_hist)) for i in (0, 1))
                metrics["vel_avg"].append({"t": round(float(t), 2), "fwd": round(avg[0], 3),
                                           "lat": round(avg[1], 3),
                                           "cmd": [float(policy.vel_cmd[0]), float(policy.vel_cmd[1])],
                                           "trunk_z": round(float(data.qpos[qpos_adr + 2]), 4),
                                           "policy": policy.current_policy})
            if start_pos is None:
                start_pos = np.array(data.qpos[qpos_adr:qpos_adr + 2])
            if step % frame_every == 0:
                p = data.qpos[qpos_adr:qpos_adr + 3]
                cam.lookat[:] = [p[0], p[1], max(0.02, float(p[2]) * 0.55)]
                renderer.update_scene(data, camera=cam)
                renderer.scene.flags[mujoco.mjtRndFlag.mjRND_SHADOW] = 0
                img = renderer.render()
                ffmpeg.stdin.write(img.tobytes())
                metrics["frames"] += 1
        ffmpeg.stdin.close()
        ffmpeg.wait()
        metrics["travelled_m"] = round(float(np.abs(data.qpos[qpos_adr:qpos_adr + 2] - start_pos).sum()), 3)
        with open(os.path.join(out_dir, f"duck_{demo}_metrics.json"), "w") as f:
            json.dump(metrics, f, indent=2)
    finally:
        if renderer is not None:
            try:
                renderer.close()
            except Exception as exc:
                print(f"renderer close warning: {exc}", flush=True)
        if ffmpeg is not None and ffmpeg.stdin is not None and not ffmpeg.stdin.closed:
            ffmpeg.stdin.close()
            ffmpeg.wait(timeout=60)
    print(f"DONE {mp4}  travelled={metrics.get('travelled_m')}", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--demo", required=True, choices=list(DEMO_DATA))
    ap.add_argument("--out", required=True)
    ap.add_argument("--width", type=int, default=1152)
    ap.add_argument("--height", type=int, default=864)
    ap.add_argument("--fps-div", type=int, default=2)
    ap.add_argument("--push", type=float, default=None)
    a = ap.parse_args()
    run(a.demo, a.out, a.width, a.height, a.fps_div, a.push)


if __name__ == "__main__":
    main()