Processors¶

class AbstractProcessor(abc.ABC):
    #: Override in subclasses to set a canonical parquet filename stem (e.g. ``"sites"``).
    #: When ``None`` (the default), ``output_name`` falls back to a snake_case of the class name.
    __output_name__: ty.ClassVar[str | None] = None

    @property
    def output_name(self) -> str:
        """Canonical name used as the parquet filename stem.

        Returns ``__output_name__`` if defined on the class, otherwise a
        snake_case of the class name (e.g. ``LicksProcessor`` → ``licks_processor``).
        """
        return self.__class__.__output_name__ or _class_name_to_snake(type(self).__name__)

    def __init__(self, dataset: Dataset, *, strict_parsing: bool = False) -> None:
        self._dataset = dataset
        self._strict_parsing = strict_parsing

    @property
    def dataset(self) -> Dataset:
        return self._dataset

    @cached_property
    def provenance(self) -> PackagingProvenance:
        """Provenance snapshot for this processor's dataset.

        Cached so that :meth:`compute` and version-check code in subclasses
        share a single :class:`~aind_behavior_vr_foraging_packaging._provenance.PackagingProvenance`
        instance rather than rebuilding it on every call.
        """

        return PackagingProvenance.build(self._dataset)

    @abc.abstractmethod
    def _compute(self) -> pd.DataFrame:
        """Compute this processor's output as a DataFrame.

        Subclasses implement this method. Callers should use :meth:`compute`,
        which wraps ``_compute`` and stamps provenance metadata into ``df.attrs``.
        """
        raise NotImplementedError

    def compute(self) -> pd.DataFrame:
        """Return the processor's output DataFrame with provenance metadata in attrs.

        Calls :meth:`_compute`, then stamps ``df.attrs`` with the session-level
        provenance keys from :class:`~aind_behavior_vr_foraging_packaging._provenance.PackagingProvenance`
        plus a processor-specific ``processor`` key (this class's name).

        Attrs already set by ``_compute`` (e.g. ``sampling_rate_hz`` from
        :class:`SniffingProcessor`) are preserved via ``setdefault``.
        """
        df = self._compute()
        for k, v in self.provenance.model_dump().items():
            df.attrs.setdefault(k, v)
        df.attrs.setdefault("processor", type(self).__name__)
        return df

    def nwbize(self, nwb_file: ty.Any) -> ty.Any:
        """Write this processor's output to *nwb_file* and return it.

        Default implementation is a no-op. Override in subclasses that have
        an NWB representation. May call ``compute()`` internally; the two
        methods are intentionally independent (no shared state).

        That independence costs a second full ``_compute()`` per session when
        both outputs are written (``--write-nwb``). Processors for which that
        is expensive decorate ``_compute`` with
        :func:`~aind_behavior_vr_foraging_packaging._base.cached_frame`, which
        removes the recomputation while preserving the no-shared-state
        guarantee — every call still hands back its own copy.
        """
        return nwb_file

    def with_strict_parsing(self, strict_parsing: bool = True) -> ty.Self:
        self._strict_parsing = strict_parsing
        return self

    @property
    def strict_parsing(self) -> bool:
        """Whether *known* data anomalies raise instead of being logged and worked around.

        The flag covers only anomalies a processor explicitly checks for and can name,
        and only where a degraded-but-meaningful output exists. The convention is::

            if <specific condition detected>:
                msg = "<what was violated>"
                if self.strict_parsing:
                    raise DatasetProcessorError(msg)
                logger.warning("%s; <what is used instead>.", msg)

        It does **not** gate general exceptions. Never write
        ``except Exception: ... if self.strict_parsing: raise`` — with the flag off (the
        default) that swallows real bugs, API drift and corrupt files as though they were
        data quirks, dropping the processor's output while the run still reports success.
        Catch only the exception types that signal an *expected* condition, narrowly —
        e.g. ``except (KeyError, FileNotFoundError)`` for a stream a given schema version
        does not declare — and let everything else propagate.

        Failures that leave nothing meaningful to emit (e.g. absent treadmill calibration,
        without which position cannot be computed at all) should raise unconditionally
        rather than consult this flag: there is no degraded output to fall back to.

        Isolating one failure from the rest of a run is the caller's job, not the flag's.
        :func:`~aind_behavior_vr_foraging_packaging.pipeline.batch.process_sessions`
        catches whatever a processor raises, so a single bad session or processor never
        aborts a batch.
        """
        return self._strict_parsing

