#!/usr/bin/env python
"""Validate a packed wds dataset dir (or several) against the loader contract.

Checks, per dataset:
  * manifest.json lists every shard-*.tar that exists, tasks.json parses
  * samples decode: cam0 JPEG opens at the declared size, meta.npz has all keys
  * dtypes/shapes: state f32[S], action_chunk f32[chunk,A], action_is_pad bool
  * normalized sanity: |mean| < 5, std in (1e-6, 50) over a sample of frames
  * task_index maps into tasks.json
Then, jointly: a few batches through wds_mix.make_mixture_loader and the batch
schema the trainer expects (incl. the "task" string -> present and str).

Usage:
    python validate_wds.py /data/wds/SetA /data/wds/SetB --batches 4
"""

from __future__ import annotations

import argparse
import io
import json
import sys
import tarfile
from pathlib import Path

import numpy as np


def check_dataset(d: Path, n_samples: int = 300) -> dict:
    errs = []
    mf = json.loads((d / "manifest.json").read_text())
    tasks = json.loads((d / "tasks.json").read_text())
    tars = sorted(d.glob("shard-*.tar"))
    if sorted(mf["shards"]) != [t.name for t in tars]:
        errs.append(f"manifest shards != on-disk tars ({len(mf['shards'])} vs {len(tars)})")
    chunk = int(mf["chunk"])

    seen, s_all, a_all = 0, [], []
    from PIL import Image

    for t in tars:
        with tarfile.open(t) as tf:
            members = tf.getmembers()
            by_key = {}
            for m in members:
                k, _, ext = m.name.partition(".")
                by_key.setdefault(k, {})[ext] = m
            for k in by_key:
                if seen >= n_samples:
                    break
                need = {"cam0.jpg", "cam1.jpg", "meta.npz"}
                raw = {m.name[len(k) + 1:]: m for m in members if m.name.startswith(k + ".")}
                if not need.issubset(raw):
                    errs.append(f"{t.name}:{k}: missing {need - set(raw)}")
                    seen += 1
                    continue
                c0 = tf.extractfile(raw["cam0.jpg"]).read()
                img = Image.open(io.BytesIO(c0))
                img.load()
                meta = np.load(io.BytesIO(tf.extractfile(raw["meta.npz"]).read()))
                for kk, dt in (("state", np.float32), ("action_chunk", np.float32),
                               ("action_is_pad", np.bool_)):
                    if kk not in meta:
                        errs.append(f"{k}: meta missing {kk}")
                    elif meta[kk].dtype != dt:
                        errs.append(f"{k}: {kk} dtype {meta[kk].dtype} != {dt}")
                ac = meta["action_chunk"]
                if ac.shape[0] != chunk:
                    errs.append(f"{k}: chunk {ac.shape[0]} != manifest {chunk}")
                if int(meta["action_dim"]) != ac.shape[1] or int(meta["state_dim"]) != meta["state"].shape[0]:
                    errs.append(f"{k}: dim fields disagree with arrays")
                if str(int(meta["task_index"])) not in tasks:
                    errs.append(f"{k}: task_index {int(meta['task_index'])} not in tasks.json")
                s_all.append(meta["state"]); a_all.append(ac[0])
                seen += 1
        if seen >= n_samples:
            break

    if s_all:
        s = np.stack(s_all); a = np.stack(a_all)
        for nm, arr in (("state", s), ("action", a)):
            m, sd = np.abs(arr.mean(0)).max(), arr.std(0)
            if m > 5:
                errs.append(f"{nm}: |mean| up to {m:.2f} — normalization looks wrong")
            if sd.max() > 50:
                errs.append(f"{nm}: std up to {sd.max():.1f} — normalization looks wrong")
    return {"dir": str(d), "samples_checked": seen, "declared": mf.get("samples"),
            "episodes": mf.get("episodes"), "errors": errs}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dirs", nargs="+", type=Path)
    ap.add_argument("--batches", type=int, default=2)
    ap.add_argument("--batch-size", type=int, default=8)
    ap.add_argument("--max-state-dim", type=int, default=256)
    ap.add_argument("--max-action-dim", type=int, default=64)
    ap.add_argument("--samples", type=int, default=300)
    args = ap.parse_args()

    bad = False
    for d in args.dirs:
        r = check_dataset(d, args.samples)
        status = "OK " if not r["errors"] else "FAIL"
        print(f"[{status}] {r['dir']}: {r['samples_checked']} samples checked "
              f"(declared {r['declared']}, {r['episodes']} eps)")
        for e in r["errors"][:20]:
            print("   -", e)
        bad = bad or bool(r["errors"])

    if args.batches > 0:
        sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
        from tinyvla.data.wds_mix import make_mixture_loader

        loader, mfs = make_mixture_loader(
            [{"dir": str(d)} for d in args.dirs],
            batch_size=args.batch_size, num_workers=0,
            max_state_dim=args.max_state_dim, max_action_dim=args.max_action_dim,
            shuffle_buffer=64, steps_per_epoch=args.batches + 1,
        )
        import torch

        want = {
            "observation.images.cam0": (torch.float32, (3, 256, 256)),
            "observation.images.cam1": (torch.float32, (3, 256, 256)),
            "camera_mask": (torch.bool, (2,)),
            "observation.state": (torch.float32, (args.max_state_dim,)),
            "action": (torch.float32, (None, args.max_action_dim)),
            "action_dim_mask": (torch.bool, (args.max_action_dim,)),
            "action_is_pad": (torch.bool, None),
            "embodiment_id": (torch.int64, ()),
            "task_index": (torch.int64, ()),
        }
        it = iter(loader)
        for b in range(args.batches):
            batch = next(it)
            for k, (dt, shp) in want.items():
                assert k in batch, f"batch missing {k}"
                assert batch[k].dtype == dt, f"{k}: dtype {batch[k].dtype} != {dt}"
                if shp:
                    tail = batch[k].shape[1:]
                    for i, s in enumerate(shp):
                        if s is not None:
                            assert tail[i] == s, f"{k}: shape {tuple(batch[k].shape)}"
            assert "task" in batch and isinstance(batch["task"][0], str), "task strings missing"
            assert not any(k.startswith("__") and torch.is_tensor(v) for k, v in batch.items())
        print(f"[OK ] mixture loader: {args.batches} batches, schema verified, "
              f"tasks e.g. {batch['task'][:2]}")

    sys.exit(1 if bad else 0)


if __name__ == "__main__":
    main()