"""
WorldMem CFG inference — block-by-block causal denoising with KV cache + PRoPE.

Delegates to WorldMemCausalDiffusionInferencePipeline, which:
  - Writes any fully-clean prefix block to the KV cache at timestep=0 first.
  - Denoises one block at a time via UniPC, each step a single forward with
    kv_cache + PRoPE recomputed for that block's viewmats.
  - Commits each block's clean K/V to cache before moving to the next, so
    later blocks attend against clean history (matching teacher-forcing
    training distribution).
  - CFG is one batched forward with batch_size = 2*B; the KV cache is also
    allocated with batch_size = 2*B. cond/uncond chunked back after each step.
"""

import argparse
import gc
import os
import sys

import torch
import torch.distributed as dist
import numpy as np
from omegaconf import OmegaConf
from torchvision.io import write_video

os.chdir(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ".")

from pipeline import WorldMemCausalDiffusionInferencePipeline
from utils.dataset import WorldMemCSVDataset
from utils.pose_utils import poses_to_viewmats
from utils.misc import set_seed

parser = argparse.ArgumentParser()
parser.add_argument("--config_path", default="configs/worldmem_diffusion_overfit.yaml")
parser.add_argument("--output_dir", default="test/res/infer_worldmem_cfg")
parser.add_argument("--checkpoint_path", default=None,
                    help="Optional generator checkpoint (model.pt)")
parser.add_argument("--data_path", default=None,
                    help="Override config.data_path (CSV) for inference test set.")
parser.add_argument("--num_samples", type=int, default=1,
                    help="Number of noise resamples per dataset entry.")
parser.add_argument("--num_inference_steps", type=int, default=50)
parser.add_argument("--guidance_scale", type=float, default=3.0,
                    help="CFG scale. <=1.0 disables CFG (single forward, B=1 cache).")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--start_index", type=int, default=0,
                    help="First dataset row (env: START_INDEX). "
                         "Ignored if --sample_indices is provided.")
parser.add_argument("--num_entries", type=int, default=None,
                    help="Number of consecutive dataset rows from --start_index (env: N). "
                         "Default: all rows from start_index to dataset end. "
                         "Ignored if --sample_indices is provided.")
parser.add_argument("--sample_indices", type=str, default=None,
                    help="Explicit dataset indices, overrides --start_index/--num_entries. "
                         "Accepts comma-separated ints and inclusive ranges, e.g. "
                         "'0,3,5' or '0-9,15,20-22', or 'all' for the full dataset.")
parser.add_argument("--num_condition_frames", type=int, default=1,
                    help="Number of clean prefix latent frames (1=first image, "
                         "3=first block when num_frame_per_block=3).")
parser.add_argument("--num_latent_frames", type=int, default=None,
                    help="Override config.image_or_video_shape[1] (latent T). "
                         "Must be a multiple of num_frame_per_block.")
parser.add_argument("--tag", type=str, default="",
                    help="Filename tag appended to outputs.")
parser.add_argument("--offload", action="store_true",
                    help="Mutually-exclusive GPU residency for generator and VAE: "
                         "during denoising the VAE lives on CPU; for VAE encode/decode "
                         "the generator lives on CPU. Trades extra H2D/D2H copies for "
                         "lower peak VRAM.")
args = parser.parse_args()

