from typing import List, Optional

from allensdk.brain_observatory.behavior.behavior_session import (
    BehaviorSession,
)
from allensdk.brain_observatory.nwb.nwb_utils import NWBWriter


class OphysNwbWriter(NWBWriter):
    """NWBWriter with additional options to modify targeted_imaging_depth."""

    def _update_session(
        self,
        lims_session: BehaviorSession,
        ophys_experiment_ids: Optional[List[int]] = None,
    ) -> BehaviorSession:
        """Call session methods to update certain values within the session.

        Should be used as part of a datarelease to resolve known data issues.

        Parameters
        ----------
        lims_session : BehaviorSession
            Input behavior session
        ophys_experiment_ids : list of int
            Subset of experiment_ids that are to be released and are in the
            same container with the experiment we are creating an NWB for.

        Returns
        -------
        modified_session : BehaviorSession
            Modified version of input session.
        """
        if ophys_experiment_ids is not None:
            lims_session.update_targeted_imaging_depth(ophys_experiment_ids)
        return lims_session
