#!/usr/bin/env python3
"""Reference loader for anchor-recorder sessions (format "anchor-recorder-session/1").

Needs Python 3.8+, ffmpeg and ffprobe on PATH. numpy is used when installed; without it the
images come back as raw bytes.

    from anchor_load import iter_stereo
    for s in iter_stereo("sessions/ue9f4c3_0002_007", pix_fmt="gray"):
        s.chunk              # "c000", "c001", ... (a new chunk follows a logged pipeline reopen)
        s.pts_us             # capture time, board CLOCK_MONOTONIC microseconds (USB arrival)
        s.left, s.right      # numpy arrays (H, W) or (H, W, 3); bytes without numpy
        s.frame              # the frames.csv row (dict of ints / None), or None if metadata was lost
        s.imu                # IMU samples carried by frames since the previous yielded pair

Command line:

    python3 anchor_load.py <session> [--limit N] [--no-decode] [--pix-fmt gray|rgb24|bgr24]

prints a JSON summary (pairs, unmatched frames, IMU samples, counter epochs, decode errors).

How pairing works (see README.txt on the card and docs/DATA-SCHEMA.md):
- MPEG-TS PTS (90 kHz) = pts_us + 1.4 s, so pts_us = round(PTS * 1e6 / 90000) - 1400000 and
  every TS frame is within 11 us of its frames.csv row. Raw TS time must never be compared to
  pts_us: 1.4 s is 42 frame periods, so a raw comparison matches the wrong frame.
- Left and right are paired by that timestamp, never by index: n restarts in every chunk and
  the L, R and frames.csv counts can differ at a truncated tail or after a power cut.
- Each decoded frame's PTS is read from ffmpeg itself (showinfo), so a frame that fails to
  decode can never shift the pairing of the frames after it.
- L.pts.txt / R.pts.txt are rounded to ~ms and are not used.

IMU (all "candidate": the camera vendor has not confirmed units, axes or timing):
- imu.csv has one row per 7-byte APP4 record. read_imu() pivots them into samples: a tag-4
  record starts a sample; tag 1 = gyro candidate, tag 2 = accel candidate, tag 3 =
  temperature candidate, all raw LSB.
- tag-4 v0 & 0xFFFF is a counter that steps 96 per sample and restarts about once a second.
  Each restart starts a new epoch; counters are comparable only within one epoch. The t48
  column is not a time and is ignored.
- A sample's carrier_pts_us is the arrival time of the frame that carried it, NOT the time
  the sample was taken.
"""
import argparse, csv, json, queue, re, subprocess, sys, threading, time
from bisect import bisect_left, bisect_right
from collections import namedtuple
from pathlib import Path

try:
    import numpy as np
except ImportError:          # stdlib fallback: frames come back as bytes
    np = None

FORMAT = "anchor-recorder-session/1"
TS_OFFSET_US = 1_400_000     # ffmpeg mpegts mux delay, measured on ue9f4c3_0002_007
MATCH_TOL_US = 11            # 90 kHz rounding is at most 5.6 us; 11 us leaves margin
Stereo = namedtuple("Stereo", "chunk pts_us left right frame imu")
_BPP = {"gray": 1, "rgb24": 3, "bgr24": 3}
_INT_COLS = ("n", "pts_us", "host_mono_ns", "jpeg_bytes", "soi", "eoi", "app4_count", "app4_len",
             "app4_counter", "app4_clock", "imu_records")


def ts_to_pts_us(ts_pts_90k):
    """MPEG-TS PTS (90 kHz ticks) -> recorder pts_us."""
    return round(int(ts_pts_90k) * 1_000_000 / 90_000) - TS_OFFSET_US


def chunks(session):
    return sorted(p for p in Path(session).glob("c[0-9][0-9][0-9]") if p.is_dir())


def segments(chunk, eye):
    return sorted((Path(chunk) / eye).glob("[0-9]*.ts"))


