"""Compute token-summed codec loss without weighting short microbatches equally."""
from __future__ import annotations
from dataclasses import dataclass
from torch import Tensor
from torch.nn import functional
from .layout import CODEBOOKS, IGNORE, VOCABULARY

@dataclass(frozen=True)
class CodecLoss:
    stream_nats: Tensor
    stream_counts: Tensor

    @property
    def numerator(self) -> Tensor:
        return self.stream_nats.sum()

    @property
    def count(self) -> Tensor:
        return self.stream_counts.sum()

def codec_loss(logits: Tensor, labels: Tensor) -> CodecLoss:
    if logits.shape[:-1] != labels.shape or logits.shape[-2:] != (CODEBOOKS, VOCABULARY):
        message = f'codec loss shape mismatch logits={tuple(logits.shape)}, labels={tuple(labels.shape)}'
        raise ValueError(message)
    valid = labels.ne(IGNORE)
    if not bool(valid.any().item()):
        message = 'codec loss has zero supervised targets'
        raise ValueError(message)
    losses = functional.cross_entropy(logits.float().reshape(-1, VOCABULARY), labels.reshape(-1), ignore_index=IGNORE, reduction='none').reshape(labels.shape)
    axes = tuple(range(labels.ndim - 1))
    return CodecLoss(losses.sum(dim=axes), valid.sum(dim=axes))