class PackagingProvenance(BaseModel):
    """Immutable provenance snapshot for one packaging run.

    All version strings are validated as semver-compatible on construction, so
    any call site is guaranteed a well-formed object.

    Attributes
    ----------
    packaging_version:
        Version of this package (``aind-behavior-vr-foraging-packaging``).
    data_contract_version:
        Version of ``aind-behavior-vr-foraging`` (the behavioural schema library).
    dataset_version:
        Version recorded in the session's ``tasklogic_input.json``.
    """

    model_config = ConfigDict(frozen=True)

    packaging_version: str
    data_contract_version: str
    dataset_version: str

    @property
    def dataset_semver(self) -> semver.Version:
        """Dataset version as a parsed :class:`semver.Version`."""
        return semver.Version.parse(self.dataset_version)

    @property
    def data_contract_semver(self) -> semver.Version:
        """Data-contract (parser) version as a parsed :class:`semver.Version`."""
        return semver.Version.parse(self.data_contract_version)

    @classmethod
    def build(cls, dataset: Dataset) -> "PackagingProvenance":
        """Construct a :class:`PackagingProvenance` from a live dataset."""
        return cls(
            packaging_version=importlib.metadata.version(_PACKAGING_PKG),
            data_contract_version=aind_behavior_vr_foraging.__semver__,
            dataset_version=str(dataset.version),
        )

