"""Encode every frame of one T-Rex episode with the frozen Wan2.2 VAE.

Replicates the baseline observation path exactly: each view is resized to
160x320 (bilinear, antialias), tiled into the DreamZero 2x2 grid
[head | right ; left | black] (320x640), resized back to 160x320 (bilinear,
as ``encode_video`` does) and encoded as a single-frame video into a
48x10x20 latent.  Output: ``latents`` [N,48,10,20] float16.
"""

from __future__ import annotations

import argparse
import time
from pathlib import Path

import numpy as np
import torch
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as TF

from fact.perception.wan_video_vae import WanVideoVAE38

VIEW_DIRS = (
    "observation.images.head_left",
    "observation.images.left_wrist",
    "observation.images.right_wrist",
)


def load_vae(path: str, device: torch.device) -> WanVideoVAE38:
    vae = WanVideoVAE38(z_dim=48, dim=160, vae_pretrained_path=path)
    state = torch.load(path, map_location="cpu")
    if isinstance(state, dict):
        for key in ("state_dict", "model", "module"):
            if isinstance(state.get(key), dict):
                state = state[key]
                break
    state = {k.removeprefix("module."): v for k, v in state.items() if torch.is_tensor(v)}
    current = vae.model.state_dict()
    compatible = {k: v for k, v in state.items() if k in current and current[k].shape == v.shape}
    coverage = sum(v.numel() for v in compatible.values()) / sum(v.numel() for v in current.values())
    if coverage < 0.99:
        raise ValueError(f"VAE coverage {coverage:.3f} from {path}")
    vae.model.load_state_dict(compatible, strict=False)
    vae = vae.to(device).eval()
    vae.mean = vae.mean.to(device)
    vae.std = vae.std.to(device)
    vae.scale = [vae.mean, vae.std.reciprocal()]
    return vae


def read_all_frames(path: Path) -> np.ndarray:
    import decord

    reader = decord.VideoReader(str(path), num_threads=2)
    return reader.get_batch(list(range(len(reader)))).asnumpy().astype(np.uint8)


def composite(views: list[np.ndarray], height: int, width: int) -> torch.Tensor:
    """3 x [N,H,W,3] uint8 -> [N,3,H,W] float in [-1,1] at the model size."""

    resized = []
    for frames in views:
        t = torch.from_numpy(frames).permute(0, 3, 1, 2)  # N,3,H,W uint8
        t = TF.resize(t, [height, width], interpolation=InterpolationMode.BILINEAR, antialias=True)
        resized.append(t)
    head, left, right = resized
    n = head.shape[0]
    grid = torch.zeros(n, 3, 2 * height, 2 * width, dtype=torch.uint8)
    grid[:, :, :height, :width] = head
    grid[:, :, height:, :width] = left
    grid[:, :, :height, width:] = right
    grid = grid.float() / 255.0 * 2 - 1
    grid = torch.nn.functional.interpolate(grid, size=(height, width), mode="bilinear", align_corners=False)
    return grid


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset-root", required=True)
    parser.add_argument("--episode", type=int, default=0)
    parser.add_argument("--vae", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--height", type=int, default=160)
    parser.add_argument("--width", type=int, default=320)
    parser.add_argument("--batch", type=int, default=16)
    args = parser.parse_args()

    device = torch.device("cuda")
    root = Path(args.dataset_root)
    views = []
    for view in VIEW_DIRS:
        path = root / "videos" / "chunk-000" / view / f"episode_{args.episode:06d}.mp4"
        frames = read_all_frames(path)
        print(f"{view}: {frames.shape}", flush=True)
        views.append(frames)
    n = min(v.shape[0] for v in views)
    views = [v[:n] for v in views]
    images = composite(views, args.height, args.width)  # N,3,H,W
    vae = load_vae(args.vae, device)
    latents = []
    start = time.time()
    with torch.no_grad():
        for i in range(0, n, args.batch):
            chunk = images[i : i + args.batch].to(device)
            video = chunk.unsqueeze(2)  # B,3,1,H,W (single-frame videos)
            z = vae.encode(video, tiled=False)  # B,48,1,h,w
            latents.append(z[:, :, 0].float().cpu())
            if (i // args.batch) % 10 == 0:
                print(f"  {i}/{n} ({time.time() - start:.1f}s)", flush=True)
    latents = torch.cat(latents).numpy().astype(np.float16)
    print("latents", latents.shape, "mean", float(latents.mean()), "std", float(latents.std()))
    np.savez_compressed(args.out, latents=latents, episode=args.episode, dataset_root=str(root))
    print("wrote", args.out)


if __name__ == "__main__":
    main()