"""PyTorch Lightning callbacks utilities."""

from __future__ import annotations

import lightning.pytorch as pl
from torch import nn

from vis4d.engine.loss_module import LossModule
from vis4d.engine.training_module import TrainingModule


def get_model(model: pl.LightningModule) -> nn.Module:
    """Get model from pl module."""
    if isinstance(model, TrainingModule):
        return model.model
    return model


def get_loss_module(loss_module: pl.LightningModule) -> LossModule:
    """Get loss_module from pl module."""
    assert hasattr(loss_module, "loss_module") and isinstance(
        loss_module.loss_module, LossModule
    ), "Loss module is not set in the training module."
    return loss_module.loss_module