sim_priors_pk/data/data_generation/observations_classes.py

from abc import ABC, abstractmethod
from typing import Callable, Optional, Tuple, List

import torch
from torch import Tensor
from torchtyping import TensorType
from sim_priors_pk.config_classes.data_config import ObservationsConfig, MetaStudyConfig
from sim_priors_pk.data.data_generation.observations_functions import fix_past_time_random_selection


def _sample_past_count_with_bias(
    low: int,
    high: int,
    *,
    generative_bias: bool,
    generator: torch.Generator,
    device: torch.device,
) -> int:
    """Sample the number of past observations under the configured bias mode."""

    if high <= 0:
        return 0

    if generative_bias:
        sample_zero = int(torch.randint(0, 2, (1,), generator=generator, device=device).item()) == 0
        if sample_zero:
            return 0

        rest_low = max(1, low)
        if rest_low > high:
            return 0
        if rest_low == high:
            return rest_low
        return int(
            torch.randint(
                rest_low,
                high + 1,
                (1,),
                generator=generator,
                device=device,
            ).item()
        )

    if low >= high:
        return int(high)

    return int(torch.randint(low, high + 1, (1,), generator=generator, device=device).item())


class ObservationStrategy(ABC):
    def __init__(self, observations_config: ObservationsConfig, meta_config: MetaStudyConfig):
        self.observations_config = observations_config
        self.meta_config = meta_config

    def _drop_non_positive_times_from_mask(self, times: Tensor, mask: Tensor) -> Tensor:
        """Optionally invalidate observations at non-positive timestamps.

        When ``drop_time_zero_observations=True`` in :class:`ObservationsConfig`,
        entries with ``time <= 0`` are excluded from downstream sampling.
        """
        if not getattr(self.observations_config, "drop_time_zero_observations", False):
            return mask
        return mask & (times > 0)

    def generate(
        self, full_simulation: Tensor, full_simulation_times: Tensor, **kwargs
    ) -> Tuple[Tensor, ...]:
        """Wrap raw generate: apply add_rem flag"""
        # call subclass raw generation
        obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._generate_raw(
            full_simulation, full_simulation_times, **kwargs
        )
        # drop remaining if not desired
        if not self.observations_config.add_rem:
            rem_sim = rem_time = rem_mask = None
        return obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask, None

    @abstractmethod
    def _generate_raw(
        self, full_simulation: Tensor, full_simulation_times: Tensor, **kwargs
    ) -> Tuple[
        Tensor,
        TensorType["B", "T_obs"],
        TensorType["B", "T_obs"],
        TensorType["B", "T_rem"],
        TensorType["B", "T_rem"],
        TensorType["B", "T_rem"],
    ]:
        """Generate observations and remaining sims raw, regardless of add_rem"""
        pass

    def get_shapes(self) -> Tuple[int, int]:
        """Wrap raw shapes: apply add_rem flag"""
        max_obs, max_rem = self._get_shapes_raw()
        if not self.observations_config.add_rem:
            max_rem = 0
        return max_obs, max_rem

    @abstractmethod
    def _get_shapes_raw(self) -> Tuple[int, int]:
        """Return max observations and max remaining assuming add_rem=True"""
        pass


