#!/usr/bin/env python3
"""
viz_predictions.py — Enhanced Web visualizer for dVLM-AD predictions.json

New Features:
- Filter by navigation command, ADE ranges, zero predictions, etc.
- Bulk export of filtered subsets
- URL state persistence for filters
- Better performance with server-side filtering

Usage:
  python viz_predictions.py --pred_json predictions.json [--port 7860]
  Then open http://localhost:7860 in a browser.

Example filters:
  http://localhost:7860?nav=GO_STRAIGHT&ade_max=2.0&hide_zero=1
"""

import argparse, json, os, re, urllib.parse
from flask import Flask, jsonify, send_file, render_template_string, abort, request

app = Flask(__name__)
ALL_PREDICTIONS = []     # Original full dataset
FILTERED_INDICES = []    # Current filtered indices into ALL_PREDICTIONS
METRICS_MAP = {}         # sample_id → {ade_3s, ade_5s, fde}
CURRENT_FILTERS = {}     # Current active filters

# ── parse ego state from question text ───────────────────────────────────────
EGO_RE = re.compile(
    r'\(t([+-]\d+\.\d+)s\)\s*\[([^\]]+)\].*?Velocity:\s*X\s*([+-]?\d+\.?\d*),\s*Y\s*([+-]?\d+\.?\d*)',
    re.DOTALL
)

def parse_ego(question):
    rows = []
    for m in EGO_RE.finditer(question):
        t = float(m.group(1))
        vx = float(m.group(3))
        vy = float(m.group(4))
        speed = (vx**2 + vy**2) ** 0.5
        rows.append({"t": t, "vx": round(vx, 2), "vy": round(vy, 2), "speed": round(speed, 2)})
    return rows

# ── filtering logic ──────────────────────────────────────────────────────────

def extract_nav_commands():
    """Extract all unique navigation commands from dataset."""
    navs = set()
    for p in ALL_PREDICTIONS:
        nav = p.get('navigation_command', '').strip()
        if nav:
            navs.add(nav)
    return sorted(list(navs))

def get_ade_range():
    """Get ADE range from dataset for filter bounds."""
    ades = []
    for p in ALL_PREDICTIONS:
        metrics = p.get('metrics') or {}
        ade = metrics.get('ade_3s') or metrics.get('ade')
        if isinstance(ade, (int, float)) and not (ade is None):
            ades.append(ade)

        # Also check METRICS_MAP
        sid = p.get('sample_id', '')
        if sid in METRICS_MAP:
            ade = METRICS_MAP[sid].get('ade_3s')
            if isinstance(ade, (int, float)) and not (ade is None):
                ades.append(ade)

    if ades:
        return min(ades), max(ades)
    return 0.0, 10.0

def is_zero_prediction(pred_traj):
    """Check if trajectory is all zeros."""
    if not pred_traj:
        return True
    return all(v == 0 for pt in pred_traj for v in (pt if isinstance(pt, (list, tuple)) else []))

def apply_filters(filters):
    """Apply filters and return list of matching indices."""
    global CURRENT_FILTERS
    CURRENT_FILTERS = dict(filters)

    indices = []
    for i, p in enumerate(ALL_PREDICTIONS):
        # Navigation filter
        if filters.get('nav') and filters['nav'] != 'ALL':
            nav = p.get('navigation_command', '').strip()
            if nav != filters['nav']:
                continue

        # ADE range filter
        if filters.get('ade_min') is not None or filters.get('ade_max') is not None:
            # Try to get ADE from metrics or METRICS_MAP
            ade = None
            metrics = p.get('metrics') or {}
            ade = metrics.get('ade_3s') or metrics.get('ade')

            sid = p.get('sample_id', '')
            if ade is None and sid in METRICS_MAP:
                ade = METRICS_MAP[sid].get('ade_3s')

            if ade is None:
                continue  # Skip samples without ADE data

            if filters.get('ade_min') is not None and ade < filters['ade_min']:
                continue
            if filters.get('ade_max') is not None and ade > filters['ade_max']:
                continue

        # Zero prediction filter
        if filters.get('hide_zero'):
            pred_traj = p.get('pred_trajectory')
            if is_zero_prediction(pred_traj):
                continue

        # Valid trajectory filter
        if filters.get('valid_only'):
            pred_traj = p.get('pred_trajectory')
            if not pred_traj or len(pred_traj) == 0:
                continue

        indices.append(i)

    return indices

