Instrument management (async)

async def abort(self) -> None:
    """Abort measurement."""
    async with self._lock():
        await create_future(self._comm.AbortAsync())

async def connect(self) -> None:
    """Connect to instrument."""
    if self.is_connected():
        return

    self._comm = await self.instrument._connect_async()

    # Disable idle messages to improve response time and reduce noise
    await create_future(self._comm.SetStatusWhenIdleAsync(False))

    firmware_warning(self._comm.Capabilities)

async def disconnect(self):
    """Disconnect from the instrument."""
    if not self.is_connected():
        return

    await create_future(self._comm.DisconnectAsync())
    self._comm.Dispose()
    del self._comm

def ensure_connection(self):
    """Raises connection error if the instrument is not connected."""
    if not self.is_connected():
        raise ConnectionError('Not connected to an instrument')

async def get_current_range(self) -> AllowedCurrentRanges:
    """Get the current range for the cell.

    Returns
    -------
    current_range: AllowedCurrentRanges
    """
    async with self._lock():
        value: PalmSens.CurrentRange = await create_future(
            self._comm.GetCurrentRangeAsync()
        )

    return cr_enum_to_string(value)

def get_estimated_duration(
    self: HasCommProtocol,
    method: PalmSens.Method | MethodTypeCompatible,
) -> float:
    """Get the estimated duration for this method.

    Parameters
    -----------
    method : MethodType
        The method to get the estimated duration for.

    Returns
    -------
    float
        Estimated duration in seconds.
    """
    self.ensure_connection()

    if not isinstance(method, PalmSens.Method):
        method = method._to_psmethod()

    capabilities = self._comm.Capabilities

    return method.GetMinimumEstimatedMeasurementDuration(capabilities)

async def get_instrument_serial(self) -> str:
    """Return instrument serial number.

    Returns
    -------
    serial : str
        Instrument serial.
    """
    async with self._lock():
        serial: PalmSens.Comm.DeviceSerialV3 = await create_future(
            self._comm.GetDeviceSerialAsync()
        )

    return serial.ToString()

async def get_potential_range(self) -> AllowedPotentialRanges:
    """Get the potential range for the cell.

    Returns
    -------
    potential_range: AllowedPotentialRanges
    """
    async with self._lock():
        # no such api: self._comm.GetPotentialRangeAsync()
        return pr_enum_to_string(self._comm.PotentialRange)

async def initialize_multiplexer(self, mux_model: int) -> int:
    """Initialize the multiplexer.

    Parameters
    ----------
    mux_model: int
        The model of the multiplexer.
        - 0 = 8 channel
        - 1 = 16 channel
        - 2 = 32 channel

    Returns
    -------
    channels : int
        Number of available multiplexes channels
    """
    async with self._lock():
        model = PalmSens.MuxModel(mux_model)

        if model == PalmSens.MuxModel.MUX8R2 and (
            self._comm.ClientConnection.GetType().Equals(
                clr.GetClrType(PalmSens.Comm.ClientConnectionPS4)
            )
            or self._comm.ClientConnection.GetType().Equals(
                clr.GetClrType(PalmSens.Comm.ClientConnectionMS)
            )
        ):
            await create_future(self._comm.ClientConnection.ReadMuxInfoAsync())

        self._comm.Capabilities.MuxModel = model

        if self._comm.Capabilities.MuxModel == PalmSens.MuxModel.MUX8:
            self._comm.Capabilities.NumMuxChannels = 8
        elif self._comm.Capabilities.MuxModel == PalmSens.MuxModel.MUX16:
            self._comm.Capabilities.NumMuxChannels = 16
        elif self._comm.Capabilities.MuxModel == PalmSens.MuxModel.MUX8R2:
            await create_future(self._comm.ClientConnection.ReadMuxInfoAsync())

    channels = self._comm.Capabilities.NumMuxChannels
    return channels

async def is_cell_on(self) -> bool:
    """Get cell status.

    Returns
    -------
    cell_on : bool
        Return true if the cell is on
    """
    async with self._lock():
        cell_on: bool = await create_future(self._comm.GetCellOnAsync())

    return cell_on

def is_connected(self) -> bool:
    """Return True if an instrument connection exists."""
    return hasattr(self, '_comm')

def is_measuring(self) -> bool:
    """Return True if device is measuring."""
    return int(self._comm.State) == CommManager.DeviceState.Measurement

