from __future__ import annotations

from collections.abc import Sequence

import torch
import torch.nn.functional as F


def death_token(vocab_size: int) -> int:
    if int(vocab_size) <= 0:
        raise ValueError(f"vocab_size must be positive, got {vocab_size}")
    return int(vocab_size) - 1


def probabilities_from_logits(
    logits: torch.Tensor,
    tau_years: float | torch.Tensor,
    *,
    dist_mode: str = "exponential",
    rho: torch.Tensor | None = None,
    death_rho: torch.Tensor | None = None,
    eps: float = 1e-8,
) -> torch.Tensor:
    """
    Convert all-future logits to tau-year event probabilities.

    Death is always treated as token vocab_size - 1. For dist_mode="mixed",
    non-death tokens use exponential hazards and death uses death_rho.
    """
    if logits.ndim != 2:
        raise ValueError(f"logits must have shape (N, V), got {tuple(logits.shape)}")
    if float(torch.as_tensor(tau_years).detach().min().cpu()) < 0:
        raise ValueError("tau_years must be non-negative")

    mode = str(dist_mode).lower()
    if mode not in {"exponential", "weibull", "mixed"}:
        raise ValueError("dist_mode must be one of: exponential, weibull, mixed")

    rate = F.softplus(logits) + float(eps)
    tau = torch.as_tensor(tau_years, dtype=rate.dtype, device=rate.device)
    if tau.ndim == 0:
        tau = tau.expand(logits.shape[0])
    if tau.ndim != 1 or tau.shape[0] != logits.shape[0]:
        raise ValueError(
            "tau_years must be a scalar or a 1D tensor with length N, got "
            f"{tuple(tau.shape)} for N={logits.shape[0]}"
        )

    if mode == "exponential":
        exposure = tau[:, None].expand_as(rate)
    elif mode == "weibull":
        if rho is None or rho.shape != logits.shape:
            raise ValueError("rho must have the same shape as logits for dist_mode='weibull'")
        exposure = torch.pow(tau[:, None].clamp_min(float(eps)), rho.to(rate.dtype))
    else:
        exposure = tau[:, None].expand_as(rate).clone()
        if death_rho is None:
            raise ValueError("death_rho is required for dist_mode='mixed'")
        death_idx = death_token(logits.shape[1])
        death_shape = tuple(death_rho.shape)
        death_rho = death_rho.to(device=rate.device, dtype=rate.dtype)
        if death_rho.ndim == 2 and death_rho.shape[1] == 1:
            death_rho = death_rho.squeeze(1)
        if death_rho.ndim != 1 or death_rho.shape[0] != logits.shape[0]:
            raise ValueError(
                "death_rho must have shape (N,) or (N, 1), got "
                f"{death_shape} for N={logits.shape[0]}"
            )
        exposure[:, death_idx] = torch.pow(tau.clamp_min(float(eps)), death_rho)

    return -torch.expm1(-rate * exposure)


def death_risk_from_probabilities(probabilities: torch.Tensor) -> torch.Tensor:
    """Return p_death(t, tau), with death fixed to token vocab_size - 1."""
    if probabilities.ndim != 2:
        raise ValueError(
            f"probabilities must have shape (N, V), got {tuple(probabilities.shape)}"
        )
    return probabilities[:, death_token(probabilities.shape[1])]


def new_disease_risk_from_probabilities(
    probabilities: torch.Tensor,
    occurred: torch.Tensor,
    disease_ids: Sequence[int],
) -> torch.Tensor:
    """
    Compute P(at least one selected disease newly occurs within tau years).

    Already occurred diseases are masked out. Death is not included here and
    should be reported separately with death_risk_from_probabilities.
    """
    if probabilities.ndim != 2 or occurred.shape != probabilities.shape:
        raise ValueError(
            "probabilities and occurred must both have shape (N, V), got "
            f"{tuple(probabilities.shape)} and {tuple(occurred.shape)}"
        )
    if not disease_ids:
        return probabilities.new_zeros(probabilities.shape[0])

    death_idx = death_token(probabilities.shape[1])
    ids = [
        idx
        for idx in dict.fromkeys(int(x) for x in disease_ids)
        if 0 <= idx < probabilities.shape[1] and idx != death_idx
    ]
    if not ids:
        return probabilities.new_zeros(probabilities.shape[0])

    idx_tensor = torch.as_tensor(ids, dtype=torch.long, device=probabilities.device)
    p = probabilities[:, idx_tensor].clamp(0.0, 1.0 - 1e-7)
    new_mask = ~occurred[:, idx_tensor].to(dtype=torch.bool)
    log_no_new = torch.log1p(-p) * new_mask.to(dtype=p.dtype)
    return -torch.expm1(log_no_new.sum(dim=1))