def _probe_size(seg):
    out = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
                          "stream=width,height", "-of", "csv=p=0", str(seg)],
                         capture_output=True, text=True).stdout.strip().splitlines()
    if not out:
        return None
    w, h = out[0].split(",")[:2]
    return int(w), int(h)


def packet_pts_us(seg):
    """pts_us of every video packet in one segment (no decode). Truncated tails are fine."""
    out = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries",
                          "packet=pts", "-of", "csv=p=0", str(seg)], capture_output=True, text=True).stdout
    vals = [x.strip().strip(",") for x in out.split()]
    return [ts_to_pts_us(v) for v in vals if v.lstrip("-").isdigit()]


def eye_pts_us(chunk, eye):
    return [t for seg in segments(chunk, eye) for t in packet_pts_us(seg)]


_SHOWINFO = re.compile(r"Parsed_showinfo.*\bn:\s*\d+\s+pts:\s*(-?\d+)")
_TB = re.compile(r"Parsed_showinfo.*config in time_base:\s*(\d+)/(\d+)")


def _decode(seg, w, h, pix_fmt, errors):
    """Yield (pts_us, image) for every frame ffmpeg can decode from one segment.

    The PTS of each output frame comes from ffmpeg's showinfo filter on stderr (in output order),
    so decode failures drop frames without shifting timestamps."""
    size = w * h * _BPP[pix_fmt]
    p = subprocess.Popen(["ffmpeg", "-hide_banner", "-nostats", "-nostdin", "-copyts", "-i", str(seg),
                          "-map", "0:v:0", "-vf", "showinfo", "-fps_mode", "passthrough",
                          "-f", "rawvideo", "-pix_fmt", pix_fmt, "pipe:1"],
                         stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    pts_q, tb = queue.Queue(), [1, 90_000]

    def read_stderr():
        for raw in p.stderr:
            ln = raw.decode("utf-8", "replace").rstrip()
            m = _TB.search(ln)
            if m:
                tb[0], tb[1] = int(m.group(1)), int(m.group(2))
                continue
            m = _SHOWINFO.search(ln)
            if m:
                pts_q.put(round(int(m.group(1)) * tb[0] * 1_000_000 / tb[1]) - TS_OFFSET_US)
            elif "showinfo" not in ln and "to muxer" not in ln and re.search(r"error|invalid|corrupt|missing|overread|concealing", ln, re.I):
                errors.append(ln)
        pts_q.put(None)

    t = threading.Thread(target=read_stderr, daemon=True)
    t.start()
    try:
        while True:
            buf = p.stdout.read(size)
            if len(buf) < size:
                break
            us = pts_q.get(timeout=60)
            if us is None:
                break
            if np is None:
                yield us, buf
            else:
                a = np.frombuffer(buf, np.uint8)
                yield us, (a.reshape(h, w) if _BPP[pix_fmt] == 1 else a.reshape(h, w, _BPP[pix_fmt]))
    finally:
        p.kill(); p.wait(); t.join(timeout=5)


def _eye(chunk, eye, pix_fmt, decode, errors):
    """Yield (pts_us, image or None) for one eye across all segments of a chunk, in time order."""
    for seg in segments(chunk, eye):
        if not decode:
            for t in sorted(packet_pts_us(seg)):
                yield t, None
            continue
        size = _probe_size(seg)
        if size is None:
            errors.append(f"{seg.name}: no video stream")
            continue
        seg_err = []
        yield from _decode(seg, size[0], size[1], pix_fmt, seg_err)
        errors.extend(f"{eye}/{seg.name}: {e}" for e in seg_err)


def read_frames(chunk):
    """frames.csv rows as dicts with int values (None for empty cells). [] if missing."""
    f = Path(chunk) / "meta" / "frames.csv"
    if not f.exists():
        return []
    rows = []
    with f.open(newline="") as fh:
        for r in csv.DictReader(fh):
            try:
                row = {k: (int(r[k]) if r.get(k) not in (None, "") else None) for k in _INT_COLS if k in r}
            except ValueError:           # a line cut short by a power loss
                continue
            if row.get("pts_us") is not None and row.get("imu_records") is not None:
                rows.append(row)
    return rows


def read_imu(chunk):
    """Pivot imu.csv (one row per 7-byte record) into samples, in arrival order.

    Each sample: carrier_n, carrier_pts_us (arrival of the carrying frame, not sample time),
    app4_counter, counter16, epoch (index of the tag-4 counter epoch within the chunk),
    gyro_candidate_lsb, accel_candidate_lsb (3-tuples of raw int16), temp_candidate_lsb (int)."""
    f = Path(chunk) / "meta" / "imu.csv"
    samples, cur, epoch, prev = [], None, -1, None
    if not f.exists():
        return samples
    with f.open(newline="") as fh:
        for r in csv.DictReader(fh):
            try:
                tag, v = int(r["tag"]), (int(r["v0"]), int(r["v1"]), int(r["v2"]))
                n, pts = int(r["n"]), int(r["pts_us"])
            except (ValueError, TypeError, KeyError):
                continue
            if tag == 4:
                c16 = v[0] & 0xFFFF
                if prev is None or c16 < prev:
                    epoch += 1
                prev = c16
                cur = {"carrier_n": n, "carrier_pts_us": pts,
                       "app4_counter": int(r["app4_counter"]) if r.get("app4_counter") else None,
                       "counter16": c16, "epoch": epoch, "gyro_candidate_lsb": None,
                       "accel_candidate_lsb": None, "temp_candidate_lsb": None}
                samples.append(cur)
            elif cur is not None:
                if tag == 1:
                    cur["gyro_candidate_lsb"] = v
                elif tag == 2:
                    cur["accel_candidate_lsb"] = v
                elif tag == 3:
                    cur["temp_candidate_lsb"] = v[0]
                else:
                    cur[f"tag{tag}_raw"] = v
    return samples


def imu_epochs(samples):
    """Summarise tag-4 counter epochs: [{epoch, samples, first_counter, last_counter,
    first_carrier_pts_us, last_carrier_pts_us, steps}] with a histogram of in-epoch steps."""
    out = {}
    for s in samples:
        e = out.get(s["epoch"])
        if e is None:
            e = out[s["epoch"]] = {"epoch": s["epoch"], "samples": 0, "first_counter": s["counter16"],
                                   "first_carrier_pts_us": s["carrier_pts_us"], "steps": {}, "_last": None}
        if e["_last"] is not None:
            d = s["counter16"] - e["_last"]
            e["steps"][d] = e["steps"].get(d, 0) + 1
        e["_last"] = s["counter16"]
        e["samples"] += 1
        e["last_counter"] = s["counter16"]
        e["last_carrier_pts_us"] = s["carrier_pts_us"]
    for e in out.values():
        e.pop("_last")
    return [out[k] for k in sorted(out)]


def _nearest(keys, t):
    i = bisect_left(keys, t)
    best = None
    for j in (i - 1, i):
        if 0 <= j < len(keys) and (best is None or abs(keys[j] - t) < abs(best - t)):
            best = keys[j]
    return best if best is not None and abs(best - t) <= MATCH_TOL_US else None


def iter_stereo(session, pix_fmt="gray", stats=None, decode=True):
    """Yield Stereo pairs for the whole session. `stats` (a dict) receives per-chunk counts:
    pairs, left_only, right_only, no_meta, imu_samples, imu_assigned, imu_epochs, decode_errors.
    decode=False pairs by packet timestamps only (fast; left/right are None)."""
    s = Path(session)
    man = json.loads((s / "manifest.json").read_text())
    if man.get("format") != FORMAT:
        print(f"anchor_load: warning: format {man.get('format')!r}, expected {FORMAT!r}", file=sys.stderr)
    if pix_fmt not in _BPP:
        raise ValueError(f"pix_fmt must be one of {sorted(_BPP)}")
    stats = stats if stats is not None else {}
    for chunk in chunks(s):
        frames = read_frames(chunk)
        fby = {r["pts_us"]: r for r in frames}
        fkeys = sorted(fby)
        imu = read_imu(chunk)
        ikeys = [x["carrier_pts_us"] for x in imu]
        errors = []
        st = stats.setdefault(chunk.name, {"pairs": 0, "left_only": 0, "right_only": 0, "no_meta": 0,
                                           "imu_samples": len(imu), "imu_assigned": 0,
                                           "imu_epochs": len({x["epoch"] for x in imu})})
        right = _eye(chunk, "R", pix_fmt, decode, errors)
        r_t, r_img = next(right, (None, None))
        prev = None
        for l_t, l_img in _eye(chunk, "L", pix_fmt, decode, errors):
            while r_t is not None and r_t < l_t - MATCH_TOL_US:        # right frame with no left
                st["right_only"] += 1
                r_t, r_img = next(right, (None, None))
            if r_t is None or abs(r_t - l_t) > MATCH_TOL_US:
                st["left_only"] += 1
                continue
            key = _nearest(fkeys, l_t)
            row = fby.get(key) if key is not None else None
            st["no_meta"] += row is None
            t = key if key is not None else l_t
            a = 0 if prev is None else bisect_right(ikeys, prev)
            b = bisect_right(ikeys, t)
            st["imu_assigned"] += max(0, b - a)
            st["pairs"] += 1
            prev = t
            yield Stereo(chunk.name, t, l_img, r_img, row, imu[a:b])
            r_t, r_img = next(right, (None, None))
        while r_t is not None:                                          # trailing right-only frames
            st["right_only"] += 1
            r_t, r_img = next(right, (None, None))
        st["decode_errors"] = errors


def main(argv=None):
    ap = argparse.ArgumentParser(description="Pair left/right frames, metadata and IMU of one session.")
    ap.add_argument("session")
    ap.add_argument("--limit", type=int, default=0, help="stop after N pairs")
    ap.add_argument("--no-decode", action="store_true", help="pair by packet timestamps only (fast)")
    ap.add_argument("--pix-fmt", default="gray", choices=sorted(_BPP))
    a = ap.parse_args(argv)
    t0, stats, n = time.time(), {}, 0
    first = None
    for x in iter_stereo(a.session, pix_fmt=a.pix_fmt, stats=stats, decode=not a.no_decode):
        n += 1
        if first is None:
            first = {"chunk": x.chunk, "pts_us": x.pts_us, "n": x.frame and x.frame["n"],
                     "left": getattr(x.left, "shape", None if x.left is None else len(x.left)),
                     "imu": x.imu[:1]}
        if a.limit and n >= a.limit:
            break
    epochs = {}
    for c in chunks(a.session):
        ep = imu_epochs(read_imu(c))
        epochs[c.name] = {"count": len(ep), "restarts": max(0, len(ep) - 1),
                          "samples_per_epoch_min": min((e["samples"] for e in ep), default=0),
                          "samples_per_epoch_max": max((e["samples"] for e in ep), default=0)}
    print(json.dumps({"session": str(a.session), "pairs": n,
                      "left_only": sum(v["left_only"] for v in stats.values()),
                      "right_only": sum(v["right_only"] for v in stats.values()),
                      "no_meta": sum(v["no_meta"] for v in stats.values()),
                      "imu_samples": sum(v["imu_samples"] for v in stats.values()),
                      "imu_assigned": sum(v["imu_assigned"] for v in stats.values()),
                      "chunks": stats, "imu_epochs": epochs, "first": first,
                      "seconds": round(time.time() - t0, 1)}, indent=1, default=str))


if __name__ == "__main__":
    main()