"""Negative binomial distribution utilities for DeepPTR."""

from __future__ import annotations

import torch
from torch import Tensor


def log_nb_positive(
    x: Tensor,
    mu: Tensor,
    theta: Tensor,
    eps: float = 1e-8,
) -> Tensor:
    """Log-likelihood of the negative binomial distribution (mean-dispersion).

    Parameterized by mean ``mu`` and inverse dispersion ``theta`` so that
    Var = mu + mu^2 / theta.  Larger theta → less overdispersion.

    Parameters
    ----------
    x : Tensor
        Observed counts (non-negative integers), shape ``(N, G)``.
    mu : Tensor
        Predicted mean, shape ``(N, G)``.
    theta : Tensor
        Inverse dispersion, shape ``(G,)`` or ``(N, G)``.
    eps : float
        Small constant for numerical stability.

    Returns
    -------
    Tensor
        Log-probability, same shape as *x*.
    """
    mu = mu.clamp(min=eps)
    theta = theta.clamp(min=eps)

    log_theta_mu = torch.log(theta + mu + eps)

    return (
        torch.lgamma(x + theta)
        - torch.lgamma(theta)
        - torch.lgamma(x + 1.0)
        + theta * (torch.log(theta + eps) - log_theta_mu)
        + x * (torch.log(mu + eps) - log_theta_mu)
    )