# ── HTML template with enhanced filtering UI ─────────────────────────────────
HTML = r"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width,initial-scale=1">
<title>dVLM-AD Prediction Viewer</title>
<style>
*{box-sizing:border-box;margin:0;padding:0}
body{font-family:'Segoe UI',system-ui,sans-serif;background:#0f1117;color:#e2e8f0;min-height:100vh}

/* Header with filters */
header{background:#1a1d2e;border-bottom:1px solid #2d3148;padding:12px 24px;position:sticky;top:0;z-index:100}
.header-top{display:flex;align-items:center;gap:16px;margin-bottom:8px}
.header-top h1{font-size:1.1rem;font-weight:700;color:#7c93f7;white-space:nowrap}
.counter{color:#94a3b8;font-size:.9rem;margin-left:auto;white-space:nowrap}

.filters{display:flex;gap:12px;align-items:center;flex-wrap:wrap;padding-top:8px;border-top:1px solid #2d3148}
.filter-group{display:flex;align-items:center;gap:6px}
.filter-group label{font-size:.8rem;color:#94a3b8;white-space:nowrap}
.filter-group select, .filter-group input{padding:4px 8px;border-radius:4px;border:1px solid #3d4268;
  background:#252840;color:#e2e8f0;font-size:.8rem}
.filter-group input[type="number"]{width:60px}
.filter-group input[type="checkbox"]{margin:0}

.filter-btn{padding:4px 12px;border-radius:4px;border:none;background:#4f5fad;color:#fff;
  cursor:pointer;font-size:.8rem;margin-left:8px}
.filter-btn:hover{background:#6272c4}
.export-btn{background:#059669;margin-left:4px}
.export-btn:hover{background:#047857}

.jump-box{display:flex;gap:6px;align-items:center;margin-left:auto}
.jump-box input{width:70px;padding:4px 8px;border-radius:6px;border:1px solid #3d4268;
  background:#252840;color:#e2e8f0;font-size:.85rem;text-align:center}
.jump-box button{padding:4px 10px;border-radius:6px;border:none;background:#4f5fad;
  color:#fff;cursor:pointer;font-size:.85rem}

.main{display:grid;grid-template-columns:1fr 1fr;gap:16px;max-width:1600px;margin:0 auto;padding:16px}
@media(max-width:900px){.main{grid-template-columns:1fr}}

.card{background:#1a1d2e;border:1px solid #2d3148;border-radius:12px;overflow:hidden}
.card-title{padding:10px 16px;font-size:.8rem;font-weight:600;text-transform:uppercase;
  letter-spacing:.05em;color:#94a3b8;border-bottom:1px solid #2d3148;
  display:flex;align-items:center;gap:8px}

/* LEFT COLUMN */
.left-col{display:flex;flex-direction:column;gap:16px}
.img-wrap{position:relative;background:#000;border-radius:0}
.img-wrap img{width:100%;display:block;object-fit:contain;max-height:340px}
.img-placeholder{width:100%;height:200px;display:flex;align-items:center;
  justify-content:center;color:#4a5568;font-size:.85rem}

.info-grid{display:grid;grid-template-columns:repeat(2,1fr);gap:10px;padding:14px}
.info-item label{display:block;font-size:.7rem;color:#64748b;text-transform:uppercase;
  letter-spacing:.05em;margin-bottom:2px}
.info-item .val{font-size:1.1rem;font-weight:700}
.nav-badge{display:inline-block;padding:3px 12px;border-radius:20px;font-size:.85rem;
  font-weight:700;letter-spacing:.03em}
.nav-straight{background:#1a3a1a;color:#4ade80;border:1px solid #166534}
.nav-left    {background:#1a2a3a;color:#60a5fa;border:1px solid #1e40af}
.nav-right   {background:#3a1a1a;color:#f87171;border:1px solid #991b1b}
.metric-ok   {color:#4ade80}
.metric-warn {color:#facc15}
.metric-bad  {color:#f87171}
.metric-zero {color:#9333ea}

.ego-table{width:100%;border-collapse:collapse;font-size:.8rem}
.ego-table th{padding:4px 8px;text-align:right;color:#64748b;font-weight:500;border-bottom:1px solid #2d3148}
.ego-table th:first-child{text-align:left}
.ego-table td{padding:4px 8px;text-align:right;border-bottom:1px solid #1e2133}
.ego-table td:first-child{text-align:left;color:#94a3b8}
.ego-table tr:last-child td{color:#7c93f7;font-weight:600}

/* RIGHT COLUMN */
.right-col{display:flex;flex-direction:column;gap:16px}
.traj-wrap{padding:10px}
canvas{display:block;width:100%;border-radius:8px;background:#0d1117}
.behavior-grid{display:grid;grid-template-columns:1fr 1fr;gap:8px;padding:12px}
.beh-item{background:#0d1117;border-radius:8px;padding:10px;text-align:center}
.beh-item label{display:block;font-size:.7rem;color:#64748b;text-transform:uppercase;margin-bottom:4px}
.beh-item .bval{font-size:.95rem;font-weight:600;color:#a78bfa}

.objects-grid{display:flex;flex-wrap:wrap;gap:6px;padding:12px}
.obj-chip{padding:3px 10px;border-radius:20px;font-size:.75rem;font-weight:500}
.obj-yes{background:#3a1a1a;color:#fca5a5;border:1px solid #7f1d1d}
.obj-no {background:#1a2a1a;color:#6b7280;border:1px solid #374151}

.explanation{padding:12px;font-size:.82rem;line-height:1.6;color:#cbd5e1;
  max-height:160px;overflow-y:auto}
.raw-out{padding:10px 12px;font-size:.75rem;font-family:monospace;
  background:#0d1117;max-height:200px;overflow:auto;line-height:1.5;
  color:#94a3b8;white-space:pre-wrap;word-break:break-all;border-radius:0 0 12px 12px}

/* PAGINATION */
.pager{display:flex;align-items:center;justify-content:center;gap:10px;
  padding:16px;background:#1a1d2e;border-top:1px solid #2d3148;position:sticky;bottom:0}
.pager button{padding:8px 20px;border-radius:8px;border:none;cursor:pointer;
  font-size:.9rem;font-weight:600;transition:.15s}
.btn-prev,.btn-next{background:#2d3148;color:#e2e8f0}
.btn-prev:hover,.btn-next:hover{background:#3d4268}
.btn-prev:disabled,.btn-next:disabled{opacity:.3;cursor:default}
.pager .page-info{font-size:.9rem;color:#64748b;min-width:120px;text-align:center}
.progress{height:3px;background:#2d3148;width:100%;max-width:200px;border-radius:2px}
.progress-fill{height:100%;background:#4f5fad;border-radius:2px;transition:.3s}

.loading{text-align:center;padding:40px;color:#4a5568;font-size:1rem}
.tag{display:inline-block;padding:1px 7px;border-radius:4px;font-size:.7rem;
  font-weight:600;vertical-align:middle;margin-left:6px}
.tag-zero{background:#4a1d7a;color:#d8b4fe}
.tag-filtered{background:#0f766e;color:#5eead4}

.filter-summary{font-size:.75rem;color:#64748b;margin-top:4px}
</style>
</head>
<body>

<header>
  <div class="header-top">
    <h1>🚗 dVLM-AD Prediction Viewer</h1>
    <div class="progress"><div class="progress-fill" id="prog"></div></div>
    <span class="counter" id="counter">— / —</span>
    <div class="jump-box">
      <input type="number" id="jumpInput" placeholder="#" min="1">
      <button onclick="jumpTo()">Go</button>
    </div>
  </div>

  <div class="filters">
    <div class="filter-group">
      <label>Navigation:</label>
      <select id="navFilter">
        <option value="ALL">All Commands</option>
        {% for nav in nav_commands %}
        <option value="{{ nav }}"{% if current_nav == nav %} selected{% endif %}>{{ nav }}</option>
        {% endfor %}
      </select>
    </div>

    <div class="filter-group">
      <label>ADE 3s:</label>
      <input type="number" id="adeMin" step="0.1" placeholder="min" value="{{ current_ade_min or '' }}">
      <span style="color:#64748b">—</span>
      <input type="number" id="adeMax" step="0.1" placeholder="max" value="{{ current_ade_max or '' }}">
    </div>

    <div class="filter-group">
      <input type="checkbox" id="hideZero"{% if current_hide_zero %} checked{% endif %}>
      <label for="hideZero">Hide Zero Pred</label>
    </div>

    <div class="filter-group">
      <input type="checkbox" id="validOnly"{% if current_valid_only %} checked{% endif %}>
      <label for="validOnly">Valid Only</label>
    </div>

    <button class="filter-btn" onclick="applyFilters()">Apply</button>
    <button class="filter-btn export-btn" onclick="exportFiltered()">Export</button>
  </div>

  <div class="filter-summary" id="filterSummary">{{ filter_summary }}</div>
</header>

<div class="main" id="main"><div class="loading">Loading…</div></div>

<div class="pager">
  <button class="btn-prev" id="prevBtn" onclick="navigate(-1)" disabled>◀ Prev</button>
  <span class="page-info" id="pageInfo"></span>
  <button class="btn-next" id="nextBtn" onclick="navigate(1)">Next ▶</button>
</div>

<script>
let idx = 0, total = {{ filtered_count }};
const originalTotal = {{ total_count }};

function esc(s){ return String(s).replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;') }

function navClass(nav){
  if(!nav) return '';
  nav = nav.toUpperCase();
  if(nav.includes('LEFT'))  return 'nav-left';
  if(nav.includes('RIGHT')) return 'nav-right';
  return 'nav-straight';
}

function metricClass(v){
  if(v===null||v===undefined) return 'metric-zero';
  if(v < 1)  return 'metric-ok';
  if(v < 3)  return 'metric-warn';
  return 'metric-bad';
}

function fmtNum(v, d=2){
  if(v===null||v===undefined||isNaN(v)) return '—';
  return (+v).toFixed(d);
}

function drawTrajectory(canvasId, predTraj, gtTraj){
  const canvas = document.getElementById(canvasId);
  if(!canvas) return;
  const dpr = window.devicePixelRatio || 1;
  const W = canvas.clientWidth || 560;
  const H = 280;
  canvas.width  = W * dpr;
  canvas.height = H * dpr;
  const ctx = canvas.getContext('2d');
  ctx.scale(dpr, dpr);

  const all = [...(gtTraj||[]), ...(predTraj||[]), [0,0]];
  const xs = all.map(p=>p[0]), ys = all.map(p=>p[1]);
  const xmin = Math.min(...xs)-2, xmax = Math.max(...xs)+2;
  const ymin = Math.min(...ys)-2, ymax = Math.max(...ys)+2;

  function toScreen(wx, wy){
    const sx = (wx - xmin)/(xmax - xmin) * (W-40) + 20;
    const sy = H - 20 - (wy - ymin)/(ymax - ymin) * (H-40);
    return [sx, sy];
  }

  ctx.fillStyle = '#0d1117';
  ctx.fillRect(0,0,W,H);

  ctx.strokeStyle = '#1e2533'; ctx.lineWidth = 1;
  for(let gx=Math.ceil(xmin); gx<=Math.floor(xmax); gx+=5){
    const [sx,] = toScreen(gx, ymin);
    ctx.beginPath(); ctx.moveTo(sx,H-20); ctx.lineTo(sx,20); ctx.stroke();
  }
  ctx.fillStyle='#4a5568'; ctx.font='9px monospace'; ctx.textAlign='center';
  for(let gx=Math.ceil(xmin/5)*5; gx<=xmax; gx+=10){
    const [sx,] = toScreen(gx,0);
    ctx.fillText(gx+'m', sx, H-6);
  }

  function drawPath(traj, color, label, dashed=false){
    if(!traj||traj.length===0) return;
    ctx.strokeStyle=color; ctx.lineWidth=2.5;
    if(dashed) ctx.setLineDash([6,3]); else ctx.setLineDash([]);
    ctx.beginPath();
    let [sx,sy] = toScreen(0,0); ctx.moveTo(sx,sy);
    traj.forEach(p=>{ [sx,sy]=toScreen(p[0],p[1]); ctx.lineTo(sx,sy); });
    ctx.stroke();
    ctx.fillStyle=color;
    traj.forEach((p,i)=>{
      [sx,sy]=toScreen(p[0],p[1]);
      ctx.beginPath(); ctx.arc(sx,sy, i===traj.length-1?5:3, 0, Math.PI*2); ctx.fill();
    });
    ctx.setLineDash([]);
  }

  drawPath(gtTraj, '#60a5fa', 'GT');
  drawPath(predTraj, '#f87171', 'Pred', true);

  const [ox,oy] = toScreen(0,0);
  ctx.fillStyle='#7c93f7'; ctx.beginPath(); ctx.arc(ox,oy,6,0,Math.PI*2); ctx.fill();
  ctx.fillStyle='#fff'; ctx.beginPath(); ctx.arc(ox,oy,2.5,0,Math.PI*2); ctx.fill();

  ctx.font='bold 11px sans-serif'; ctx.textAlign='left';
  ctx.fillStyle='#60a5fa'; ctx.fillRect(16,14,18,3); ctx.fillText('GT',38,19);
  ctx.fillStyle='#f87171'; ctx.fillRect(70,14,18,3); ctx.fillText('Pred',92,19);
}

async function loadSample(i){
  document.getElementById('main').innerHTML = '<div class="loading">Loading…</div>';
  const res = await fetch(`/api/sample/${i}`);
  if(!res.ok){ document.getElementById('main').innerHTML='<div class="loading">Error loading sample.</div>'; return; }
  const d = await res.json();

  const isZero = d.is_zero_pred;
  const nav = d.navigation_command||'';
  const ego = d.ego_states||[];
  const cur = ego.length ? ego[ego.length-1] : null;
  const parsed = d.model_output_parsed||{};
  const behavior = parsed.future_meta_behavior||{};
  const objs = parsed.critical_objects||{};
  const ade3 = d.metrics?.ade_3s ?? d.metrics?.ade;
  const ade5 = d.metrics?.ade_5s;
  const fde  = d.metrics?.fde;

  const egoRows = ego.map(e=>`
    <tr><td>${e.t>=0?'+':''}${e.t.toFixed(1)}s</td>
    <td>${fmtNum(e.vx)}</td><td>${fmtNum(e.vy)}</td><td>${fmtNum(e.speed)}</td></tr>
  `).join('');

  const chipHtml = Object.entries(objs).map(([k,v])=>`
    <span class="obj-chip ${v==='yes'?'obj-yes':'obj-no'}">${k.replace(/_/g,' ')}</span>
  `).join('');

  document.getElementById('main').innerHTML = `
  <div class="left-col">
    <div class="card">
      <div class="card-title">📷 Camera View
        ${isZero ? '<span class="tag tag-zero">ZERO PRED</span>' : ''}
        ${total < originalTotal ? '<span class="tag tag-filtered">FILTERED</span>' : ''}
      </div>
      <div class="img-wrap">
        ${d.image_full_path
          ? `<img src="/image?path=${encodeURIComponent(d.image_full_path)}" alt="scene">`
          : '<div class="img-placeholder">Image not available</div>'}
      </div>
      <div class="info-grid">
        <div class="info-item">
          <label>Navigation</label>
          <div class="val"><span class="nav-badge ${navClass(nav)}">${nav||'—'}</span></div>
        </div>
        <div class="info-item">
          <label>Sample ID</label>
          <div class="val" style="font-size:.75rem;color:#94a3b8;word-break:break-all">${d.sample_id}</div>
        </div>
        <div class="info-item">
          <label>Current Speed</label>
          <div class="val" style="color:#7c93f7">${cur ? fmtNum(cur.speed) : '—'} <span style="font-size:.8rem;color:#64748b">m/s</span></div>
        </div>
        <div class="info-item">
          <label>Vel (Vx / Vy)</label>
          <div class="val" style="font-size:.9rem;color:#94a3b8">${cur ? fmtNum(cur.vx)+' / '+fmtNum(cur.vy) : '—'} <span style="font-size:.75rem;color:#64748b">m/s</span></div>
        </div>
      </div>
    </div>

    <div class="card">
      <div class="card-title">📈 Ego State History</div>
      <div style="padding:10px 12px 12px">
        <table class="ego-table">
          <thead><tr><th>Time</th><th>Vx (m/s)</th><th>Vy (m/s)</th><th>Speed</th></tr></thead>
          <tbody>${egoRows}</tbody>
        </table>
      </div>
    </div>

    <div class="card">
      <div class="card-title">🎯 Metrics (JMT 4Hz)</div>
      <div class="info-grid">
        <div class="info-item">
          <label>ADE @ 3s</label>
          <div class="val ${metricClass(ade3)}">${fmtNum(ade3)} <span style="font-size:.8rem;color:#64748b">m</span></div>
        </div>
        <div class="info-item">
          <label>ADE @ 5s</label>
          <div class="val ${metricClass(ade5)}">${fmtNum(ade5)} <span style="font-size:.8rem;color:#64748b">m</span></div>
        </div>
        <div class="info-item" style="grid-column:span 2">
          <label>FDE @ 5s</label>
          <div class="val ${metricClass(fde)}">${fmtNum(fde)} <span style="font-size:.8rem;color:#64748b">m</span></div>
        </div>
      </div>
    </div>
  </div>

  <div class="right-col">
    <div class="card">
      <div class="card-title">🛣️ Trajectory (forward=x, left=y)</div>
      <div class="traj-wrap">
        <canvas id="trajCanvas" style="width:100%;height:280px"></canvas>
      </div>
      <div class="behavior-grid">
        <div class="beh-item">
          <label>Longitudinal</label>
          <div class="bval">${esc(behavior.longitudinal||'—')}</div>
        </div>
        <div class="beh-item">
          <label>Lateral</label>
          <div class="bval">${esc(behavior.lateral||'—')}</div>
        </div>
      </div>
    </div>

    <div class="card">
      <div class="card-title">🔍 Critical Objects</div>
      <div class="objects-grid">${chipHtml||'<span style="color:#4a5568;padding:4px">No data</span>'}</div>
    </div>

    <div class="card">
      <div class="card-title">💬 Explanation</div>
      <div class="explanation">${esc(parsed.explanation||d.model_output_raw||'—')}</div>
    </div>

    <div class="card">
      <div class="card-title">📝 Raw Model Output</div>
      <div class="raw-out">${esc(d.model_output_raw||'—')}</div>
    </div>
  </div>`;

  requestAnimationFrame(()=>{
    drawTrajectory('trajCanvas', d.pred_trajectory_1hz, d.gt_trajectory_1hz);
  });

  document.getElementById('counter').textContent = `${i+1} / ${total}`;
  document.getElementById('pageInfo').textContent = `${i+1} / ${total}`;
  document.getElementById('prog').style.width = `${100*(i+1)/total}%`;
  document.getElementById('prevBtn').disabled = (i===0);
  document.getElementById('nextBtn').disabled = (i===total-1);
  document.getElementById('jumpInput').value = i+1;
}

function navigate(dir){
  idx = Math.max(0, Math.min(total-1, idx+dir));
  loadSample(idx);
}

function jumpTo(){
  const v = parseInt(document.getElementById('jumpInput').value);
  if(v>=1 && v<=total){ idx=v-1; loadSample(idx); }
}

function applyFilters(){
  const nav = document.getElementById('navFilter').value;
  const adeMin = document.getElementById('adeMin').value;
  const adeMax = document.getElementById('adeMax').value;
  const hideZero = document.getElementById('hideZero').checked;
  const validOnly = document.getElementById('validOnly').checked;

  const params = new URLSearchParams();
  if(nav && nav !== 'ALL') params.set('nav', nav);
  if(adeMin) params.set('ade_min', adeMin);
  if(adeMax) params.set('ade_max', adeMax);
  if(hideZero) params.set('hide_zero', '1');
  if(validOnly) params.set('valid_only', '1');

  window.location.href = '/?' + params.toString();
}

async function exportFiltered(){
  const res = await fetch('/api/export');
  if(!res.ok) return alert('Export failed');
  const blob = await res.blob();
  const url = URL.createObjectURL(blob);
  const a = document.createElement('a');
  a.href = url;
  a.download = 'filtered_predictions.json';
  a.click();
  URL.revokeObjectURL(url);
}

document.addEventListener('keydown', e=>{
  if(e.target.tagName==='INPUT') return;
  if(e.key==='ArrowLeft'||e.key==='a')  navigate(-1);
  if(e.key==='ArrowRight'||e.key==='d') navigate(1);
});

// Load URL parameters into filter controls
const urlParams = new URLSearchParams(window.location.search);
if(urlParams.get('nav')) document.getElementById('navFilter').value = urlParams.get('nav');
if(urlParams.get('ade_min')) document.getElementById('adeMin').value = urlParams.get('ade_min');
if(urlParams.get('ade_max')) document.getElementById('adeMax').value = urlParams.get('ade_max');
if(urlParams.get('hide_zero')) document.getElementById('hideZero').checked = true;
if(urlParams.get('valid_only')) document.getElementById('validOnly').checked = true;

// Initial load
if(total > 0) loadSample(0);
else document.getElementById('main').innerHTML = '<div class="loading">No samples match current filters.</div>';
</script>
</body>
</html>"""

# ── API endpoints ─────────────────────────────────────────────────────────────

@app.route('/')
def index():
    """Main page with filtering UI."""

    # Parse query parameters for filters
    filters = {}
    if request.args.get('nav'):
        filters['nav'] = request.args.get('nav')
    if request.args.get('ade_min'):
        try:
            filters['ade_min'] = float(request.args.get('ade_min'))
        except:
            pass
    if request.args.get('ade_max'):
        try:
            filters['ade_max'] = float(request.args.get('ade_max'))
        except:
            pass
    if request.args.get('hide_zero'):
        filters['hide_zero'] = True
    if request.args.get('valid_only'):
        filters['valid_only'] = True

    # Apply filters
    global FILTERED_INDICES
    FILTERED_INDICES = apply_filters(filters)

    # Generate filter summary
    filter_parts = []
    if filters.get('nav') and filters['nav'] != 'ALL':
        filter_parts.append(f"nav={filters['nav']}")
    if filters.get('ade_min') is not None:
        filter_parts.append(f"ade≥{filters['ade_min']}")
    if filters.get('ade_max') is not None:
        filter_parts.append(f"ade≤{filters['ade_max']}")
    if filters.get('hide_zero'):
        filter_parts.append("no-zero")
    if filters.get('valid_only'):
        filter_parts.append("valid-only")

    filter_summary = f"Showing {len(FILTERED_INDICES)} / {len(ALL_PREDICTIONS)} samples"
    if filter_parts:
        filter_summary += f" ({', '.join(filter_parts)})"

    nav_commands = extract_nav_commands()

    return render_template_string(HTML,
                                  nav_commands=nav_commands,
                                  current_nav=filters.get('nav'),
                                  current_ade_min=filters.get('ade_min'),
                                  current_ade_max=filters.get('ade_max'),
                                  current_hide_zero=filters.get('hide_zero'),
                                  current_valid_only=filters.get('valid_only'),
                                  filter_summary=filter_summary,
                                  filtered_count=len(FILTERED_INDICES),
                                  total_count=len(ALL_PREDICTIONS))

@app.route('/api/sample/<int:i>')
def api_sample(i):
    """Get filtered sample by index."""
    if i < 0 or i >= len(FILTERED_INDICES):
        abort(404)

    # Get actual prediction from filtered indices
    actual_idx = FILTERED_INDICES[i]
    p = ALL_PREDICTIONS[actual_idx]

    ego = parse_ego(p.get('question', ''))

    # Down-sample GT from 4Hz (20pts) to 1Hz (5pts) at t=1,2,3,4,5s → indices 3,7,11,15,19
    gt_20  = p.get('gt_trajectory') or []
    gt_1hz = [gt_20[j] for j in [3, 7, 11, 15, 19] if j < len(gt_20)]
    pred_1hz = p.get('pred_trajectory') or []

    # JMT metrics from pre-computed eval_detailed.json or fallback
    sid = p.get('sample_id', '')
    if sid in METRICS_MAP:
        metrics = METRICS_MAP[sid]
    else:
        metrics = p.get('metrics') or {}

    # model_output_parsed may be stored as JSON string → ensure it's dict
    mop = p.get('model_output_parsed')
    if isinstance(mop, str):
        try:
            mop = json.loads(mop)
        except:
            mop = None

    is_zero = is_zero_prediction(pred_1hz)

    out = {
        "sample_id":           sid,
        "image_full_path":     p.get('image_full_path'),
        "navigation_command":  p.get('navigation_command'),
        "ego_states":          ego,
        "pred_trajectory_1hz": pred_1hz,
        "gt_trajectory_1hz":   gt_1hz,
        "gt_trajectory":       gt_20,
        "model_output_raw":    p.get('model_output_raw'),
        "model_output_parsed": mop,
        "metrics":             metrics,
        "is_zero_pred":        is_zero,
        "original_index":      actual_idx,  # For debugging
    }
    return jsonify(out)

@app.route('/api/export')
def export_filtered():
    """Export current filtered subset as JSON."""
    filtered_predictions = [ALL_PREDICTIONS[i] for i in FILTERED_INDICES]

    # Create export object with metadata
    export_obj = {
        "metadata": {
            "total_samples": len(ALL_PREDICTIONS),
            "filtered_samples": len(filtered_predictions),
            "filters": CURRENT_FILTERS,
            "exported_at": __import__('datetime').datetime.now().isoformat()
        },
        "predictions": filtered_predictions
    }

    # Stream JSON response
    from flask import Response
    import json as json_module

    def generate():
        yield json_module.dumps(export_obj, ensure_ascii=False, indent=2)

    return Response(generate(),
                   mimetype='application/json',
                   headers={'Content-Disposition': 'attachment; filename=filtered_predictions.json'})

@app.route('/image')
def serve_image():
    """Serve image files."""
    path = request.args.get('path', '')
    if not path or not os.path.isfile(path):
        abort(404)
    return send_file(path, mimetype='image/jpeg')

# ── entry point ───────────────────────────────────────────────────────────────

def main():
    global ALL_PREDICTIONS, METRICS_MAP, FILTERED_INDICES
    ap = argparse.ArgumentParser()
    ap.add_argument('--pred_json', required=True)
    ap.add_argument('--eval_detailed', default=None,
                    help='eval_detailed.json from evaluate_predictions.py (JMT metrics)')
    ap.add_argument('--port', type=int, default=7860)
    ap.add_argument('--host', default='0.0.0.0')
    args = ap.parse_args()

    print(f'Loading {args.pred_json} …')
    with open(args.pred_json) as f:
        obj = json.load(f)
    ALL_PREDICTIONS = obj['predictions'] if isinstance(obj, dict) else obj
    print(f'Loaded {len(ALL_PREDICTIONS)} samples.')

    # Initialize with no filters (show all)
    FILTERED_INDICES = list(range(len(ALL_PREDICTIONS)))

    # Load detailed metrics if available
    auto_detailed = os.path.join(os.path.dirname(os.path.abspath(args.pred_json)),
                                 'eval_detailed.json')
    detailed_path = args.eval_detailed or (auto_detailed if os.path.exists(auto_detailed) else None)
    if detailed_path and os.path.exists(detailed_path):
        with open(detailed_path) as f:
            for row in json.load(f):
                METRICS_MAP[row['sample_id']] = {
                    'ade_3s': row.get('ade_3s'),
                    'ade_5s': row.get('ade_5s'),
                    'fde':    row.get('fde'),
                }
        print(f'Loaded JMT metrics for {len(METRICS_MAP)} samples from {detailed_path}')
    else:
        print('No eval_detailed.json found — metrics will use raw per-sample values.')

    # Print available navigation commands for user reference
    nav_commands = extract_nav_commands()
    ade_min, ade_max = get_ade_range()
    print(f'Navigation commands: {nav_commands}')
    print(f'ADE range: {ade_min:.2f} - {ade_max:.2f}')

    print(f'\nOpen http://localhost:{args.port}')
    print('Filter examples:')
    print(f'  http://localhost:{args.port}?nav=GO_STRAIGHT')
    print(f'  http://localhost:{args.port}?nav=GO_LEFT&ade_max=2.0')
    print(f'  http://localhost:{args.port}?hide_zero=1&valid_only=1')
    print()

    app.run(host=args.host, port=args.port, debug=False)

if __name__ == '__main__':
    main()