from __future__ import annotations

from typing import Iterable

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


PAD_IDX = 0
CHECKUP_IDX = 1
NO_EVENT_IDX = 2


def _make_ignore_mask(
    vocab_size: int,
    ignored_idx: Iterable[int],
    device: torch.device,
) -> torch.Tensor:
    ignore_mask = torch.zeros(vocab_size, dtype=torch.bool, device=device)
    for idx in ignored_idx:
        idx = int(idx)
        if 0 <= idx < vocab_size:
            ignore_mask[idx] = True
    return ignore_mask


def _valid_vocab_mask(
    vocab_size: int,
    ignored_idx: Iterable[int],
    device: torch.device,
) -> torch.Tensor:
    return ~_make_ignore_mask(vocab_size, ignored_idx, device)


def _zero_loss_like(logits: torch.Tensor) -> torch.Tensor:
    return logits.sum() * 0.0


class Delphi2MLoss(nn.Module):
    """Next-token plus exponential time-to-next-token supervision."""

    def __init__(
        self,
        t_min: float = 1.0 / 365.25,
        ignored_tokens: Iterable[int] | None = None,
        exclude_ignored_from_intensity: bool = True,
        max_exp_input: float = 60.0,
        ce_weight: float = 1.0,
        time_weight: float = 1.0,
    ):
        super().__init__()
        self.t_min = float(t_min)
        self.ignored_tokens = (
            [PAD_IDX, CHECKUP_IDX]
            if ignored_tokens is None
            else [int(x) for x in ignored_tokens]
        )
        self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity)
        self.max_exp_input = float(max_exp_input)
        self.ce_weight = float(ce_weight)
        self.time_weight = float(time_weight)

    def forward(
        self,
        logits: torch.Tensor,
        target_events: torch.Tensor,
        target_times: torch.Tensor,
        current_times: torch.Tensor,
        padding_mask: torch.Tensor,
        return_components: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
        if logits.dim() != 3:
            raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}")
        bsz, seq_len, vocab_size = logits.shape

        expected = (bsz, seq_len)
        if target_events.shape != expected:
            raise ValueError(f"target_events must be {expected}, got {tuple(target_events.shape)}")
        if target_times.shape != expected:
            raise ValueError(f"target_times must be {expected}, got {tuple(target_times.shape)}")
        if current_times.shape != expected:
            raise ValueError(f"current_times must be {expected}, got {tuple(current_times.shape)}")
        if padding_mask.shape != expected:
            raise ValueError(f"padding_mask must be {expected}, got {tuple(padding_mask.shape)}")

        valid_mask = padding_mask.bool()
        for idx in self.ignored_tokens:
            valid_mask = valid_mask & (target_events != int(idx))
        valid_mask = valid_mask & (target_events > PAD_IDX)

        if not valid_mask.any():
            total_loss = _zero_loss_like(logits)
            if return_components:
                return total_loss, {
                    "ce": total_loss.detach(),
                    "time": total_loss.detach(),
                    "total": total_loss.detach(),
                }
            return total_loss

        logits_valid = logits[valid_mask]
        target_events_valid = target_events[valid_mask]
        target_times_valid = target_times[valid_mask]
        current_times_valid = current_times[valid_mask]

        logits_safe = torch.nan_to_num(
            logits_valid,
            nan=0.0,
            posinf=self.max_exp_input,
            neginf=-self.max_exp_input,
        )

        loss_ce = F.cross_entropy(
            logits_safe,
            target_events_valid,
            reduction="mean",
        )

        logits_for_lse = logits_safe
        if self.exclude_ignored_from_intensity:
            ignore_mask = _make_ignore_mask(vocab_size, self.ignored_tokens, logits.device)
            logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))

        log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
        log_lambda_total = -torch.log(torch.exp(-log_lambda_total) + self.t_min)

        dt = torch.clamp(target_times_valid - current_times_valid, min=self.t_min)
        log_dt_inv = -torch.log(dt + self.t_min)
        loss_dt = -(
            log_lambda_total
            - torch.exp(
                torch.clamp(log_lambda_total - log_dt_inv, max=self.max_exp_input)
            )
        )
        loss_dt = loss_dt.mean()

        total_loss = self.ce_weight * loss_ce + self.time_weight * loss_dt
        if return_components:
            return total_loss, {
                "ce": loss_ce.detach(),
                "time": loss_dt.detach(),
                "total": total_loss.detach(),
            }
        return total_loss


