#!/usr/bin/env python
"""Extract Wan2.2 VAE video latents + UMT5 text embeddings for a LIBERO LeRobot dataset.

VLA-vs-WAM study, WAM data pipeline. LingBot-VA trains on pre-extracted video latents, not raw
frames (see wan_va/dataset/lerobot_latent_dataset.py). Robbyant released these for LIBERO-Long
only; this script reproduces the *exact* per-camera `.pth` format so we can train the WAM on
libero_spatial / libero_object / libero_goal too.

Output layout (mirrors the released libero-long-lerobot):
  <root>/latents/chunk-000/<cam_key>/episode_{idx:06d}_{start}_{end}.pth
Each .pth is a dict:
  latent (F*H*W, 48) bf16, latent_num_frames/height/width, video_num_frames/height/width,
  text_emb (512, 4096) bf16, text, frame_ids (np.int), start_frame, end_frame, fps, ori_fps
Also writes <root>/empty_emb.pt (UMT5 embedding of "") for classifier-free guidance.

Wan2.2-TI2V VAE: 16x spatial, 4x temporal; latents normalized as (z - latents_mean)/latents_std.
Video frames are truncated to the largest (4k+1) count so temporal compression is exact.

Usage:
  # Verify the pipeline reproduces the released libero-long episode-0 latent, then stop:
  python extract_latents.py --verify_episode0

  # Extract a whole dataset:
  python extract_latents.py \
      --lerobot_root data/libero-spatial-lerobot \
      --wan_dir ckpts/Wan2.2-TI2V-5B-Diffusers \
      --cam_keys observation.images.agentview_rgb observation.images.eye_in_hand_rgb
"""

import argparse
import importlib.machinery
import json
import sys
import types
from pathlib import Path

import numpy as np
import torch


def _stub_flash_attn():
    """Allow importing diffusers without a built flash-attn -- but prefer the real package.

    Bug fixed here (hit on job 446): the original version always stubbed both module names
    unconditionally, even when a real, working flash-attn was installed. That stub only defines
    flash_attn_func, not flash_attn_varlen_func -- so as soon as diffusers' own internal
    `from flash_attn import flash_attn_func, flash_attn_varlen_func` (attention_dispatch.py, hit
    just from `from diffusers import AutoencoderKLWan`, nothing to do with the transformer/
    flash-attn kernel this script never uses) ran, it got the stub instead of the real package
    and crashed with ImportError. Now we only stub a name if it genuinely isn't importable.
    """
    for name in ("flash_attn", "flash_attn_interface"):
        if name in sys.modules:
            continue
        try:
            importlib.import_module(name)
            continue
        except ImportError:
            pass
        mod = types.ModuleType(name)
        mod.flash_attn_func = lambda *a, **k: None
        mod.__spec__ = importlib.machinery.ModuleSpec(name, None)
        mod.__version__ = "0.0.0-stub"
        sys.modules[name] = mod


def largest_4k_plus_1(n: int) -> int:
    """Largest m <= n with m % 4 == 1 (so temporal compression (m-1)/4+1 is exact)."""
    m = n - ((n - 1) % 4)
    return max(m, 1)


def read_video_frames(mp4_path, target_hw):
    """Return uint8 array [T, H, W, 3] resized to target_hw."""
    import av
    import cv2
    container = av.open(str(mp4_path))
    frames = [f.to_ndarray(format="rgb24") for f in container.decode(video=0)]
    container.close()
    frames = np.stack(frames)  # [T,H,W,3]
    H, W = target_hw
    if frames.shape[1] != H or frames.shape[2] != W:
        frames = np.stack([cv2.resize(f, (W, H), interpolation=cv2.INTER_AREA) for f in frames])
    return frames


class Encoders:
    def __init__(self, wan_dir, device="cuda", dtype=torch.bfloat16):
        from diffusers import AutoencoderKLWan
        from transformers import UMT5EncoderModel, T5TokenizerFast
        self.device, self.dtype = device, dtype
        self.vae = AutoencoderKLWan.from_pretrained(wan_dir, subfolder="vae", torch_dtype=torch.float32).to(device).eval()
        self.text_encoder = UMT5EncoderModel.from_pretrained(wan_dir, subfolder="text_encoder", torch_dtype=dtype).to(device).eval()
        self.tokenizer = T5TokenizerFast.from_pretrained(wan_dir, subfolder="tokenizer")
        cfg = self.vae.config
        self.mean = torch.tensor(cfg.latents_mean, dtype=torch.float32, device=device).view(1, -1, 1, 1, 1)
        self.std = torch.tensor(cfg.latents_std, dtype=torch.float32, device=device).view(1, -1, 1, 1, 1)
        self.text_max_len = 512

    @torch.no_grad()
    def encode_video(self, frames_uint8):
        """frames_uint8 [T,H,W,3] -> normalized latent [C,f,h,w] float32, plus dims."""
        T = largest_4k_plus_1(frames_uint8.shape[0])
        frames = frames_uint8[:T]
        # [T,H,W,3] uint8 -> [1,3,T,H,W] in [-1,1]
        x = torch.from_numpy(frames).to(self.device).float().div(127.5).sub(1.0)
        x = x.permute(3, 0, 1, 2).unsqueeze(0)  # [1,3,T,H,W]
        z = self.vae.encode(x).latent_dist.mode()          # [1,48,f,h,w]
        z = (z - self.mean) / self.std                      # Wan normalization
        z = z[0]                                            # [48,f,h,w]
        return z.cpu(), frames_uint8.shape[0], frames.shape[1], frames.shape[2]

    @torch.no_grad()
    def encode_text(self, text):
        tok = self.tokenizer([text], padding="max_length", max_length=self.text_max_len,
                             truncation=True, return_tensors="pt").to(self.device)
        emb = self.text_encoder(input_ids=tok.input_ids, attention_mask=tok.attention_mask).last_hidden_state
        return emb[0].to(self.dtype).cpu()                  # [512,4096]


