# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Knowledge Distillation recipe for next-token prediction with NeMo-AutoModel.

This recipe fine-tunes a *student* model using the logits of a frozen *teacher* model. It
extends ``FinetuneRecipeForNextTokenPrediction`` adding:

1. teacher_model – an additional HF/NeMo model loaded in ``eval`` mode
2. kd_loss_fn     – KL-divergence between temperature-scaled distributions
3. kd_ratio       – linear mix between CE loss and KD loss

The training loop is copied from the parent class but the loss becomes:
    loss = (1-kd_ratio) * ce_loss + kd_ratio * kd_loss

The file exposes ``KnowledgeDistillationRecipeForNextTokenPrediction`` and a
``main`` entry-point so it can be launched exactly the same way as other
recipes:

    python -m torch.distributed.run --nproc-per-node=8 \
        nemo_automodel/recipes/llm/knowledge_distillation.py \
        -c examples/llm/llama_3_2_1b_kd.yaml
"""

from __future__ import annotations

import logging
import time
from contextlib import nullcontext
from typing import Any, Dict, Optional

import torch
import wandb
from torchao.float8 import precompute_float8_dynamic_scale_for_fsdp

from nemo_automodel._transformers.auto_tokenizer import NeMoAutoTokenizer
from nemo_automodel.components.config._arg_parser import parse_args_and_load_config
from nemo_automodel.components.distributed.cp_utils import make_cp_batch_and_ctx
from nemo_automodel.components.distributed.utils import get_sync_ctx
from nemo_automodel.components.loggers.metric_logger import MetricsSample
from nemo_automodel.components.loss.linear_ce import FusedLinearCrossEntropy
from nemo_automodel.components.training.rng import ScopedRNG
from nemo_automodel.components.training.utils import ScopedModuleOffloading, count_tail_padding
from nemo_automodel.recipes.llm.train_ft import (
    TrainFinetuneRecipeForNextTokenPrediction,
    calculate_loss,
)

logger = logging.getLogger(__name__)


def _build_kd_loss_fn(cfg_kd):
    if cfg_kd is None:
        logger.info("No KD loss function provided, using KLDivLoss")
        return torch.nn.KLDivLoss(reduction="batchmean")
    return cfg_kd.instantiate()


def _build_teacher_model(
    cfg_teacher,
    seed,
    has_packed_sequence,
    device_mesh=None,
    moe_mesh=None,
    distributed_config=None,
    device=None,
):
    """Build and initialize the teacher model for knowledge distillation.

    Uses the same infrastructure as student model (NeMoAutoModelForCausalLM) but without
    PEFT, FP8, or QAT since the teacher should be frozen in full precision.

    Args:
        cfg_teacher: Configuration for teacher model instantiation.
        seed: Random seed for reproducibility.
        has_packed_sequence: Whether using packed sequences.
        device_mesh: Device mesh for distributed training.
        moe_mesh: MOE mesh for expert parallelism.
        distributed_config: Strategy-specific distributed config.
        device: Device to place the teacher model on.

    Returns:
        The frozen teacher model ready for inference.

    Note:
        The `offload_teacher_model` config option is not supported with this approach.
        Device placement is handled internally by NeMoAutoModelForCausalLM infrastructure.
    """

    assert cfg_teacher is not None, "`teacher_model` section missing from YAML config"
    logger.info("Instantiating teacher model")

    # Build teacher model using the same infrastructure as student
    # but without PEFT/FP8/QAT (teacher should be frozen in full precision)
    with ScopedRNG(seed=seed, ranked=True):
        kwargs: Dict[str, Any] = {
            "has_packed_sequence": has_packed_sequence,
            "device_mesh": device_mesh,
            "moe_mesh": moe_mesh,
            "distributed_config": distributed_config,
        }

        teacher_model = cfg_teacher.instantiate(**kwargs)

        # Ensure the teacher model is on the correct device
        teacher_model = teacher_model.to(device)

        # Set teacher to eval mode and freeze parameters
        teacher_model.eval()
        for p in teacher_model.parameters():
            p.requires_grad_(False)

        return teacher_model


def _verify_tokenizer_compatibility(student_cfg, teacher_cfg, trust_remote_code=True):
    if student_cfg is None or teacher_cfg is None:
        raise ValueError("Student and teacher model configs are required")
    student_tokenizer = NeMoAutoTokenizer.from_pretrained(
        student_cfg.pretrained_model_name_or_path, trust_remote_code=trust_remote_code
    )
    teacher_tokenizer = NeMoAutoTokenizer.from_pretrained(
        teacher_cfg.pretrained_model_name_or_path, trust_remote_code=trust_remote_code
    )
    if student_tokenizer.vocab_size != teacher_tokenizer.vocab_size:
        raise ValueError(
            "Student and teacher tokenizers have different vocab sizes; Support will be added in the future"
        )
    if student_tokenizer.pad_token != teacher_tokenizer.pad_token:
        raise ValueError("Student and teacher tokenizers have different pad tokens")
    del student_tokenizer, teacher_tokenizer


class KnowledgeDistillationRecipeForNextTokenPrediction(TrainFinetuneRecipeForNextTokenPrediction):
    """Fine-tune a student model via knowledge distillation."""

    def setup(self):  # noqa: C901 – same complexity as parent
        """Build student & teacher, dataloaders, optimizers, etc."""
        # Right now, we only support tokenizer compatibility for the same tokenizer.
        # We will add support for different tokenizers in the future.
        _verify_tokenizer_compatibility(self.cfg.get("model", None), self.cfg.get("teacher_model", None))

        # Let the parent class build *everything* for the student first
        super().setup()
        if self.pp_enabled:
            raise ValueError("Pipeline parallelism support will be added in the future for knowledge distillation")

        self._offload_teacher_model = self.cfg.get("offload_teacher_model", False)
        teacher_device = self.dist_env.device if not self._offload_teacher_model else "cpu"

        self.teacher_model = _build_teacher_model(
            cfg_teacher=self.cfg.get("teacher_model", None),
            seed=self.cfg.get("seed", 42),
            has_packed_sequence=self.cfg.get("packed_sequence.packed_sequence_size", 0) > 0,
            device_mesh=self.device_mesh,
            moe_mesh=self.moe_mesh,
            distributed_config=self.distributed_config,
            device=teacher_device,
        )
        logger.info("Teacher Model: " + str(self.teacher_model))
        # KD
        self.kd_loss_fn = _build_kd_loss_fn(self.cfg.get("kd_loss_fn", None))
        self.kd_ratio: float = float(self.cfg.get("kd_ratio", 0.5))
        logger.info("KD Loss config: " + str(self.cfg.get("kd_loss_fn", None)))
        temperature = getattr(self.kd_loss_fn, "temperature", "N/A")
        logger.info(f"Knowledge-distillation enabled: ratio={self.kd_ratio}, T={temperature}")

        # Buffers for logging
        self._kd_loss_buffer = []
        self._ce_loss_buffer = []

    #  Override the forward backward step to inject KD loss
    def _forward_backward_step(
        self,
        idx,
        batch,
        *,
        num_label_tokens,
        num_batches,
        is_train: bool = True,
    ):
        """Override the forward backward step to include knowledge distillation loss."""
        batch = {k: v.to(self.dist_env.device, non_blocking=True) for k, v in batch.items()}
        labels = batch.pop("labels")
        train_ctx, batch = make_cp_batch_and_ctx(self.device_mesh, batch, labels)

        model = self.model_parts[0]
        sync_ctx = (
            get_sync_ctx(
                model,
                idx == num_batches - 1,
                defer_fsdp_grad_sync=getattr(self.distributed_config, "defer_fsdp_grad_sync", True),
            )
            if is_train
            else nullcontext()
        )
        with train_ctx(), sync_ctx:
            # No grad for teacher forward
            with (
                ScopedModuleOffloading(self.teacher_model, enabled=self._offload_teacher_model),
                torch.inference_mode(),
            ):
                teacher_logits = self.teacher_model(**batch)
                teacher_logits = getattr(teacher_logits, "logits", teacher_logits).detach().clone()

            # Student forward
            student_keep_last = isinstance(self.loss_fn, FusedLinearCrossEntropy)
            if student_keep_last:
                # Student forward keeping only last token logits to match loss_fn
                student_out = model(logits_to_keep=1, **batch)
            else:
                student_out = model(**batch)

            student_logits = getattr(student_out, "logits", student_out)  # shape (B, S, V)
            # Cross-entropy loss against true labels (same as parent)
            ce_loss = calculate_loss(
                self.loss_fn,
                logits=student_logits,
                labels=labels,
                model=model,
                hidden_states=student_out.hidden_states[-1] if "hidden_states" in student_out else None,
                num_label_tokens=num_label_tokens,
            )
            # Reminder: kd_loss is normalized by num_label_tokens,
            # which typically is larger than the number of labels in this batch,
            # because it contains the total number of labels for all batches contained
            # in one optimization step (grad_acc_steps = gbs / mbs).
            kd_loss = self.kd_loss_fn(
                student_logits,
                teacher_logits,
                labels,
                num_batch_labels=num_label_tokens,
            )
            local_loss = (1.0 - self.kd_ratio) * ce_loss + self.kd_ratio * kd_loss
            if is_train:
                (local_loss * self._get_dp_group_size(include_cp=True)).backward()
            # return the losses for logging
            detached_local = local_loss.detach().clone()
            return detached_local, kd_loss.detach().clone(), ce_loss.detach().clone()

    def _run_train_optim_step(self, batches, max_grad_norm: Optional[float] = None):
        """Execute a single training step.

        Args:
            batches: List of batches of training data.
            max_grad_norm: Gradient clipping norm. Optional, if None will not clip gradients.
        """

        num_label_tokens = torch.tensor(
            sum((batch["labels"] != -100).sum().item() for batch in batches), dtype=torch.long
        )
        num_label_tokens = self._dp_allreduce(num_label_tokens).item()
        loss_buffer = []

        # number of tokens in the batch, excluding any tail padding.
        num_tokens_in_batch = torch.tensor(
            sum(batch["labels"].numel() - count_tail_padding(batch["labels"]) for batch in batches),
            dtype=torch.long,
        )
        num_tokens_in_batch = self._dp_allreduce(num_tokens_in_batch).item()
        num_batches = len(batches)
        for i, batch in enumerate(batches):
            local_loss, kd_loss, ce_loss = self._forward_backward_step(
                i, batch, num_label_tokens=num_label_tokens, num_batches=num_batches
            )
            loss_buffer.append(local_loss)
            self._ce_loss_buffer.append(ce_loss)
            self._kd_loss_buffer.append(kd_loss)

        grad_norm = 0
        # Clip gradients **after** any rescaling.
        # TODO(@boxiangw): Fix TP gradient clipping
        if max_grad_norm is not None:
            if not self.device_mesh or self.device_mesh["tp"].size() == 1:
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    [p for p in self.model_parts[0].parameters() if p.requires_grad], max_grad_norm
                )
                if hasattr(grad_norm, "full_tensor"):
                    grad_norm = grad_norm.full_tensor()  # collect the summed grad norm across ranks

            if isinstance(grad_norm, torch.Tensor):
                grad_norm = grad_norm.item()

        self.checkpointer.maybe_wait_for_staging()
        for opt in self.optimizer:
            opt.step()
            opt.zero_grad()

        if self.lr_scheduler is not None:
            for scheduler in self.lr_scheduler:
                scheduler.step(1)

        # Precompute FP8 scales
        fp8_config = self.cfg.get("fp8", None)
        if (
            fp8_config is not None
            and fp8_config.get("enabled", False)
            and fp8_config.get("precompute_float8_dynamic_scale_for_fsdp", False)
            and not self.pp_enabled
            and self.device_mesh is not None
            and self.device_mesh["dp_shard"].size() > 1
        ):
            precompute_float8_dynamic_scale_for_fsdp(self.model_parts[0])

        # Note(MegatronFSDP): Need to call these functions for MegatronFSDP if not using latest api
        # self.model_parts[0].install_optimized_model_weights()
        # self.model_parts[0].zero_grad_buffer()

        t = time.perf_counter()
        time_delta = t - self.timestamp
        self.timestamp = t
        tps = num_tokens_in_batch / time_delta
        reporting_loss = torch.sum(torch.stack(loss_buffer))
        reporting_loss = self._dp_allreduce(reporting_loss, include_cp=True)
        reporting_loss = reporting_loss.cpu().item()
        # fix reporting_loss, tps across ranks

        ce_loss = self._dp_allreduce(torch.stack(self._ce_loss_buffer).sum(), include_cp=True).item()
        kd_loss = self._dp_allreduce(torch.stack(self._kd_loss_buffer).sum(), include_cp=True).item()
        # Clear buffers for next step
        self._ce_loss_buffer.clear()
        self._kd_loss_buffer.clear()

        # return reporting_loss, grad_norm, tps, num_tokens_in_batch, num_label_tokens
        return MetricsSample(
            step=self.step_scheduler.step,
            epoch=self.step_scheduler.epoch,
            metrics={
                "loss": reporting_loss,
                "ce_loss": ce_loss,
                "kd_loss": kd_loss,
                "grad_norm": grad_norm,
                "lr": self.optimizer[0].param_groups[0]["lr"],
                "mem": torch.cuda.max_memory_allocated() / 1024**3,
                "tps": tps,
                "tps_per_gpu": tps / max(self._get_dp_group_size(), 1),
                "num_tokens_per_step": num_tokens_in_batch,
                "num_label_tokens": num_label_tokens,
                "kd_ratio": self.kd_ratio,
                "temperature": getattr(self.kd_loss_fn, "temperature", float("nan")),
            },
        )

    @torch.no_grad()
    def _run_validation_epoch(self, val_dataloader):
        """Run one pass over `self.val_dataloader`."""
        if self.pp_enabled:
            logger.warning("Validation is not supported for pipeline parallelism")
            return

        with ScopedRNG(seed=1, ranked=True):
            for mp in self.model_parts:
                mp.eval()

            total_loss = torch.tensor(0.0, dtype=torch.float32, device=self.dist_env.device)
            ce_loss = torch.tensor(0.0, dtype=torch.float32, device=self.dist_env.device)
            kd_loss = torch.tensor(0.0, dtype=torch.float32, device=self.dist_env.device)
            total_num_label_tokens = 0

            for batch in val_dataloader:
                num_label_tokens = (batch["labels"] != -100).sum().item()
                local_loss, _kd_loss, _ce_loss = self._forward_backward_step(
                    0,
                    batch,
                    num_label_tokens=num_label_tokens,
                    num_batches=1,
                    is_train=False,
                )
                total_num_label_tokens += num_label_tokens
                ce_loss += _ce_loss
                kd_loss += _kd_loss
                total_loss += local_loss

        total_loss = self._dp_allreduce(total_loss, include_cp=True).item()
        ce_loss = self._dp_allreduce(ce_loss, include_cp=True).item()
        kd_loss = self._dp_allreduce(kd_loss, include_cp=True).item()
        total_num_label_tokens = self._dp_allreduce(torch.tensor(total_num_label_tokens, dtype=torch.long)).item()

        val_loss = total_loss / max(total_num_label_tokens, 1e-8)
        return MetricsSample(
            step=self.step_scheduler.step,
            epoch=self.step_scheduler.epoch,
            metrics={
                "val_loss": val_loss,
                "ce_loss": ce_loss,
                "kd_loss": kd_loss,
                "lr": self.optimizer[0].param_groups[0]["lr"],
                "num_label_tokens": total_num_label_tokens,
                "mem": torch.cuda.max_memory_allocated() / 1024**3,
            },
        )

    def log_val_metrics(self, val_name, log_data, metric_logger=None):
        if not self.dist_env.is_main or log_data is None:
            return

        if wandb.run is not None:
            wandb.log(log_data.to_dict() | {"val_name": val_name}, step=log_data.step)

        # JSONL validation log
        if not metric_logger is None:
            metric_logger.log(log_data)

        # assumes all model parts' optimizers have the same learning rate
        logging.info(
            "[val] {} | step {} | epoch {} | loss {:.4f} | ce_loss {:.4f} | kd_loss {:.4f} | lr {:.2e} | num_label_tokens {}".format(
                val_name,
                log_data.step,
                log_data.epoch,
                log_data.metrics["val_loss"],
                log_data.metrics["ce_loss"],
                log_data.metrics["kd_loss"],
                log_data.metrics["lr"],
                log_data.metrics["num_label_tokens"],
            )
        )

    def log_train_metrics(self, log_data) -> float:
        """Log metrics to wandb and other loggers.

        Args:
            log_data: MetricsSample object, containing:
                step: int, the current step.
                epoch: int, the current epoch.
                metrics: Dict[str, float], containing:
                    "loss": Training loss.
                    "grad_norm": Grad norm from the training step.
                    "lr": Learning rate.
                    "mem": Memory allocated.
                    "tps": Tokens per second.
                    "tps_per_gpu": Tokens per second per GPU.
                    "num_label_tokens": Number of label tokens.
        """
        if not self.dist_env.is_main:
            return

        # Log to remote services (WandB) according to step_scheduler frequency
        if self.step_scheduler.is_remote_logging_step:
            if wandb.run is not None:
                wandb.log(log_data.to_dict(), step=log_data.step)

        logging.info(
            "step {} | epoch {} | "
            "loss {:.4f} | ce_loss {:.4f} | kd_loss {:.4f} | "
            "lr {:.2e} | tps {:.2f} | kd_ratio {:.2f} | temperature {:.2f}".format(
                log_data.step,
                log_data.epoch,
                log_data.metrics["loss"],
                log_data.metrics["ce_loss"],
                log_data.metrics["kd_loss"],
                log_data.metrics["lr"],
                log_data.metrics["tps"],
                log_data.metrics["kd_ratio"],
                log_data.metrics["temperature"],
            )
        )
        torch.cuda.reset_peak_memory_stats()


# Entry point
def main(config_path="examples/llm_kd/llama3_2/llama3_2_1b_kd.yaml"):
    """Run the KD recipe from CLI or directly."""
    cfg = parse_args_and_load_config(config_path)
    trainer = KnowledgeDistillationRecipeForNextTokenPrediction(cfg)
    trainer.setup()
    trainer.run_train_validation_loop()


if __name__ == "__main__":  # pragma: no cover
    main()