class UniqueTimeSetExponentialLoss(nn.Module):
    """Next distinct timestamp event-set supervision with sum reduction."""

    def __init__(
        self,
        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
        t_min: float = 1.0 / 365.25,
        max_exp_input: float = 60.0,
        exclude_ignored_from_intensity: bool = True,
    ):
        super().__init__()
        self.ignored_idx = [int(x) for x in ignored_idx]
        self.t_min = float(t_min)
        self.max_exp_input = float(max_exp_input)
        self.exclude_ignored_from_intensity = bool(exclude_ignored_from_intensity)

    def forward(
        self,
        logits: torch.Tensor,
        target_multi_hot: torch.Tensor,
        target_dt_unique: torch.Tensor,
        readout_mask: torch.Tensor,
        return_components: bool = False,
    ) -> torch.Tensor | tuple[torch.Tensor, dict[str, torch.Tensor]]:
        if logits.dim() != 3:
            raise ValueError(f"logits must be (B, L, K), got {tuple(logits.shape)}")
        bsz, seq_len, vocab_size = logits.shape

        if target_multi_hot.shape != (bsz, seq_len, vocab_size):
            raise ValueError(
                "target_multi_hot must match logits shape, "
                f"got {tuple(target_multi_hot.shape)} vs {tuple(logits.shape)}"
            )
        if target_dt_unique.shape != (bsz, seq_len):
            raise ValueError(
                f"target_dt_unique must be {(bsz, seq_len)}, got {tuple(target_dt_unique.shape)}"
            )
        if readout_mask.shape != (bsz, seq_len):
            raise ValueError(f"readout_mask must be {(bsz, seq_len)}, got {tuple(readout_mask.shape)}")

        ignore_mask = _make_ignore_mask(vocab_size, self.ignored_idx, logits.device)

        num_targets = target_multi_hot[:, :, ~ignore_mask].sum(dim=-1)
        valid_mask = readout_mask.bool() & (num_targets > 0)

        if not valid_mask.any():
            total_loss = _zero_loss_like(logits)
            if return_components:
                return total_loss, {
                    "observed": total_loss.detach(),
                    "penalty": total_loss.detach(),
                    "total": total_loss.detach(),
                }
            return total_loss

        logits_safe = torch.nan_to_num(
            logits[valid_mask],
            nan=0.0,
            posinf=self.max_exp_input,
            neginf=-self.max_exp_input,
        )
        target_valid = target_multi_hot[valid_mask].to(logits_safe.dtype)
        target_valid[:, ignore_mask] = 0.0

        observed_term = (logits_safe * target_valid).sum(dim=-1)
        penalty_scale = target_valid.sum(dim=-1)

        logits_for_lse = logits_safe
        if self.exclude_ignored_from_intensity:
            logits_for_lse = logits_safe.masked_fill(ignore_mask.unsqueeze(0), float("-inf"))

        dt_clamped = torch.clamp(target_dt_unique[valid_mask], min=self.t_min)
        log_lambda_total = torch.logsumexp(logits_for_lse, dim=-1)
        log_penalty = log_lambda_total + dt_clamped.log()
        penalty = torch.exp(torch.clamp(log_penalty, max=self.max_exp_input))

        observed_loss = -observed_term
        penalty_loss = penalty_scale * penalty
        total_loss = (observed_loss + penalty_loss).mean()

        if return_components:
            return total_loss, {
                "observed": observed_loss.mean().detach(),
                "penalty": penalty_loss.mean().detach(),
                "total": total_loss.detach(),
            }
        return total_loss


class ExponentialLoss(nn.Module):
    """Query-conditioned all-future-event exponential point-process loss."""

    def __init__(
        self,
        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
        eps: float = 1e-8,
    ):
        super().__init__()
        self.ignored_idx = tuple(int(i) for i in ignored_idx)
        self.eps = eps

    def forward(
        self,
        logits: torch.Tensor,
        targets: torch.Tensor,
        exposure: torch.Tensor,
    ) -> torch.Tensor:
        _, vocab_size = logits.shape
        rate = F.softplus(logits) + self.eps
        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)

        penalty = exposure.to(rate.dtype) * rate[:, valid_vocab].sum(dim=-1)
        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
        for idx in self.ignored_idx:
            target_valid &= targets != idx

        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
        observed = rate.log().gather(1, safe_targets)
        observed = (observed * target_valid.to(rate.dtype)).sum(dim=-1)
        return (-observed + penalty).mean()


