"""Audit every episode of a T-Rex-format LeRobot v2.1 dataset.

Usage: python scripts/data/audit_episode_quality.py --dataset-root data/trex_full
Contact = fingertip |(fx,fy,fz)| minus the per-finger whole-episode median baseline (if >0.5 N) exceeds 0.5 N,
the same definition as data/sharpa/tools/compare_difficulty.py. Writes <root>/meta/episode_quality.json.
"""
import json, glob, os, sys, numpy as np, pyarrow.parquet as pq, time
import argparse
ap=argparse.ArgumentParser(description="Per-episode quality audit for T-Rex-format LeRobot datasets (duration, bimanual motion, contact ratio, force baseline).")
ap.add_argument("--dataset-root",required=True); ap.add_argument("--out",default=None,help="json output (default <root>/meta/episode_quality.json)")
args=ap.parse_args(); root=args.dataset_root
eps={e["episode_index"]:e for e in map(json.loads, open(f"{root}/meta/episodes.jsonl"))}
F_THR=0.5   # N above per-finger baseline = contact
rows=[]; t0=time.time()
files=sorted(glob.glob(f"{root}/data/chunk-*/episode_*.parquet"))
for i,f in enumerate(files):
    t=pq.read_table(f, columns=["episode_index","timestamp","observation.state","observation.tactile_force"])
    ei=int(t["episode_index"][0].as_py())
    ts=np.array(t["timestamp"].to_pylist(),dtype=np.float64)
    st=np.stack(t["observation.state"].to_numpy(zero_copy_only=False)).astype(np.float32)
    tac=np.stack(t["observation.tactile_force"].to_numpy(zero_copy_only=False)).astype(np.float32).reshape(-1,10,6)
    fn=np.linalg.norm(tac[:,:,:3],axis=-1)          # [T,10]
    med=np.median(fn,0); base=np.where(med>0.5, med, 0.0)   # remove un-zeroed baseline per finger (same as compare_difficulty.py)
    fn=fn-base
    contact=fn>F_THR
    l=contact[:,:5].any(1); r=contact[:,5:].any(1)
    T=len(ts); dur=float(ts[-1]-ts[0]) if T>1 else 0.0
    def path(a): return float(np.abs(np.diff(a,axis=0)).sum())
    def events(m):
        on=0; gap=99; prev=False
        for v in m:
            if v and not prev and gap>9: on+=1
            gap=0 if v else gap+1; prev=v
        return on
    dq=np.abs(np.diff(st,axis=0)); idle=float((dq.max(1)<1e-3).mean()) if T>1 else 1.0
    rows.append(dict(ep=ei, task=eps[ei]["tasks"][0], frames=T, dur=dur,
        contact=float((l|r).mean()), both=float((l&r).mean()), left=float(l.mean()), right=float(r.mean()),
        larm=path(st[:,0:7]), rarm=path(st[:,29:36]), lhand=path(st[:,7:29]), rhand=path(st[:,36:58]),
        fp99=float(np.percentile(fn,99)), fmax=float(fn.max()), fingers=int((contact.mean(0)>0.05).sum()),
        events=events(l)+events(r), idle=idle, baseline_max=float(base.max()),
        ts_mono=bool(np.all(np.diff(ts)>0)), dt_med=float(np.median(np.diff(ts))) if T>1 else 0))
    if i%500==0: print(f"{i}/{len(files)} {time.time()-t0:.0f}s", flush=True)
out=args.out or f"{root}/meta/episode_quality.json"; json.dump(rows, open(out,"w")); print("wrote",out)
d=np.array([r["dur"] for r in rows]); print("durations: min %.1f median %.1f p90 %.1f max %.1f; n>=30s: %d; 30-120s: %d"%(d.min(),np.median(d),np.percentile(d,90),d.max(),(d>=30).sum(),((d>=30)&(d<=120)).sum()))
cand=[r for r in rows if 30<=r["dur"]<=120 and r["larm"]>1.0 and r["rarm"]>1.0 and r["left"]>0.1 and r["right"]>0.1]
print("candidates (30-120 s, both arms move >1 rad, each hand in contact >10%):", len(cand))
cand.sort(key=lambda r: -(r["contact"]*0.6+r["both"]*0.4))
for r in cand[:15]:
    print(f"  ep {r['ep']:5d} {r['dur']:6.1f}s contact {r['contact']:.2f} both {r['both']:.2f} L {r['left']:.2f} R {r['right']:.2f} larm {r['larm']:.1f} rarm {r['rarm']:.1f} fingers {r['fingers']} fp99 {r['fp99']:.1f}N | {r['task'][:80]}")
import statistics
print("dataset medians: contact %.2f both %.2f events %d baseline_max %.1fN"%(statistics.median(r['contact'] for r in rows),statistics.median(r['both'] for r in rows),statistics.median(r['events'] for r in rows),statistics.median(r['baseline_max'] for r in rows)))
print("episodes with baseline offset >0.5N on some finger:", sum(r['baseline_max']>0.5 for r in rows))
# second ranking: prefer 40-90 s, high contact & both, many fingers
cand2=[r for r in cand if 40<=r['dur']<=90 and r['fingers']>=6 and r['events']>=6]
cand2.sort(key=lambda r: -(r['contact']+r['both']+0.02*r['fingers']))
print("shortlist (40-90 s, >=6 fingers, >=6 events):", len(cand2))
for r in cand2[:12]:
    print(f"  ep {r['ep']:5d} {r['dur']:6.1f}s contact {r['contact']:.2f} both {r['both']:.2f} L {r['left']:.2f} R {r['right']:.2f} ev {r['events']:3d} idle {r['idle']:.2f} fingers {r['fingers']} fp99 {r['fp99']:.1f}N base {r['baseline_max']:.1f} | {r['task']}")
print("DONE")