"""
NB-Transformer Package

A PyTorch Lightning-based implementation of transformers for fast Negative Binomial GLM 
parameter estimation - a modern replacement for DESeq2 statistical analysis.

The package provides attention-based models that learn to estimate parameters of NB GLM
models from variable-length sets of observations, providing 14.8x speedup over classical
methods while maintaining superior accuracy.

Main components:
- DispersionTransformer: Fast NB GLM parameter estimation (mu, beta, alpha) 
- PairSetTransformer: Base transformer model for pair-set tasks
- SyntheticNBGLMDataset: Online synthetic data generation for NB GLM
- DispersionLightningModule: PyTorch Lightning training module
- Statistical inference utilities for p-values and confidence intervals
"""

from .model import PairSetTransformer, DispersionTransformer
from .dataset import SyntheticNBGLMDataset, create_dataloaders
from .utils import (
    normalize_data,
    denormalize_data,
    compute_rmse,
    compute_mae,
    EarlyStopping,
    mean_pooling,
    masked_mean_pooling,
    pad_sequences,
    create_padding_mask
)
from .inference import (
    compute_fisher_weights,
    compute_standard_errors,
    compute_wald_statistics,
    compute_nb_glm_inference,
    validate_calibration,
    summarize_calibration_results,
    load_pretrained_model,
    quick_inference_example
)
from .method_of_moments import (
    MethodOfMomentsEstimator,
    estimate_nb_glm_parameters,
    estimate_batch_parameters,
    estimate_batch_parameters_vectorized,
    MoMEstimator,
    estimate_parameters
)
__version__ = "1.0.0"
__author__ = "Valentine Svensson"
__email__ = "valentine.svensson@gmail.com"

__all__ = [
    "PairSetTransformer",
    "DispersionTransformer", 
    "SyntheticNBGLMDataset",
    "create_dataloaders",
    "normalize_data",
    "denormalize_data",
    "compute_rmse", 
    "compute_mae",
    "EarlyStopping",
    "mean_pooling",
    "masked_mean_pooling",
    "pad_sequences",
    "create_padding_mask",
    "compute_fisher_weights",
    "compute_standard_errors", 
    "compute_wald_statistics",
    "compute_nb_glm_inference",
    "validate_calibration",
    "summarize_calibration_results",
    "load_pretrained_model",
    "quick_inference_example",
    "MethodOfMomentsEstimator",
    "estimate_nb_glm_parameters",
    "estimate_batch_parameters",
    "estimate_batch_parameters_vectorized",
    "MoMEstimator",
    "estimate_parameters"
]