class WeibullLoss(nn.Module):
    """Query-conditioned all-future-event Weibull point-process loss."""

    def __init__(
        self,
        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
        eps: float = 1e-8,
    ):
        super().__init__()
        self.ignored_idx = tuple(int(i) for i in ignored_idx)
        self.eps = eps

    def forward(
        self,
        logits: torch.Tensor,
        weibull_rho: torch.Tensor,
        targets: torch.Tensor,
        dt: torch.Tensor,
        exposure: torch.Tensor,
    ) -> torch.Tensor:
        _, vocab_size = logits.shape
        if weibull_rho is None:
            raise ValueError("weibull_rho is required for WeibullLoss")
        if weibull_rho.shape != logits.shape:
            raise ValueError(
                "weibull_rho must have the same shape as logits. "
                f"Got logits={tuple(logits.shape)}, weibull_rho={tuple(weibull_rho.shape)}"
            )

        dtype = logits.dtype
        rate = F.softplus(logits) + self.eps
        rho = weibull_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)
        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)

        t_exp = exposure.to(dtype).clamp_min(self.eps).unsqueeze(1)
        penalty = (rate * torch.pow(t_exp, rho))[:, valid_vocab].sum(dim=-1)

        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
        for idx in self.ignored_idx:
            target_valid &= targets != idx

        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
        target_rate = rate.gather(1, safe_targets)
        target_rho = rho.gather(1, safe_targets)
        target_dt = dt.to(dtype).clamp_min(self.eps)
        log_intensity = (
            target_rate.log()
            + target_rho.log()
            + (target_rho - 1.0) * target_dt.log()
        )
        observed = (log_intensity * target_valid.to(dtype)).sum(dim=-1)
        return (-observed + penalty).mean()


class MixedLoss(nn.Module):
    """Exponential diseases plus one Weibull death endpoint."""

    def __init__(
        self,
        death_idx: int,
        ignored_idx: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
        eps: float = 1e-8,
    ):
        super().__init__()
        self.death_idx = int(death_idx)
        self.ignored_idx = tuple(int(i) for i in ignored_idx)
        self.eps = eps

    def forward(
        self,
        logits: torch.Tensor,
        death_rho: torch.Tensor,
        targets: torch.Tensor,
        dt: torch.Tensor,
        exposure: torch.Tensor,
    ) -> torch.Tensor:
        _, vocab_size = logits.shape
        dtype = logits.dtype
        rate = F.softplus(logits) + self.eps

        if death_rho.dim() == 2:
            death_rho = death_rho.squeeze(-1)
        death_rho = death_rho.to(device=logits.device, dtype=dtype).clamp_min(self.eps)

        valid_vocab = _valid_vocab_mask(vocab_size, self.ignored_idx, logits.device)
        valid_disease_vocab = valid_vocab.clone()
        valid_disease_vocab[self.death_idx] = False

        t_exp = exposure.to(dtype).clamp_min(self.eps)
        disease_penalty = t_exp * rate[:, valid_disease_vocab].sum(dim=-1)
        death_rate = rate[:, self.death_idx]
        death_penalty = death_rate * torch.pow(t_exp, death_rho)
        penalty = disease_penalty + death_penalty

        target_valid = torch.ones_like(targets, dtype=torch.bool, device=logits.device)
        for idx in self.ignored_idx:
            target_valid &= targets != idx

        disease_event_mask = target_valid & (targets != self.death_idx)
        safe_targets = targets.clamp(min=0, max=vocab_size - 1)
        disease_log_rate = rate.log().gather(1, safe_targets)
        observed_disease = (disease_log_rate * disease_event_mask.to(dtype)).sum(dim=-1)

        death_event_mask = target_valid & (targets == self.death_idx)
        death_observed = death_event_mask.any(dim=1)
        death_dt = (dt.to(dtype).clamp_min(self.eps) * death_event_mask.to(dtype)).sum(dim=1)
        death_log_intensity = (
            death_rate.log()
            + death_rho.log()
            + (death_rho - 1.0) * death_dt.clamp_min(self.eps).log()
        )
        observed_death = death_log_intensity * death_observed.to(dtype)

        return (-observed_disease - observed_death + penalty).mean()


def build_loss(name: str, **kwargs) -> nn.Module:
    name = name.lower()
    if name in {"delphi2m", "d2m", "next_token"}:
        return Delphi2MLoss(**kwargs)
    if name in {"uts", "unique_time_set", "unique_time_exponential"}:
        return UniqueTimeSetExponentialLoss(**kwargs)
    if name in {"exponential", "query_exponential"}:
        return ExponentialLoss(**kwargs)
    if name in {"weibull", "query_weibull"}:
        return WeibullLoss(**kwargs)
    if name in {"mixed", "query_mixed"}:
        return MixedLoss(**kwargs)
    raise ValueError(
        f"Unknown loss {name!r}. Available: delphi2m, uts, exponential, weibull, mixed."
    )
