# targets.py
from __future__ import annotations

from dataclasses import dataclass
import numpy as np


PAD_IDX = 0
CHECKUP_IDX = 1
NO_EVENT_IDX = 2
DAYS_PER_YEAR = 365.25


@dataclass(frozen=True)
class NextTokenTargets:
    """
    Delphi2M-style next-token supervision targets.

    Shapes:
        input_events:      (L,)
        input_times_years: (L,)
        target_events:     (L,)
        target_times_years:(L,)

    where L = N - 1.
    """
    input_events: np.ndarray
    input_times_years: np.ndarray
    target_events: np.ndarray
    target_times_years: np.ndarray


def _as_numpy_1d(
    x: np.ndarray,
    name: str,
    dtype: np.dtype | type | None = None,
) -> np.ndarray:
    arr = np.asarray(x)
    if arr.ndim != 1:
        raise ValueError(f"{name} must be 1D, got shape {arr.shape}")
    if dtype is not None:
        arr = arr.astype(dtype)
    return arr


def validate_event_sequence(
    labels: np.ndarray,
    times_days: np.ndarray,
    *,
    require_sorted: bool = True,
) -> None:
    """
    Validate one patient's event sequence.

    labels:
        1D integer label ids.

    times_days:
        1D event times in days.

    require_sorted:
        If True, raises when times_days is not non-decreasing.
    """
    labels = _as_numpy_1d(labels, "labels")
    times_days = _as_numpy_1d(times_days, "times_days")

    if len(labels) != len(times_days):
        raise ValueError(
            f"labels and times_days must have the same length, "
            f"got {len(labels)} and {len(times_days)}"
        )

    if len(labels) == 0:
        raise ValueError("Empty event sequence is not valid.")

    if np.any(labels < 0):
        raise ValueError("labels contains negative ids.")

    if not np.all(np.isfinite(times_days)):
        raise ValueError("times_days contains non-finite values.")

    if require_sorted and len(times_days) > 1:
        if np.any(np.diff(times_days) < 0):
            raise ValueError("times_days must be non-decreasing.")


def build_next_token_targets(
    labels: np.ndarray,
    times_days: np.ndarray,
    *,
    require_sorted: bool = True,
) -> NextTokenTargets:
    """
    Build Delphi2M-style autoregressive next-token targets.

    Given full sequence:

        labels:     [x0, x1, x2, ..., xN-1]
        times_days: [t0, t1, t2, ..., tN-1]

    returns:

        input_events:       [x0, x1, ..., xN-2]
        input_times_years:  [t0, t1, ..., tN-2] / 365.25
        target_events:      [x1, x2, ..., xN-1]
        target_times_years: [t1, t2, ..., tN-1] / 365.25

    This function does not ignore PAD/CHECKUP/NO_EVENT. Ignoring belongs to
    the loss function because different objectives may use different ignore ids.
    """
    labels = _as_numpy_1d(labels, "labels", np.int64)
    times_days = _as_numpy_1d(times_days, "times_days", np.float32)
    validate_event_sequence(labels, times_days, require_sorted=require_sorted)

    if len(labels) < 2:
        raise ValueError(
            "Need at least two events to build next-token targets."
        )

    input_events = labels[:-1].astype(np.int64)
    input_times_years = (times_days[:-1] / DAYS_PER_YEAR).astype(np.float32)

    target_events = labels[1:].astype(np.int64)
    target_times_years = (times_days[1:] / DAYS_PER_YEAR).astype(np.float32)

    return NextTokenTargets(
        input_events=input_events,
        input_times_years=input_times_years,
        target_events=target_events,
        target_times_years=target_times_years,
    )