# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/losses/01_class_balanced_cross_entropy_loss.ipynb.

# %% auto 0
__all__ = ['ClassBalancedCrossEntropyLossConfig', 'ClassBalancedCrossEntropyLoss']

# %% ../../nbs/losses/01_class_balanced_cross_entropy_loss.ipynb 2
from collections import Counter

import torch
from torch import nn
from torch.nn import functional as F

from ..docstrings import populate_docstring
from ..utils.custom_base_model import CustomBaseModel, Field

# %% ../../nbs/losses/01_class_balanced_cross_entropy_loss.ipynb 4
class ClassBalancedCrossEntropyLossConfig(CustomBaseModel):
    num_classes: int = Field(..., description="Number of classes to weight cross entropy loss.")
    ema_decay: float = Field(
        0.99, description="Exponential moving average decay. By default 0.99 is used which has a half life of ~69 steps"
    )

# %% ../../nbs/losses/01_class_balanced_cross_entropy_loss.ipynb 6
class ClassBalancedCrossEntropyLoss(nn.Module):
    @populate_docstring
    def __init__(self, config: ClassBalancedCrossEntropyLossConfig = {}, **kwargs):
        """Class-balanced cross-entropy loss with running prevalence estimation.

        This loss reweights the standard multi-class cross-entropy by the inverse of
        the observed class prevalences in the training data. Class prevalences are
        estimated online via an exponential moving average (EMA) using class counts
        from the incoming targets.

        Notes:
            - Targets must be discrete integer class indices in [0, num_classes-1].
              Probabilistic/soft labels are not supported.
            - For classes that haven't been observed yet, their prevalence is treated
              as NaN and replaced in the weight vector by the mean of observed weights
              (then clamped within 3 standard deviations to avoid extreme values).

        Args:
            config: {CONFIG_INSTANCE_DOC}
            **kwargs: {CONFIG_KWARGS_DOC}
        """
        super().__init__()

        self.config = ClassBalancedCrossEntropyLossConfig.model_validate(config | kwargs)

        # Save class prevalence percentages to weight the cross entropy loss
        self.class_prevalences = [None] * self.config.num_classes
        # Initialized with None as we don't know the initial class prevalence estimates.

    @torch.no_grad()
    def update_class_prevalences(self, target: torch.Tensor, ignore_index: int = -100):
        """Update the running class-prevalence estimates from a target tensor.

        The method counts class occurrences in the provided target, converts counts
        to per-batch prevalences, then updates the internal EMA-tracked prevalence
        vector for each class.

        Args:
            target: A tensor of integer class indices with any shape, typically
                (N,) or (N, ...) for segmentation. Values outside
                [0, num_classes-1] are ignored.
            ignore_index: A class index to ignore during updates.
        """
        # Count the number of times each class is encountered
        class_counts = Counter(target.flatten().int().tolist())

        # Remove ignore_index if present
        class_counts.pop(ignore_index, None)

        # Remove those counts that exceed the number of classes that are being tracked
        class_counts = {k: v for k, v in class_counts.items() if k < self.config.num_classes}

        # Calculate prevalence of each class
        total_count = sum(class_counts.values())
        new_prevalences = (
            {class_id: count / total_count for class_id, count in class_counts.items()} if total_count > 0 else {}
        )

        # Update current prevalences using EMA to allow for distribution shift in data
        decay = self.config.ema_decay
        for i in range(self.config.num_classes):
            # when encountered for the first time, assume each class was equally prevalent at the start
            if self.class_prevalences[i] is None and i in new_prevalences:
                self.class_prevalences[i] = 1 / self.config.num_classes

            # update to the new value using EMA
            if self.class_prevalences[i] is not None:
                self.class_prevalences[i] = self.class_prevalences[i] * decay + new_prevalences.get(i, 0.0) * (
                    1 - decay
                )

    @torch.no_grad()
    def get_class_prevalences(self, device=torch.device("cpu")) -> torch.Tensor:
        """Return the current vector of class prevalences as a tensor.

        For classes that haven't been observed yet, the corresponding entry will be
        NaN. This method does not perform any imputation or normalization beyond
        returning the current EMA state.

        Args:
            device: The device on which to place the returned tensor.

        Returns:
            Tensor of shape (num_classes,) with dtype float32 containing per-class
            prevalence estimates in [0, 1] or NaN for unseen classes.
        """

        # Get class prevalences and replace None with nan
        class_prevalences = [
            prevalence if prevalence is not None else torch.nan for prevalence in self.class_prevalences
        ]

        # Calculate weights as inverse of class prevalences
        class_prevalences = torch.tensor(class_prevalences, dtype=torch.float32, device=device)

        return class_prevalences

    @torch.no_grad()
    def get_class_weights(self, device=torch.device("cpu")) -> torch.Tensor:
        """Compute per-class weights as the inverse of prevalences with safeguards.

        Steps:
            1) Convert current EMA prevalences to a tensor with NaNs for unseen classes.
            2) Take the inverse to obtain raw weights (higher weight for rarer classes).
            3) Replace NaNs by the mean of observed weights to avoid biasing toward
               unseen classes, then clamp to mean ± 3·std to prevent extreme values.
            4) Renormalize weights to sum to num_classes (so the average weight is 1).

        Args:
            device: The device on which to place the returned tensor.

        Returns:
            Tensor of shape (num_classes,) containing normalized class weights.
        """
        # Get class prevalences
        class_prevalences = self.get_class_prevalences(device)

        # If there were no class prevalences, then all weights will be nan; handle this
        if class_prevalences.isnan().all():
            class_prevalences = class_prevalences.nan_to_num(1.0)

        # Calculate weights as inverse of class prevalences
        weights = 1 / class_prevalences

        # Substitute nan values with mean and clamp weights to a limit
        # Assumption: all classes are visited at least once in the dataset
        mu, std = weights.nanmean(), weights[~weights.isnan()].std()
        weights = weights.nan_to_num(nan=mu, posinf=mu, neginf=mu)
        weights = weights.clamp(mu - 3 * std, mu + 3 * std)

        # Add failsafe just in case nans are still present
        weights = weights.nan_to_num(1.0)

        # Normalize weights to sum to self.config.num_classes
        weights = self.config.num_classes * weights / weights.sum()

        return weights

    def forward(
        self,
        input: torch.Tensor,
        target: torch.Tensor,
        update_class_prevalences: bool = True,
        return_class_weights: bool = False,
        *args,
        **kwargs
    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
        """Compute class-balanced cross entropy.

        This updates the internal class-prevalence EMA using the provided targets,
        then computes cross-entropy with a weight vector derived from the current
        prevalences.

        Args:
            input: Logits of shape (N, C, ...) where C == num_classes. Any extra spatial
                dimensions (e.g., H, W, D) are supported as long as ``target`` is broadcastable
                to the same non-channel shape expected by torch.nn.functional.cross_entropy.
            target: Integer class indices with shape matching input without the channel
                dimension, e.g., (N, ...) with values in [0, C-1].
            update_class_prevalences: If True, update the internal class prevalence estimates
                using the provided targets.
            return_class_weights: If True, also return the per-class weight tensor used
                for this call.
            *args, **kwargs: Additional keyword args forwarded to F.cross_entropy
                (e.g., reduction='mean').

        Returns:
            If return_class_weights is False: a scalar tensor loss.
            If True: a tuple (loss, class_weights).
        """
        # Update class prevalences
        if update_class_prevalences:
            self.update_class_prevalences(target, kwargs.get("ignore_index", -100))

        # Get class weights
        class_weights = self.get_class_weights(input.device)

        # Calculate loss
        loss = F.cross_entropy(input, target, weight=class_weights, *args, **kwargs)

        if return_class_weights:
            return loss, class_weights
        return loss
