import importlib.metadata

from .base import (
    LightningIRClassFactory,
    LightningIRConfig,
    LightningIRConfigClassFactory,
    LightningIRModel,
    LightningIRModelClassFactory,
    LightningIRModule,
    LightningIROutput,
    LightningIRTokenizer,
    LightningIRTokenizerClassFactory,
)
from .bi_encoder import (
    BiEncoderConfig,
    BiEncoderEmbedding,
    BiEncoderModel,
    BiEncoderModule,
    BiEncoderOutput,
    BiEncoderTokenizer,
    MultiVectorBiEncoderConfig,
    MultiVectorBiEncoderModel,
    SingleVectorBiEncoderConfig,
    SingleVectorBiEncoderModel,
)
from .callbacks import IndexCallback, RankCallback, RegisterLocalDatasetCallback, ReRankCallback, SearchCallback
from .cross_encoder import (
    CrossEncoderConfig,
    CrossEncoderModel,
    CrossEncoderModule,
    CrossEncoderOutput,
    CrossEncoderTokenizer,
)
from .data import (
    DocDataset,
    DocSample,
    IndexBatch,
    LightningIRDataModule,
    QueryDataset,
    QuerySample,
    RankBatch,
    RankSample,
    RunDataset,
    SearchBatch,
    TrainBatch,
    TupleDataset,
)
from .data.external_datasets import _register_external_datasets
from .loss import (
    ApproxMRR,
    ApproxNDCG,
    ApproxRankMSE,
    ConstantMarginMSE,
    ContrastiveLocalLoss,
    FLOPSRegularization,
    InBatchCrossEntropy,
    InfoNCE,
    KLDivergence,
    L1Regularization,
    L2Regularization,
    PearsonCorrelation,
    RankNet,
    ScoreBasedInBatchCrossEntropy,
    ScoreBasedInBatchLossFunction,
    SupervisedMarginMSE,
)
from .main import LightningIRTrainer, LightningIRWandbLogger
from .models.register_external_models import _register_external_models
from .models.register_internal_models import _register_internal_models
from .retrieve import (
    FaissFlatIndexConfig,
    FaissFlatIndexer,
    FaissIVFIndexConfig,
    FaissIVFIndexer,
    FaissIVFPQIndexConfig,
    FaissIVFPQIndexer,
    FaissPQIndexConfig,
    FaissPQIndexer,
    FaissSearchConfig,
    FaissSearcher,
    IndexConfig,
    Indexer,
    SearchConfig,
    Searcher,
    SeismicIndexConfig,
    SeismicIndexer,
    SeismicSearchConfig,
    SeismicSearcher,
    TorchDenseIndexConfig,
    TorchDenseIndexer,
    TorchDenseSearchConfig,
    TorchDenseSearcher,
    TorchSparseIndexConfig,
    TorchSparseIndexer,
    TorchSparseSearchConfig,
    TorchSparseSearcher,
)
from .schedulers import (
    ConstantLRSchedulerWithLinearWarmup,
    GenericConstantSchedulerWithLinearWarmup,
    GenericConstantSchedulerWithQuadraticWarmup,
    GenericLinearSchedulerWithLinearWarmup,
    LinearLRSchedulerWithLinearWarmup,
    WarmupLRScheduler,
)

__version__ = importlib.metadata.version("lightning-ir")

_register_internal_models()
_register_external_models()
_register_external_datasets()


__all__ = [
    "ApproxMRR",
    "ApproxNDCG",
    "ApproxRankMSE",
    "BiEncoderConfig",
    "BiEncoderEmbedding",
    "BiEncoderModel",
    "BiEncoderModule",
    "BiEncoderOutput",
    "BiEncoderTokenizer",
    "ConstantLRSchedulerWithLinearWarmup",
    "ConstantMarginMSE",
    "ContrastiveLocalLoss",
    "CrossEncoderConfig",
    "CrossEncoderModel",
    "CrossEncoderModule",
    "CrossEncoderOutput",
    "CrossEncoderTokenizer",
    "DocDataset",
    "DocSample",
    "FaissFlatIndexConfig",
    "FaissFlatIndexer",
    "FaissIVFIndexConfig",
    "FaissIVFIndexer",
    "FaissIVFPQIndexConfig",
    "FaissIVFPQIndexer",
    "FaissPQIndexConfig",
    "FaissPQIndexer",
    "FaissSearchConfig",
    "FaissSearcher",
    "FLOPSRegularization",
    "GenericConstantSchedulerWithLinearWarmup",
    "GenericConstantSchedulerWithQuadraticWarmup",
    "GenericLinearSchedulerWithLinearWarmup",
    "InBatchCrossEntropy",
    "IndexBatch",
    "IndexCallback",
    "IndexConfig",
    "Indexer",
    "InfoNCE",
    "KLDivergence",
    "L1Regularization",
    "L2Regularization",
    "LightningIRClassFactory",
    "LightningIRConfig",
    "LightningIRConfigClassFactory",
    "LightningIRDataModule",
    "LightningIRModel",
    "LightningIRModelClassFactory",
    "LightningIRModule",
    "LightningIROutput",
    "LightningIRTokenizer",
    "LightningIRTokenizerClassFactory",
    "LightningIRTrainer",
    "LightningIRWandbLogger",
    "LinearLRSchedulerWithLinearWarmup",
    "MultiVectorBiEncoderConfig",
    "MultiVectorBiEncoderModel",
    "PearsonCorrelation",
    "QueryDataset",
    "QuerySample",
    "RankBatch",
    "RankCallback",
    "RankNet",
    "RankSample",
    "RegisterLocalDatasetCallback",
    "ReRankCallback",
    "RunDataset",
    "ScoreBasedInBatchCrossEntropy",
    "ScoreBasedInBatchLossFunction",
    "SearchBatch",
    "SearchCallback",
    "SearchConfig",
    "Searcher",
    "SeismicIndexConfig",
    "SeismicIndexer",
    "SeismicSearchConfig",
    "SeismicSearcher",
    "SingleVectorBiEncoderConfig",
    "SingleVectorBiEncoderModel",
    "SupervisedMarginMSE",
    "TorchDenseIndexConfig",
    "TorchDenseIndexer",
    "TorchDenseSearchConfig",
    "TorchDenseSearcher",
    "TorchSparseIndexConfig",
    "TorchSparseIndexer",
    "TorchSparseSearchConfig",
    "TorchSparseSearcher",
    "TrainBatch",
    "TupleDataset",
    "WarmupLRScheduler",
]
