"""Model loading, validation cases, camera poses and per-forward timing for the
HY-WorldPlay base.

Run inside the ``hy_worldplay_july20260720`` environment with the upstream clone
on ``sys.path`` (``enter_repo`` does that).  Weights come from the
HY-WorldPlay-DEV-Predictor project's HF cache; the validation cases (25 VBench
I2V images cropped to 832x480, 4 camera actions each) come from that project's
``hyworldplay_predictor_vbench_val25_test50`` dataset.
"""

import json
import os
import sys
import time
from types import SimpleNamespace

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
REPO = os.path.join(ROOT, "repos/HY-WorldPlay")
DEV = "/local/zoubin/cz/projects/HY-WorldPlay-DEV-Predictor"
HF = "/local/zoubin/cz/checkpoints/hy_worldplay/huggingface/hub"
BASE_MODEL = os.path.join(
    HF, "models--tencent--HunyuanVideo-1.5/snapshots/9b49404b3f5df2a8f0b31df27a0c7ab872e7b038")
ACTION_CKPT = os.path.join(
    HF, "models--tencent--HY-WorldPlay/snapshots/f4c29235647707b571479a69b569e4166f9f5bf8/"
        "ar_distilled_action_model/diffusion_pytorch_model.safetensors")
DATASET = os.path.join(DEV, "datasets/hyworldplay_predictor_vbench_val25_test50")

# The four bidirectional camera actions of the validation protocol, for 32 latents
# (125 frames).  ``pose_for_latents`` rescales them for shorter search videos.
ACTIONS = (("w_s", "w", "s"), ("a_d", "a", "d"), ("up_down", "up", "down"),
           ("left_right", "left", "right"))

NUM_STEPS = 4
CHUNK_LATENTS = 4


def enter_repo():
    if REPO not in sys.path:
        sys.path.insert(0, REPO)
    if ROOT not in sys.path:
        sys.path.insert(0, ROOT)
    # 125-frame videos sit at ~71 GB with the stock pipeline; without expandable
    # segments the allocator fragments another ~9 GB and the cache banks do not fit.
    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
    os.environ.setdefault("HF_HUB_OFFLINE", "1")
    os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")


def init_state():
    """Single-process inference state identical to the DEV project's ``init_state``."""
    from hyvideo.commons.infer_state import initialize_infer_state
    initialize_infer_state(SimpleNamespace(
        sage_blocks_range="0-53", use_sageattn=False, enable_torch_compile=False,
        use_fp8_gemm=False, quant_type="fp8-per-block", include_patterns="double_blocks",
        use_vae_parallel=False))


def load_pipeline(device=None):
    import torch
    from hyvideo.pipelines.worldplay_video_pipeline import HunyuanVideo_1_5_Pipeline

    init_state()
    pipe = HunyuanVideo_1_5_Pipeline.create_pipeline(
        pretrained_model_name_or_path=BASE_MODEL,
        transformer_version="480p_i2v",
        enable_offloading=True,        # text encoder / VAE live on CPU between uses
        enable_group_offloading=False,
        create_sr_pipeline=False,
        force_sparse_attn=False,
        transformer_dtype=torch.bfloat16,
        action_ckpt=ACTION_CKPT,
    )
    pipe.transformer.eval().requires_grad_(False)
    return pipe


def pose_for_latents(action, num_latents):
    """``"w-15,s-16"`` for 32 latents; the first move gets floor((n-1)/2) latents,
    the reverse move the rest, so the video still turns around half-way."""
    _, fwd, back = action
    n = num_latents - 1
    return f"{fwd}-{n // 2},{back}-{n - n // 2}"


def load_cases(split="validation"):
    """The 25 validation (or 50 test) cases with image paths on this machine."""
    rows = []
    with open(os.path.join(DATASET, "cases.jsonl")) as f:
        for line in f:
            r = json.loads(line)
            if r["split"] != split:
                continue
            r["image_path"] = os.path.join(DATASET, split, "input_images",
                                           f"case_{r['split_case_id']:04d}.jpg")
            rows.append(r)
    rows.sort(key=lambda r: r["split_case_id"])
    return rows


