# SPDX-License-Identifier: Apache-2.0
"""DMD2 distillation method (algorithm layer)."""

from __future__ import annotations

from typing import Any, Literal

import torch
import torch.nn.functional as F

from fastvideo.train.methods.base import TrainingMethod, LogScalar
from fastvideo.train.models.base import ModelBase
from fastvideo.train.utils.optimizer import (
    build_optimizer_and_scheduler, )
from fastvideo.train.utils.config import (
    get_optional_float,
    get_optional_int,
    parse_betas,
)


class DMD2Method(TrainingMethod):
    """DMD2 distillation algorithm (method layer).

    Owns role model instances directly:
    - ``self.student`` — trainable student :class:`ModelBase`
    - ``self.teacher`` — frozen teacher :class:`ModelBase`
    - ``self.critic`` — trainable critic :class:`ModelBase`
    """

    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("DMD2Method requires role 'student'")
        if "teacher" not in role_models:
            raise ValueError("DMD2Method requires role 'teacher'")
        if "critic" not in role_models:
            raise ValueError("DMD2Method requires role 'critic'")

        self.teacher = role_models["teacher"]
        self.critic = role_models["critic"]

        if not self.student._trainable:
            raise ValueError("DMD2Method requires student to be trainable")
        if self.teacher._trainable:
            raise ValueError("DMD2Method requires teacher to be "
                             "non-trainable")
        if not self.critic._trainable:
            raise ValueError("DMD2Method requires critic to be trainable")
        self._cfg_uncond = self._parse_cfg_uncond()
        self._rollout_mode = self._parse_rollout_mode()
        self._denoising_step_list: torch.Tensor | None = (None)

        # Initialize preprocessors on student.
        self.student.init_preprocessors(self.training_config)

        self._init_optimizers_and_schedulers()

    @property
    def _optimizer_dict(self, ) -> dict[str, torch.optim.Optimizer]:
        return {
            "student": self._student_optimizer,
            "critic": self._critic_optimizer,
        }

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

    # TrainingMethod override: single_train_step
    def single_train_step(
        self,
        batch: dict[str, Any],
        iteration: int,
    ) -> tuple[
            dict[str, torch.Tensor],
            dict[str, Any],
            dict[str, LogScalar],
    ]:
        latents_source: Literal["data", "zeros"] = "data"
        if self._rollout_mode == "simulate":
            latents_source = "zeros"

        training_batch = self.student.prepare_batch(
            batch,
            generator=self.cuda_generator,
            latents_source=latents_source,
        )

        update_student = self._should_update_student(iteration)

        generator_loss = torch.zeros(
            (),
            device=training_batch.latents.device,
            dtype=training_batch.latents.dtype,
        )
        student_ctx = None
        if update_student:
            generator_pred_x0 = self._student_rollout(training_batch, with_grad=True)
            student_ctx = (
                training_batch.timesteps,
                training_batch.attn_metadata_vsa,
            )
            generator_loss = self._dmd_loss(generator_pred_x0, training_batch)

        (
            fake_score_loss,
            critic_ctx,
            critic_outputs,
        ) = self._critic_flow_matching_loss(training_batch)

        total_loss = generator_loss + fake_score_loss
        loss_map = {
            "total_loss": total_loss,
            "generator_loss": generator_loss,
            "fake_score_loss": fake_score_loss,
        }

        outputs: dict[str, Any] = dict(critic_outputs)
        outputs["_fv_backward"] = {
            "update_student": update_student,
            "student_ctx": student_ctx,
            "critic_ctx": critic_ctx,
        }
        metrics: dict[str, LogScalar] = {"update_student": float(update_student)}
        return loss_map, outputs, metrics

    # TrainingMethod override: backward
    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))
        backward_ctx = outputs.get("_fv_backward")
        if not isinstance(backward_ctx, dict):
            super().backward(
                loss_map,
                outputs,
                grad_accum_rounds=grad_accum_rounds,
            )
            return

        update_student = bool(backward_ctx.get("update_student", False))
        if update_student:
            student_ctx = backward_ctx.get("student_ctx")
            if student_ctx is None:
                raise RuntimeError("Missing student backward context")
            self.student.backward(
                loss_map["generator_loss"],
                student_ctx,
                grad_accum_rounds=grad_accum_rounds,
            )

        critic_ctx = backward_ctx.get("critic_ctx")
        if critic_ctx is None:
            raise RuntimeError("Missing critic backward context")
        self.critic.backward(
            loss_map["fake_score_loss"],
            critic_ctx,
            grad_accum_rounds=grad_accum_rounds,
        )

    # TrainingMethod override: get_optimizers
    def get_optimizers(
        self,
        iteration: int,
    ) -> list[torch.optim.Optimizer]:
        optimizers: list[torch.optim.Optimizer] = []
        optimizers.append(self._critic_optimizer)
        if self._should_update_student(iteration):
            optimizers.append(self._student_optimizer)
        return optimizers

    # TrainingMethod override: get_lr_schedulers
    def get_lr_schedulers(
        self,
        iteration: int,
    ) -> list[Any]:
        schedulers: list[Any] = []
        schedulers.append(self._critic_lr_scheduler)
        if self._should_update_student(iteration):
            schedulers.append(self._student_lr_scheduler)
        return schedulers

    # TrainingMethod override: get_grad_clip_targets
    def get_grad_clip_targets(
        self,
        iteration: int,
    ) -> dict[str, torch.nn.Module]:
        targets: dict[str, torch.nn.Module] = {}
        if self._should_update_student(iteration):
            targets["student"] = (self.student.transformer)
        targets["critic"] = self.critic.transformer
        return targets

    def _parse_rollout_mode(self, ) -> Literal["simulate", "data_latent"]:
        raw = self.method_config.get("rollout_mode", None)
        if raw is None:
            raise ValueError("method_config.rollout_mode must be set "
                             "for DMD2")
        if not isinstance(raw, str):
            raise ValueError("method_config.rollout_mode must be a "
                             "string, "
                             f"got {type(raw).__name__}")
        mode = raw.strip().lower()
        if mode in ("simulate", "sim"):
            return "simulate"
        if mode in ("data_latent", "data", "vae_latent"):
            return "data_latent"
        raise ValueError("method_config.rollout_mode must be one of "
                         "{simulate, data_latent}, got "
                         f"{raw!r}")

    def _parse_cfg_uncond(self, ) -> dict[str, Any] | None:
        raw = self.method_config.get("cfg_uncond", None)
        if raw is None:
            return None
        if not isinstance(raw, dict):
            raise ValueError("method_config.cfg_uncond must be a dict "
                             f"when set, got {type(raw).__name__}")

        cfg: dict[str, Any] = dict(raw)

        on_missing_raw = cfg.get("on_missing", "error")
        if on_missing_raw is None:
            on_missing_raw = "error"
        if not isinstance(on_missing_raw, str):
            raise ValueError("method_config.cfg_uncond.on_missing must "
                             "be a string, got "
                             f"{type(on_missing_raw).__name__}")
        on_missing = on_missing_raw.strip().lower()
        if on_missing not in {"error", "ignore"}:
            raise ValueError("method_config.cfg_uncond.on_missing must "
                             "be one of {error, ignore}, got "
                             f"{on_missing_raw!r}")
        cfg["on_missing"] = on_missing

        for channel, policy_raw in list(cfg.items()):
            if channel == "on_missing":
                continue
            if policy_raw is None:
                continue
            if not isinstance(policy_raw, str):
                raise ValueError("method_config.cfg_uncond values must "
                                 "be strings, got "
                                 f"{channel}="
                                 f"{type(policy_raw).__name__}")
            policy = policy_raw.strip().lower()
            allowed = {"keep", "zero", "drop"}
            if channel == "text":
                allowed = {*allowed, "negative_prompt"}
            if policy not in allowed:
                raise ValueError("method_config.cfg_uncond values must "
                                 "be one of "
                                 f"{sorted(allowed)}, got "
                                 f"{channel}={policy_raw!r}")
            cfg[channel] = policy

        return cfg

    def _init_optimizers_and_schedulers(self) -> None:
        tc = self.training_config

        # Student optimizer/scheduler.
        student_lr = float(tc.optimizer.learning_rate)
        student_betas = tc.optimizer.betas
        student_sched = str(tc.optimizer.lr_scheduler)
        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=student_betas,
            scheduler_name=student_sched,
        )

        # Critic optimizer/scheduler — must be set in
        # method config.
        critic_lr_raw = get_optional_float(
            self.method_config,
            "fake_score_learning_rate",
            where="method.fake_score_learning_rate",
        )
        if critic_lr_raw is None or critic_lr_raw == 0.0:
            raise ValueError("method.fake_score_learning_rate must "
                             "be set to a positive value")
        critic_lr = float(critic_lr_raw)

        critic_betas_raw = self.method_config.get("fake_score_betas", None)
        if critic_betas_raw is None:
            raise ValueError("method.fake_score_betas must be set "
                             "(e.g. [0.0, 0.999])")
        critic_betas = parse_betas(
            critic_betas_raw,
            where="method.fake_score_betas",
        )

        critic_sched_raw = self.method_config.get("fake_score_lr_scheduler", None)
        if critic_sched_raw is None:
            raise ValueError("method.fake_score_lr_scheduler must "
                             "be set (e.g. 'constant')")
        critic_sched = str(critic_sched_raw)
        critic_params = [p for p in self.critic.transformer.parameters() if p.requires_grad]
        (
            self._critic_optimizer,
            self._critic_lr_scheduler,
        ) = build_optimizer_and_scheduler(
            params=critic_params,
            optimizer_config=tc.optimizer,
            loop_config=tc.loop,
            learning_rate=critic_lr,
            betas=critic_betas,
            scheduler_name=critic_sched,
        )

    def _should_update_student(
        self,
        iteration: int,
    ) -> bool:
        interval = get_optional_int(
            self.method_config,
            "generator_update_interval",
            where="method.generator_update_interval",
        )
        if interval is None:
            interval = 1
        if interval <= 0:
            return True
        return iteration % interval == 0

    def _get_denoising_step_list(
        self,
        device: torch.device,
    ) -> torch.Tensor:
        if (self._denoising_step_list is not None and self._denoising_step_list.device == device):
            return self._denoising_step_list

        raw = self.method_config.get("dmd_denoising_steps", None)
        if not isinstance(raw, list) or not raw:
            raise ValueError("method_config.dmd_denoising_steps must "
                             "be set for DMD2 distillation")

        steps = torch.tensor(
            [int(s) for s in raw],
            dtype=torch.long,
            device=device,
        )

        warp = self.method_config.get("warp_denoising_step", None)
        if warp is None:
            warp = False
        if bool(warp):
            timesteps = torch.cat((
                self.student.noise_scheduler.timesteps.to("cpu"),
                torch.tensor([0], dtype=torch.float32),
            )).to(device)
            steps = timesteps[1000 - steps]

        self._denoising_step_list = steps
        return steps

    def _sample_rollout_timestep(
        self,
        device: torch.device,
    ) -> torch.Tensor:
        step_list = self._get_denoising_step_list(device)
        index = torch.randint(
            0,
            len(step_list),
            [1],
            device=device,
            dtype=torch.long,
            generator=self.cuda_generator,
        )
        return step_list[index]

    def _student_rollout(
        self,
        batch: Any,
        *,
        with_grad: bool,
    ) -> torch.Tensor:
        latents = batch.latents
        device = latents.device
        dtype = latents.dtype
        step_list = self._get_denoising_step_list(device)

        if self._rollout_mode != "simulate":
            timestep = self._sample_rollout_timestep(device)
            noise = torch.randn(
                latents.shape,
                device=device,
                dtype=dtype,
                generator=self.cuda_generator,
            )
            noisy_latents = self.student.add_noise(latents, noise, timestep)
            pred_x0 = self.student.predict_x0(
                noisy_latents,
                timestep,
                batch,
                conditional=True,
                cfg_uncond=self._cfg_uncond,
                attn_kind="vsa",
            )
            batch.dmd_latent_vis_dict["generator_timestep"] = timestep
            return pred_x0

        target_timestep_idx = torch.randint(
            0,
            len(step_list),
            [1],
            device=device,
            dtype=torch.long,
            generator=self.cuda_generator,
        )
        target_timestep_idx_int = int(target_timestep_idx.item())
        target_timestep = step_list[target_timestep_idx]

        current_noise_latents = torch.randn(
            latents.shape,
            device=device,
            dtype=dtype,
            generator=self.cuda_generator,
        )
        current_noise_latents_copy = (current_noise_latents.clone())

        max_target_idx = len(step_list) - 1
        noise_latents: list[torch.Tensor] = []
        noise_latent_index = target_timestep_idx_int - 1

        if max_target_idx > 0:
            with torch.no_grad():
                for step_idx in range(max_target_idx):
                    current_timestep = step_list[step_idx]
                    current_timestep_tensor = (current_timestep * torch.ones(
                        1,
                        device=device,
                        dtype=torch.long,
                    ))

                    pred_clean = self.student.predict_x0(
                        current_noise_latents,
                        current_timestep_tensor,
                        batch,
                        conditional=True,
                        cfg_uncond=self._cfg_uncond,
                        attn_kind="vsa",
                    )

                    next_timestep = step_list[step_idx + 1]
                    next_timestep_tensor = (next_timestep * torch.ones(
                        1,
                        device=device,
                        dtype=torch.long,
                    ))
                    noise = torch.randn(
                        latents.shape,
                        device=device,
                        dtype=pred_clean.dtype,
                        generator=self.cuda_generator,
                    )
                    current_noise_latents = (self.student.add_noise(
                        pred_clean,
                        noise,
                        next_timestep_tensor,
                    ))
                    noise_latents.append(current_noise_latents.clone())

        if noise_latent_index >= 0:
            if noise_latent_index >= len(noise_latents):
                raise RuntimeError("noise_latent_index is out of bounds")
            noisy_input = noise_latents[noise_latent_index]
        else:
            noisy_input = current_noise_latents_copy

        if with_grad:
            pred_x0 = self.student.predict_x0(
                noisy_input,
                target_timestep,
                batch,
                conditional=True,
                cfg_uncond=self._cfg_uncond,
                attn_kind="vsa",
            )
        else:
            with torch.no_grad():
                pred_x0 = self.student.predict_x0(
                    noisy_input,
                    target_timestep,
                    batch,
                    conditional=True,
                    cfg_uncond=self._cfg_uncond,
                    attn_kind="vsa",
                )

        batch.dmd_latent_vis_dict["generator_timestep"] = target_timestep.float().detach()
        return pred_x0

    def _critic_flow_matching_loss(
        self,
        batch: Any,
    ) -> tuple[torch.Tensor, Any, dict[str, Any]]:
        with torch.no_grad():
            generator_pred_x0 = self._student_rollout(batch, with_grad=False)

        device = generator_pred_x0.device
        fake_score_timestep = torch.randint(
            0,
            int(self.student.num_train_timesteps),
            [1],
            device=device,
            dtype=torch.long,
            generator=self.cuda_generator,
        )
        fake_score_timestep = (self.student.shift_and_clamp_timestep(fake_score_timestep))

        noise = torch.randn(
            generator_pred_x0.shape,
            device=device,
            dtype=generator_pred_x0.dtype,
            generator=self.cuda_generator,
        )
        noisy_x0 = self.student.add_noise(generator_pred_x0, noise, fake_score_timestep)

        pred_noise = self.critic.predict_noise(
            noisy_x0,
            fake_score_timestep,
            batch,
            conditional=True,
            cfg_uncond=self._cfg_uncond,
            attn_kind="dense",
        )
        target = noise - generator_pred_x0
        flow_matching_loss = torch.mean((pred_noise - target)**2)

        batch.fake_score_latent_vis_dict = {
            "generator_pred_video": generator_pred_x0,
            "fake_score_timestep": fake_score_timestep,
        }
        outputs = {"fake_score_latent_vis_dict": (batch.fake_score_latent_vis_dict)}
        return (
            flow_matching_loss,
            (batch.timesteps, batch.attn_metadata),
            outputs,
        )

    def _dmd_loss(
        self,
        generator_pred_x0: torch.Tensor,
        batch: Any,
    ) -> torch.Tensor:
        guidance_scale = get_optional_float(
            self.method_config,
            "real_score_guidance_scale",
            where="method.real_score_guidance_scale",
        )
        if guidance_scale is None:
            guidance_scale = 1.0
        device = generator_pred_x0.device

        with torch.no_grad():
            timestep = torch.randint(
                0,
                int(self.student.num_train_timesteps),
                [1],
                device=device,
                dtype=torch.long,
                generator=self.cuda_generator,
            )
            timestep = (self.student.shift_and_clamp_timestep(timestep))

            noise = torch.randn(
                generator_pred_x0.shape,
                device=device,
                dtype=generator_pred_x0.dtype,
                generator=self.cuda_generator,
            )
            noisy_latents = self.student.add_noise(generator_pred_x0, noise, timestep)

            faker_x0 = self.critic.predict_x0(
                noisy_latents,
                timestep,
                batch,
                conditional=True,
                cfg_uncond=self._cfg_uncond,
                attn_kind="dense",
            )
            real_cond_x0 = self.teacher.predict_x0(
                noisy_latents,
                timestep,
                batch,
                conditional=True,
                cfg_uncond=self._cfg_uncond,
                attn_kind="dense",
            )
            real_uncond_x0 = self.teacher.predict_x0(
                noisy_latents,
                timestep,
                batch,
                conditional=False,
                cfg_uncond=self._cfg_uncond,
                attn_kind="dense",
            )
            real_cfg_x0 = real_uncond_x0 + (real_cond_x0 - real_uncond_x0) * guidance_scale

            denom = torch.abs(generator_pred_x0 - real_cfg_x0).mean()
            grad = (faker_x0 - real_cfg_x0) / denom
            grad = torch.nan_to_num(grad)

        loss = 0.5 * F.mse_loss(
            generator_pred_x0.float(),
            (generator_pred_x0.float() - grad.float()).detach(),
        )
        return loss