# targets.py
from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

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


@dataclass(frozen=True)
class UniqueTimeSetTargets:
    """
    Unique-time set supervision targets.

    Shapes:
        readout_mask:      (L,)
        target_dt_unique:  (L,)
        target_multi_hot:  (L, vocab_size)

    where L = N - 1.

    Only group-end positions can have readout_mask=True.
    target_dt_unique is measured in years.
    """
    readout_mask: np.ndarray
    target_dt_unique: np.ndarray
    target_multi_hot: np.ndarray


@dataclass(frozen=True)
class TargetPack:
    """
    Combined target package for one patient sequence.

    Contains both next-token targets and unique-time-set targets.
    The training pipeline decides which one to use.
    """
    next_token: NextTokenTargets
    unique_time_set: UniqueTimeSetTargets


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,
    )


def build_unique_time_set_targets(
    labels: np.ndarray,
    times_days: np.ndarray,
    *,
    vocab_size: int,
    ignored_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
    require_sorted: bool = True,
) -> UniqueTimeSetTargets:
    """
    Build next-unique-time set targets.

    This is the target construction used by your UTS / default mode.

    For each input position i:
        - only if i is the last token of its timestamp group;
        - find the next distinct timestamp group;
        - target is the set of valid event labels at that next timestamp.

    Example:

        t=49: X
        t=50: A, B, C
        t=51: D, E

    Supervises:

        X@49           -> {A, B, C}@50
        group_end@50   -> {D, E}@51

    It does NOT supervise:

        A@50 -> B@50
        B@50 -> C@50

    Parameters
    ----------
    labels:
        Full event sequence labels, shape (N,).

    times_days:
        Full event sequence times in days, shape (N,).

    vocab_size:
        Size of output vocabulary.

    ignored_target_ids:
        Label ids that should not enter target_multi_hot.
        Usually:
            no no-event:      {0, 1}
            with no-event:    {0, 1, 2}
        For UTS, I recommend ignoring <NO_EVENT> unless explicitly testing it
        as an event target.

    Returns
    -------
    UniqueTimeSetTargets
    """
    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 vocab_size <= 0:
        raise ValueError(f"vocab_size must be positive, got {vocab_size}")

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

    input_len = len(labels) - 1

    readout_mask = np.zeros(input_len, dtype=bool)
    target_dt_unique = np.zeros(input_len, dtype=np.float32)
    target_multi_hot = np.zeros((input_len, vocab_size), dtype=bool)

    ignored = {int(x) for x in ignored_target_ids}

    unique_times = np.unique(times_days)
    time_to_group_idx = {t: i for i, t in enumerate(unique_times)}
    group_indices = np.array([time_to_group_idx[t]
                             for t in times_days], dtype=np.int64)

    for i in range(input_len):
        current_group = group_indices[i]

        is_last_in_group = (
            i == input_len - 1
            or group_indices[i + 1] != current_group
        )
        if not is_last_in_group:
            continue

        next_group_idx = current_group + 1
        if next_group_idx >= len(unique_times):
            continue

        next_time = unique_times[next_group_idx]
        next_labels = labels[group_indices == next_group_idx]

        valid_next_labels: list[int] = []
        for lab in next_labels:
            lab_int = int(lab)
            if lab_int in ignored:
                continue
            if lab_int < 0 or lab_int >= vocab_size:
                continue
            valid_next_labels.append(lab_int)

        # If next timestamp contains only technical tokens, do not supervise UTS.
        if len(valid_next_labels) == 0:
            continue

        readout_mask[i] = True
        target_dt_unique[i] = float(next_time - times_days[i]) / DAYS_PER_YEAR
        target_multi_hot[i, valid_next_labels] = True

    return UniqueTimeSetTargets(
        readout_mask=readout_mask,
        target_dt_unique=target_dt_unique.astype(np.float32),
        target_multi_hot=target_multi_hot,
    )


def build_all_targets(
    labels: np.ndarray,
    times_days: np.ndarray,
    *,
    vocab_size: int,
    ignored_uts_target_ids: Iterable[int] = (PAD_IDX, CHECKUP_IDX),
    require_sorted: bool = True,
) -> TargetPack:
    """
    Build both next-token targets and unique-time-set targets for one patient.

    This is the function dataset.py should usually call during initialization.

    The dataset can then store:
        event_seq           = target_pack.next_token.input_events
        time_seq            = target_pack.next_token.input_times_years

        target_event_seq    = target_pack.next_token.target_events
        target_time_seq     = target_pack.next_token.target_times_years

        readout_mask        = target_pack.unique_time_set.readout_mask
        target_dt_unique    = target_pack.unique_time_set.target_dt_unique
        target_multi_hot    = target_pack.unique_time_set.target_multi_hot
    """
    next_token = build_next_token_targets(
        labels=labels,
        times_days=times_days,
        require_sorted=require_sorted,
    )

    unique_time_set = build_unique_time_set_targets(
        labels=labels,
        times_days=times_days,
        vocab_size=vocab_size,
        ignored_target_ids=ignored_uts_target_ids,
        require_sorted=require_sorted,
    )

    return TargetPack(
        next_token=next_token,
        unique_time_set=unique_time_set,
    )


def get_group_end_mask_from_times(
    times_days: np.ndarray,
    *,
    input_len: int | None = None,
) -> np.ndarray:
    """
    Convenience utility for debugging.

    Returns a bool mask indicating the last token of each same-time group
    within the input sequence.

    If input_len is None, uses len(times_days) - 1, matching model input length.
    """
    times_days = _as_numpy_1d(times_days, "times_days", np.float32)

    if input_len is None:
        input_len = len(times_days) - 1

    if input_len < 0 or input_len > len(times_days):
        raise ValueError(
            f"Invalid input_len={input_len} for sequence length {len(times_days)}"
        )

    out = np.zeros(input_len, dtype=bool)

    for i in range(input_len):
        is_last_in_group = (
            i == input_len - 1
            or times_days[i + 1] != times_days[i]
        )
        out[i] = is_last_in_group

    return out


def summarize_targets(
    target_pack: TargetPack,
) -> dict[str, int | float]:
    """
    Small debugging helper for logging.
    """
    nt = target_pack.next_token
    uts = target_pack.unique_time_set

    n_tokens = int(len(nt.input_events))
    n_readout = int(uts.readout_mask.sum())
    n_positive_labels = int(uts.target_multi_hot.sum())

    mean_set_size = (
        float(n_positive_labels / n_readout)
        if n_readout > 0
        else 0.0
    )

    return {
        "n_input_tokens": n_tokens,
        "n_uts_readouts": n_readout,
        "n_uts_positive_labels": n_positive_labels,
        "mean_uts_set_size": mean_set_size,
    }