"""Build a self-contained HTML viewer for a decoded raiden recording.

Reads:
  <session>/robot_data.npz              — proprio timelines
  <session>/mp4/<cam>.mp4               — decoded camera videos (per pyzed decode_svo2.py)
  <session>/mp4/<cam>.timestamps.json   — per-frame ns timestamps (from decode_svo2.py)

Writes:
  <session>/viewer.html                 — 2x2 video grid + synced proprio plots
  <session>/viewer_data.json            — proprio arrays + per-cam frame timestamps

The HTML uses Plotly for the joint/gripper plots and native <video> elements
for the 4 camera streams. All videos share a single scrub timeline — the
scene camera is the "master" and the others seek via linear-time mapping
(their per-frame timestamps get baked into the js so the sync is exact).
"""
import argparse
import json
from pathlib import Path

import numpy as np


HTML = """<!doctype html>
<html>
<head>
<meta charset='utf-8'>
<title>Raiden recording — {session}</title>
<script src='https://cdn.plot.ly/plotly-3.0.1.min.js'></script>
<style>
  body {{ font-family: -apple-system, sans-serif; background: #111; color: #ddd; margin: 12px; }}
  h1 {{ font-size: 16px; margin: 0 0 8px 0; color: #bbb; }}
  .grid {{ display: grid; grid-template-columns: 1fr 1fr; gap: 6px; max-width: 1400px; }}
  .cam {{ background: #000; position: relative; }}
  .cam label {{ position: absolute; top: 4px; left: 6px;
                background: rgba(0,0,0,0.55); padding: 2px 6px; font-size: 12px;
                border-radius: 3px; }}
  video, img {{ width: 100%; display: block; background: #000; }}
  .controls {{ margin: 10px 0; }}
  input[type=range] {{ width: 100%; }}
  .plots {{ max-width: 1400px; }}
  .plotrow {{ display: grid; grid-template-columns: 1fr 1fr; gap: 6px; }}
</style>
</head>
<body>
<h1>{session}  ·  {n_frames} proprio frames  ·  {duration:.1f}s @ {hz:.0f} Hz</h1>

<div class='grid'>
  <div class='cam'><label>scene_camera</label>
    <img id='i0' src='frames/scene_camera_first.jpg?raw&v={build_ts}'
      onerror="document.getElementById('info').textContent += 'FAIL: ' + this.currentSrc + '\\n'"></div>
  <div class='cam'><label>ego_camera</label>
    <img id='i1' src='frames/ego_camera_first.jpg?raw&v={build_ts}'
      onerror="document.getElementById('info').textContent += 'FAIL: ' + this.currentSrc + '\\n'"></div>
  <div class='cam'><label>left_wrist_camera</label>
    <img id='i2' src='frames/left_wrist_camera_first.jpg?raw&v={build_ts}'
      onerror="document.getElementById('info').textContent += 'FAIL: ' + this.currentSrc + '\\n'"></div>
  <div class='cam'><label>right_wrist_camera</label>
    <img id='i3' src='frames/right_wrist_camera_first.jpg?raw&v={build_ts}'
      onerror="document.getElementById('info').textContent += 'FAIL: ' + this.currentSrc + '\\n'"></div>
</div>

<div class='controls'>
  <button id='playbtn'>▶ play</button>
  <span id='clock' style='margin-left:12px; font-family: monospace;'>0.00 / {duration:.2f} s</span>
  <input type='range' id='scrub' min='0' max='{duration}' step='0.01' value='0'>
  <div id='info' style='font-family:monospace; color:#f80; margin-top:4px; white-space: pre;'></div>
  <div style='font-family:monospace; margin-top:8px; color:#8ac;'>Diagnostic: top-level jpg test →
    <img src='test_top.jpg?raw&v={build_ts}' style='height:60px; vertical-align:middle; background:#600;'
         onerror="this.style.background='#f00'; this.alt='FAIL'"
         onload="this.style.background='#060'"></div>
</div>

<div class='plots'>
  <div class='plotrow'>
    <div id='plot_left_joints'></div>
    <div id='plot_right_joints'></div>
  </div>
  <div class='plotrow'>
    <div id='plot_left_grip'></div>
    <div id='plot_right_grip'></div>
  </div>
</div>

<script>
const CAM_ORDER = ['scene_camera', 'ego_camera', 'left_wrist_camera', 'right_wrist_camera'];
const videos = [0,1,2,3].map(i => document.getElementById('v' + i));
// Cache-bust: append a build stamp so browsers re-fetch mp4s after we
// re-encode. Without this, the previous mp4v-encoded videos linger in
// cache and render as black frames.
const CB = '?raw&v={build_ts}';
CAM_ORDER.forEach((n, i) => {{
  const v = videos[i];
  v.src = 'mp4/' + n + '.mp4' + CB;
  v.load();
  v.addEventListener('error', () => {{
    document.getElementById('info').textContent =
      'ERR: video ' + n + ' failed to load: ' + (v.error ? v.error.message : '?');
  }});
  v.addEventListener('loadedmetadata', () => {{
    console.log(n, 'metadata loaded  dur=', v.duration, 's');
  }});
}});

fetch('viewer_data.json').then(r => r.json()).then(D => {{
  const dur = D.duration_s;
  // Plots
  const t_s = D.robot_t_s;
  const layout = (title) => ({{
    title: {{text: title, font: {{color:'#ddd', size:12}}}},
    paper_bgcolor: '#111', plot_bgcolor: '#111',
    margin: {{t:24, r:10, b:24, l:36}},
    height: 240,
    xaxis: {{color:'#aaa', gridcolor:'#222'}},
    yaxis: {{color:'#aaa', gridcolor:'#222'}},
    showlegend: false,
    shapes: [{{type:'line', x0:0, x1:0, y0:0, y1:1, yref:'paper',
              line:{{color:'#fa0', width:2}}, name:'cursor'}}],
  }});

  const jointColors = ['#e66','#e96','#dd6','#9d6','#6dd','#69e'];
  const mkTraces = (arr) => arr[0].map((_, j) => ({{
    x: t_s, y: arr.map(row => row[j]),
    type: 'scattergl', mode: 'lines',
    line: {{color: jointColors[j % jointColors.length], width: 1}},
    name: 'j' + j,
  }}));

  Plotly.newPlot('plot_left_joints', mkTraces(D.follower_l_joint_pos), layout('left joints (rad)'));
  Plotly.newPlot('plot_right_joints', mkTraces(D.follower_r_joint_pos), layout('right joints (rad)'));
  Plotly.newPlot('plot_left_grip', [{{
    x: t_s, y: D.follower_l_gripper_pos, mode:'lines',
    line:{{color:'#4bf', width:1.5}}}}], layout('left gripper'));
  Plotly.newPlot('plot_right_grip', [{{
    x: t_s, y: D.follower_r_gripper_pos, mode:'lines',
    line:{{color:'#fb4', width:1.5}}}}], layout('right gripper'));

  // Playback state
  let playing = false;
  const scrub = document.getElementById('scrub');
  const clock = document.getElementById('clock');
  const playbtn = document.getElementById('playbtn');

  function seekAll(t) {{
    // Map t_s (session-relative) to each video's currentTime.
    // Each cam has its own per-frame timestamps; find the frame whose ts is
    // nearest to (t0_ns + t_s*1e9) and set currentTime to (frame_idx / fps).
    for (let i = 0; i < 4; i++) {{
      const cam = CAM_ORDER[i];
      const info = D.cam_ts[cam];
      const target_ns = D.robot_t0_ns + t * 1e9;
      // Binary search for nearest ts index
      let lo = 0, hi = info.timestamps_ns.length - 1;
      while (lo < hi) {{
        const mid = (lo + hi) >> 1;
        if (info.timestamps_ns[mid] < target_ns) lo = mid + 1; else hi = mid;
      }}
      // Set to (idx / fps). This may differ from the exact stamp but is
      // close enough for scrubbing; browser video decoding snaps to keyframe.
      videos[i].currentTime = lo / info.fps;
    }}
    // Update cursor lines on each plot
    const updateCursor = (plotId) => {{
      const gd = document.getElementById(plotId);
      Plotly.relayout(gd, {{'shapes[0].x0': t, 'shapes[0].x1': t}});
    }};
    ['plot_left_joints','plot_right_joints','plot_left_grip','plot_right_grip']
        .forEach(updateCursor);
    clock.textContent = t.toFixed(2) + ' / ' + dur.toFixed(2) + ' s';
  }}

  // Seek only the SECONDARY videos + plot cursors + scrub label. Never
  // reseeks the master (scene_camera) during playback — that would feed
  // back into its own currentTime and freeze the video.
  function updateFromMaster(t) {{
    const target_ns = D.robot_t0_ns + t * 1e9;
    for (let i = 1; i < 4; i++) {{
      const info = D.cam_ts[CAM_ORDER[i]];
      let lo = 0, hi = info.timestamps_ns.length - 1;
      while (lo < hi) {{
        const mid = (lo + hi) >> 1;
        if (info.timestamps_ns[mid] < target_ns) lo = mid + 1; else hi = mid;
      }}
      const desired = lo / info.fps;
      // Only seek if drift > 100ms to avoid feedback thrash.
      if (Math.abs(videos[i].currentTime - desired) > 0.1) {{
        videos[i].currentTime = desired;
      }}
    }}
    ['plot_left_joints','plot_right_joints','plot_left_grip','plot_right_grip']
      .forEach(id => Plotly.relayout(document.getElementById(id),
                                     {{'shapes[0].x0': t, 'shapes[0].x1': t}}));
    clock.textContent = t.toFixed(2) + ' / ' + dur.toFixed(2) + ' s';
  }}

  scrub.addEventListener('input', () => {{
    playing = false; playbtn.textContent = '▶ play';
    videos.forEach(v => v.pause());
    seekAll(parseFloat(scrub.value));  // full sync for manual scrub
  }});

  playbtn.addEventListener('click', () => {{
    playing = !playing;
    playbtn.textContent = playing ? '⏸ pause' : '▶ play';
    if (playing) {{
      if (parseFloat(scrub.value) >= dur - 0.05) {{
        scrub.value = 0; seekAll(0);
      }}
      videos.forEach(v => v.play());
      // Tick reads master.currentTime → session t → syncs SECONDARIES only.
      const tick = () => {{
        if (!playing) return;
        const info = D.cam_ts.scene_camera;
        const frame_idx = Math.floor(videos[0].currentTime * info.fps);
        const idx = Math.max(0, Math.min(frame_idx, info.timestamps_ns.length - 1));
        const t = (info.timestamps_ns[idx] - D.robot_t0_ns) / 1e9;
        scrub.value = Math.max(0, Math.min(t, dur));
        updateFromMaster(parseFloat(scrub.value));
        if (t >= dur) {{ playing = false; playbtn.textContent = '▶ play'; }}
        else requestAnimationFrame(tick);
      }};
      requestAnimationFrame(tick);
    }} else {{
      videos.forEach(v => v.pause());
    }}
  }});

  // Init
  seekAll(0);
}});
</script>
</body>
</html>
"""


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("session_dir", type=Path)
    args = ap.parse_args()
    d = args.session_dir

    print(f"Loading {d}/robot_data.npz")
    z = np.load(d / "robot_data.npz", allow_pickle=True)
    t_ns = z["timestamps"].astype(np.int64)
    t0_ns = int(t_ns[0])
    dur_s = float((t_ns[-1] - t_ns[0]) / 1e9)
    t_s = ((t_ns - t0_ns) / 1e9).tolist()
    n = len(t_ns)

    cam_ts = {}
    for cam in ("scene_camera", "ego_camera", "left_wrist_camera", "right_wrist_camera"):
        p = d / "mp4" / f"{cam}.timestamps.json"
        if p.exists():
            cam_ts[cam] = json.load(p.open())
            print(f"  {cam}: {cam_ts[cam]['n_frames']} frames @ {cam_ts[cam]['fps']} Hz")
        else:
            print(f"  ⚠ {p} missing")

    data = {
        "session": d.name,
        "robot_t0_ns": t0_ns,
        "robot_t_s": t_s,
        "duration_s": dur_s,
        "follower_l_joint_pos": z["follower_l_joint_pos"].tolist(),
        "follower_r_joint_pos": z["follower_r_joint_pos"].tolist(),
        "follower_l_gripper_pos": z["follower_l_gripper_pos"][:, 0].tolist(),
        "follower_r_gripper_pos": z["follower_r_gripper_pos"][:, 0].tolist(),
        "cam_ts": cam_ts,
    }
    (d / "viewer_data.json").write_text(json.dumps(data))
    print(f"  wrote viewer_data.json ({(d / 'viewer_data.json').stat().st_size / 1e6:.1f} MB)")

    import time as _t
    build_ts = int(_t.time())
    html = HTML.format(session=d.name, duration=dur_s, n_frames=n, hz=n / dur_s,
                       build_ts=build_ts)
    (d / "viewer.html").write_text(html)
    print(f"  wrote viewer.html")


if __name__ == "__main__":
    main()