async def measure(
    self,
    method: MethodTypeCompatible,
    *,
    callback: Callback | CallbackEIS | None = None,
    stream: Path | str | None = None,
    sync_event: asyncio.Event | None = None,
):
    """Start measurement using given method parameters.

    Parameters
    ----------
    method: MethodType
        Method parameters for measurement.
    callback: Callback, optional
        If specified, call this function on every new set of data points.
        New data points are batched, and contain all points since the last
        time it was called. Each point is an instance of `ps.data.CallbackData`
        for non-impedimetric or `ps.data.CallbackDataEIS`.
        for impedimetric measurments.
    stream: Path | str | None
        If defined, stream data directly to this file in JSON Lines text format
        (https://jsonlines.org). This option is useful for long-term measurements.
        In case of a PC crash or power outage, the most recent measurement data will
        still be available.
    sync_event: asyncio.Event
        Event for hardware synchronization. Do not use directly.
        Instead, initiate hardware sync via `InstrumentPoolAsync.measure()`.
    """
    self.ensure_connection()
    self.validate_method(method)  # type: ignore

    measurement_manager = MeasurementManagerAsync(comm=self._comm)

    return await measurement_manager.measure(
        method,
        callback=callback,
        stream=stream,
        sync_event=sync_event,
        listeners=self._listeners,
    )

def on(
    self,
    event: AllowedEvents,
    callback: Callable[..., None],
) -> EventHandle:
    """Register a callback to invoke for the specified event.

    Parameters
    ----------
    event : AllowedEvents
        Name of the event to subscribe to. For most events this simply
        appends callback to the internal listeners dictionary.
    callback : Callable[..., None]
        Function that will be invoked when the event is triggered.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.

    Notes
    -----
    * ``'receive_message'`` and ``'receive_status'`` are handled specially:
      they create instances of :class:`EventHandleReceiveMessage` or
      :class:`EventHandleStatus`, respectively, which attach the callback
      directly to the underlying communication layer.
    """
    if event == 'receive_message':
        return EventHandleReceiveMessage(emitter=self, event=event, callback=callback)

    elif event == 'receive_status':
        return EventHandleStatus(emitter=self, event=event, callback=callback)

    self._listeners[event].append(callback)
    return EventHandle(emitter=self, event=event, callback=callback)

