# SPDX-License-Identifier: Apache-2.0
"""Knowledge Distillation method for ODE-init training.

Trains a student model with MSE loss to reproduce a teacher model's
multi-step ODE denoising trajectories. The resulting checkpoint
(exported via dcp_to_diffusers) serves as the ``ode_init`` weight
initialization for downstream Self-Forcing training.

Teacher path generation is cached to disk so it only runs once.
Interrupted generation resumes from the last completed sample.

Typical YAML::

    models:
      student:
        _target_: fastvideo.train.models.wan.WanModel
        init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
        trainable: true
      teacher:           # omit once cache is complete
        _target_: fastvideo.train.models.wan.WanModel
        init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers
        trainable: false
        disable_custom_init_weights: true

    method:
      _target_: fastvideo.train.methods.knowledge_distillation.kd.KDMethod
      teacher_path_cache: /data/kd_cache/wan14b_4step
      t_list: [999, 937, 833, 624, 0]   # integer timesteps
      student_sample_steps: 4
      teacher_guidance_scale: 1.0
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

from tqdm import tqdm

import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.utils.data import Dataset

from fastvideo.dataset.parquet_dataset_map_style import DP_SP_BatchSampler
from fastvideo.distributed import (
    get_sp_world_size,
    get_world_group,
    get_world_rank,
    get_world_size,
    get_sp_group,
)
from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (
    SelfForcingFlowMatchScheduler, )
from fastvideo.models.utils import pred_noise_to_pred_video
from fastvideo.train.methods.base import LogScalar, TrainingMethod
from fastvideo.train.models.base import ModelBase
from fastvideo.train.utils.optimizer import build_optimizer_and_scheduler

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Cache helpers
# ---------------------------------------------------------------------------

_COMPLETE_SENTINEL = "COMPLETE"
_METADATA_FILE = "metadata.json"


class _KDPathCache:
    """Utility class for KD teacher-path on-disk cache management.

    Cache layout::

        <cache_dir>/
        ├── metadata.json       # config written before generation starts
        ├── samples/
        │   ├── 00000000.pt     # one file per source dataset sample
        │   ├── 00000001.pt
        │   └── ...
        └── COMPLETE            # sentinel written only when all samples done

    Each ``.pt`` file contains a dict::

        {
            "trajectory_latents":   Tensor[S, T, C, H, W],  # all S teacher steps
            "trajectory_timesteps": Tensor[S],               # int64
            "real":                 Tensor[T, C, H, W],
            "text_embedding":       Tensor[L, D],
            "text_attention_mask":  Tensor[L],
        }

    The full teacher trajectory is stored so that ``t_list`` can be changed
    without regenerating the cache.  Subsampling to the student timesteps
    happens at training time.
    ``real`` is the teacher's final clean x0 prediction after all ODE steps.
    All latents are in permuted ``[T, C, H, W]`` format (not ``[C, T, H, W]``).
    """

    @staticmethod
    def is_complete(cache_dir: str) -> bool:
        return (Path(cache_dir) / _COMPLETE_SENTINEL).exists()

    @staticmethod
    def validate_or_create_metadata(
        cache_dir: str,
        teacher_id: str,
        teacher_inference_steps: int,
        total_samples: int,
        source_data_path: str,
    ) -> None:
        """Write metadata if absent; raise if present but mismatched."""
        meta_path = Path(cache_dir) / _METADATA_FILE
        if get_world_rank() == 0:
            Path(cache_dir).mkdir(parents=True, exist_ok=True)
            if meta_path.exists():
                with open(meta_path) as f:
                    stored = json.load(f)
                if stored.get("teacher") != teacher_id:
                    raise ValueError(f"Cache teacher {stored['teacher']!r} != config "
                                     f"teacher {teacher_id!r}. Delete {cache_dir} "
                                     "and re-run.")
                if stored.get("teacher_inference_steps") != teacher_inference_steps:
                    raise ValueError(f"Cache teacher_inference_steps "
                                     f"{stored['teacher_inference_steps']} != config "
                                     f"{teacher_inference_steps}. Delete {cache_dir} "
                                     "and re-run.")
            else:
                meta: dict[str, Any] = {
                    "total_samples": total_samples,
                    "teacher": teacher_id,
                    "teacher_inference_steps": teacher_inference_steps,
                    "source_data_path": source_data_path,
                }
                with open(meta_path, "w") as f:
                    json.dump(meta, f, indent=2)
        if dist.is_initialized():
            dist.barrier()

    @staticmethod
    def find_missing(cache_dir: str, total_samples: int) -> list[int]:
        """Return sorted list of sample indices not yet on disk."""
        samples_dir = Path(cache_dir) / "samples"
        samples_dir.mkdir(parents=True, exist_ok=True)
        existing = {int(f.stem) for f in samples_dir.glob("*.pt")}
        return [i for i in range(total_samples) if i not in existing]

    @staticmethod
    def mark_complete(cache_dir: str) -> None:
        """Write COMPLETE sentinel (rank 0 only) and barrier."""
        if get_world_rank() == 0:
            (Path(cache_dir) / _COMPLETE_SENTINEL).touch()
        if dist.is_initialized():
            dist.barrier()


# ---------------------------------------------------------------------------
# Dataset for reading cached paths at training time
# ---------------------------------------------------------------------------


class _KDPathDataset(Dataset):
    """Reads per-sample ``.pt`` files from a completed KD cache.

    Returns dicts with keys:
        - ``path``                 [num_steps, T, C, H, W]
        - ``real``                 [T, C, H, W]
        - ``text_embedding``       [L, D]
        - ``text_attention_mask``  [L]
        - ``path_timesteps``       [num_steps]  int64
    """

    def __init__(self, cache_dir: str) -> None:
        self._samples_dir = Path(cache_dir) / "samples"
        meta_path = Path(cache_dir) / _METADATA_FILE
        with open(meta_path) as f:
            meta = json.load(f)
        self._total: int = int(meta["total_samples"])

    def __len__(self) -> int:
        return self._total

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        path = self._samples_dir / f"{idx:08d}.pt"
        return torch.load(path, weights_only=True, map_location="cpu")


# ---------------------------------------------------------------------------
# Lazy dataloader wrapper
# ---------------------------------------------------------------------------


class _KDDataLoaderWrapper:
    """Wraps the parquet dataloader initially; activates to KD cache after generation.

    The builder captures ``student.dataloader`` after ``KDMethod.__init__``
    runs. At that point this wrapper is the stored value.  The trainer
    calls ``on_train_start()`` before iterating, so by the time
    ``__iter__`` is called the wrapper has been activated and yields
    from the KD cache dataloader.
    """

    def __init__(self, source_loader: Any) -> None:
        self._source = source_loader
        self._active: Any = None

    def activate(
        self,
        cache_dir: str,
        batch_size: int,
        num_workers: int,
        seed: int,
    ) -> None:
        from torchdata.stateful_dataloader import StatefulDataLoader
        from fastvideo.platforms import current_platform

        dataset = _KDPathDataset(cache_dir)
        sampler = DP_SP_BatchSampler(
            batch_size=batch_size,
            dataset_size=len(dataset),
            num_sp_groups=get_world_size() // get_sp_world_size(),
            sp_world_size=get_sp_world_size(),
            global_rank=get_world_rank(),
            drop_last=True,
            seed=seed,
        )
        self._active = StatefulDataLoader(
            dataset,
            batch_sampler=sampler,
            collate_fn=_kd_collate,
            num_workers=num_workers,
            pin_memory=True,
            pin_memory_device=current_platform.device_name,
            persistent_workers=num_workers > 0,
        )

    def __iter__(self):
        if self._active is None:
            raise RuntimeError("_KDDataLoaderWrapper has not been "
                               "activated. Ensure on_train_start() "
                               "completed successfully.")
        return iter(self._active)

    def state_dict(self):
        if self._active is not None and hasattr(self._active, "state_dict"):
            return self._active.state_dict()
        return {}

    def load_state_dict(self, state: dict) -> None:
        if self._active is not None and hasattr(self._active, "load_state_dict"):
            self._active.load_state_dict(state)


def _kd_collate(samples: list[dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
    """Stack a list of per-sample dicts into a batched dict."""
    keys = samples[0].keys()
    return {k: torch.stack([s[k] for s in samples]) for k in keys}


# ---------------------------------------------------------------------------
# KD Method
# ---------------------------------------------------------------------------


class KDMethod(TrainingMethod):
    """Knowledge Distillation training method.

    Trains the student with MSE loss on teacher ODE trajectories cached
    to ``method_config.teacher_path_cache``.

    Roles:
        - ``student`` (required, trainable): the model being distilled.
        - ``teacher`` (optional, non-trainable): used to generate the cache
          on first run; freed from GPU memory afterwards.

    If the cache is incomplete and no teacher is configured, an error
    is raised at the start of training.
    """

    def __init__(
        self,
        *,
        cfg: Any,
        role_models: dict[str, ModelBase],
    ) -> None:
        super().__init__(cfg=cfg, role_models=role_models)

        if "student" not in role_models:
            raise ValueError("KDMethod requires role 'student'")
        if not self.student._trainable:
            raise ValueError("KDMethod requires student to be trainable")

        mcfg = self.method_config

        # --- Parse method config ---
        raw_t_list = mcfg.get("t_list")
        if not isinstance(raw_t_list, list) or not raw_t_list:
            raise ValueError("method_config.t_list must be a non-empty list "
                             "of integer timestep values, e.g. "
                             "[999, 937, 833, 624, 0]")
        self._t_list: list[int] = [int(t) for t in raw_t_list]

        raw_steps = mcfg.get("student_sample_steps")
        if raw_steps is None:
            raw_steps = len(self._t_list) - 1
        self._num_steps: int = int(raw_steps)
        if len(self._t_list) != self._num_steps + 1:
            raise ValueError(f"len(t_list)={len(self._t_list)} must equal "
                             f"student_sample_steps+1={self._num_steps + 1}")

        cache_dir = mcfg.get("teacher_path_cache")
        if not cache_dir:
            raise ValueError("method_config.teacher_path_cache must be set")
        self._cache_dir: str = str(cache_dir)

        self._teacher_guidance_scale: float = float(mcfg.get("teacher_guidance_scale", 1.0))
        self._teacher_inference_steps: int = int(mcfg.get("teacher_inference_steps", 48))

        # --- Optional teacher ---
        self.teacher: ModelBase | None = role_models.get("teacher")
        if self.teacher is not None and getattr(self.teacher, "_trainable", False):
            raise ValueError("KDMethod requires teacher to be non-trainable "
                             "(set trainable: false in models.teacher)")

        # --- Build parquet dataloader via student.init_preprocessors ---
        self.student.init_preprocessors(self.training_config)

        # Wrap the parquet dataloader so we can swap it after generation.
        self._source_loader = self.student.dataloader
        self._kd_wrapper = _KDDataLoaderWrapper(self._source_loader)
        self.student.dataloader = self._kd_wrapper  # builder captures this

        # --- Build SelfForcingFlowMatchScheduler for sigma lookups ---
        # num_inference_steps=1000 gives a dense grid accurate for any t.
        tc = self.training_config
        self._flow_shift = float(getattr(tc.pipeline_config, "flow_shift", 0.0) or 0.0)
        self._sf_scheduler = SelfForcingFlowMatchScheduler(
            num_inference_steps=1000,
            num_train_timesteps=int(self.student.num_train_timesteps),
            shift=self._flow_shift,
            sigma_min=0.0,
            extra_one_step=True,
            training=False,
        )

        # --- Student optimizer / scheduler (same as FineTuneMethod) ---
        self._init_optimizers_and_schedulers()

    # ------------------------------------------------------------------
    # TrainingMethod abstract implementations
    # ------------------------------------------------------------------

    @property
    def _optimizer_dict(self) -> dict[str, Any]:
        return {"student": self._student_optimizer}

    @property
    def _lr_scheduler_dict(self) -> dict[str, Any]:
        return {"student": self._student_lr_scheduler}

    def get_optimizers(self, iteration: int) -> list[torch.optim.Optimizer]:
        del iteration
        return [self._student_optimizer]

    def get_lr_schedulers(self, iteration: int) -> list[Any]:
        del iteration
        return [self._student_lr_scheduler]

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def on_train_start(self) -> None:
        super().on_train_start()  # sets CUDA RNG, calls student.on_train_start()

        if not _KDPathCache.is_complete(self._cache_dir):
            if self.teacher is None:
                raise RuntimeError(f"teacher_path_cache at {self._cache_dir!r} is "
                                   "incomplete and no teacher model is configured. "
                                   "Add a 'teacher' entry under 'models:' in the YAML "
                                   "or provide a complete cache directory.")

            # Prepare teacher for inference (world_group needed for
            # negative conditioning setup).
            self.teacher.world_group = get_world_group()
            self.teacher.sp_group = get_sp_group()
            self.teacher.on_train_start()

            # Share the student's VAE with the teacher for latent normalization.
            # The teacher doesn't call init_preprocessors() (non-trainable),
            # so self.teacher.vae is None. Both models use the same Wan VAE
            # constants (latents_mean / latents_std), so sharing is safe.
            if getattr(self.teacher, "vae", None) is None:
                self.teacher.vae = self.student.vae

            source_dataset = self._source_loader.dataset
            total = sum(source_dataset.lengths)
            teacher_id = getattr(self.teacher, "_init_from", "unknown")
            data_path = str(getattr(self.training_config.data, "data_path", ""))

            _KDPathCache.validate_or_create_metadata(
                self._cache_dir,
                teacher_id,
                self._teacher_inference_steps,
                total,
                data_path,
            )

            missing = _KDPathCache.find_missing(self._cache_dir, total)
            if missing:
                logger.info(
                    "Generating KD teacher paths: %d/%d samples missing. "
                    "Cache: %s",
                    len(missing),
                    total,
                    self._cache_dir,
                )
                self._generate_paths(source_dataset, missing, total)

            _KDPathCache.mark_complete(self._cache_dir)
            logger.info("KD cache complete: %s", self._cache_dir)

        # Free teacher from GPU memory — not needed after cache is built.
        if self.teacher is not None:
            del self.teacher
            self.teacher = None
            self._role_models.pop("teacher", None)
            if "teacher" in self.role_modules:
                del self.role_modules["teacher"]
            torch.cuda.empty_cache()
            if dist.is_initialized():
                dist.barrier()

        # Activate the KD cache dataloader.
        tc = self.training_config
        self._kd_wrapper.activate(
            cache_dir=self._cache_dir,
            batch_size=int(tc.data.train_batch_size),
            num_workers=int(tc.data.dataloader_num_workers),
            seed=int(tc.data.seed),
        )

        # Validate that t_list is a subset of the cached trajectory timesteps
        # and build a lookup table used in single_train_step.
        sample0 = _KDPathDataset(self._cache_dir)[0]
        traj_ts = sample0["trajectory_timesteps"].tolist()
        t_to_idx: dict[int, int] = {int(t): i for i, t in enumerate(traj_ts)}
        missing_ts = [t for t in self._t_list[:-1] if t not in t_to_idx]
        if missing_ts:
            raise ValueError(f"t_list timesteps {missing_ts} are not present in the cached "
                             f"trajectory timesteps.\nCached: {traj_ts}\n"
                             f"Adjust t_list to use a subset of the cached timesteps.")
        self._t_to_traj_idx: dict[int, int] = {t: t_to_idx[t] for t in self._t_list[:-1]}
        logger.info(
            "t_list → trajectory indices: %s",
            {t: self._t_to_traj_idx[t]
             for t in self._t_list[:-1]},
        )

    # ------------------------------------------------------------------
    # Path generation
    # ------------------------------------------------------------------

    def _generate_paths(
        self,
        source_dataset: Any,
        missing: list[int],
        total: int,
    ) -> None:
        """Generate and cache teacher ODE paths for all missing indices.

        All ranks participate in each teacher forward pass (required by
        FSDP). Only rank 0 writes files. A dist.barrier() follows each
        save to keep ranks in sync.
        """
        missing_set = set(missing)
        samples_dir = Path(self._cache_dir) / "samples"
        samples_dir.mkdir(parents=True, exist_ok=True)
        device = self.student.device

        pbar = tqdm(
            total=len(missing),
            desc="Generating KD cache",
            disable=get_world_rank() != 0,
            unit="sample",
        )
        for global_idx in range(total):
            if global_idx not in missing_set:
                continue  # all ranks skip together — no FSDP divergence

            # Load single sample by stable global index.
            raw = source_dataset.__getitems__([global_idx])
            raw = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in raw.items()}

            with torch.no_grad():
                traj_latents, traj_timesteps, real = self._teacher_ode_rollout(raw)
                # traj_latents:   [S, T, C, H, W]
                # traj_timesteps: [S]  int64
                # real:           [T, C, H, W]

            if get_world_rank() == 0:
                out = {
                    "trajectory_latents": traj_latents.cpu(),
                    "trajectory_timesteps": traj_timesteps.cpu(),
                    "real": real.cpu(),
                    "text_embedding": raw["text_embedding"][0].cpu(),
                    "text_attention_mask": raw["text_attention_mask"][0].cpu(),
                }
                torch.save(out, samples_dir / f"{global_idx:08d}.pt")
                pbar.update(1)

            if dist.is_initialized():
                dist.barrier()

        pbar.close()

    def _teacher_ode_rollout(
        self,
        raw_batch: dict[str, Any],
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Run high-quality teacher ODE; return full trajectory.

        The teacher runs ``teacher_inference_steps`` ODE steps (e.g. 48),
        saving every intermediate state.  The full trajectory is returned so
        the cache is independent of ``t_list`` — subsampling happens at
        training time.

        Returns:
            trajectory_latents:   ``[S, T, C, H, W]`` — state before each step
            trajectory_timesteps: ``[S]`` int64 — teacher timestep at each state
            real:                 ``[T, C, H, W]`` — teacher's final clean x0
        """
        assert self.teacher is not None
        device = self.student.device
        # bf16 is required by FlashAttention and is consistent with training
        # precision. Models load as float32 but inference always runs in bf16.
        dtype = torch.bfloat16

        # prepare_batch runs text encoding; latents: [1, T, C, H, W]  (B=1)
        training_batch = self.teacher.prepare_batch(
            raw_batch,
            generator=self.cuda_generator,
            latents_source="data",
        )
        latents = training_batch.latents.to(dtype)
        B, T = latents.shape[:2]

        # Build teacher's full ODE schedule.
        teacher_sched = SelfForcingFlowMatchScheduler(
            num_inference_steps=self._teacher_inference_steps,
            num_train_timesteps=int(self.student.num_train_timesteps),
            shift=self._flow_shift,
            sigma_min=0.0,
            extra_one_step=True,
            training=False,
        )
        teacher_t_list = [int(t.item()) for t in teacher_sched.timesteps]

        # Start from pure noise at the highest teacher timestep.
        noise = torch.randn(
            latents.shape,
            device=device,
            dtype=dtype,
            generator=self.cuda_generator,
        )
        t0 = teacher_t_list[0]
        t0_flat = torch.full((B * T, ), t0, device=device, dtype=torch.float32)
        x = self._sf_scheduler.add_noise(
            latents.flatten(0, 1),
            noise.flatten(0, 1),
            t0_flat,
        ).unflatten(0, (B, T)).to(dtype)  # [B, T, C, H, W]; scheduler may upcast to fp32

        # Run full teacher ODE, recording every state before each step.
        traj_states: list[torch.Tensor] = []  # each [T, C, H, W]
        pred_x0: torch.Tensor | None = None
        teacher_next = teacher_t_list[1:] + [0]

        for t_cur, t_next in zip(teacher_t_list, teacher_next, strict=True):
            traj_states.append(x.squeeze(0).clone())

            timestep_b = torch.full((B, ), t_cur, device=device, dtype=torch.float32)
            training_batch.timesteps = timestep_b

            pred_noise = self.teacher.predict_noise(
                x,
                timestep_b,
                training_batch,
                conditional=True,
                attn_kind="dense",
            )
            timestep_bt = timestep_b.unsqueeze(1).expand(B, T)
            pred_x0 = pred_noise_to_pred_video(
                pred_noise=pred_noise.flatten(0, 1),
                noise_input_latent=x.flatten(0, 1),
                timestep=timestep_bt,
                scheduler=self._sf_scheduler,
            ).unflatten(0, (B, T))

            if t_next > 0:
                sigma_cur = self._timestep_to_sigma(t_cur, B * T, device)
                sigma_next = self._timestep_to_sigma(t_next, B * T, device)
                x_flat = x.flatten(0, 1)
                p_flat = pred_x0.flatten(0, 1)
                eps = ((x_flat - (1.0 - sigma_cur) * p_flat) / sigma_cur.clamp_min(1e-8))
                x = ((1.0 - sigma_next) * p_flat + sigma_next * eps).unflatten(0, (B, T)).to(dtype)

        assert pred_x0 is not None, "teacher_t_list must have at least one step"

        trajectory_latents = torch.stack(traj_states)  # [S, T, C, H, W]
        trajectory_timesteps = torch.tensor(teacher_t_list, dtype=torch.long)
        real = pred_x0.squeeze(0)  # [T, C, H, W]
        return trajectory_latents, trajectory_timesteps, real

    def _timestep_to_sigma(
        self,
        timestep_val: int,
        numel: int,
        device: torch.device,
    ) -> torch.Tensor:
        """Return sigma for a scalar timestep value, broadcast to [numel, 1, 1, 1]."""
        sigmas = self._sf_scheduler.sigmas.to(device=device, dtype=torch.float32)
        timesteps = self._sf_scheduler.timesteps.to(device=device, dtype=torch.float32)
        t = torch.tensor([float(timestep_val)], device=device)
        idx = torch.argmin((timesteps - t).abs())
        return sigmas[idx].expand(numel).reshape(numel, 1, 1, 1)

    # ------------------------------------------------------------------
    # Training step
    # ------------------------------------------------------------------

    def single_train_step(
        self,
        batch: dict[str, Any],
        iteration: int,
        *,
        current_vsa_sparsity: float = 0.0,
    ) -> tuple[
            dict[str, torch.Tensor],
            dict[str, Any],
            dict[str, LogScalar],
    ]:
        del iteration
        # batch keys: trajectory_latents [B, S, T, C, H, W],
        #             trajectory_timesteps [B, S],
        #             real [B, T, C, H, W],
        #             text_embedding [B, L, D], text_attention_mask [B, L]

        device = self.student.device
        dtype = batch["real"].dtype

        # Randomly select which student denoising step to train on.
        step_i = int(torch.randint(
            0,
            self._num_steps,
            (1, ),
            generator=self.cuda_generator,
        ).item())

        t = self._t_list[step_i]
        traj_idx = self._t_to_traj_idx[t]
        noisy_input = batch["trajectory_latents"][:, traj_idx].to(device, dtype=dtype)  # [B, T, C, H, W]
        target_x0 = batch["real"].to(device, dtype=dtype)  # [B, T, C, H, W]
        B, T = noisy_input.shape[:2]

        # Build a proxy batch so prepare_batch sets up text conditioning and
        # attention metadata. We pass "real" as vae_latent (permuted back to
        # [B, C, T, H, W] format that prepare_batch expects), then override
        # the noisy input and timestep below.
        proxy_batch = {
            "vae_latent": target_x0.permute(0, 2, 1, 3, 4),  # [B, C, T, H, W]
            "text_embedding": batch["text_embedding"].to(device, dtype=dtype),
            "text_attention_mask": batch["text_attention_mask"].to(device, dtype=dtype),
        }
        training_batch = self.student.prepare_batch(
            proxy_batch,
            generator=self.cuda_generator,
            current_vsa_sparsity=current_vsa_sparsity,
            latents_source="data",
        )

        # Override timestep with the actual KD timestep.
        timestep_b = torch.full((B, ), float(t), device=device, dtype=torch.float32)
        training_batch.timesteps = timestep_b

        # Student forward: predict noise from the cached noisy latent.
        pred_noise = self.student.predict_noise(
            noisy_input,
            timestep_b,
            training_batch,
            conditional=True,
            attn_kind="dense",
        )  # [B, T, C, H, W]

        # Convert predicted noise to predicted x0 using _sf_scheduler.
        timestep_bt = timestep_b.unsqueeze(1).expand(B, T)
        pred_x0 = pred_noise_to_pred_video(
            pred_noise=pred_noise.flatten(0, 1),
            noise_input_latent=noisy_input.flatten(0, 1),
            timestep=timestep_bt,
            scheduler=self._sf_scheduler,
        ).unflatten(0, (B, T))  # [B, T, C, H, W]

        loss = 0.5 * F.mse_loss(pred_x0.float(), target_x0.float())

        loss_map: dict[str, torch.Tensor] = {
            "total_loss": loss,
            "kd_loss": loss,
        }
        outputs: dict[str, Any] = {
            "_fv_backward": (
                training_batch.timesteps,
                training_batch.attn_metadata,
            )
        }
        metrics: dict[str, LogScalar] = {"kd_step_idx": float(step_i)}
        return loss_map, outputs, metrics

    def backward(
        self,
        loss_map: dict[str, torch.Tensor],
        outputs: dict[str, Any],
        *,
        grad_accum_rounds: int = 1,
    ) -> None:
        grad_accum_rounds = max(1, int(grad_accum_rounds))
        ctx = outputs.get("_fv_backward")
        if ctx is None:
            super().backward(
                loss_map,
                outputs,
                grad_accum_rounds=grad_accum_rounds,
            )
            return
        self.student.backward(
            loss_map["total_loss"],
            ctx,
            grad_accum_rounds=grad_accum_rounds,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _init_optimizers_and_schedulers(self) -> None:
        tc = self.training_config
        student_lr = float(tc.optimizer.learning_rate)
        if student_lr <= 0.0:
            raise ValueError("training.optimizer.learning_rate must be > 0")

        student_params = [p for p in self.student.transformer.parameters() if p.requires_grad]
        (
            self._student_optimizer,
            self._student_lr_scheduler,
        ) = build_optimizer_and_scheduler(
            params=student_params,
            optimizer_config=tc.optimizer,
            loop_config=tc.loop,
            learning_rate=student_lr,
            betas=tc.optimizer.betas,
            scheduler_name=str(tc.optimizer.lr_scheduler),
        )


# ---------------------------------------------------------------------------
# Causal KD Method
# ---------------------------------------------------------------------------


class KDCausalMethod(KDMethod):
    """KD for causal Wan: per-frame block-quantized timestep sampling.

    Identical to :class:`KDMethod` except ``single_train_step`` samples
    a **per-frame** denoising step index (block-quantized to groups of
    ``num_frames_per_block`` frames) instead of one index per batch.
    This matches the legacy ``ODEInitTrainingPipeline`` training scheme
    required by causal / streaming student models.

    Additional YAML field under ``method``::

        num_frames_per_block: 3   # frames sharing the same noise level
    """

    def __init__(
        self,
        *,
        cfg: Any,
        role_models: dict[str, ModelBase],
    ) -> None:
        super().__init__(cfg=cfg, role_models=role_models)
        self._num_frames_per_block: int = int(self.method_config.get("num_frames_per_block", 3))
        if self._num_frames_per_block < 1:
            raise ValueError("num_frames_per_block must be >= 1")

    def single_train_step(
        self,
        batch: dict[str, Any],
        iteration: int,
        *,
        current_vsa_sparsity: float = 0.0,
    ) -> tuple[
            dict[str, torch.Tensor],
            dict[str, Any],
            dict[str, LogScalar],
    ]:
        del iteration
        device = self.student.device
        dtype = batch["real"].dtype

        # trajectory_latents: [B, S, T, C, H, W]
        B, _S, T, C, H, W = batch["trajectory_latents"].shape
        K = self._num_steps  # number of student steps (excludes t=0)

        # Gather the K relevant trajectory steps for t_list[:-1].
        traj_indices = torch.tensor(
            [self._t_to_traj_idx[t] for t in self._t_list[:-1]],
            dtype=torch.long,
            device=device,
        )
        relevant = batch["trajectory_latents"].to(device, dtype=dtype)[:, traj_indices]  # [B, K, T, C, H, W]

        # Sample per-frame step indices, block-quantized.
        indexes = self._sample_per_frame_step_idx(B, T, K, device)  # [B, T] in [0, K)

        # Gather noisy_input per frame: [B, T, C, H, W]
        noisy_input = torch.gather(
            relevant,
            dim=1,
            index=indexes[:, None, :, None, None, None].expand(B, 1, T, C, H, W),
        ).squeeze(1)

        target_x0 = batch["real"].to(device, dtype=dtype)  # [B, T, C, H, W]

        # Per-frame timestep [B, T]
        t_list_tensor = torch.tensor(self._t_list[:-1], dtype=torch.float32, device=device)
        timestep_per_frame = t_list_tensor[indexes]  # [B, T]

        proxy_batch = {
            "vae_latent": target_x0.permute(0, 2, 1, 3, 4),
            "text_embedding": batch["text_embedding"].to(device, dtype=dtype),
            "text_attention_mask": batch["text_attention_mask"].to(device, dtype=dtype),
        }
        training_batch = self.student.prepare_batch(
            proxy_batch,
            generator=self.cuda_generator,
            current_vsa_sparsity=current_vsa_sparsity,
            latents_source="data",
        )
        training_batch.timesteps = timestep_per_frame  # [B, T]

        pred_noise = self.student.predict_noise(
            noisy_input,
            timestep_per_frame,
            training_batch,
            conditional=True,
            attn_kind="dense",
        )  # [B, T, C, H, W]

        pred_x0 = pred_noise_to_pred_video(
            pred_noise=pred_noise.flatten(0, 1),
            noise_input_latent=noisy_input.flatten(0, 1),
            timestep=timestep_per_frame.flatten(0, 1),
            scheduler=self._sf_scheduler,
        ).unflatten(0, (B, T))  # [B, T, C, H, W]

        loss = 0.5 * F.mse_loss(pred_x0.float(), target_x0.float())

        loss_map: dict[str, torch.Tensor] = {
            "total_loss": loss,
            "kd_loss": loss,
        }
        outputs: dict[str, Any] = {
            "_fv_backward": (
                training_batch.timesteps,
                training_batch.attn_metadata,
            )
        }
        metrics: dict[str, LogScalar] = {
            "kd_mean_step_idx": float(indexes.float().mean().item()),
        }
        return loss_map, outputs, metrics

    def _sample_per_frame_step_idx(
        self,
        B: int,
        T: int,
        K: int,
        device: torch.device,
    ) -> torch.Tensor:
        """Sample per-frame step indices, block-quantized.

        Each group of ``num_frames_per_block`` consecutive frames shares
        the same randomly sampled step index in ``[0, K)``.

        Returns:
            ``[B, T]`` long tensor with values in ``[0, K)``.
        """
        n_blocks = (T + self._num_frames_per_block - 1) // self._num_frames_per_block
        block_idx = torch.randint(
            0,
            K,
            (B, n_blocks),
            generator=self.cuda_generator,
            device=device,
        )  # [B, n_blocks]
        # Expand each block index to num_frames_per_block frames, slice to T.
        per_frame = block_idx.repeat_interleave(self._num_frames_per_block, dim=1)[:, :T]
        return per_frame  # [B, T]