def make_latent_dict(z_cfhw, video_num_frames, vheight, vwidth, text_emb, text, start, end, fps, ori_fps):
    C, f, h, w = z_cfhw.shape
    latent_flat = z_cfhw.permute(1, 2, 3, 0).reshape(f * h * w, C).contiguous().to(torch.bfloat16)  # (f h w) c
    used = largest_4k_plus_1(video_num_frames)
    return {
        "latent": latent_flat,
        "latent_num_frames": int(f), "latent_height": int(h), "latent_width": int(w),
        "video_num_frames": int(used), "video_height": int(vheight), "video_width": int(vwidth),
        "text_emb": text_emb, "text": text,
        "frame_ids": np.arange(used, dtype=np.int64),
        "start_frame": int(start), "end_frame": int(end),
        "fps": int(fps), "ori_fps": int(ori_fps),
    }


def verify_episode0(args):
    """Extract libero-long episode-0 and compare to the released latent to confirm conventions."""
    _stub_flash_attn()
    root = Path(args.verify_root)
    ref = torch.load(root / "latents/chunk-000/observation.images.agentview_rgb/episode_000000_0_388.pth",
                     weights_only=False)
    enc = Encoders(args.wan_dir, device=args.device)
    frames = read_video_frames(root / "videos/chunk-000/observation.images.agentview_rgb/episode_000000.mp4",
                               (ref["video_height"], ref["video_width"]))
    z, vnf, vh, vw = enc.encode_video(frames)
    print(f"[verify] my latent shape (C,f,h,w)={tuple(z.shape)} ; released flat={tuple(ref['latent'].shape)} "
          f"num_frames={ref['latent_num_frames']} h={ref['latent_height']} w={ref['latent_width']}")
    mine = z.permute(1, 2, 3, 0).reshape(-1, z.shape[0]).float()
    rel = ref["latent"].float()
    if mine.shape == rel.shape:
        err = (mine - rel).abs()
        corr = torch.corrcoef(torch.stack([mine.flatten(), rel.flatten()]))[0, 1]
        print(f"[verify] shape MATCH | mean|abs err|={err.mean():.4f} max={err.max():.4f} corr={corr:.4f}")
        print("[verify] -> PASS (normalization/convention confirmed)" if corr > 0.98
              else "[verify] -> MISMATCH: inspect VAE version / normalization / frame sampling")
    else:
        print("[verify] shape MISMATCH -> check spatial/temporal compression or frame truncation")


def extract_dataset(args):
    _stub_flash_attn()
    root = Path(args.lerobot_root)
    meta = [json.loads(l) for l in (root / "meta/episodes.jsonl").read_text().splitlines()]
    info = json.load(open(root / "meta/info.json"))
    fps = int(info.get("fps", 60))
    enc = Encoders(args.wan_dir, device=args.device)

    # empty text embedding for CFG
    torch.save(enc.encode_text(""), root / "empty_emb.pt")
    print(f"[extract] wrote {root/'empty_emb.pt'}")

    for ep in meta:
        idx = ep["episode_index"]
        chunk = idx // int(info.get("chunks_size", 1000))
        for seg in ep["action_config"]:
            start, end, text = seg["start_frame"], seg["end_frame"], seg["action_text"]
            text_emb = enc.encode_text(text)
            for cam in args.cam_keys:
                mp4 = root / f"videos/chunk-{chunk:03d}/{cam}/episode_{idx:06d}.mp4"
                frames = read_video_frames(mp4, (info["features"][cam]["shape"][0], info["features"][cam]["shape"][1]))
                frames = frames[start:end]
                z, vnf, vh, vw = enc.encode_video(frames)
                d = make_latent_dict(z, frames.shape[0], vh, vw, text_emb, text, start, end, fps, fps)
                out = root / f"latents/chunk-{chunk:03d}/{cam}/episode_{idx:06d}_{start}_{end}.pth"
                out.parent.mkdir(parents=True, exist_ok=True)
                torch.save(d, out)
        print(f"[extract] episode {idx} done")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--wan_dir", default="ckpts/Wan2.2-TI2V-5B-Diffusers")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--verify_episode0", action="store_true",
                    help="Reproduce released libero-long episode-0 latent and compare (no dataset written)")
    ap.add_argument("--verify_root", default="data/libero-long-lerobot")
    ap.add_argument("--lerobot_root", default=None)
    ap.add_argument("--cam_keys", nargs="+",
                    default=["observation.images.agentview_rgb", "observation.images.eye_in_hand_rgb"])
    args = ap.parse_args()
    if args.verify_episode0:
        verify_episode0(args)
    else:
        assert args.lerobot_root, "--lerobot_root required for extraction"
        extract_dataset(args)


if __name__ == "__main__":
    main()