"""
torch-supercluster
==================
SuperCluster: Learning Prediction Prototypes via Target-Guided Cross-Attention Clustering.

Quick start
-----------
Binary classification::

    from supercluster import SuperCluster, center_diversity_loss
    import torch, torch.nn as nn

    model = SuperCluster(input_dim=10, embedding_dim=64, num_clusters=8)
    logits, cluster_weights = model(x)
    loss = nn.BCEWithLogitsLoss()(logits, y) + 0.1 * center_diversity_loss(model.centers)

7-class classification::

    model = SuperCluster(input_dim=54, embedding_dim=128,
                         num_clusters=8, num_classes=7)
    logits, cluster_weights = model(x)
    loss = nn.CrossEntropyLoss()(logits, y_long)

Regression::

    model = SuperCluster(input_dim=10, embedding_dim=64, num_clusters=5)
    preds, cluster_weights = model(x)
    loss = nn.MSELoss()(preds, y)
"""

from supercluster.model import (
    SuperCluster,
    CrossTransformerLayer,
    CrossAttention,
    MLPBlock,
    # backward-compatibility alias
    SupervisedClusteringModelV2,
)
from supercluster.losses import center_diversity_loss, effective_prototype_count

__version__ = "0.1.0"
__author__ = "Aaron John Danielson"
__license__ = "MIT"

__all__ = [
    "SuperCluster",
    "center_diversity_loss",
    "effective_prototype_count",
    # building blocks (for advanced users)
    "CrossTransformerLayer",
    "CrossAttention",
    "MLPBlock",
    # backward compat
    "SupervisedClusteringModelV2",
]
