# BADAS Training Package
"""
Training modules for video collision anticipation.

Usage:
    # Standard training
    python -m badas.training.lightning_training --config configs/exp_combined_1to1.yml

    # Knowledge distillation
    python -m badas.training.distillation --config configs/distillation_vitl_to_vitb.yml
"""

from badas.training.config import (
    Config,
    ModelConfig,
    DataConfig,
    TrainingConfig,
    SystemConfig,
    LoggingConfig,
    get_config,
)
from badas.training.lightning_module import (
    VideoLightningModule,
    CollisionAnticipationLoss,
    FocalLoss,
    NegativeMarginLoss,
)
from badas.training.dataset import VideoDataset
from badas.training.lightning_data_module import VideoDataModule
from badas.training.distillation import (
    DistillationConfig,
    FullDistillationConfig,
    DistillationModule,
    DistillationLoss,
    VideoMAEv2Classifier,
)

__all__ = [
    # Config
    "Config",
    "ModelConfig",
    "DataConfig",
    "TrainingConfig",
    "SystemConfig",
    "LoggingConfig",
    "get_config",
    # Training
    "VideoLightningModule",
    "CollisionAnticipationLoss",
    "FocalLoss",
    "NegativeMarginLoss",
    # Distillation
    "DistillationConfig",
    "FullDistillationConfig",
    "DistillationModule",
    "DistillationLoss",
    "VideoMAEv2Classifier",
    # Data
    "VideoDataset",
    "VideoDataModule",
]
