#!/usr/bin/env python3
"""Minimal example: load TLT12 from its WebDataset shards and inspect a batch.

    pip install -r requirements.txt
    python example_load.py --shards shards/train
    python example_load.py --shards shards/train --gif sequence.gif

The dataset ships as plain tar archives, so nothing here is magic: each sample
is one traversal sequence -- 32 JPEG frames plus a few .npy arrays and a small
JSON blob. `wds_dataset.py` decodes that into tensors.
"""
import argparse
import glob
import json
import os

import torch

from wds_dataset import TraversalWebDataset, make_loader


def describe(split_dir):
    """Print the split's dataset_info.json header."""
    # Shards may be plain or gzipped; webdataset reads both.
    shards = sorted(glob.glob(os.path.join(split_dir, "*.tar")) +
                    glob.glob(os.path.join(split_dir, "*.tar.gz")))
    if not shards:
        raise SystemExit(f"No .tar or .tar.gz shards found in {split_dir}")
    ds = TraversalWebDataset(shards, info_dir=split_dir, shuffle=0,
                             shardshuffle=0)
    print(f"split        : {split_dir}")
    print(f"shards       : {len(shards)}")
    print(f"sequences    : {ds.n_sequences}")
    print(f"frames/seq   : {ds.n_frames}")
    print(f"image size   : {ds.image_size}px")
    print(f"latent names : {ds.latent_names}")
    return ds


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--shards", default="shards/train",
                    help="Directory holding *.tar and dataset_info.json "
                         "(default: shards/train)")
    ap.add_argument("--batch-size", type=int, default=4)
    ap.add_argument("--num-workers", type=int, default=2)
    ap.add_argument("--gif", default=None,
                    help="Optional path to write the first sequence as an "
                         "animated GIF")
    args = ap.parse_args()

    ds = describe(args.shards)

    # Human-readable category names; the shards carry only synset IDs.
    with open(os.path.join(os.path.dirname(os.path.abspath(__file__)),
                           "synset_labels.json")) as fh:
        synsets = json.load(fh)["synsets"]

    # ── one batch ────────────────────────────────────────────────────────────
    loader = make_loader(ds, batch_size=args.batch_size,
                         num_workers=args.num_workers)
    batch = next(iter(loader))

    print("\nOne batch:")
    print(f"  frames               {tuple(batch['frames'].shape)}  "
          f"{batch['frames'].dtype}  range "
          f"[{batch['frames'].min():.2f}, {batch['frames'].max():.2f}]")
    print(f"  latents              {tuple(batch['latents'].shape)}")
    print(f"  base_latent          {tuple(batch['base_latent'].shape)}")
    print(f"  traversal_velocities {tuple(batch['traversal_velocities'].shape)}")

    print("\nPer-sequence:")
    for i in range(len(batch["synset_id"])):
        sid = batch["synset_id"][i]
        entry = synsets.get(sid, {})
        varying = [ds.latent_names[k] for k in batch["traversal_factors"][i]]
        print(f"  [{i}] {entry.get('name', '?'):<11s} ({sid})  "
              f"class={entry.get('class_index', '?'):<3}  "
              f"varying: {', '.join(varying) or '(none)'}")

    # Classification target, if you want one: synset -> contiguous class index.
    labels = torch.tensor([synsets[s]["class_index"] for s in batch["synset_id"]])
    print(f"\n  class labels {labels.tolist()}")

    # ── optional GIF of the first sequence ───────────────────────────────────
    if args.gif:
        from PIL import Image
        seq = (batch["frames"][0] * 255).byte().permute(0, 2, 3, 1).numpy()
        imgs = [Image.fromarray(f) for f in seq]
        imgs[0].save(args.gif, save_all=True, append_images=imgs[1:],
                     duration=1000 // 12, loop=0)
        print(f"\nWrote {args.gif}  ({len(imgs)} frames)")


if __name__ == "__main__":
    main()