"""Neural network modules for scDistill.

Two-Phase Architecture:
- Encoder: X → Z (latent representation)
- Decoder: Z → X_corrected (batch-corrected expression via NB)

Training:
- Phase 1: PCA + Teacher (Harmony) → Z* (batch-corrected target)
- Phase 2a: Encoder Training (X → Z*)
- Phase 2b: Decoder Training (Z* → X, with batch conditioning)
"""

from .encoder import Encoder
from .decoder import Decoder
from .training import (
    get_device,
    SparseAwareDataset,
    nb_nll,
    train_encoder,
    train_decoder,
)
from .inference import get_corrected_expression

__all__ = [
    # Core modules
    "Encoder",
    "Decoder",
    # Utilities
    "get_device",
    "SparseAwareDataset",
    # Training functions
    "nb_nll",
    "train_encoder",
    "train_decoder",
    # Inference functions
    "get_corrected_expression",
]
