#!/usr/bin/env python
"""Pack one shard-builder output dir into self-contained webdataset tars.

Each wds sample carries EVERYTHING training needs — no sidecar lookups in the
loader, which is what makes the mixture loader trivial and robust:

    __key__     "{dsname}_{ep:06d}_{fr:05d}"
    cam0.jpg    primary camera JPEG (as produced by build_shards*)
    cam1.jpg    wrist camera JPEG, or b"" for single-camera datasets
    meta.npz    state        float32[S]   normalized (dataset mean/std)
                action_chunk float32[50,A] normalized FUTURE actions from this
                             step; episode tail padded by repeating the last row
                action_is_pad bool[50]
                action_dim / state_dim / task_index / embodiment_id  int32

Normalization happens HERE, at pack time, with the per-dataset stats the shard
builder wrote — the loader never touches stats. The stats used are copied into
the output manifest so unnormalization at eval time uses the same numbers.

JPEG bytes are copied verbatim (no re-encode); tars are uncompressed because
the payload already is.

Usage:
    python pack_wds.py --src /workspace/data/CanToDrawer \
        --out /workspace/data/wds/CanToDrawer --chunk 50 --maxcount 3000
"""

from __future__ import annotations

import argparse
import io
import json
from pathlib import Path

import numpy as np
import pyarrow.parquet as pq


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", type=Path, required=True, help="shard-builder output dir")
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--name", default=None, help="dataset name for keys (default: src dirname)")
    ap.add_argument("--chunk", type=int, default=50)
    ap.add_argument("--maxcount", type=int, default=3000)
    ap.add_argument("--embodiment-id", type=int, default=0)
    args = ap.parse_args()
    name = args.name or args.src.name

    import webdataset as wds

    manifest = json.loads(next(args.src.glob("manifest-part*.json")).read_text())
    stats = json.loads(next(args.src.glob("stats-part*.json")).read_text())
    s_mean = np.asarray(stats["observation.state"]["mean"], np.float32)
    s_std = np.clip(np.asarray(stats["observation.state"]["std"], np.float32), 1e-6, None)
    a_mean = np.asarray(stats["action"]["mean"], np.float32)
    a_std = np.clip(np.asarray(stats["action"]["std"], np.float32), 1e-6, None)

    # frames table -> per-episode contiguous state/action, RAM-resident
    tabs = [pq.read_table(p) for p in sorted(args.src.glob("frames-part*.parquet"))]
    import pyarrow as pa

    t = pa.concat_tables(tabs) if len(tabs) > 1 else tabs[0]
    ep = np.asarray(t["episode_index"], np.int64)
    fr = np.asarray(t["frame_index"], np.int64)
    order = np.lexsort((fr, ep))
    state = np.asarray(t["state"].to_numpy(zero_copy_only=False).tolist(), np.float32)[order]
    action = np.asarray(t["action"].to_numpy(zero_copy_only=False).tolist(), np.float32)[order]
    task_i = np.asarray(t["task_index"], np.int32)[order]
    ep, fr = ep[order], fr[order]
    uniq, starts = np.unique(ep, return_index=True)
    ep_start = dict(zip(uniq.tolist(), starts.tolist()))
    ep_len = dict(zip(uniq.tolist(), np.diff(np.append(starts, len(ep))).tolist()))

    state = (state - s_mean) / s_std
    action = (action - a_mean) / a_std

    S, A, CH = state.shape[1], action.shape[1], args.chunk
    args.out.mkdir(parents=True, exist_ok=True)
    sink = wds.ShardWriter(str(args.out / "shard-%05d.tar"), maxcount=args.maxcount, verbose=0)
    n = 0
    for shard in sorted((args.src / "shards").glob("shard-*.parquet")):
        pf = pq.ParquetFile(shard)
        for batch in pf.iter_batches(batch_size=512):
            d = batch.to_pydict()
            for e, f, c0, c1 in zip(d["episode_index"], d["frame_index"], d["cam0"], d["cam1"]):
                e, f = int(e), int(f)
                base, ln = ep_start[e], ep_len[e]
                i = base + f
                avail = min(CH, ln - f)
                chunk = np.empty((CH, A), np.float32)
                chunk[:avail] = action[i : i + avail]
                if avail < CH:
                    chunk[avail:] = action[i + avail - 1]
                is_pad = np.zeros(CH, bool)
                is_pad[avail:] = True
                buf = io.BytesIO()
                np.savez(buf, state=state[i], action_chunk=chunk, action_is_pad=is_pad,
                         action_dim=np.int32(A), state_dim=np.int32(S),
                         task_index=task_i[i], embodiment_id=np.int32(args.embodiment_id))
                sink.write({"__key__": f"{name}_{e:06d}_{f:05d}",
                            "cam0.jpg": c0, "cam1.jpg": bytes(c1 or b""),
                            "meta.npz": buf.getvalue()})
                n += 1
    sink.close()

    tasks = {}
    for p in sorted(args.src.glob("tasks-part*.json")):
        tasks.update(json.loads(p.read_text()))
    (args.out / "tasks.json").write_text(json.dumps(tasks, ensure_ascii=False))
    (args.out / "manifest.json").write_text(json.dumps({
        **manifest, "name": name, "samples": n, "chunk": CH,
        "embodiment_id": args.embodiment_id,
        "norm_stats": stats,  # exactly what was applied, for unnormalization
        "shards": sorted(p.name for p in args.out.glob("shard-*.tar")),
    }, indent=2))
    print(f"{name}: {n:,} семплов -> {len(list(args.out.glob('shard-*.tar')))} тарболов в {args.out}")


if __name__ == "__main__":
    main()