# ─── Distributed (optional, via torchrun) ────────────────────────────────────
if "LOCAL_RANK" in os.environ:
    dist.init_process_group(backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = dist.get_world_size()
    torch.cuda.set_device(local_rank)
    device = torch.device(f"cuda:{local_rank}")
else:
    local_rank = 0
    world_size = 1
    device = torch.device("cuda")

set_seed(args.seed + local_rank)
dtype = torch.bfloat16

# ─── Config ──────────────────────────────────────────────────────────────────
config = OmegaConf.load(args.config_path)
default_config = OmegaConf.load("configs/default_config.yaml")
config = OmegaConf.merge(default_config, config)
config.use_text_encoder = False
if args.data_path is not None:
    print(f"Overriding config.data_path: {config.data_path} -> {args.data_path}")
    config.data_path = args.data_path

if args.num_latent_frames is not None:
    nfb = int(config.num_frame_per_block)
    assert args.num_latent_frames % nfb == 0, (
        f"--num_latent_frames={args.num_latent_frames} must be a multiple of "
        f"num_frame_per_block={nfb}")
    print(f"Overriding latent T: {config.image_or_video_shape[1]} -> {args.num_latent_frames}")
    config.image_or_video_shape[1] = args.num_latent_frames
    config.num_frames = args.num_latent_frames

num_latent_frames = config.image_or_video_shape[1]
latent_C, latent_H, latent_W = config.image_or_video_shape[2:]
n_cond = int(args.num_condition_frames)
assert 1 <= n_cond < num_latent_frames, \
    f"num_condition_frames must be in [1, {num_latent_frames - 1}], got {n_cond}"

guidance_scale = float(args.guidance_scale)
use_cfg = guidance_scale > 1.0
if local_rank == 0:
    os.makedirs(args.output_dir, exist_ok=True)
if dist.is_initialized():
    dist.barrier()
print(f"Config: T={num_latent_frames}  latent {latent_H}x{latent_W}  "
      f"nfb={config.num_frame_per_block}  "
      f"guidance_scale={guidance_scale} (cfg={use_cfg})  "
      f"n_cond={n_cond}")

# ─── Pipeline ────────────────────────────────────────────────────────────────
torch.set_grad_enabled(False)
pipeline = WorldMemCausalDiffusionInferencePipeline(config, device=device)
pipeline.generator = pipeline.generator.to(device=device, dtype=dtype)
pipeline.generator.eval()
pipeline.vae = pipeline.vae.to(device=device, dtype=dtype)


# ─── Optional CPU/GPU swap for generator/VAE ────────────────────────────────
# When --offload is set, only one of {generator, VAE} lives on the GPU at any
# time. The "owner" variable below tracks which model currently holds GPU
# residency so we can no-op redundant swaps. Idle model is parked on CPU.
_GPU_OWNER = "both"  # "generator" | "vae" | "both" | "none"


def _swap_to(target: str):
    """Make ``target`` the sole GPU resident; park the other on CPU.

    target ∈ {"generator", "vae"}. No-op when --offload is disabled.
    Frees pipeline.kv_cache before parking the generator and runs an
    explicit gc + empty_cache so VRAM is reclaimed before the next load.
    """
    global _GPU_OWNER
    if not args.offload:
        return
    if _GPU_OWNER == target:
        return
    if target == "vae":
        # Park generator on CPU. Drop transient GPU state (kv_cache) first
        # so the move actually frees memory rather than fragmenting it.
        if pipeline.kv_cache is not None:
            pipeline.kv_cache = None
        if pipeline.crossattn_cache is not None:
            pipeline.crossattn_cache = None
        pipeline.generator.to("cpu")
        gc.collect()
        torch.cuda.empty_cache()
        pipeline.vae.to(device=device, dtype=dtype)
    elif target == "generator":
        pipeline.vae.to("cpu")
        gc.collect()
        torch.cuda.empty_cache()
        pipeline.generator.to(device=device, dtype=dtype)
    else:
        raise ValueError(f"Unknown swap target: {target}")
    _GPU_OWNER = target


if args.offload:
    # Initial state: VAE on GPU (first thing we do per-sample is encode).
    pipeline.generator.to("cpu")
    gc.collect()
    torch.cuda.empty_cache()
    _GPU_OWNER = "vae"
    print("Offload enabled: generator/VAE are mutually exclusive on GPU.")

if args.checkpoint_path:
    weight_tag = "generator"
    print(f"Loading checkpoint: {args.checkpoint_path} ({weight_tag})")
    ckpt = torch.load(args.checkpoint_path, map_location="cpu", weights_only=False)
    sd = ckpt.get("generator", ckpt.get("model", ckpt))
    sample_key = next(iter(sd.keys()))
    target = pipeline.generator if sample_key.startswith("model.") else pipeline.generator.model
    missing, unexpected = target.load_state_dict(sd, strict=False)
    if missing:
        print(f"  missing keys ({len(missing)}): {missing[:5]}{'...' if len(missing) > 5 else ''}")
    if unexpected:
        print(f"  unexpected keys ({len(unexpected)}): {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}")
else:
    print("No checkpoint — using pretrained WAN backbone + random action encoder.")

# ─── Dataset ─────────────────────────────────────────────────────────────────
dataset = WorldMemCSVDataset(
    config.data_path,
    num_latent_frames=num_latent_frames,
    sample_strategy="first",
)
print(f"Dataset size: {len(dataset)}")

def _parse_sample_indices(spec: str, dataset_size: int) -> list:
    """Parse '0,3,5' / '0-9,15,20-22' / 'all' into a deduplicated sorted index list."""
    if spec.strip().lower() == "all":
        return list(range(dataset_size))
    out = set()
    for token in spec.split(","):
        token = token.strip()
        if not token:
            continue
        if "-" in token:
            lo, hi = token.split("-", 1)
            lo, hi = int(lo), int(hi)
            assert lo <= hi, f"bad range '{token}': lo > hi"
            out.update(range(lo, hi + 1))
        else:
            out.add(int(token))
    indices = sorted(out)
    for i in indices:
        assert 0 <= i < dataset_size, (
            f"index {i} out of range [0, {dataset_size}) for --sample_indices='{spec}'")
    return indices


if args.sample_indices is not None:
    sample_indices = _parse_sample_indices(args.sample_indices, len(dataset))
    assert sample_indices, f"--sample_indices='{args.sample_indices}' parsed to empty list"
    print(f"Using explicit --sample_indices ({len(sample_indices)}): {sample_indices}")
else:
    assert args.start_index >= 0, f"--start_index must be >= 0, got {args.start_index}"
    if args.num_entries is not None:
        assert args.num_entries > 0, f"--num_entries must be > 0, got {args.num_entries}"
        end_index = args.start_index + args.num_entries
        sample_indices = list(range(args.start_index, min(end_index, len(dataset))))
        if end_index > len(dataset):
            print(f"Warning: start_index={args.start_index} + num_entries={args.num_entries} "
                  f"exceeds dataset size {len(dataset)}; using indices "
                  f"{args.start_index}..{len(dataset) - 1}")
    else:
        sample_indices = list(range(args.start_index, len(dataset)))

if world_size > 1:
    sample_indices = sample_indices[local_rank::world_size]
    print(f"[rank {local_rank}/{world_size}] shard indices: {sample_indices}")
elif sample_indices:
    print(f"Will run on dataset indices: {sample_indices}")

def run_one_sample(ds_idx: int):
    print(f"\n=== Dataset index {ds_idx} ===")
    rng_state = np.random.get_state()
    np.random.seed(args.seed)
    sample = dataset[ds_idx]
    np.random.set_state(rng_state)

    frames_raw = sample["frames"].unsqueeze(0).to(device=device, dtype=dtype)   # (1,3,T_vid,H,W)
    actions = sample["actions"].unsqueeze(0).to(device=device, dtype=dtype)     # (1,T_vid,25)
    poses = sample["poses"].unsqueeze(0).to(device=device, dtype=dtype)         # (1,T_lat,5)

    # VAE encode: needs VAE on GPU.
    _swap_to("vae")
    gt_latent = pipeline.vae.encode_to_latent(frames_raw).to(device=device, dtype=dtype)
    viewmats = poses_to_viewmats(poses)                                         # (1,T,4,4)

    for sample_idx in range(args.num_samples):
        print(f"  noise sample {sample_idx + 1}/{args.num_samples} (prefix={n_cond}) ...")
        torch.manual_seed(args.seed + sample_idx)

        noise = torch.randn(1, num_latent_frames, latent_C, latent_H, latent_W,
                            device=device, dtype=dtype)
        # Denoising: needs generator on GPU; VAE parked on CPU.
        _swap_to("generator")
        with torch.no_grad():
            video, latents = pipeline.inference(
                noise=noise,
                actions=actions,
                viewmats=viewmats,
                initial_latent=gt_latent,
                num_condition_frames=n_cond,
                guidance_scale=guidance_scale,
                sampling_steps=args.num_inference_steps,
                return_latents=True,
                skip_vae_decode=args.offload,
            )

        # Decode: park generator, bring VAE back to GPU.
        # In offload mode, `video` is None — we decode externally now that the
        # VAE owns the GPU. In non-offload mode the decode already happened.
        _swap_to("vae")
        if args.offload:
            video = pipeline.vae.decode_to_pixel(latents.to(dtype=dtype))
            video = (video * 0.5 + 0.5).clamp(0, 1)

        gt_video = pipeline.vae.decode_to_pixel(gt_latent.to(dtype=dtype))
        gt_video = (gt_video * 0.5 + 0.5).clamp(0, 1)

        gen_uint8 = (video[0].permute(0, 2, 3, 1).cpu().float() * 255).byte()
        gt_uint8 = (gt_video[0].permute(0, 2, 3, 1).cpu().float() * 255).byte()
        side_by_side = torch.cat([gt_uint8, gen_uint8], dim=2)

        tag = f"_{args.tag}" if args.tag else ""
        out_path = os.path.join(
            args.output_dir,
            f"worldmem_s{ds_idx}_n{sample_idx}_cond{n_cond}_cfg{guidance_scale}{tag}.mp4")
        write_video(out_path, side_by_side, fps=10)

        out_dir = os.path.join(args.output_dir, 'pred')
        os.makedirs(out_dir, exist_ok=True)
        out_gen_path = os.path.join(
            out_dir,
            f"worldmem_s{ds_idx}_n{sample_idx}_cond{n_cond}_cfg{guidance_scale}{tag}.mp4")
        write_video(out_gen_path, gen_uint8, fps=10)

        out_dir = os.path.join(args.output_dir, 'gt')
        os.makedirs(out_dir, exist_ok=True)
        out_gt_path = os.path.join(
            out_dir,
            f"worldmem_s{ds_idx}_n{sample_idx}_cond{n_cond}_cfg{guidance_scale}{tag}.mp4")
        write_video(out_gt_path, gt_uint8, fps=10)

        print(f"  Saved {out_path}  (left=GT | right=generated)")

        # Release per-sample GPU residents and defrag the allocator before the
        # next iteration re-allocates the KV cache inside pipeline.inference().
        del video, latents, gt_video, gen_uint8, gt_uint8, side_by_side
        gc.collect()
        torch.cuda.empty_cache()


if not sample_indices:
    print(f"[rank {local_rank}] no samples in this shard, skipping.")
else:
    for ds_idx in sample_indices:
        run_one_sample(ds_idx)

if dist.is_initialized():
    dist.barrier()
if local_rank == 0:
    print("\nDone.")