class SiteTableProcessor(AbstractProcessor):
    __output_name__ = "sites"

    def __init__(self, dataset: contraqctor.contract.Dataset, *, strict_parsing: bool = False) -> None:
        super().__init__(dataset, strict_parsing=strict_parsing)

        if self.provenance.dataset_semver != self.provenance.data_contract_semver:
            logger.warning(
                "Dataset version %s does not match parser version %s",
                self.provenance.dataset_semver,
                self.provenance.data_contract_semver,
            )
        self.rig_configuration = self._ensure_json_not_pydantic(
            self.dataset["Behavior"]["InputSchemas"]["Rig"].load().data
        )

    @staticmethod
    def _ensure_json_not_pydantic(d: t.Any) -> dict:
        if isinstance(d, BaseModel):
            return d.model_dump()
        return d

    @staticmethod
    def _parse_speaker_choice_feedback(dataset: contraqctor.contract.Dataset) -> pd.DataFrame:
        speaker_choice = dataset.at("Behavior").at("HarpBehavior").load().at("PwmStart").load().data
        speaker_choice = speaker_choice[(speaker_choice["MessageType"] == "WRITE") & (speaker_choice["PwmDO2"])]
        return speaker_choice

    @staticmethod
    def _parse_odor_onset(dataset: contraqctor.contract.Dataset) -> pd.Series:
        odor_onset = dataset.at("Behavior").at("HarpOlfactometer").load().at("EndValveState").load().data
        odor_onset = odor_onset[odor_onset["MessageType"] == "WRITE"]["EndValve0"]
        odor_onset = odor_onset[(odor_onset) & (~odor_onset.shift(1, fill_value=False))]
        return odor_onset

    @staticmethod
    def _parse_continuous_patch_state(dataset: contraqctor.contract.Dataset) -> pd.DataFrame:
        patches_state = dataset.at("Behavior").at("SoftwareEvents").at("PatchState").load().data
        expanded = pd.json_normalize(patches_state["data"])
        expanded.index = patches_state.index
        patches_state = patches_state.join(expanded)
        return patches_state

    def _parse_patch_state_at_reward(self, dataset: contraqctor.contract.Dataset) -> pd.DataFrame:
        if self.provenance.dataset_semver < semver.Version(major=0, minor=6, patch=0):
            raise DatasetProcessorError("PatchStateAtReward is only available in dataset version 0.6.0 and above")
        # TODO this is likely something we want to overload for 0.5.x to work.
        patches_state_at_reward = dataset.at("Behavior").at("SoftwareEvents").at("PatchStateAtReward").load().data
        expanded = pd.json_normalize(patches_state_at_reward["data"])
        expanded.index = patches_state_at_reward.index
        patches_state_at_reward = patches_state_at_reward.join(expanded)
        return patches_state_at_reward

    @staticmethod
    def _parse_wait_reward_outcome(dataset: contraqctor.contract.Dataset) -> pd.Series:
        try:
            return dataset.at("Behavior").at("SoftwareEvents").at("WaitRewardOutcome").load().data
        except FileNotFoundError:
            return pd.Series(dtype=bool)

    @staticmethod
    def _as_dict(d: contraqctor.contract.DataStream | PydanticModel | BaseModel | dict) -> dict:
        if isinstance(d, (PydanticModel, contraqctor.contract.DataStream)):
            d = t.cast(BaseModel | dict, d.data)
        if isinstance(d, dict):
            return d
        if isinstance(d, BaseModel):
            return d.model_dump()
        else:
            raise TypeError(f"Cannot convert type {type(d)} to dict")

    @staticmethod
    def _parse_friction(dataset: contraqctor.contract.Dataset) -> pd.Series:
        d = dataset.at("Behavior").at("HarpTreadmill").at("BrakeCurrentSetPoint").load().data
        return d.loc[d["MessageType"] == "WRITE", "BrakeCurrentSetPoint"]

    @staticmethod
    def _parse_is_stopped(dataset: contraqctor.contract.Dataset) -> pd.DataFrame | None:
        return dataset.at("Behavior").at("OperationControl").at("IsStopped").load().data

    def _parse_velocity(self, dataset: contraqctor.contract.Dataset) -> pd.Series | None:
        rig_config = self._ensure_json_not_pydantic(self.rig_configuration)
        return PositionAndVelocityProcessor.compute_position_and_velocity_from_treadmill(dataset, rig_config)[
            "velocity"
        ]

    def _get_olfactometer_channel_count(self, dataset: contraqctor.contract.Dataset) -> int:
        extra_olfs = getattr(self.rig_configuration, "harp_olfactometer_extension", None)
        n_extra_channels = 4 * len(extra_olfs) if extra_olfs is not None else 0
        return (
            3 + n_extra_channels
        )  # The channel 3 is always used as carrier, therefore only 3 odor channels are available.

    def _process_odor_concentration(self, odor_specification: BaseModel | dict | None, n_channels: int) -> list[float]:

        concentration = [0.0] * n_channels
        if odor_specification is None:
            return concentration
        odor_specification = self._ensure_json_not_pydantic(odor_specification)
        return TypeAdapter(OdorMixture).validate_python(odor_specification)

    @staticmethod
    def _load_blocks(dataset: contraqctor.contract.Dataset) -> pd.DataFrame:
        blocks = t.cast(pd.DataFrame, dataset.at("Behavior").at("SoftwareEvents").at("Block").load().data)
        blocks["block_count"] = range(len(blocks))
        return blocks

    def _select_patch_state_at_reward(
        self, candidates: pd.DataFrame, *, choice_time: float, site_start: float, site_stop: float
    ) -> pd.DataFrame:
        """Narrow the in-site patch-reward states down to the one belonging to this site.

        From dataset 0.6.0 onwards the experimenter's manual rewards are logged separately, as
        ``ForceGiveReward``, so only the site's own earned reward can land in the interval and
        more than one row means the parse is wrong. Overridden by
        :class:`~._legacy_site_table.LegacySiteTableProcessor`, where manual and earned rewards
        share a single logging path and cannot be told apart from the record alone.
        """
        assert len(candidates) <= 1, (
            f"Site interval [{site_start}, {site_stop}) booked {len(candidates)} rewards to this patch, "
            "expected at most one. A site can only be rewarded once, and from dataset 0.6.0 the "
            "experimenter's manual water is logged separately as ForceGiveReward, so it cannot account "
            f"for the extras. Reward times: {list(candidates.index)}"
        )
        return candidates

    def _select_reward_onset_time(
        self,
        site_water_delivery: pd.Series,
        reward_metadata_sliced: pd.DataFrame,
        *,
        site_start: float,
        choice_time: float,
    ) -> float:
        """Pick the valve opening that delivered this site's earned reward.

        ``site_water_delivery`` is guaranteed non-empty by the caller. See
        :meth:`_select_patch_state_at_reward` for why the legacy processor overrides this.
        """
        if len(reward_metadata_sliced) > 1:
            closest_index = site_water_delivery.index.get_indexer(pd.Index([site_start]), method="nearest")[0]
            return t.cast(float, site_water_delivery.index[closest_index])
        return t.cast(float, site_water_delivery.index[0])

    def _select_choice_feedback(
        self,
        candidates: pd.DataFrame,
        *,
        site_index: int,
        t_start: float,
        t_end: float,
    ) -> pd.DataFrame:
        """Validate and return the choice-feedback events for a single site interval.

        The base implementation enforces that at most one choice-cue event falls inside
        the interval — a violation indicates a parse or hardware problem and raises.
        Override in subclasses (e.g. :class:`~._legacy_site_table.LegacySiteTableProcessor`)
        to handle known hardware quirks gracefully.
        """
        assert len(candidates) <= 1, (
            f"Multiple speaker choices in site interval [{t_start}, {t_end}); "
            f"site index {site_index}, timestamps: {candidates.index.tolist()}"
        )
        return candidates

    def process_to_sites(self) -> list[Site]:
        """
        Processes sites, patches, and blocks from the dataset and merges them.
        Returns a DataFrame with merged information.
        """
        dataset = self.dataset
        odor_sites = t.cast(pd.DataFrame, dataset.at("Behavior").at("SoftwareEvents").at("ActiveSite").load().data)
        patches = t.cast(pd.DataFrame, dataset.at("Behavior").at("SoftwareEvents").at("ActivePatch").load().data)
        patches["patch_count"] = range(len(patches))
        blocks = self._load_blocks(dataset)

        # Merge nearest patch (backward in time)
        merged = pd.merge_asof(
            odor_sites.sort_index(),
            patches[["data", "patch_count"]].rename(columns={"data": "patch_data"}).sort_index(),
            left_index=True,
            right_index=True,
            direction="backward",
            suffixes=("", "_patch"),
        )
        merged["patch_index"] = merged["patch_data"].apply(lambda d: d["state_index"])

        # Merge nearest block (backward in time)
        merged = pd.merge_asof(
            merged.sort_index(),
            blocks[["block_count"]].sort_index(),
            left_index=True,
            right_index=True,
            direction="backward",
        )

        choice_feedback = self._parse_speaker_choice_feedback(dataset)
        water_delivery = parse_water_delivery(dataset)
        reward_metadata = parse_reward_metadata(dataset)

        # In theory we only need the metadata, but since we are aligning
        # temporally later, we must have access to the hardware-aligned times.
        manual_water_delivery = parse_manual_water_delivery(dataset)

        odor_onset = self._parse_odor_onset(dataset)
        patch_state_at_reward = self._parse_patch_state_at_reward(dataset)
        friction = self._parse_friction(dataset)
        olfactometer_channel_count = self._get_olfactometer_channel_count(dataset)
        wait_reward_outcome = self._parse_wait_reward_outcome(dataset)
        is_stopped = self._parse_is_stopped(dataset)
        velocity = self._parse_velocity(dataset)

        # Precompute all site indices
        merged["site_label"] = merged["data"].apply(lambda d: d["label"])
        merged["patch_label"] = merged["patch_data"].apply(lambda d: d["label"])

        # Site-level indices
        merged["_site_index_in_patch"] = merged.groupby("patch_count").cumcount()
        merged["_site_index_in_block"] = merged.groupby("block_count").cumcount()
        merged["_site_index_by_type"] = merged.groupby("site_label").cumcount()
        merged["_site_index_in_patch_by_type"] = merged.groupby(["patch_count", "site_label"]).cumcount()
        merged["_site_index_in_block_by_type"] = merged.groupby(["block_count", "site_label"]).cumcount()

        # Patch-level indices (computed on patches, then mapped back to sites via patch_count)
        patches_with_blocks = pd.merge_asof(
            patches.sort_index(),
            blocks[["block_count"]].sort_index(),
            left_index=True,
            right_index=True,
            direction="backward",
        )
        patches_with_blocks["patch_label"] = patches_with_blocks["data"].apply(lambda d: d["label"])
        patches_with_blocks["_patch_index_in_block"] = patches_with_blocks.groupby("block_count").cumcount()
        patches_with_blocks["_patch_index_by_type"] = patches_with_blocks.groupby("patch_label").cumcount()
        patches_with_blocks["_patch_index_in_block_by_type"] = patches_with_blocks.groupby(
            ["block_count", "patch_label"]
        ).cumcount()
        merged = merged.join(
            patches_with_blocks.set_index("patch_count")[
                ["_patch_index_in_block", "_patch_index_by_type", "_patch_index_in_block_by_type"]
            ],
            on="patch_count",
        )

        # Only mutable states that requires trial-based
        current_friction = 0  # Keeps track of the last known friction. Sites with null friction will not update this.

        sites: list[Site] = []
        # We reject the last site because it may not have completed and would require custom logic to handle
        for i in range(len(merged) - 1):
            # We generally assume that all relevant events happen within the software-event derived timestamp intervals
            # Note this may not always be true depending on system jitter, but it is generally a safe assumption.
            # If you find edge cases where this is not true, submit an issue so we can investigate and improve the parser.

            this_timestamp = t.cast(float, merged.index[i])
            next_timestamp = t.cast(float, merged.index[i + 1])

            this_site = merged.iloc[i]["data"]
            this_patch = merged.iloc[i]["patch_data"]

            site_choice_feedback = slice_by_index(choice_feedback, this_timestamp, next_timestamp)
            site_choice_feedback = self._select_choice_feedback(
                site_choice_feedback, site_index=i, t_start=this_timestamp, t_end=next_timestamp
            )

            choice_time: float = (
                t.cast(float, site_choice_feedback.index[0]) if not site_choice_feedback.empty else np.nan
            )

            site_odor_onset = slice_by_index(odor_onset, this_timestamp, next_timestamp)
            site_force_reward = slice_by_index(manual_water_delivery, this_timestamp, next_timestamp)

            this_friction = slice_by_index(friction, this_timestamp, next_timestamp)
            if not this_friction.empty:
                current_friction = this_friction.values[-1]

            site_patch_state_at_reward = slice_by_index(patch_state_at_reward, this_timestamp, next_timestamp)
            site_patch_state_at_reward = site_patch_state_at_reward[
                site_patch_state_at_reward["PatchId"] == merged.iloc[i]["patch_index"]
            ]
            site_patch_state_at_reward = self._select_patch_state_at_reward(
                site_patch_state_at_reward,
                choice_time=choice_time,
                site_start=this_timestamp,
                site_stop=next_timestamp,
            )

            ##
            row = merged.iloc[i]

            # Compute last_stop_time, last_stop_duration, velocity_at_last_stop
            # We skip calculation if no choice_time was found or IsStopped data is unavailable
            site_stop_time: float = np.nan
            site_stop_duration: float = np.nan
            site_velocity_at_stop: float = np.nan
            if not np.isnan(choice_time) and is_stopped is not None:
                site_is_stopped = slice_by_index(is_stopped, this_timestamp, choice_time, end_inclusive=True)
                stops_before_choice = site_is_stopped[site_is_stopped["IsStopped"]]
                if stops_before_choice.empty:
                    msg = f"Choice occurred at {choice_time} but no IsStopped=True event found in site interval [{this_timestamp}, {next_timestamp})"
                    if self.strict_parsing:
                        raise DatasetProcessorError(msg)
                    else:
                        logger.warning(msg + ". Falling back to global search.")
                        stops_before_choice = is_stopped[is_stopped["IsStopped"] & (is_stopped.index <= choice_time)]
                    if stops_before_choice.empty:
                        raise DatasetProcessorError(
                            f"Choice occurred at {choice_time} but no IsStopped=True event found before choice time"
                        )

                site_stop_time = (
                    t.cast(float, stops_before_choice.index[-1]) if not stops_before_choice.empty else np.nan
                )
                site_stop_duration = choice_time - site_stop_time
                if velocity is not None:
                    closest_ts = get_closest_from_timestamp(np.array([site_stop_time]), velocity, search_mode="closest")
                    site_velocity_at_stop = float(velocity[closest_ts[0]])

            if site_odor_onset.empty and this_site["odor_specification"] is not None:
                # Sometimes the timestamp for the odor onset arrives slightly before the site. We should investigate
                # but for now we just log a warning and use the site onset instead after checking if this is the issue
                odor_onset_before_site = odor_onset[
                    (odor_onset.index < this_timestamp) & (odor_onset.index >= this_timestamp - 0.002)
                ]  # we use a 2ms conservative window
                if odor_onset_before_site.empty:
                    if self.strict_parsing:
                        raise DatasetProcessorError("No odor onset found in site interval")
                    else:
                        logger.warning("No odor onset found in site interval")
                        odor_onset_time = np.nan
                else:
                    logger.warning("Odor onset found slightly (<2ms) before site interval, using site onset instead")
                    odor_onset_time = this_timestamp
            else:
                # we always take the first odor onset in case animal goes in and out
                odor_onset_time = t.cast(float, site_odor_onset.index[0]) if not site_odor_onset.empty else np.nan

            site_water_delivery = slice_by_index(water_delivery, this_timestamp, next_timestamp)
            reward_metadata_sliced = slice_by_index(reward_metadata, this_timestamp, next_timestamp)
            if reward_metadata_sliced.empty or bool(reward_metadata_sliced["data"].fillna(0).eq(0).all()):
                # Note: for None or 0 reward metadata there won't be a hardware water delivery event
                # However, if the experimenter manually triggered a reward around this time, we should not count that
                # as a reward for this site either, so we make an explicit decision to set reward_onset_time to nan
                reward_onset_time = np.nan
            else:
                if len(site_water_delivery) == 0:
                    if self.strict_parsing:
                        raise DatasetProcessorError(
                            "Valid reward metadata found but no water delivery in site interval"
                        )
                    else:
                        logger.error("Valid reward metadata found but no water delivery in site interval")
                        reward_onset_time = np.nan
                else:
                    reward_onset_time = self._select_reward_onset_time(
                        site_water_delivery,
                        reward_metadata_sliced,
                        site_start=this_timestamp,
                        choice_time=choice_time,
                    )

            wait_reward_outcome_sliced = slice_by_index(wait_reward_outcome, this_timestamp, next_timestamp)
            has_waited_reward_delay = (
                wait_reward_outcome_sliced.iloc[0]["data"]["IsSuccessfulWait"]
                if not wait_reward_outcome_sliced.empty
                else None
            )

            site = Site(
                start_time=this_timestamp,
                stop_time=next_timestamp,
                start_position=this_site["start_position"],
                length=this_site["length"],
                site_label=str(this_site["label"]),
                friction=current_friction,
                patch_label=str(this_patch["label"]),
                odor_concentration=self._process_odor_concentration(
                    this_patch["odor_specification"], olfactometer_channel_count
                ),
                patch_index=row["patch_count"],
                patch_index_in_block=row["_patch_index_in_block"],
                patch_index_by_type=row["_patch_index_by_type"],
                patch_index_in_block_by_type=row["_patch_index_in_block_by_type"],
                site_index=i,
                site_index_in_patch=row["_site_index_in_patch"],
                site_index_in_block=row["_site_index_in_block"],
                site_index_by_type=row["_site_index_by_type"],
                site_index_in_patch_by_type=row["_site_index_in_patch_by_type"],
                site_index_in_block_by_type=row["_site_index_in_block_by_type"],
                odor_onset_time=odor_onset_time,
                reward_onset_time=reward_onset_time,
                reward_amount=np.nan
                if site_patch_state_at_reward.empty
                else site_patch_state_at_reward.iloc[0]["Amount"],
                reward_probability=np.nan
                if site_patch_state_at_reward.empty
                else site_patch_state_at_reward.iloc[0]["Probability"],
                reward_available=np.nan
                if site_patch_state_at_reward.empty
                else site_patch_state_at_reward.iloc[0]["Available"],
                has_reward=not np.isnan(reward_onset_time),
                has_forced_rewards=not site_force_reward.empty,
                choice_cue_time=choice_time,
                has_choice=not site_choice_feedback.empty,
                reward_delay_duration=reward_onset_time - choice_time
                if reward_onset_time is not np.nan and choice_time is not None
                else np.nan,
                has_waited_reward_delay=has_waited_reward_delay,
                last_stop_time=None if np.isnan(site_stop_time) else site_stop_time,
                last_stop_duration=None if np.isnan(site_stop_duration) else site_stop_duration,
                velocity_at_last_stop=None if np.isnan(site_velocity_at_stop) else site_velocity_at_stop,
                block_index=row["block_count"],
            )
            sites.append(site)
        return sites

    @cached_frame
    def _compute(self) -> pd.DataFrame:
        """Returns site table as a DataFrame with one row per site."""
        sites = self.process_to_sites()
        return pd.DataFrame([s.model_dump() for s in sites])

    def nwbize(self, nwb_file: t.Any) -> t.Any:
        """Add sites to *nwb_file* from compute() output."""
        df = self.compute()
        for col in df.columns:
            if col in ("start_time", "stop_time"):
                continue
            nwb_file.add_trial_column(name=col, description=col)
        for _, row in df.iterrows():
            trial = {k: (np.nan if v is None else v) for k, v in row.to_dict().items()}
            nwb_file.add_trial(**trial)
        return nwb_file