# Core shared modules for BADAS
# These modules are used by both training and inference

from badas.core.base import VideoClassifierBase
from badas.core.vjepa2 import VJEPA2Classifier
from badas.core.videomae import VideoMAEv2Classifier
from badas.core.registry import (
    register_model,
    get_model_class,
    create_model,
    detect_model_type,
    list_registered_models,
)
from badas.core.modules import (
    LinearHead,
    MLPHead,
    AttentionHead,
    TemporalProcessor,
    MeanPooling,
    MaxPooling,
    LSTMProcessor,
    AttentionProcessor,
    AttentiveProbe,
)
from badas.core.preprocessing import (
    IMAGENET_MEAN,
    IMAGENET_STD,
    calculate_frame_indices,
    sample_frames_uniform,
    resize_frame,
    preprocess_frame_cpu,
    preprocess_single_frame,
    normalize_tensor,
    VideoPreprocessor,
    preprocess_video_frames,
    preprocess_frames_batch,
    # Debug utilities
    denormalize_tensor,
    save_tensor_as_video,
    save_tensor_debug,
)


__all__ = [
    # Base class
    "VideoClassifierBase",
    # Model implementations
    "VJEPA2Classifier",
    "VideoMAEv2Classifier",
    # Registry
    "register_model",
    "get_model_class",
    "create_model",
    "detect_model_type",
    "list_registered_models",
    # Classification heads
    "LinearHead",
    "MLPHead",
    "AttentionHead",
    # Temporal processors
    "TemporalProcessor",
    "MeanPooling",
    "MaxPooling",
    "LSTMProcessor",
    "AttentionProcessor",
    "AttentiveProbe",
    # Preprocessing
    "IMAGENET_MEAN",
    "IMAGENET_STD",
    "calculate_frame_indices",
    "sample_frames_uniform",
    "resize_frame",
    "preprocess_frame_cpu",
    "preprocess_single_frame",
    "normalize_tensor",
    "VideoPreprocessor",
    "preprocess_video_frames",
    "preprocess_frames_batch",
    # Debug utilities
    "denormalize_tensor",
    "save_tensor_as_video",
    "save_tensor_debug",
]