class ForwardTimer:
    """Wraps ``pipe.transformer.forward`` with CUDA events and classifies every call:
    text-cache pass, per-chunk context pass (``cache_vision=True``), or a denoising
    forward.  Only the denoising forwards make up ``denoise_dit_ms``; the context
    passes are the analogue of Self-Forcing's KV-refresh pass and are reported
    separately, never added in.  Also tells the cache controller where it is."""

    def __init__(self, transformer, controller=None, num_steps=NUM_STEPS, first_chunk_steps=None):
        import torch
        self.transformer = transformer
        self.controller = controller
        self.num_steps = num_steps
        # chunk 0 may run a different number of denoising steps than the rest
        self.first_chunk_steps = first_chunk_steps or num_steps
        self.torch = torch
        self.reset()

    def reset(self):
        self.events = {"denoise": [], "context": [], "text": []}
        self.denoise_calls = 0
        if self.controller is not None:
            self.controller.reset_video()

    def __enter__(self):
        self._orig = self.transformer.forward
        timer = self

        def forward(*args, **kw):
            torch = timer.torch
            if kw.get("ar_txt_inference"):
                kind = "text"
            elif kw.get("cache_vision"):
                kind = "context"
            else:
                kind = "denoise"
            if kind == "denoise" and timer.controller is not None:
                if timer.denoise_calls < timer.first_chunk_steps:
                    block, step = 0, timer.denoise_calls
                else:
                    rest = timer.denoise_calls - timer.first_chunk_steps
                    block, step = 1 + rest // timer.num_steps, rest % timer.num_steps
                if step == 0:
                    timer.controller.begin_chunk(block)
                timer.controller.denoise_step(step)
            s = torch.cuda.Event(enable_timing=True)
            e = torch.cuda.Event(enable_timing=True)
            s.record()
            out = timer._orig(*args, **kw)
            e.record()
            timer.events[kind].append((s, e))
            if kind == "denoise":
                timer.denoise_calls += 1
                if timer.controller is not None:
                    timer.controller.end_step()
            return out

        self.transformer.forward = forward
        return self

    def __exit__(self, *exc):
        self.transformer.forward = self._orig

    def summary(self):
        self.torch.cuda.synchronize()
        per = {k: [s.elapsed_time(e) for s, e in v] for k, v in self.events.items()}
        return {
            "denoise_dit_ms": sum(per["denoise"]),
            "per_step_denoise_ms": per["denoise"],
            "num_denoise_forwards": len(per["denoise"]),
            "context_kv_dit_ms": sum(per["context"]),
            "num_context_forwards": len(per["context"]),
            "text_cache_ms": sum(per["text"]),
        }


class _ChunkTimesteps:
    """Stands in for the pipeline's ``timesteps`` tensor so that chunk 0 iterates a
    longer schedule than the following chunks.  ``_ar_rollout_inner`` only iterates
    it (once per chunk), reads ``.dtype`` and ``len()``; the scheduler is re-pointed
    at the matching grid at the start of every chunk (the pipeline itself calls
    ``set_timesteps(num_inference_steps)`` just before, so this always wins)."""

    def __init__(self, scheduler, first_n, n, device, extra):
        self.scheduler, self.first_n, self.n, self.device, self.extra = scheduler, first_n, n, device, extra
        self.calls = 0
        self.dtype = scheduler.timesteps.dtype

    def __len__(self):
        return self.n

    def __iter__(self):
        k = self.first_n if self.calls == 0 else self.n
        self.calls += 1
        self.scheduler.set_timesteps(k, device=self.device, **self.extra)
        return iter(self.scheduler.timesteps)


def generate(pipe, case, action, controller=None, num_frames=125, seed=0, decode=True,
             num_steps=NUM_STEPS, first_chunk_steps=None):
    """One video.  Returns (video [1,C,T,H,W] in [0,1] or None, latents, timing).
    ``first_chunk_steps`` gives chunk 0 its own step count (the others run ``num_steps``)."""
    import torch
    from hyvideo.generate import pose_to_input
    import hyvideo.pipelines.worldplay_video_pipeline as wp

    patched = first_chunk_steps is not None and first_chunk_steps != num_steps
    orig_retrieve = wp.retrieve_timesteps

    def retrieve_first_chunk(scheduler, n, device=None, timesteps=None, sigmas=None, **kw):
        ts, n_out = orig_retrieve(scheduler, n, device, timesteps=timesteps, sigmas=sigmas, **kw)
        return _ChunkTimesteps(scheduler, first_chunk_steps, n_out, device, kw), n_out

    num_latents = (num_frames - 1) // 4 + 1
    assert num_latents % CHUNK_LATENTS == 0
    pose = pose_for_latents(action, num_latents)
    viewmats, Ks, act = pose_to_input(pose, num_latents)

    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    if patched:
        wp.retrieve_timesteps = retrieve_first_chunk
    try:
      with ForwardTimer(pipe.transformer, controller, num_steps=num_steps,
                        first_chunk_steps=first_chunk_steps) as timer, torch.no_grad():
        wall0 = time.time()
        out = pipe(
            enable_sr=False, prompt=case["caption"], aspect_ratio="16:9",
            num_inference_steps=num_steps, video_length=num_frames, negative_prompt="",
            seed=seed, output_type="pt" if decode else "latent", prompt_rewrite=False,
            return_pre_sr_video=True, reference_image=case["image_path"],
            viewmats=viewmats.unsqueeze(0), Ks=Ks.unsqueeze(0), action=act.unsqueeze(0),
            few_step=True, chunk_latent_frames=CHUNK_LATENTS, model_type="ar",
            user_height=480, user_width=832, transformer_resident_ar_rollout=True,
        )
        wall = time.time() - wall0
    finally:
        if patched:
            wp.retrieve_timesteps = orig_retrieve
    timing = timer.summary()
    timing["wall_s"] = wall
    timing["peak_mem_gib"] = torch.cuda.max_memory_allocated() / 1024 ** 3
    video = out.videos if decode else None
    latents = None if decode else out.videos
    return video, latents, timing