from __future__ import annotations

from collections.abc import Sequence
from typing import Any, overload

import numpy as np

try:
    import torch
    import torch.nn.functional as F
except ModuleNotFoundError:  # pragma: no cover - optional for numpy-only use
    torch = None
    F = None


ArrayLike = Any


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 _infer_vocab_size(x: ArrayLike, vocab_size: int | None) -> int:
    if x.ndim != 2:
        raise ValueError(f"Expected a 2D array/tensor with shape (N, V), got {tuple(x.shape)}")
    inferred = int(x.shape[1])
    if vocab_size is None:
        return inferred
    if int(vocab_size) != inferred:
        raise ValueError(f"vocab_size={vocab_size} does not match input width {inferred}")
    return int(vocab_size)


def _normalize_disease_ids(
    disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None,
    *,
    vocab_size: int,
    excluded_token_ids: Sequence[int],
) -> list[int]:
    death_idx = _death_token(vocab_size)
    excluded = {
        int(idx)
        for idx in excluded_token_ids
        if 0 <= int(idx) < vocab_size
    }
    excluded.add(death_idx)
    if disease_ids is None:
        return [idx for idx in range(vocab_size) if idx not in excluded]

    if torch is not None and isinstance(disease_ids, torch.Tensor):
        raw = disease_ids.detach().cpu().reshape(-1).tolist()
    else:
        raw = np.asarray(disease_ids).reshape(-1).tolist()

    out: list[int] = []
    seen: set[int] = set()
    for value in raw:
        idx = int(value)
        if idx < 0 or idx >= vocab_size:
            raise ValueError(f"disease id {idx} is outside [0, {vocab_size})")
        if idx in excluded:
            continue
        if idx not in seen:
            seen.add(idx)
            out.append(idx)
    return out


@overload
def future_event_free_survival_from_probabilities(
    probabilities: torch.Tensor,
    occurred: torch.Tensor,
    disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
    *,
    vocab_size: int | None = None,
    excluded_token_ids: Sequence[int] = (0, 1, 2),
    eps: float = 1e-7,
) -> torch.Tensor:
    ...


@overload
def future_event_free_survival_from_probabilities(
    probabilities: np.ndarray,
    occurred: np.ndarray,
    disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
    *,
    vocab_size: int | None = None,
    excluded_token_ids: Sequence[int] = (0, 1, 2),
    eps: float = 1e-7,
) -> np.ndarray:
    ...


def future_event_free_survival_from_probabilities(
    probabilities: ArrayLike,
    occurred: ArrayLike,
    disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
    *,
    vocab_size: int | None = None,
    excluded_token_ids: Sequence[int] = (0, 1, 2),
    eps: float = 1e-7,
) -> ArrayLike:
    """
    Compute P(alive and no new selected disease in the next tau years).

    Parameters
    ----------
    probabilities:
        Matrix with shape (N, V). Entry (i, d) is p_d(t, tau), the model's
        future first-occurrence probability for token d over the chosen tau.
        The death probability is always read from token V - 1.

    occurred:
        Boolean matrix with shape (N, V). Entry (i, d) is True if disease d has
        already occurred at or before query time t. Already occurred diseases do
        not contribute to "new disease" risk.

    disease_ids:
        Optional subset of disease tokens. If None, all non-death tokens are
        included except excluded_token_ids. If provided, death and excluded
        tokens are ignored here and death is still handled separately as
        survival.

    vocab_size:
        Optional vocabulary size. If omitted, inferred from probabilities.

    excluded_token_ids:
        Technical tokens to exclude from "new disease" calculations. Defaults
        to (0, 1, 2), matching PAD, CHECKUP, and NO_EVENT.

    Returns
    -------
    Array/tensor with shape (N,):
        Approximate probability of being alive and having no newly occurring
        disease among the selected disease tokens over the same tau horizon.
    """
    vocab_size = _infer_vocab_size(probabilities, vocab_size)
    death_idx = _death_token(vocab_size)
    selected = _normalize_disease_ids(
        disease_ids,
        vocab_size=vocab_size,
        excluded_token_ids=excluded_token_ids,
    )

    if tuple(occurred.shape) != tuple(probabilities.shape):
        raise ValueError(
            "occurred must have the same shape as probabilities, got "
            f"{tuple(occurred.shape)} vs {tuple(probabilities.shape)}"
        )

    if torch is not None and isinstance(probabilities, torch.Tensor):
        probs = probabilities.clamp(min=0.0, max=1.0 - float(eps))
        occurred_bool = occurred.to(device=probs.device, dtype=torch.bool)
        log_survival = torch.log1p(-probs[:, death_idx])
        if selected:
            ids = torch.as_tensor(selected, dtype=torch.long, device=probs.device)
            new_mask = ~occurred_bool[:, ids]
            log_no_new = torch.log1p(-probs[:, ids]) * new_mask.to(probs.dtype)
            log_survival = log_survival + log_no_new.sum(dim=1)
        return torch.exp(log_survival)

    probs_np = np.clip(np.asarray(probabilities), 0.0, 1.0 - float(eps))
    occurred_bool_np = np.asarray(occurred, dtype=bool)
    log_survival_np = np.log1p(-probs_np[:, death_idx])
    if selected:
        selected_arr = np.asarray(selected, dtype=np.int64)
        new_mask_np = ~occurred_bool_np[:, selected_arr]
        log_no_new_np = np.log1p(-probs_np[:, selected_arr]) * new_mask_np
        log_survival_np = log_survival_np + log_no_new_np.sum(axis=1)
    return np.exp(log_survival_np)


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.

    The death token is always treated as vocab_size - 1. For dist_mode="mixed",
    non-death tokens use exponential hazards and the death token uses
    death_rho. For dist_mode="weibull", rho must have the same shape as logits.
    """
    if torch is None or F is None:
        raise ImportError("probabilities_from_logits requires PyTorch.")
    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 future_event_free_survival_from_logits(
    logits: torch.Tensor,
    occurred: torch.Tensor,
    tau_years: float | torch.Tensor,
    disease_ids: Sequence[int] | np.ndarray | torch.Tensor | None = None,
    *,
    dist_mode: str = "exponential",
    rho: torch.Tensor | None = None,
    death_rho: torch.Tensor | None = None,
    eps: float = 1e-8,
) -> torch.Tensor:
    """
    Convenience wrapper for computing future event-free survival from logits.

    Returns P(alive and no new selected disease in the next tau years), with
    death fixed to token vocab_size - 1.
    """
    probabilities = probabilities_from_logits(
        logits=logits,
        tau_years=tau_years,
        dist_mode=dist_mode,
        rho=rho,
        death_rho=death_rho,
        eps=eps,
    )
    return future_event_free_survival_from_probabilities(
        probabilities=probabilities,
        occurred=occurred,
        disease_ids=disease_ids,
        vocab_size=logits.shape[1],
        excluded_token_ids=excluded_token_ids,
    )