def on_curve_begin(self, callback: Callable[[Curve], None]) -> EventHandle:
    """Register a callback to invoke at the start of a new curve.

    For EIS use `on_eis_data_start`.

    Parameters
    ----------
    callback : Callable[[Curve]]
        The function to call when triggered.
        Passes [pypalmsens.data.Curve][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('curve_begin', callback=callback)

def on_curve_end(self, callback: Callable[[Curve], None]) -> EventHandle:
    """Register a callback to invoke at the end of a curve.

    For EIS use `on_eis_data_end`.

    Parameters
    ----------
    callback : Callable[[Curve]]
        The function to call when triggered.
        Passes [pypalmsens.data.Curve][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('curve_end', callback=callback)

def on_curve_new_data(self, callback: Callable[[CallbackData], None]) -> EventHandle:
    """Register a callback to invoke when new data are received

    Note that the data are batched depending on available resources.

    For EIS use `on_eis_new_data`.

    Parameters
    ----------
    callback : Callable[[CallbackData]]
        The function to call when triggered.
        Passes [pypalmsens.data.CallbackData][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('curve_new_data', callback=callback)

def on_eis_data_begin(self, callback: Callable[[EISData], None]) -> EventHandle:
    """Register a callback to invoke at the start of a new EIS data set.

    Parameters
    ----------
    callback : Callable[[EISData]]
        The function to call when triggered.
        Passes [pypalmsens.data.EISData][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('eis_data_begin', callback=callback)

def on_eis_data_end(self, callback: Callable[[], None]) -> EventHandle:
    """Register a callback to invoke at the end of an EIS data set.

    Parameters
    ----------
    callback : Callable
        The function to call when triggered.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('eis_data_end', callback=callback)

def on_eis_new_data(self, callback: Callable[[CallbackDataEIS], None]) -> EventHandle:
    """Register a callback to invoke when new eis data are received.

    Data points are batched depending on available resources.

    Parameters
    ----------
    callback : Callable[[CallbackDataEIS]]
        The function to call when triggered.
        Passes [pypalmsens.data.CallbackDataEIS][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('eis_new_data', callback=callback)

def on_error(self, callback: Callable[..., None]) -> EventHandle:
    """Register a callback to invoke when an error occurs during a measurement.

    These errors can be a connection or communication error.

    Parameters
    ----------
    callback : Callable
        The function to call when triggered.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('error', callback=callback)

def on_measurement_begin(self, callback: Callable[[Measurement], None]) -> EventHandle:
    """Register a callback to invoke at the start of a measurement.

    Parameters
    ----------
    callback : Callable
        The function to call when triggered.
        Passes [pypalmsens.data.Measurement][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('measurement_begin', callback=callback)

def on_measurement_end(self, callback: Callable[[Measurement], None]) -> EventHandle:
    """Register a callback to invoke at the end of a measurement.

    Parameters
    ----------
    callback : Callable
        The function to call when triggered.
        Passes [pypalmsens.data.Measurement][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('measurement_end', callback=callback)

def on_measurement_setup(self, callback: Callable[[], None]) -> EventHandle:
    """Register a callback to invoke before the measurement starts.

    Use this to set up file resources, database connections, etc.

    Parameters
    ----------
    callback : Callable
        The function to call when triggered.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('measurement_setup', callback=callback)

def on_measurement_teardown(self, callback: Callable[[], None]) -> EventHandle:
    """Register a callback to invoke after the measurement ends.

    The measurement ends when it finnished successfully or after an error occurs.
    Use this to close files or clean up resources.

    Parameters
    ----------
    callback : Callable
        The function to call when triggered.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('measurement_teardown', callback=callback)

def on_receive_message(self, callback: Callable[[str], None], /):
    """Register a callback for when a new message is received.

    The callback will be invoked, for example, when a method is started,
    or when `send_string` is called in MethodSCRIPT.

    Parameters
    ----------
    callback : callable[[str]]
        The function to call when triggered.
        Passes [str][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('receive_message', callback=callback)

def on_receive_status(self, callback: Callable[[Status], None], /):
    """Register a callback for idle status update events.

    Requires active event loop (i.e. async only).

    The callback will be invoked whenever the instrument sends
    updated current/potential values during idle state or pretreatment phases.
    The update frequency varies per device.

    Parameters
    ----------
    callback : callable[[Status]]
        The function to call when triggered.
        Passes [pypalmsens.data.Status][] as argument to the callback.

    Returns
    -------
    EventHandle
        Handle that can be used to cancel the subscription.
    """
    return self.on('receive_status', callback=callback)

async def query(self, command: str, delay: float | None = None) -> str:
    """Send a command using the communication protocol and return its response.

    This is a method for direct communication with the instrument.
    It writes the command to the device, waits for completion,
    reads the full response, and returns it as a string.

    For commands that run for a long time (e.g. scripts), this
    method will block until the script completes or times out.

    See also [pypalmsens.CommProtocolAsync][].

    Parameters
    ----------
    command : str
        The command to send (e.g., 'i' to get the serial number).
        If `command` does not end with `'\\n'`, one is automatically added.
    delay : float, optional
        Pause (in seconds) between read attempts. Defaults to `self.delay`.

    Returns
    -------
    response : str
        The complete response from the device.
    """
    if not isinstance(self._comm.ClientConnection, PalmSens.Comm.ClientConnectionMS):
        raise TypeError(
            'The Communication Protocol is only supported on MethodSCRIPT devices.'
        )

    emit_idle_messages: bool

    async with self._lock():
        # this temporarily turns off idle messages to reduce cross-talk
        if emit_idle_messages := await create_future(self._comm.GetStatusWhenIdleAsync()):
            await create_future(self._comm.SetStatusWhenIdleAsync(False))

        comm = CommProtocolAsync(self.instrument)

        try:
            response = await comm.query(command, delay=delay)
        finally:
            if emit_idle_messages:
                await create_future(self._comm.SetStatusWhenIdleAsync(emit_idle_messages))

    return response

async def read_current(self) -> float:
    """Read the current in µA.

    Returns
    -------
    current : float
        Current in µA.
    """
    async with self._lock():
        current: float = await create_future(self._comm.GetCurrentAsync())

    return single_to_double(current)

async def read_potential(self) -> float:
    """Read the potential in V.

    Returns
    -------
    potential : float
        Potential in V.
    """

    async with self._lock():
        potential: float = await create_future(self._comm.GetPotentialAsync())

    return single_to_double(potential)

async def set_cell(self, cell_on: bool) -> None:
    """Turn the cell on or off.

    Parameters
    ----------
    cell_on : bool
        If true, turn on the cell
    """
    async with self._lock():
        await create_future(self._comm.SetCellOnAsync(cell_on))

async def set_current_range(self, current_range: AllowedCurrentRanges):
    """Set the current range for the cell.

    Parameters
    ----------
    current_range: AllowedCurrentRanges
        Set the current range as a string.
        See `pypalmsens.settings.AllowedCurrentRanges` for options.
    """
    async with self._lock():
        await create_future(
            self._comm.SetCurrentRangeAsync(cr_string_to_enum(current_range))
        )

async def set_multiplexer_channel(self, channel: int):
    """Sets the multiplexer channel.

    Parameters
    ----------
    channel : int
        Index of the channel to set.
    """
    async with self._lock():
        await create_future(self._comm.ClientConnection.SetMuxChannelAsync(channel))

async def set_mux8r2_settings(
    self,
    connect_sense_to_working_electrode: bool = False,
    combine_reference_and_counter_electrodes: bool = False,
    use_channel_1_reference_and_counter_electrodes: bool = False,
    set_unselected_channel_working_electrode: int = 0,
):
    """Set the settings for the Mux8R2 multiplexer.

    Parameters
    ---------
    connect_sense_to_working_electrode: float
        Connect the sense electrode to the working electrode. Default is False.
    combine_reference_and_counter_electrodes: float
        Combine the reference and counter electrodes. Default is False.
    use_channel_1_reference_and_counter_electrodes: float
        Use channel 1 reference and counter electrodes for all working electrodes. Default is False.
    set_unselected_channel_working_electrode: float
        Set the unselected channel working electrode to disconnected/floating (0), ground (1), or standby potential (2). Default is 0.
    """
    self.ensure_connection()

    if self._comm.Capabilities.MuxModel != PalmSens.MuxModel.MUX8R2:
        raise ValueError(
            f"Incompatible mux model: {self._comm.Capabilities.MuxModel}, expected 'MUXR2'."
        )

    mux_settings = PalmSens.Method.MuxSettings(False)
    mux_settings.ConnSEWE = connect_sense_to_working_electrode
    mux_settings.ConnectCERE = combine_reference_and_counter_electrodes
    mux_settings.CommonCERE = use_channel_1_reference_and_counter_electrodes
    mux_settings.UnselWE = PalmSens.Method.MuxSettings.UnselWESetting(
        set_unselected_channel_working_electrode
    )

    async with self._lock():
        await create_future(
            self._comm.ClientConnection.SetMuxSettingsAsync(MuxType(1), mux_settings)
        )

async def set_potential(self, potential: float) -> None:
    """Set the potential of the cell.

    Parameters
    ----------
    potential : float
        Potential in V
    """
    async with self._lock():
        await create_future(self._comm.SetPotentialAsync(potential))

async def set_potential_range(self, potential_range: AllowedPotentialRanges):
    """Set the potential range for the cell.

    Parameters
    ----------
    potential_range: AllowedPotentialRanges
        Set the potential range as a string.
        See `pypalmsens.settings.AllowedPotentialRanges` for options.
    """
    async with self._lock():
        await create_future(
            self._comm.SetPotentialRangeAsync(pr_string_to_enum(potential_range))
        )

def status(self) -> Status:
    """Get status.

    Sets device 'StatusWhenIdle' flag on device, which tells it
    to periodically send an updated status message.
    """
    self.ensure_connection()

    if not self._comm.StatusWhenIdle:
        self._comm.StatusWhenIdle = True

        while not (status := self._comm.get_Status()):
            time.sleep(0.1)

    else:
        status = self._comm.get_Status()

    return Status(
        status,
        device_state=str(self._comm.get_State()),  # type:ignore
    )

def supported_applied_current_ranges(self: HasCommProtocol) -> list[AllowedCurrentRanges]:
    """List applied current ranges supported by this device.

    Returns
    -------
    current_ranges: list[AllowedCurrentRanges]
        List of supported current ranges.
    """
    self.ensure_connection()
    return CapabilitiesInterface(comm=self._comm).supported_applied_current_ranges

def supported_bipot_current_ranges(self: HasCommProtocol) -> list[AllowedCurrentRanges]:
    """List bipot current ranges supported by this device.

    Returns
    -------
    current_ranges: list[AllowedCurrentRanges]
        List of supported current ranges.
    """
    self.ensure_connection()
    return CapabilitiesInterface(comm=self._comm).supported_bipot_current_ranges

def supported_current_ranges(self: HasCommProtocol) -> list[AllowedCurrentRanges]:
    """List current ranges supported by this device.

    Returns
    -------
    current_ranges: list[AllowedCurrentRanges]
        List of supported current ranges.
    """
    self.ensure_connection()
    return CapabilitiesInterface(comm=self._comm).supported_current_ranges

def supported_methods(self: HasCommProtocol) -> list[AllowedMethods]:
    """List methods supported by this device.

    Returns
    -------
    methods: list[AllowedMethods]
        List of supported methods.
    """
    self.ensure_connection()
    return CapabilitiesInterface(comm=self._comm).supported_methods

def supported_potential_ranges(self: HasCommProtocol) -> list[AllowedPotentialRanges]:
    """List applied potential ranges supported by this device.

    Returns
    -------
    potential_ranges: list[AllowedPotentialRanges]
        List of supported potential ranges.
    """
    self.ensure_connection()
    return CapabilitiesInterface(comm=self._comm).supported_potential_ranges

def validate_method(
    self: HasCommProtocol,
    method: MethodTypeCompatible,
):
    """Validate method.

    Raise ValueError if the method cannot be validated.

    Parameters
    -----------
    method: MethodType
        The method to validate.
    """
    self.ensure_connection()

    capabilities = self._comm.Capabilities

    psmethod = method._to_psmethod()
    errors = psmethod.Validate(capabilities)

    if any(error.IsFatal for error in errors):
        message = '\n'.join([error.Message for error in errors])
        raise MethodIncompatibleError(f'Method not compatible:\n{message}')

async def wait_digital_trigger(self, wait_for_high: bool) -> None:
    """Wait for digital trigger.

    Parameters
    ----------
    wait_for_high: bool
        Wait for digital line high before starting
    """
    async with self._lock():
        while True:
            if await create_future(self._comm.DigitalLineD0Async()) == wait_for_high:
                break
            await asyncio.sleep(0.05)

async def add(self, manager: InstrumentManagerAsync) -> None:
    """Open and add manager to the pool.

    Parameters
    ----------
    manager : InstrumentManagerAsync
        Instance of an instrument manager.
    """
    await manager.connect()
    self.managers.append(manager)

async def connect(self, attempts: int = 1) -> None:
    """Connect all instrument managers in the pool.

    Parameters
    ----------
    attempts: int, optional
        Number of attempts to establish connection.
        Use this if you experience connection issues via USB.
    """
    tasks = [manager.connect() for manager in self.managers]

    try:
        _ = await asyncio.gather(*tasks)
    except Exception:
        if attempts <= 1:
            raise
        await asyncio.sleep(0.5)
    else:
        return

    await self.connect(attempts=attempts - 1)

async def disconnect(self) -> None:
    """Disconnect all instrument managers in the pool."""
    tasks = [manager.disconnect() for manager in self.managers]
    await asyncio.gather(*tasks)

def is_connected(self) -> bool:
    """Return true if all managers in the pool are connected."""
    return all(manager.is_connected() for manager in self.managers)

def is_disconnected(self) -> bool:
    """Return true if all managers in the pool are disconnected."""
    return not any(manager.is_connected() for manager in self.managers)

async def measure(
    self,
    method: MethodType,
    callback: Sequence[Callback | CallbackEIS] | Callback | CallbackEIS | None = None,
    **kwargs,
) -> list[Measurement]:
    """Concurrently start measurement on all managers in the pool.

    For hardware synchronization, set `.general.use_hardware_sync` on the method.
    For MethodSCRIPT, use 'set_channel_sync 1'.

    In addition, the pool must contain:
    - channels from a single multichannel instrument only
    - the first channel of the multichannel instrument
    - at least two channels

    All instruments are prepared and put in a waiting state.
    The measurements are started via a hardware sync trigger on channel 1.

    Parameters
    ----------
    method : MethodType
        Method parameters for measurement.
    callback : list[Callback] | Callback | CallbackEIS | None
        If specified, call these functions/this function on every new set of data points.
        New data points are batched, and contain all points since the last
        time it was called.

        Specify a sequence of callbacks to set a different function for every channel.
        The number of callbacks must match the number of channels.

        Specify a single callback to set the same function to all channels.
    **kwargs
        These keyword parameters are passed to the measure function.
    """
    tasks: list[Awaitable[Measurement]] = []

    callbacks: Sequence[Callback | CallbackEIS | None]

    if isinstance(callback, Sequence):
        if len(callback) != len(self.managers):
            raise IndexError('Number of callbacks does not match number of channels.')
        callbacks = callback
    else:
        callbacks = [callback or None for _ in self.managers]

    if method._use_hardware_sync:
        return await self._measure_hw_sync(method, callbacks=callbacks)

    for manager, _callback in zip(self.managers, callbacks):
        tasks.append(manager.measure(method, callback=_callback, **kwargs))

    results = await asyncio.gather(*tasks)
    return results

async def remove(self, manager: InstrumentManagerAsync) -> None:
    """Close and remove manager from pool.

    Parameters
    ----------
    manager : InstrumentManagerAsync
        Instance of an instrument manager.
    """
    self.managers.remove(manager)
    _ = await manager.disconnect()