class PKPeakHalfLifeStrategy(ObservationStrategy):
    """Observation strategy tailored to pharmacokinetic (PK) curves.

    The strategy samples observations around the absorption peak and along the
    elimination phase of a PK simulation. It uses a canonical grid composed of
    four segments:

    1. Several points before the peak that are proportional to the configured
       peak time.
    2. The peak itself.
    3. Several points after the peak spaced by multiples of the provided
       half-life.
    4. Optional remainder points that are handed back to the caller when
       ``add_rem`` is enabled.

    For **synthetic simulations**, the strategy still uses this canonical grid
    and nearest-neighbour alignment.

    For **empirical data**, measurements are treated as already canonical:

    * No canonical time grid construction.
    * No time normalisation or template matching.
    * No interpolation or re-scaling to canonical coordinates.

    Empirical sequences are only padded / truncated to the internal capacity
    implied by :class:`ObservationsConfig` and :class:`MetaStudyConfig`, and
    then passed through the same past/future splitting logic.

    Past/future splitting
    ----------------------
    When ``split_past_future=True``, the canonical sequence for each row is
    split into:

    * a *past* observation block of fixed width (``max_obs``), and
    * an optional *remainder* block of width (``max_rem``).

    In the default mode (no fixed past selection), the number of past points
    is sampled according to ``generative_bias``:

    * ``False`` samples in ``[min_past, max_past]``.
    * ``True`` samples exactly ``0`` with probability 0.5 and, otherwise,
      samples uniformly in ``[max(1, min_past), max_past]``.

    Under ``generative_bias=False``, **short sequences** receive a special treatment: when
    the number of valid canonical points is less than or equal to the
    observation capacity, *all* valid points are placed in the observation
    block and none are shifted into the remainder.

    Fixed past selection
    --------------------
    Calling :meth:`fix_past_selection(k)` activates a strict mode in which
    the strategy tries to expose exactly ``k`` earliest valid timestamps as
    "past" for each series, subject to the following structural limits:

    1. The number of real data points available in the series.
    2. The observation capacity dictated by :meth:`_get_shapes_raw`.

    Concretely, for each row:

    * Let ``k`` be the fixed past count.
    * Let ``total_valid`` be the number of valid canonical points.
    * Let ``past_required = min(k, total_valid)``.

    The observation block receives

    * ``obs_count = min(past_required, max_obs)`` earliest valid points.

    If ``past_required > obs_count`` (for example because ``k`` exceeds the
    number of observation slots), the remaining required past events
    ``past_required - obs_count`` are the *first entries* in the remainder
    block (subject to the remainder capacity). This guarantees that, as long
    as data and shapes allow, the first ``k`` valid timestamps appear in
    ``obs`` + ``rem`` before any later timestamps.

    Calling :meth:`release_past_selection()` returns to the default stochastic
    behaviour governed by ``min_past``/``max_past``.
    """

    _PEAK_PHASE_MULTIPLIERS = (0.1, 0.2, 0.5, 0.8)
    _POST_PEAK_HALF_LIFE_MULTIPLIERS = (
        0.25,
        0.50,
        1.00,
        2.00,
        4.00,
        6.00,
        8.00,
        9.00,
        14.0,
        19.0,
        30.0,
    )
    _RAW_CANONICAL_POINTS = len(_PEAK_PHASE_MULTIPLIERS) + 1 + len(_POST_PEAK_HALF_LIFE_MULTIPLIERS)

    def __init__(
        self, observations_config: ObservationsConfig, meta_config: MetaStudyConfig
    ) -> None:
        super().__init__(observations_config, meta_config)
        self.max_num_obs = observations_config.max_num_obs
        self.split_past_future = observations_config.split_past_future
        self.min_past = observations_config.min_past
        self.max_past = observations_config.max_past
        self.generative_bias = observations_config.generative_bias
        # None → default random selection. When set, the strategy enforces a
        # strict fixed-past semantics as documented above.
        self._fixed_past_obs_count: Optional[int] = None

    def fix_past_selection(self, obs_count: int) -> None:
        """Activate strict ``k``-past behaviour.

        When this mode is active and ``split_past_future=True``, every call to
        :meth:`generate` or :meth:`generate_empirical` will:

        * expose up to ``obs_count`` earliest valid timestamps in the
          observation block, bounded by the available data and the observation
          capacity;
        * place any additional required past events (when ``obs_count`` is
          larger than the observation capacity) at the *front* of the remainder
          block (when a remainder is present).

        The strategy is allowed to allocate fewer than ``obs_count`` past
        events only when:

        * the series contains fewer real data points than ``obs_count``, or
        * the observation/remainder shapes leave insufficient slots.

        In all other cases the earliest valid timestamps are allocated in the
        order: observation block first, then remainder.
        """

        if not self.split_past_future:
            # No split → fixed past count is meaningless.
            return

        if obs_count < self.min_past or obs_count > self.max_past:
            raise ValueError(
                "Fixed past observation count must lie within the configured min/max bounds."
            )
        self._fixed_past_obs_count = int(obs_count)

    def release_past_selection(self) -> None:
        """Return to the default random past selection behaviour."""
        self._fixed_past_obs_count = None

    @classmethod
    def _build_canonical_grid(
        cls,
        *,
        t_peak: float,
        t_half: float,
        device: torch.device,
        dtype: torch.dtype,
    ) -> Tensor:
        """Construct the canonical grid for a single simulation.

        The grid covers the pre-peak, peak and post-peak regime of the curve by
        scaling two fundamental quantities supplied at runtime: the time of the
        peak concentration ``t_peak`` and the half-life ``t_half``. Both values
        are expected to be expressed in the same units as the simulation time
        axis.
        """
        before_peak = [mult * t_peak for mult in cls._PEAK_PHASE_MULTIPLIERS]
        after_peak = [t_peak + mult * t_half for mult in cls._POST_PEAK_HALF_LIFE_MULTIPLIERS]
        values = before_peak + [t_peak] + after_peak
        return torch.tensor(values, device=device, dtype=dtype)

    def _canonical_grid_capacity(self) -> int:
        """Return the number of canonical grid points available.

        The capacity is the minimum between the simulator resolution and the
        theoretical number of canonical points. This ensures that the
        observation tensors never attempt to gather indices outside the
        original simulation.
        """
        time_steps = getattr(self.meta_config, "time_num_steps", self.max_num_obs)
        return max(
            0,
            min(int(self.max_num_obs), int(time_steps), self._RAW_CANONICAL_POINTS),
        )

    def _get_shapes_raw(self) -> Tuple[int, int]:
        """Compute the maximum number of observation and remainder slots.

        Returns
        -------
        max_obs, max_rem : int, int
            * ``max_obs`` – maximum number of observation time-steps.
            * ``max_rem`` – maximum number of remainder time-steps when
              ``add_rem`` is enabled.

        Raises
        ------
        ValueError
            If a past/future split is requested but the canonical capacity
            cannot satisfy the configured ``min_past`` requirement.
        """
        canonical_cap = self._canonical_grid_capacity()
        if canonical_cap == 0:
            return 0, 0

        if self.split_past_future:
            if canonical_cap < self.min_past:
                raise ValueError("Canonical grid capacity is smaller than the configured min_past")
            max_obs = min(self.max_past, canonical_cap)
            max_rem = max(0, canonical_cap - self.min_past)
        else:
            max_obs = canonical_cap
            max_rem = canonical_cap

        return max_obs, max_rem

    @staticmethod
    def _deduplicate_sorted_indices(
        idx: Tensor, valid_mask: Optional[Tensor] = None
    ) -> Tuple[Tensor, Tensor]:
        """Collapse repeated gather indices while preserving alignment.

        ``idx`` is expected to be monotonically increasing. Consecutive
        duplicates are collapsed into a single entry at the front of the tensor
        and the corresponding ``valid_mask`` entries are shifted accordingly.
        """
        if valid_mask is None:
            valid_mask = torch.ones_like(idx, dtype=torch.bool)

        if idx.numel() <= 1:
            return idx, valid_mask

        duplicate_mask = torch.zeros_like(idx, dtype=torch.bool)
        duplicate_mask[1:] = idx[1:] == idx[:-1]

        if not duplicate_mask.any():
            return idx, valid_mask

        unique_mask = ~duplicate_mask
        kept_idx = idx[unique_mask]
        duplicate_idx = idx[duplicate_mask]

        padded_idx = torch.empty_like(idx)
        padded_idx[: kept_idx.numel()] = kept_idx
        padded_idx[kept_idx.numel() :] = duplicate_idx

        kept_valid = valid_mask[unique_mask]
        padded_mask = torch.zeros_like(valid_mask)
        padded_mask[: kept_valid.numel()] = kept_valid

        return padded_idx, padded_mask

    def _assemble_from_canonical(
        self,
        canonical_vals: Tensor,
        canonical_times: Tensor,
        canonical_mask: Tensor,
        *,
        generator: Optional[torch.Generator] = None,
    ) -> Tuple[Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
        """Convert canonical tensors into output observations.

        The canonical representation stores **all** admissible samples for a
        batch element. This helper slices the canonical tensors into the
        "past" observations that will be returned to the caller and (when
        requested) the "future" remainder.

        Allocation invariants
        ---------------------
        For each batch row:

        * Let ``valid_idx`` be the indices where ``canonical_mask`` is True,
          sorted in ascending order.
        * The observation block always receives the **earliest**
          ``obs_count`` indices from ``valid_idx``.
        * The remainder block (when present) receives later indices only; it
          never contains timestamps that precede those in the observation block.
        * Under ``generative_bias=False``, short sequences
          (``total_valid <= max_obs``) keep all valid points in the
          observation block and do not shift points to the remainder.

        When :meth:`fix_past_selection(k)` is active, we define::

            past_required = min(k, total_valid)

        and allocate:

        * ``obs_count = min(past_required, max_obs)`` to the observation
          block; and
        * any surplus past events ``past_required - obs_count`` at the **front**
          of the remainder block (subject to the remainder capacity), followed
          by any truly future points.

        Releasing the fixed selection returns to the stochastic behaviour
        controlled by ``generative_bias``.
        """
        max_obs, max_rem = self._get_shapes_raw()
        device = canonical_vals.device
        dtype = canonical_vals.dtype
        batch, _ = canonical_vals.shape

        obs_out = torch.zeros(batch, max_obs, dtype=dtype, device=device)
        obs_time = torch.zeros_like(obs_out)
        obs_mask = torch.zeros(batch, max_obs, dtype=torch.bool, device=device)

        rem_sim = rem_time = rem_mask = None
        if max_rem > 0:
            rem_sim = torch.zeros(batch, max_rem, dtype=dtype, device=device)
            rem_time = torch.zeros_like(rem_sim)
            rem_mask = torch.zeros(batch, max_rem, dtype=torch.bool, device=device)

        gen = generator if generator is not None else torch.default_generator

        for row in range(batch):
            valid_idx = canonical_mask[row].nonzero(as_tuple=True)[0]
            total_valid = int(valid_idx.numel())
            if total_valid == 0:
                continue

            fixed_k = self._fixed_past_obs_count if self.split_past_future else None

            # ------------------------------------------------------------------
            # 1) Decide obs_count
            # ------------------------------------------------------------------
            if self.split_past_future and fixed_k is not None:
                # Strict fixed-past semantics. Structural limits:
                #   - real data (total_valid)
                #   - observation capacity (max_obs)
                past_required = min(fixed_k, total_valid)
                obs_capacity = min(max_obs, total_valid)
                obs_count = min(past_required, obs_capacity)
            else:
                # Default stochastic behaviour; the short-series fix is kept
                # for the non-biased mode only.
                if self.split_past_future:
                    low = min(self.min_past, total_valid)
                    high = min(self.max_past, total_valid)

                    sampled = _sample_past_count_with_bias(
                        low=low,
                        high=high,
                        generative_bias=self.generative_bias,
                        generator=gen,
                        device=device,
                    )

                    if (not self.generative_bias) and total_valid <= max_obs:
                        # Short-series fix: never push valid points into the
                        # remainder just to satisfy a random split.
                        obs_count = total_valid
                    else:
                        obs_count = min(sampled, max_obs)
                else:
                    obs_count = min(total_valid, max_obs)

            # Safety clamp.
            obs_count = max(0, min(obs_count, min(max_obs, total_valid)))

            # ------------------------------------------------------------------
            # 2) Fill observation block (earliest obs_count indices)
            # ------------------------------------------------------------------
            if obs_count > 0:
                take = valid_idx[:obs_count]
                obs_out[row, :obs_count] = canonical_vals[row, take]
                obs_time[row, :obs_count] = canonical_times[row, take]
                obs_mask[row, :obs_count] = True

            # ------------------------------------------------------------------
            # 3) Fill remainder block (if enabled)
            # ------------------------------------------------------------------
            if rem_sim is not None:
                if self.split_past_future and fixed_k is not None:
                    # Remaining required past events plus genuine future.
                    past_required = min(fixed_k, total_valid)
                    # indices that are still part of the fixed past window
                    # but did not fit into the observation block
                    extra_past_idx = valid_idx[obs_count:past_required]
                    future_idx = valid_idx[past_required:]

                    candidates: List[Tensor] = []
                    if extra_past_idx.numel() > 0:
                        candidates.append(extra_past_idx)
                    if future_idx.numel() > 0:
                        candidates.append(future_idx)
                    if candidates:
                        remainder_candidates = torch.cat(candidates, dim=0)
                    else:
                        remainder_candidates = valid_idx.new_empty((0,), dtype=valid_idx.dtype)
                else:
                    # Default behaviour: everything after the obs window.
                    remainder_candidates = valid_idx[obs_count:]

                rem_count = min(int(remainder_candidates.numel()), max_rem)
                if rem_count > 0:
                    rem_idx = remainder_candidates[:rem_count]
                    rem_sim[row, :rem_count] = canonical_vals[row, rem_idx]
                    rem_time[row, :rem_count] = canonical_times[row, rem_idx]
                    rem_mask[row, :rem_count] = True

        return obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask

    def _align_simulation_to_canonical(
        self,
        full_simulation: Tensor,
        full_simulation_times: Tensor,
        *,
        time_scales: Tensor,
        num_obs_sampler: Optional[Callable[[int], Tensor]] = None,
    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """Gather canonical samples from a simulated PK curve.

        Synthetic behaviour is unchanged compared to the original strategy:
        we build a canonical grid, snap it to the nearest simulation times and
        optionally subsample points via ``num_obs_sampler``.
        """
        device = full_simulation.device
        dtype = full_simulation.dtype
        batch, _ = full_simulation.shape
        time_steps = int(full_simulation_times.size(1))

        # DataLoader workers may receive empty row slices (B=0). In that case
        # there is no reference timeline to align against; return an empty
        # canonical block and let _assemble_from_canonical create [B, *] outputs.
        if batch == 0 or time_steps == 0:
            zero = torch.zeros(batch, 0, dtype=dtype, device=device)
            mask = torch.zeros(batch, 0, dtype=torch.bool, device=device)
            return zero, zero, mask, time_scales.clone()

        canonical_cap = self._canonical_grid_capacity()
        if canonical_cap == 0:
            zero = torch.zeros(batch, 0, dtype=dtype, device=device)
            mask = torch.zeros(batch, 0, dtype=torch.bool, device=device)
            return zero, zero, mask, time_scales.clone()

        grid = self._build_canonical_grid(
            t_peak=time_scales[0].item(),
            t_half=time_scales[1].item(),
            device=device,
            dtype=dtype,
        )[:canonical_cap]

        ref_times = full_simulation_times[0]
        min_time = ref_times.min()
        max_time = ref_times.max()
        grid_valid_mask = (grid >= min_time) & (grid <= max_time)

        idx = torch.cdist(grid[:, None], ref_times[:, None]).argmin(dim=1)
        idx, order = idx.sort()
        grid_valid_mask = grid_valid_mask[order]
        idx, grid_valid_mask = self._deduplicate_sorted_indices(idx, grid_valid_mask)

        gather_idx = idx[None, :].expand(batch, -1)
        batch_idx = torch.arange(batch, device=device)[:, None]

        canonical_vals = full_simulation[batch_idx, gather_idx]
        canonical_times = full_simulation_times[batch_idx, gather_idx]

        invalid_slots = ~grid_valid_mask
        if invalid_slots.any():
            canonical_vals[:, invalid_slots] = 0
            canonical_times[:, invalid_slots] = 0

        if num_obs_sampler is None:
            total_counts = torch.full((batch,), canonical_cap, dtype=torch.long, device=device)
        else:
            sampled = num_obs_sampler(batch).to(device=device).long()
            total_counts = sampled.clamp(min=0, max=canonical_cap)

        max_valid = int(grid_valid_mask.sum().item())
        if max_valid == 0:
            total_counts.zero_()
        else:
            total_counts.clamp_(max=max_valid)

        valid_order = grid_valid_mask.long().cumsum(dim=0) - 1
        valid_order = torch.where(
            grid_valid_mask,
            valid_order,
            torch.full_like(valid_order, -1, dtype=valid_order.dtype),
        )
        canonical_mask = grid_valid_mask[None, :] & (valid_order[None, :] < total_counts[:, None])
        canonical_mask = self._drop_non_positive_times_from_mask(canonical_times, canonical_mask)

        return canonical_vals, canonical_times, canonical_mask, time_scales.clone()

    def _align_empirical_to_canonical(
        self,
        empirical_obs: Tensor,
        empirical_times: Tensor,
        empirical_mask: Tensor,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """(Legacy) Project empirical observations onto the canonical grid.

        This method is retained for backward compatibility but is **not** used
        by :meth:`generate_empirical`, which now treats empirical data as
        already canonical. New code should avoid calling this helper.
        """
        device = empirical_obs.device
        dtype = empirical_obs.dtype
        batch, _ = empirical_obs.shape
        canonical_cap = self._canonical_grid_capacity()

        canonical_vals = torch.zeros(batch, canonical_cap, dtype=dtype, device=device)
        canonical_times = torch.zeros_like(canonical_vals)
        canonical_mask = torch.zeros(batch, canonical_cap, dtype=torch.bool, device=device)

        if canonical_cap == 0:
            return canonical_vals, canonical_times, canonical_mask

        for row in range(batch):
            valid_idx = empirical_mask[row].nonzero(as_tuple=True)[0]
            if valid_idx.numel() == 0:
                continue

            obs_row = empirical_obs[row, valid_idx]
            time_row = empirical_times[row, valid_idx]
            max_time = torch.maximum(time_row.max(), torch.tensor(1.0, device=device))
            norm_time = time_row / max_time

            peak_idx = obs_row.argmax().item()
            t_peak = norm_time[peak_idx].item()
            post_times = norm_time[peak_idx:]
            post_obs = obs_row[peak_idx:]
            half_level = obs_row[peak_idx] / 2
            below_half = (post_obs <= half_level).nonzero(as_tuple=True)[0]
            if below_half.numel() == 0:
                half_time = post_times[-1].item()
            else:
                half_time = post_times[below_half[0]].item()
            t_half = max(half_time - t_peak, 1e-3)

            grid = self._build_canonical_grid(
                t_peak=t_peak if t_peak > 0 else 1e-3,
                t_half=t_half,
                device=device,
                dtype=dtype,
            )[:canonical_cap].clamp(max=1.0)

            actual_grid = grid * max_time
            distances = torch.cdist(actual_grid[:, None], time_row[:, None])
            nearest = distances.argmin(dim=1)

            usable = min(time_row.numel(), grid.numel())
            if usable == 0:
                continue

            canonical_vals[row, :usable] = obs_row[nearest[:usable]]
            canonical_times[row, :usable] = time_row[nearest[:usable]]
            canonical_mask[row, :usable] = True

        canonical_mask = self._drop_non_positive_times_from_mask(canonical_times, canonical_mask)

        return canonical_vals, canonical_times, canonical_mask

    def _prepare_empirical_as_canonical(
        self,
        empirical_obs: Tensor,
        empirical_times: Tensor,
        empirical_mask: Tensor,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """Treat empirical observations as already canonical.

        This helper:

        * does **not** build any canonical grid;
        * does **not** normalise or re-scale time;
        * simply copies valid empirical points in their original order into
          fixed-size tensors, padding with zeros / False as needed.

        The resulting tensors have width equal to the canonical capacity so
        that they can be passed to :meth:`_assemble_from_canonical`.
        """
        device = empirical_obs.device
        dtype = empirical_obs.dtype
        batch, _ = empirical_obs.shape
        canonical_cap = self._canonical_grid_capacity()

        canonical_vals = torch.zeros(batch, canonical_cap, dtype=dtype, device=device)
        canonical_times = torch.zeros_like(canonical_vals)
        canonical_mask = torch.zeros(batch, canonical_cap, dtype=torch.bool, device=device)

        if canonical_cap == 0:
            return canonical_vals, canonical_times, canonical_mask

        for row in range(batch):
            valid_idx = empirical_mask[row].nonzero(as_tuple=True)[0]
            if valid_idx.numel() == 0:
                continue

            take_count = min(int(valid_idx.numel()), canonical_cap)
            take_idx = valid_idx[:take_count]

            canonical_vals[row, :take_count] = empirical_obs[row, take_idx]
            canonical_times[row, :take_count] = empirical_times[row, take_idx]
            canonical_mask[row, :take_count] = True

        canonical_mask = self._drop_non_positive_times_from_mask(canonical_times, canonical_mask)

        return canonical_vals, canonical_times, canonical_mask

    def _generate_raw(
        self, full_simulation: Tensor, full_simulation_times: Tensor, **kwargs
    ) -> Tuple[
        Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Tensor
    ]:
        """Deterministic canonical PK sampling for synthetic simulations."""
        time_scales: Optional[Tensor] = kwargs.get("time_scales")
        if time_scales is None:
            raise ValueError("time_scales must be provided for PKPeakHalfLifeStrategy")

        canonical_vals, canonical_times, canonical_mask, rescaled = (
            self._align_simulation_to_canonical(
                full_simulation,
                full_simulation_times,
                time_scales=time_scales,
                num_obs_sampler=kwargs.get("num_obs_sampler"),
            )
        )

        obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._assemble_from_canonical(
            canonical_vals,
            canonical_times,
            canonical_mask,
            generator=kwargs.get("generator"),
        )

        return obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask, rescaled

    def _generate_random(
        self,
        full_simulation: Tensor,
        full_simulation_times: Tensor,
        *,
        time_scales: Tensor,
        generator: Optional[torch.Generator] = None,
    ) -> Tuple[
        Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Tensor
    ]:
        """Randomised variant of canonical observation generation.

        The pre- and post-peak segments are sampled from uniform distributions
        bounded by the canonical limits. This keeps the semantic meaning of the
        selected points while injecting stochasticity that can improve
        robustness during training.
        """
        device, dtype = full_simulation.device, full_simulation.dtype
        batch = full_simulation.size(0)
        time_steps = int(full_simulation_times.size(1))
        if batch == 0 or time_steps == 0:
            canonical_vals = torch.zeros(batch, 0, dtype=dtype, device=device)
            canonical_times = torch.zeros(batch, 0, dtype=dtype, device=device)
            canonical_mask = torch.zeros(batch, 0, dtype=torch.bool, device=device)
            obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._assemble_from_canonical(
                canonical_vals, canonical_times, canonical_mask, generator=generator
            )
            return obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask, time_scales.clone()
        t_peak, t_half = time_scales[0].item(), time_scales[1].item()

        n_pre = len(self._PEAK_PHASE_MULTIPLIERS)
        n_post = len(self._POST_PEAK_HALF_LIFE_MULTIPLIERS)

        # Uniform samples before peak
        pre_times = torch.rand(n_pre, device=device, dtype=dtype) * t_peak
        # Always include the peak
        peak_time = torch.tensor([t_peak], device=device, dtype=dtype)
        # Uniform samples after peak
        post_times = []
        for mult in self._POST_PEAK_HALF_LIFE_MULTIPLIERS:
            t_end = t_peak + mult * t_half
            t_rand = torch.empty(1, device=device, dtype=dtype).uniform_(t_peak, t_end)
            post_times.append(t_rand)
        post_times = torch.cat(post_times, dim=0)

        # Truncate to canonical capacity
        grid = torch.cat([pre_times, peak_time, post_times], dim=0)
        canonical_cap = self._canonical_grid_capacity()
        grid = grid[:canonical_cap]

        # Map grid to nearest simulation points
        ref_times = full_simulation_times[0]
        idx = torch.cdist(grid[:, None], ref_times[:, None]).argmin(dim=1)
        idx, _ = idx.sort()
        valid_mask = torch.ones_like(idx, dtype=torch.bool)
        idx, valid_mask = self._deduplicate_sorted_indices(idx, valid_mask)
        gather_idx = idx[None, :].expand(batch, -1)
        batch_idx = torch.arange(batch, device=device)[:, None]

        canonical_vals = full_simulation[batch_idx, gather_idx]
        canonical_times = full_simulation_times[batch_idx, gather_idx]
        invalid_slots = ~valid_mask
        if invalid_slots.any():
            canonical_vals[:, invalid_slots] = 0
            canonical_times[:, invalid_slots] = 0

        canonical_mask = valid_mask[None, :].expand(batch, -1).clone()
        canonical_mask = self._drop_non_positive_times_from_mask(canonical_times, canonical_mask)

        obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._assemble_from_canonical(
            canonical_vals, canonical_times, canonical_mask, generator=generator
        )
        return obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask, time_scales.clone()

    def generate(
        self,
        full_simulation: Tensor,
        full_simulation_times: Tensor,
        **kwargs,
    ) -> Tuple[
        Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Tensor
    ]:
        """Generate PK observations for synthetic simulations.

        With probability ``randomize_prob`` (default 0.5) the method delegates
        to :meth:`_generate_random`; otherwise the deterministic
        :meth:`_generate_raw` path is taken. Setting ``deterministic_only=True``
        forces the deterministic branch. Both paths require ``time_scales`` and
        honour the ``add_rem`` flag.
        """
        time_scales: Optional[Tensor] = kwargs.get("time_scales")
        if time_scales is None:
            raise ValueError("time_scales must be provided for PKPeakHalfLifeStrategy")

        deterministic_only = kwargs.pop("deterministic_only", False)

        use_random = False
        if not deterministic_only:
            use_random = torch.rand(()) < getattr(self, "randomize_prob", 0.5)

        if use_random:
            obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask, rescaled = self._generate_random(
                full_simulation,
                full_simulation_times,
                time_scales=time_scales,
                generator=kwargs.get("generator"),
            )
        else:
            obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask, rescaled = self._generate_raw(
                full_simulation,
                full_simulation_times,
                **kwargs,
            )

        if not self.observations_config.add_rem:
            rem_sim = rem_time = rem_mask = None

        return obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask, rescaled

    def generate_empirical(
        self,
        empirical_obs: Tensor,
        empirical_times: Tensor,
        empirical_mask: Tensor,
        *,
        generator: Optional[torch.Generator] = None,
    ) -> Tuple[Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
        """Generate observations from empirical data.

        Empirical measurements are assumed to already live on their correct
        time grid. This routine:

        * does **not** perform canonical alignment or time normalisation;
        * only pads / truncates sequences to match the internal capacity;
        * applies past/future splitting via :meth:`_assemble_from_canonical`
          using the configuration in :class:`ObservationsConfig`.

        Synthetic simulations keep using the canonical alignment path.
        """
        canonical_vals, canonical_times, canonical_mask = self._prepare_empirical_as_canonical(
            empirical_obs,
            empirical_times,
            empirical_mask,
        )

        obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._assemble_from_canonical(
            canonical_vals,
            canonical_times,
            canonical_mask,
            generator=generator,
        )

        if not self.observations_config.add_rem:
            rem_sim = rem_time = rem_mask = None

        return obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask


class PKPeakHalfLifeStrategyOld(ObservationStrategy):
    """Observation strategy tailored to pharmacokinetic (PK) curves.

    The strategy samples observations around the absorption peak and along the
    elimination phase of a PK simulation.  It uses a canonical grid composed of
    four segments:

    1. Several points before the peak that are proportional to the configured
       peak time.
    2. The peak itself.
    3. Several points after the peak spaced by multiples of the provided
       half-life.
    4. Optional remainder points that are handed back to the caller when
       ``add_rem`` is enabled.

    The resulting observation tensor can be optionally split into "past" and
    "future" observations according to :class:`ObservationsConfig`.

    Parameters
    ----------
    observations_config:
        Simulation-level configuration that defines sampling constraints such
        as ``max_num_obs`` or the minimum/maximum number of "past" points when
        a split is requested.
    meta_config:
        Meta-study configuration.  Only the ``time_num_steps`` attribute is
        used and allows clamping the canonical grid to the resolution of the
        simulator.
    """

    _PEAK_PHASE_MULTIPLIERS = (0.1, 0.2, 0.5, 0.8)
    _POST_PEAK_HALF_LIFE_MULTIPLIERS = (
        0.25,
        0.50,
        1.00,
        2.00,
        4.00,
        6.00,
        8.00,
        9.00,
        14.0,
        19.0,
        30.0,
    )
    _RAW_CANONICAL_POINTS = len(_PEAK_PHASE_MULTIPLIERS) + 1 + len(_POST_PEAK_HALF_LIFE_MULTIPLIERS)

    def __init__(
        self, observations_config: ObservationsConfig, meta_config: MetaStudyConfig
    ) -> None:
        super().__init__(observations_config, meta_config)
        self.max_num_obs = observations_config.max_num_obs
        self.split_past_future = observations_config.split_past_future
        self.min_past = observations_config.min_past
        self.max_past = observations_config.max_past
        self.generative_bias = observations_config.generative_bias
        # ``None`` indicates that the number of past observations should be
        # sampled according to the standard strategy.  When populated it forces
        # :meth:`_assemble_from_canonical` to always select the provided number
        # of past observations (within the valid range).
        self._fixed_past_obs_count: Optional[int] = None

    def fix_past_selection(self, obs_count: int) -> None:
        """Force the past observation count to ``obs_count`` when splitting.

        The override is only applied when ``split_past_future`` is enabled.  The
        provided ``obs_count`` must fall within ``[min_past, max_past]``.
        """

        if not self.split_past_future:
            return

        if obs_count < self.min_past or obs_count > self.max_past:
            raise ValueError(
                "Fixed past observation count must lie within the configured min/max bounds."
            )
        self._fixed_past_obs_count = int(obs_count)

    def release_past_selection(self) -> None:
        """Return to the default random past selection behaviour."""

        self._fixed_past_obs_count = None

    @classmethod
    def _build_canonical_grid(
        cls,
        *,
        t_peak: float,
        t_half: float,
        device: torch.device,
        dtype: torch.dtype,
    ) -> Tensor:
        """Construct the canonical grid for a single simulation.

        The grid covers the pre-peak, peak and post-peak regime of the curve by
        scaling two fundamental quantities supplied at runtime: the time of the
        peak concentration ``t_peak`` and the half-life ``t_half``.  Both values
        are expected to be expressed in the same units as the simulation time
        axis.

        Parameters
        ----------
        t_peak:
            Estimated time of the concentration peak.
        t_half:
            Estimated half-life used to position post-peak points.
        device, dtype:
            Torch device and dtype for the returned tensor so that it matches
            the simulation tensors that will be gathered later on.

        Returns
        -------
        torch.Tensor
            One-dimensional tensor containing monotonically increasing times
            representing the canonical sampling grid.
        """
        before_peak = [mult * t_peak for mult in cls._PEAK_PHASE_MULTIPLIERS]
        after_peak = [t_peak + mult * t_half for mult in cls._POST_PEAK_HALF_LIFE_MULTIPLIERS]
        values = before_peak + [t_peak] + after_peak
        return torch.tensor(values, device=device, dtype=dtype)

    def _canonical_grid_capacity(self) -> int:
        """Return the number of canonical grid points available.

        The capacity is the minimum between the simulator resolution and the
        theoretical number of canonical points.  This ensures that the
        observation tensors never attempt to gather indices outside the
        original simulation.

        Returns
        -------
        int
            Maximum number of grid points that can be sampled for each
            simulation in the batch.
        """
        time_steps = getattr(self.meta_config, "time_num_steps", self.max_num_obs)
        return max(
            0,
            min(int(self.max_num_obs), int(time_steps), self._RAW_CANONICAL_POINTS),
        )

    def _get_shapes_raw(self) -> Tuple[int, int]:
        """Compute the maximum number of observation and remainder slots.

        The method applies the canonical grid capacity alongside the
        ``split_past_future`` configuration to decide how many points can be
        surfaced directly as observations and how many should be exposed as
        "remaining" (future) points.

        Returns
        -------
        tuple[int, int]
            The first entry is the maximum number of observations.  The second
            entry is the maximum number of remaining observations when
            ``add_rem`` is enabled.

        Raises
        ------
        ValueError
            If a past/future split is requested but the canonical capacity
            cannot satisfy the configured ``min_past`` requirement.
        """
        canonical_cap = self._canonical_grid_capacity()
        if canonical_cap == 0:
            return 0, 0

        if self.split_past_future:
            if canonical_cap < self.min_past:
                raise ValueError("Canonical grid capacity is smaller than the configured min_past")
            max_obs = min(self.max_past, canonical_cap)
            max_rem = max(0, canonical_cap - self.min_past)
        else:
            max_obs = canonical_cap
            max_rem = canonical_cap

        return max_obs, max_rem

    @staticmethod
    def _deduplicate_sorted_indices(
        idx: Tensor, valid_mask: Optional[Tensor] = None
    ) -> Tuple[Tensor, Tensor]:
        """Collapse repeated gather indices while preserving alignment."""

        if valid_mask is None:
            valid_mask = torch.ones_like(idx, dtype=torch.bool)

        if idx.numel() <= 1:
            return idx, valid_mask

        duplicate_mask = torch.zeros_like(idx, dtype=torch.bool)
        duplicate_mask[1:] = idx[1:] == idx[:-1]

        if not duplicate_mask.any():
            return idx, valid_mask

        unique_mask = ~duplicate_mask
        kept_idx = idx[unique_mask]
        duplicate_idx = idx[duplicate_mask]

        padded_idx = torch.empty_like(idx)
        padded_idx[: kept_idx.numel()] = kept_idx
        padded_idx[kept_idx.numel() :] = duplicate_idx

        kept_valid = valid_mask[unique_mask]
        padded_mask = torch.zeros_like(valid_mask)
        padded_mask[: kept_valid.numel()] = kept_valid

        return padded_idx, padded_mask

    def _assemble_from_canonical(
        self,
        canonical_vals: Tensor,
        canonical_times: Tensor,
        canonical_mask: Tensor,
        *,
        generator: Optional[torch.Generator] = None,
    ) -> Tuple[Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
        """Convert canonical tensors into output observations.

        The canonical representation stores **all** admissible samples for a
        batch element.  This helper slices the canonical tensors into the
        "past" observations that will be returned to the caller and (when
        requested) the "future" remainder.  The selection proceeds row by row:

        1. ``canonical_mask`` is inspected to identify the indices that contain
           valid information.  These are the only points that may be surfaced.
        2. When ``split_past_future`` is ``False`` every valid point is treated
           as part of the observation history up to the configured capacity.
        3. Otherwise we randomly draw ``obs_count`` between ``min_past`` and
           ``max_past`` (capped by the number of valid canonical entries).  The
           first ``obs_count`` indices become past observations while the
           remaining valid points are placed in the remainder tensors.

        Parameters
        ----------
        canonical_vals, canonical_times:
            Tensors produced by aligning the simulation or empirical data to
            the canonical grid.
        canonical_mask:
            Boolean tensor marking valid entries for each batch element.
        generator:
            Optional random generator used when sampling ``obs_count`` in
            split-past/future mode.

        Returns
        -------
        tuple of tensors
            Observation and remaining tensors matching the shapes dictated by
            :meth:`_get_shapes_raw`.  All tensors share the same device and
            dtype as the inputs.  ``None`` is returned for remainder tensors
            when the capacity is zero.
        """
        max_obs, max_rem = self._get_shapes_raw()
        device = canonical_vals.device
        dtype = canonical_vals.dtype
        batch, _ = canonical_vals.shape

        obs_out = torch.zeros(batch, max_obs, dtype=dtype, device=device)
        obs_time = torch.zeros_like(obs_out)
        obs_mask = torch.zeros(batch, max_obs, dtype=torch.bool, device=device)

        rem_sim = rem_time = rem_mask = None
        if max_rem > 0:
            rem_sim = torch.zeros(batch, max_rem, dtype=dtype, device=device)
            rem_time = torch.zeros_like(rem_sim)
            rem_mask = torch.zeros(batch, max_rem, dtype=torch.bool, device=device)

        gen = generator if generator is not None else torch.default_generator

        for row in range(batch):
            valid_idx = canonical_mask[row].nonzero(as_tuple=True)[0]
            total_valid = valid_idx.numel()
            if total_valid == 0:
                continue

            if self.split_past_future:
                low = min(self.min_past, total_valid)
                high = min(self.max_past, total_valid)
                if self._fixed_past_obs_count is not None:
                    obs_count = min(self._fixed_past_obs_count, total_valid)
                else:
                    obs_count = _sample_past_count_with_bias(
                        low=low,
                        high=high,
                        generative_bias=self.generative_bias,
                        generator=gen,
                        device=device,
                    )
                obs_count = min(obs_count, max_obs)
            else:
                obs_count = min(total_valid, max_obs)

            if obs_count > 0:
                take = valid_idx[:obs_count]
                obs_out[row, :obs_count] = canonical_vals[row, take]
                obs_time[row, :obs_count] = canonical_times[row, take]
                obs_mask[row, :obs_count] = True

            if rem_sim is not None:
                rem_candidates = valid_idx[obs_count:]
                rem_count = min(rem_candidates.numel(), max_rem)
                if rem_count > 0:
                    rem_idx = rem_candidates[:rem_count]
                    rem_sim[row, :rem_count] = canonical_vals[row, rem_idx]
                    rem_time[row, :rem_count] = canonical_times[row, rem_idx]
                    rem_mask[row, :rem_count] = True

        return obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask

    def _align_simulation_to_canonical(
        self,
        full_simulation: Tensor,
        full_simulation_times: Tensor,
        *,
        time_scales: Tensor,
        num_obs_sampler: Optional[Callable[[int], Tensor]] = None,
    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """Gather the canonical samples from a simulated PK curve.

        The routine creates the canonical grid described in the configuration
        (using the provided ``time_scales``) and then performs a nearest-neighbour
        lookup on the simulated trajectory.  Each grid location picks the
        closest time point from the reference simulation (the first batch row);
        the same indices are applied to every batch element so that values and
        times remain aligned across the batch.  ``num_obs_sampler`` can further
        prune the resulting grid by specifying how many of those canonical
        points should remain valid for each row.

        Parameters
        ----------
        full_simulation, full_simulation_times:
            Batched tensors representing the simulated concentration curve and
            its time axis.
        time_scales:
            Two-element tensor with ``t_peak`` and ``t_half`` scaling factors.
        num_obs_sampler:
            Optional callable that samples how many canonical points should be
            retained for each batch element.

        Returns
        -------
        tuple of torch.Tensor
            The canonical values, their corresponding times, a boolean mask of
            valid entries and the (cloned) ``time_scales`` tensor.  When the
            canonical capacity is zero, zero-sized tensors are returned for the
            first three entries.
        """
        device = full_simulation.device
        dtype = full_simulation.dtype
        batch, _ = full_simulation.shape
        time_steps = int(full_simulation_times.size(1))

        # Empty worker slices (B=0) and zero-step trajectories are valid edge
        # cases; return empty canonical tensors and keep shape assembly
        # delegated to _assemble_from_canonical.
        if batch == 0 or time_steps == 0:
            zero = torch.zeros(batch, 0, dtype=dtype, device=device)
            mask = torch.zeros(batch, 0, dtype=torch.bool, device=device)
            return zero, zero, mask, time_scales.clone()

        canonical_cap = self._canonical_grid_capacity()
        if canonical_cap == 0:
            zero = torch.zeros(batch, 0, dtype=dtype, device=device)
            mask = torch.zeros(batch, 0, dtype=torch.bool, device=device)
            return zero, zero, mask, time_scales.clone()

        grid = self._build_canonical_grid(
            t_peak=time_scales[0].item(),
            t_half=time_scales[1].item(),
            device=device,
            dtype=dtype,
        )[:canonical_cap]

        ref_times = full_simulation_times[0]
        min_time = ref_times.min()
        max_time = ref_times.max()
        grid_valid_mask = (grid >= min_time) & (grid <= max_time)

        idx = torch.cdist(grid[:, None], ref_times[:, None]).argmin(dim=1)
        idx, order = idx.sort()
        grid_valid_mask = grid_valid_mask[order]
        idx, grid_valid_mask = self._deduplicate_sorted_indices(idx, grid_valid_mask)

        gather_idx = idx[None, :].expand(batch, -1)
        batch_idx = torch.arange(batch, device=device)[:, None]

        canonical_vals = full_simulation[batch_idx, gather_idx]
        canonical_times = full_simulation_times[batch_idx, gather_idx]

        invalid_slots = ~grid_valid_mask
        if invalid_slots.any():
            canonical_vals[:, invalid_slots] = 0
            canonical_times[:, invalid_slots] = 0

        if num_obs_sampler is None:
            total_counts = torch.full((batch,), canonical_cap, dtype=torch.long, device=device)
        else:
            sampled = num_obs_sampler(batch).to(device=device).long()
            total_counts = sampled.clamp(min=0, max=canonical_cap)

        max_valid = int(grid_valid_mask.sum().item())
        if max_valid == 0:
            total_counts.zero_()
        else:
            total_counts.clamp_(max=max_valid)

        valid_order = grid_valid_mask.long().cumsum(dim=0) - 1
        valid_order = torch.where(
            grid_valid_mask,
            valid_order,
            torch.full_like(valid_order, -1, dtype=valid_order.dtype),
        )
        canonical_mask = grid_valid_mask[None, :] & (valid_order[None, :] < total_counts[:, None])
        canonical_mask = self._drop_non_positive_times_from_mask(canonical_times, canonical_mask)

        return canonical_vals, canonical_times, canonical_mask, time_scales.clone()

    def _align_empirical_to_canonical(
        self,
        empirical_obs: Tensor,
        empirical_times: Tensor,
        empirical_mask: Tensor,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """Project empirical observations onto the canonical grid.

        The projection normalises the empirical time axis to estimate the peak
        and half-life from the data itself.  This allows harmonising real
        measurements with the canonical layout used during simulation-driven
        training.

        Parameters
        ----------
        empirical_obs, empirical_times, empirical_mask:
            Batched tensors storing empirical observations, the corresponding
            time stamps and a mask of valid entries.

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
            Canonical values, times and boolean masks aligned to the canonical
            sampling scheme.
        """
        device = empirical_obs.device
        dtype = empirical_obs.dtype
        batch, _ = empirical_obs.shape
        canonical_cap = self._canonical_grid_capacity()

        canonical_vals = torch.zeros(batch, canonical_cap, dtype=dtype, device=device)
        canonical_times = torch.zeros_like(canonical_vals)
        canonical_mask = torch.zeros(batch, canonical_cap, dtype=torch.bool, device=device)

        if canonical_cap == 0:
            return canonical_vals, canonical_times, canonical_mask

        for row in range(batch):
            valid_idx = empirical_mask[row].nonzero(as_tuple=True)[0]
            if valid_idx.numel() == 0:
                continue

            obs_row = empirical_obs[row, valid_idx]
            time_row = empirical_times[row, valid_idx]
            max_time = torch.maximum(time_row.max(), torch.tensor(1.0, device=device))
            norm_time = time_row / max_time

            peak_idx = obs_row.argmax().item()
            t_peak = norm_time[peak_idx].item()
            post_times = norm_time[peak_idx:]
            post_obs = obs_row[peak_idx:]
            half_level = obs_row[peak_idx] / 2
            below_half = (post_obs <= half_level).nonzero(as_tuple=True)[0]
            if below_half.numel() == 0:
                half_time = post_times[-1].item()
            else:
                half_time = post_times[below_half[0]].item()
            t_half = max(half_time - t_peak, 1e-3)

            grid = self._build_canonical_grid(
                t_peak=t_peak if t_peak > 0 else 1e-3,
                t_half=t_half,
                device=device,
                dtype=dtype,
            )[:canonical_cap].clamp(max=1.0)

            actual_grid = grid * max_time
            distances = torch.cdist(actual_grid[:, None], time_row[:, None])
            nearest = distances.argmin(dim=1)

            usable = min(time_row.numel(), grid.numel())
            if usable == 0:
                continue

            canonical_vals[row, :usable] = obs_row[nearest[:usable]]
            canonical_times[row, :usable] = time_row[nearest[:usable]]
            canonical_mask[row, :usable] = True

        canonical_mask = self._drop_non_positive_times_from_mask(canonical_times, canonical_mask)

        return canonical_vals, canonical_times, canonical_mask

    def _generate_raw(
        self, full_simulation: Tensor, full_simulation_times: Tensor, **kwargs
    ) -> Tuple[
        Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Tensor
    ]:
        time_scales: Optional[Tensor] = kwargs.get("time_scales")
        if time_scales is None:
            raise ValueError("time_scales must be provided for PKPeakHalfLifeStrategy")

        canonical_vals, canonical_times, canonical_mask, rescaled = (
            self._align_simulation_to_canonical(
                full_simulation,
                full_simulation_times,
                time_scales=time_scales,
                num_obs_sampler=kwargs.get("num_obs_sampler"),
            )
        )

        obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._assemble_from_canonical(
            canonical_vals,
            canonical_times,
            canonical_mask,
            generator=kwargs.get("generator"),
        )

        return obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask, rescaled

    def _generate_random(
        self,
        full_simulation: Tensor,
        full_simulation_times: Tensor,
        *,
        time_scales: Tensor,
        generator: Optional[torch.Generator] = None,
    ) -> Tuple[
        Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Tensor
    ]:
        """Randomized variant of canonical observation generation.

        Instead of fixed multipliers, the pre- and post-peak segments are
        sampled from uniform distributions bounded by the canonical limits.
        This keeps the semantic meaning of the selected points while injecting
        stochasticity that improves robustness when training amortised
        inference models.
        """
        device, dtype = full_simulation.device, full_simulation.dtype
        batch = full_simulation.size(0)
        time_steps = int(full_simulation_times.size(1))
        if batch == 0 or time_steps == 0:
            canonical_vals = torch.zeros(batch, 0, dtype=dtype, device=device)
            canonical_times = torch.zeros(batch, 0, dtype=dtype, device=device)
            canonical_mask = torch.zeros(batch, 0, dtype=torch.bool, device=device)
            obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._assemble_from_canonical(
                canonical_vals, canonical_times, canonical_mask, generator=generator
            )
            return obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask, time_scales.clone()
        t_peak, t_half = time_scales[0].item(), time_scales[1].item()

        n_pre = len(self._PEAK_PHASE_MULTIPLIERS)
        n_post = len(self._POST_PEAK_HALF_LIFE_MULTIPLIERS)

        # Uniform samples before peak
        pre_times = torch.rand(n_pre, device=device, dtype=dtype) * t_peak
        # Always include the peak
        peak_time = torch.tensor([t_peak], device=device, dtype=dtype)
        # Uniform samples after peak
        post_times = []
        for mult in self._POST_PEAK_HALF_LIFE_MULTIPLIERS:
            t_end = t_peak + mult * t_half
            t_rand = torch.empty(1, device=device, dtype=dtype).uniform_(t_peak, t_end)
            post_times.append(t_rand)
        post_times = torch.cat(post_times, dim=0)

        # Truncate to canonical capacity
        grid = torch.cat([pre_times, peak_time, post_times], dim=0)
        canonical_cap = self._canonical_grid_capacity()
        grid = grid[:canonical_cap]

        # Map grid to nearest simulation points
        ref_times = full_simulation_times[0]
        idx = torch.cdist(grid[:, None], ref_times[:, None]).argmin(dim=1)
        idx, _ = idx.sort()
        valid_mask = torch.ones_like(idx, dtype=torch.bool)
        idx, valid_mask = self._deduplicate_sorted_indices(idx, valid_mask)
        gather_idx = idx[None, :].expand(batch, -1)
        batch_idx = torch.arange(batch, device=device)[:, None]

        canonical_vals = full_simulation[batch_idx, gather_idx]
        canonical_times = full_simulation_times[batch_idx, gather_idx]
        invalid_slots = ~valid_mask
        if invalid_slots.any():
            canonical_vals[:, invalid_slots] = 0
            canonical_times[:, invalid_slots] = 0

        canonical_mask = valid_mask[None, :].expand(batch, -1).clone()
        canonical_mask = self._drop_non_positive_times_from_mask(canonical_times, canonical_mask)

        obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._assemble_from_canonical(
            canonical_vals, canonical_times, canonical_mask, generator=generator
        )
        return obs_out, obs_time, obs_mask, rem_sim, rem_time, rem_mask, time_scales.clone()

    def generate(
        self,
        full_simulation: Tensor,
        full_simulation_times: Tensor,
        **kwargs,
    ) -> Tuple[
        Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor], Tensor
    ]:
        """Generate PK observations using canonical or randomized schedules.

        With probability ``randomize_prob`` (default 0.5) the method delegates
        to :meth:`_generate_random`; otherwise the deterministic
        :meth:`_generate_raw` path is taken.  Setting the keyword argument
        ``deterministic_only=True`` forces the deterministic branch regardless
        of the random draw.  Both paths require the caller to provide
        ``time_scales`` specifying the peak and half-life.  The method honours
        the ``add_rem`` flag by optionally returning remainder tensors.
        """
        time_scales: Optional[Tensor] = kwargs.get("time_scales")
        if time_scales is None:
            raise ValueError("time_scales must be provided for PKPeakHalfLifeStrategy")

        deterministic_only = kwargs.pop("deterministic_only", False)

        use_random = False
        if not deterministic_only:
            use_random = torch.rand(()) < getattr(self, "randomize_prob", 0.5)

        if use_random:
            obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask, rescaled = self._generate_random(
                full_simulation,
                full_simulation_times,
                time_scales=time_scales,
                generator=kwargs.get("generator"),
            )
        else:
            obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask, rescaled = self._generate_raw(
                full_simulation,
                full_simulation_times,
                **kwargs,
            )

        if not self.observations_config.add_rem:
            rem_sim = rem_time = rem_mask = None

        return obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask, rescaled

    def generate_empirical(
        self,
        empirical_obs: Tensor,
        empirical_times: Tensor,
        empirical_mask: Tensor,
        *,
        generator: Optional[torch.Generator] = None,
    ) -> Tuple[Tensor, Tensor, Tensor, Optional[Tensor], Optional[Tensor], Optional[Tensor]]:
        canonical_vals, canonical_times, canonical_mask = self._align_empirical_to_canonical(
            empirical_obs,
            empirical_times,
            empirical_mask,
        )

        obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask = self._assemble_from_canonical(
            canonical_vals,
            canonical_times,
            canonical_mask,
            generator=generator,
        )

        if not self.observations_config.add_rem:
            rem_sim = rem_time = rem_mask = None

        return obs, obs_time, obs_mask, rem_sim, rem_time, rem_mask


class FixPastTimeRandomSelectionStrategy(ObservationStrategy):
    """Randomly sample observations and split with fixed-capacity past/future slots.

    For ``split_past_future=True`` this strategy enforces the contract:
    ``obs_capacity=max_past`` and ``rem_capacity=max_num_obs-max_past``
    (subject to ``fixed_M_max=min(max_num_obs, time_num_steps)``).
    """

    def __init__(self, config: ObservationsConfig, meta_config: MetaStudyConfig):
        super().__init__(config, meta_config)
        time_steps = getattr(meta_config, "time_num_steps", config.max_num_obs)
        self.fixed_M_max = min(config.max_num_obs, time_steps)
        self.split_past_future = config.split_past_future
        self.max_past = config.max_past
        self.min_past = config.min_past
        self.generative_bias = config.generative_bias
        self.boundary_ratio = getattr(config, "past_time_ratio", 0.1)

    def _generate_raw(self, full_simulation: Tensor, full_simulation_times: Tensor, **kwargs):
        return fix_past_time_random_selection(
            full_simulation=full_simulation,
            full_simulation_times=full_simulation_times,
            boundary_ratio=self.boundary_ratio,
            fixed_M_max=self.fixed_M_max,
            num_obs_sampler=kwargs.get("num_obs_sampler", None),
            generator=kwargs.get("generator", None),
        )

    def _get_shapes_raw(self) -> Tuple[int, int]:
        """Return fixed-capacity shapes for random split outputs.

        With ``split_past_future=True``:
        - ``max_obs`` is bounded by ``max_past``
        - ``max_rem`` is bounded by ``max_num_obs - max_past``
        """
        if self.split_past_future:
            if self.min_past is None or self.max_past is None:
                raise ValueError(
                    "min_past and max_past must be specified when split_past_future=True"
                )
            if self.fixed_M_max < self.min_past:
                raise ValueError("fixed_M_max is smaller than the configured min_past")
            max_obs = min(self.max_past, self.fixed_M_max)
            max_rem = max(0, self.fixed_M_max - self.max_past)
        else:
            max_obs = self.fixed_M_max
            max_rem = self.fixed_M_max

        return max_obs, max_rem

    def _split_by_boundary(
        self,
        obs: TensorType["B", "M"],
        obs_time: TensorType["B", "M"],
        obs_mask: TensorType["B", "M"],
        *,
        generator: Optional[torch.Generator] = None,
    ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
        """Split sampled observations into strict past and future blocks.

        The split is boundary-based and strict:
        - Past block samples ``k`` points from ``time <= boundary`` candidates,
          where ``k`` follows ``min_past``/``max_past`` (and ``generative_bias``),
          capped by available candidates and ``K_max``.
        - When ``k > 0``, remainder receives up to ``R_cap`` points sampled
          from ``time > boundary`` only (strict future).
        - When ``k == 0``, boundary splitting is ignored for remainder and
          points are sampled from all valid candidates.

        Extra past/future candidates are ignored, and missing entries are
        padded by zeros with mask=False.
        """
        B, M = obs.shape
        # K_max: capacity of the past block [B, K_max]
        K_max = min(int(self.max_past), int(M))
        K_min = min(int(self.min_past), K_max)
        # R_cap: fixed capacity of the remainder block [B, R_cap]
        R_cap = max(0, int(M) - K_max)

        boundary = self.meta_config.time_stop * self.boundary_ratio
        gen = generator if generator is not None else torch.default_generator

        past_obs = torch.zeros(B, K_max, dtype=obs.dtype, device=obs.device)
        past_time = torch.zeros_like(past_obs)
        past_mask = torch.zeros(B, K_max, dtype=torch.bool, device=obs.device)

        rem_obs = torch.zeros(B, R_cap, dtype=obs.dtype, device=obs.device)
        rem_time = torch.zeros_like(rem_obs)
        rem_mask = torch.zeros(B, R_cap, dtype=torch.bool, device=obs.device)

        for b in range(B):
            valid_idx = obs_mask[b].nonzero(as_tuple=True)[0]
            past_candidates = valid_idx[obs_time[b, valid_idx] <= boundary]
            future_candidates = valid_idx[obs_time[b, valid_idx] > boundary]

            if past_candidates.numel() > 1:
                order = torch.argsort(obs_time[b, past_candidates])
                past_candidates = past_candidates[order]
            if future_candidates.numel() > 1:
                order = torch.argsort(obs_time[b, future_candidates])
                future_candidates = future_candidates[order]

            # Past is sampled uniformly without replacement from pre-boundary points.
            k_high = min(K_max, int(past_candidates.numel()))
            k_low = min(K_min, k_high)
            k = _sample_past_count_with_bias(
                low=int(k_low),
                high=int(k_high),
                generative_bias=self.generative_bias,
                generator=gen,
                device=obs.device,
            )
            if k > 0 and past_candidates.numel() > 0:
                chosen_offsets = torch.randperm(
                    past_candidates.numel(),
                    generator=gen,
                    device=obs.device,
                )[:k]
                chosen_past = past_candidates[chosen_offsets]
                chosen_order = torch.argsort(obs_time[b, chosen_past])
                chosen_past = chosen_past[chosen_order]
            else:
                chosen_past = past_candidates[:0]

            num_past = chosen_past.numel()
            if num_past > 0:
                past_obs[b, :num_past] = obs[b, chosen_past]
                past_time[b, :num_past] = obs_time[b, chosen_past]
                past_mask[b, :num_past] = True

            # If no past point is selected, allow remainder sampling across the
            # whole valid domain. Otherwise keep strict future-only remainder.
            rem_pool = valid_idx if num_past == 0 else future_candidates
            if rem_pool.numel() > 1:
                order = torch.argsort(obs_time[b, rem_pool])
                rem_pool = rem_pool[order]

            if R_cap <= 0 or rem_pool.numel() == 0:
                chosen_rem = rem_pool[:0]
            elif rem_pool.numel() <= R_cap:
                chosen_rem = rem_pool
            else:
                chosen_offsets = torch.randperm(
                    rem_pool.numel(),
                    generator=gen,
                    device=obs.device,
                )[:R_cap]
                chosen_rem = rem_pool[chosen_offsets]
                chosen_order = torch.argsort(obs_time[b, chosen_rem])
                chosen_rem = chosen_rem[chosen_order]

            r = chosen_rem.numel()
            if r > 0:
                rem_obs[b, :r] = obs[b, chosen_rem]
                rem_time[b, :r] = obs_time[b, chosen_rem]
                rem_mask[b, :r] = True

        return past_obs, past_time, past_mask, rem_obs, rem_time, rem_mask

    def generate(
        self, full_simulation: Tensor, full_simulation_times: Tensor, **kwargs
    ) -> Tuple[Tensor, ...]:
        obs, obs_time, obs_mask, _, _, _ = self._generate_raw(
            full_simulation, full_simulation_times, **kwargs
        )
        obs_mask = self._drop_non_positive_times_from_mask(obs_time, obs_mask)

        if self.split_past_future:
            out = self._split_by_boundary(
                obs,
                obs_time,
                obs_mask,
                generator=kwargs.get("generator", None),
            )
        else:
            past_obs, past_time, past_mask = obs, obs_time, obs_mask
            rem_obs = rem_time = rem_mask = None
            out = (past_obs, past_time, past_mask, rem_obs, rem_time, rem_mask)

        if not self.observations_config.add_rem:
            out = out[:3] + (None, None, None)

        return (*out, None)


class ObservationStrategyFactory:
    @staticmethod
    def from_config(
        obs_config: ObservationsConfig, meta_config: MetaStudyConfig
    ) -> ObservationStrategy:
        # Legacy compatibility:
        # - omitted ``type`` defaults via dataclass to ``pk_peak_half_life``
        # - explicit YAML ``type: null`` is loaded as ``None`` and also falls
        #   back to ``pk_peak_half_life``
        strategy_type = getattr(obs_config, "type", None)
        if strategy_type is None:
            normalized_type = "pk_peak_half_life"
        elif isinstance(strategy_type, str):
            stripped = strategy_type.strip()
            if stripped == "" or stripped.lower() in {"null", "none"}:
                normalized_type = "pk_peak_half_life"
            else:
                normalized_type = stripped.lower()
        else:
            normalized_type = str(strategy_type).strip().lower()

        if normalized_type in {
            "observations_pk_peak_halflife",
            "pk_peak_half_life",
        }:
            return PKPeakHalfLifeStrategy(obs_config, meta_config)
        if normalized_type in {
            "fix_past_time_random_selection",
            "random",
        }:
            return FixPastTimeRandomSelectionStrategy(obs_config, meta_config)
        raise ValueError(f"Unknown observation type: {strategy_type}")