import os
import torch
from transformers import Trainer, TrainerCallback, TrainerControl, TrainerState


class LlavaNextTrainer(Trainer):
    def __init__(
        self, *args, projector_lr=None, vision_lr=None, llm_lr=None, lm_head_lr=None, **kwargs
    ):
        super().__init__(*args, **kwargs)
        self.projector_lr = projector_lr
        self.vision_lr = vision_lr
        self.llm_lr = llm_lr
        self.lm_head_lr = lm_head_lr

    def _build_param_groups(self, model):
        groups = []
        seen = set()

        def add_param_group(params, lr):
            trainable_params = [p for p in params if p.requires_grad]
            if not trainable_params:
                return

            groups.append(
                {"params": trainable_params, "lr": lr, "weight_decay": self.args.weight_decay}
            )
            for p in trainable_params:
                seen.add(id(p))

        if hasattr(model, "multi_modal_projector"):
            lr = self.projector_lr if self.projector_lr is not None else self.args.learning_rate
            add_param_group(model.multi_modal_projector.parameters(), lr)
        if hasattr(model, "vision_tower"):
            lr = self.vision_lr if self.vision_lr is not None else self.args.learning_rate
            add_param_group(model.vision_tower.parameters(), lr)
        if hasattr(model, "language_model"):
            lr = self.llm_lr if self.llm_lr is not None else self.args.learning_rate
            add_param_group(model.language_model.parameters(), lr)
        if hasattr(model, "lm_head"):
            lr = self.lm_head_lr if self.lm_head_lr is not None else self.args.learning_rate
            add_param_group(model.lm_head.parameters(), lr)

        other_trainable_params = [
            p for p in model.parameters() if p.requires_grad and id(p) not in seen
        ]
        add_param_group(other_trainable_params, self.args.learning_rate)

        return groups

    def create_optimizer(self):
        if self.optimizer is not None:
            return

        optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(self.args)
        param_groups = self._build_param_groups(self.model)
        self.optimizer = optimizer_cls(param_groups, **optimizer_kwargs)


class GpuStatsCallback(TrainerCallback):
    def __init__(self):
        self.rank = int(os.environ.get("RANK", "0"))
        self.local_rank = int(os.environ.get("LOCAL_RANK", "0"))
        self.world_size = int(os.environ.get("WORLD_SIZE", "1"))

    def on_log(self, args, state: TrainerState, control: TrainerControl, **kwargs):
        if not torch.cuda.is_available():
            return

        dev = torch.cuda.current_device()
        stats = {
            "gpu/local_rank": self.local_rank,
            "gpu/world_rank": self.rank,
            "gpu/device_index": dev,
            "gpu/mem_alloc_MB": torch.cuda.memory_allocated(dev) / (1024**2),
            "gpu/mem_reserved_MB": torch.cuda.memory_reserved(dev) / (1024**2),
            "gpu/max_mem_alloc_MB": torch.cuda.max_memory_allocated(dev) / (1024**2),
            "gpu/max_mem_reserved_MB": torch.cuda.max_memory_reserved(dev) / (1024**2),
        }
        try:
            import wandb

            if args.report_to and "wandb" in args.report_to:
                wandb.log(stats, step=int(state.global_step))
        except Exception:
            pass
