# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

from __future__ import annotations
from typing import Callable, TYPE_CHECKING
import wandb

import torch

from lipforcing.callbacks.callback import Callback
import lipforcing.utils.logging_utils as logger
from lipforcing.utils.basic_utils import get_batch_size_total
from lipforcing.utils.distributed import is_rank0

if TYPE_CHECKING:
    from lipforcing.methods import FastGenModel
    from lipforcing.configs.config import BaseConfig


class CTScheduleCallback(Callback):
    config: "BaseConfig"

    def __init__(
        self,
        q: float = 2.0,
        ratio_limit: float = 0.999,
        kimg_per_stage: int = 12500,
        batch_size: int = 1,
    ):
        self.q = q
        self.ratio_limit = ratio_limit
        self.kimg_per_stage = kimg_per_stage
        self.batch_size = batch_size

        self.stage = 0
        self.ratio = 0.0

    def _get_cur_stage(self, model, iteration):
        # Start from the saved iteration of the first-stage model in TCM
        if hasattr(model, "resume_iter"):
            assert isinstance(model.resume_iter, int)
            iteration = iteration + model.resume_iter

        batch_size = self.batch_size
        if hasattr(self, "config"):
            # override the batch_size using self.config
            batch_size = get_batch_size_total(self.config)

        cur_nimg = iteration * batch_size
        stage = cur_nimg // (self.kimg_per_stage * 1000)
        return stage, cur_nimg

    def _update_schedule(self, stage):
        self.stage = stage
        self.ratio = 1 - 1 / self.q ** (stage + 1)
        if self.ratio > self.ratio_limit:
            logger.info(f"Clipping ratio from {self.ratio} -> {self.ratio_limit}")
            self.ratio = self.ratio_limit

    def on_train_begin(self, model: FastGenModel, iteration: int = 0) -> None:
        stage, _ = self._get_cur_stage(model, iteration)
        self._update_schedule(stage)
        setattr(model, "ratio", self.ratio)

    def on_training_step_end(
        self,
        model: FastGenModel,
        data_batch: dict[str, torch.Tensor],
        output_batch: dict[str, torch.Tensor | Callable],
        loss_dict: dict[str, torch.Tensor],
        iteration: int = 0,
    ) -> None:
        del data_batch, output_batch, loss_dict
        new_stage, cur_nimg = self._get_cur_stage(model, iteration)
        if new_stage > self.stage:
            self._update_schedule(new_stage)
            setattr(model, "ratio", self.ratio)

        if hasattr(self, "config"):
            # only wandb log when config exists
            if iteration % self.config.trainer.logging_iter == 0 and is_rank0():
                if wandb.run:
                    wandb.log({"ct_schedule/kimg": cur_nimg / 1e3, "ct_schedule/ratio": self.ratio}, step=iteration)