"""Stimulation domain models and enums."""

from __future__ import annotations

from enum import IntEnum

from pydantic import BaseModel, ConfigDict

# ===================================================================
# Stimulation enums
# ===================================================================


class StimShape(IntEnum):
    """Stimulation waveform shape. Values match proto ``allego.StimShape``."""

    BIPHASIC = 0
    BIPHASIC_INTERPHASE_DELAY = 1
    TRIPHASIC = 2
    MONOPHASIC = 3


class StimPolarity(IntEnum):
    """Stimulation polarity. Values match proto ``allego.StimPolarity``."""

    CATHODIC_FIRST = 0
    ANODIC_FIRST = 1


class StimTriggerEdgeOrLevel(IntEnum):
    """Stimulation trigger edge/level mode. Values match proto ``allego.StimTriggerEdgeOrLevel``."""

    EDGE = 0
    LEVEL = 1


class StimTriggerHighOrLow(IntEnum):
    """Stimulation trigger high/low. Values match proto ``allego.StimTriggerHighOrLow``."""

    HIGH = 0
    LOW = 1


class StimPulseMode(IntEnum):
    """Stimulation pulse mode. Values match proto ``allego.StimPulseOrTrain``."""

    SINGLE_PULSE = 0
    PULSE_TRAIN = 1


class StimKeypressIndex(IntEnum):
    KEYPRESS_1 = 1
    KEYPRESS_2 = 2
    KEYPRESS_3 = 3
    KEYPRESS_4 = 4
    KEYPRESS_5 = 5
    KEYPRESS_6 = 6
    KEYPRESS_7 = 7
    KEYPRESS_8 = 8


class StimStep(IntEnum):
    """Stimulation step size. Values match proto ``allego.StimStep``."""

    STIM_STEP_SIZE_MIN = 0
    STIM_STEP_SIZE_10_NA = 1
    STIM_STEP_SIZE_20_NA = 2
    STIM_STEP_SIZE_50_NA = 3
    STIM_STEP_SIZE_100_NA = 4
    STIM_STEP_SIZE_200_NA = 5
    STIM_STEP_SIZE_500_NA = 6
    STIM_STEP_SIZE_1_UA = 7
    STIM_STEP_SIZE_2_UA = 8
    STIM_STEP_SIZE_5_UA = 9
    STIM_STEP_SIZE_10_UA = 10
    STIM_STEP_SIZE_MAX = 11


class StimParams(BaseModel):
    """Stimulation parameters for a single channel.

    All timing values are in microseconds (μs). Current amplitudes are in μA
    for amplifier channels and V for DAC channels.

    Attributes:
        stim_sys_chan_idx: System channel index (ntvAbsKeyIdx).
        stim_shape: Stimulation waveform shape. References StimShape. Defaults to StimShape.BIPHASIC.
        stim_polarity: Stimulation polarity. References StimPolarity. Defaults to StimPolarity.CATHODIC_FIRST.
        first_phase_duration: Duration of the first phase in μs. Defaults to 100.0.
        second_phase_duration: Duration of the second phase in μs. Defaults to 100.0.
        interphase_delay: Delay between phases in μs. Defaults to 0.0.
        first_phase_amplitude: Amplitude of the first phase (μA for amp, V for DAC). Defaults to 0.1.
        second_phase_amplitude: Amplitude of the second phase (μA for amp, V for DAC). Defaults to 0.0.
        baseline_voltage: Baseline voltage in V. Defaults to 0.0.
        trigger_edge_or_level: Trigger mode. References StimTriggerEdgeOrLevel. Defaults to StimTriggerEdgeOrLevel.EDGE.
        trigger_high_or_low: Trigger polarity. References StimTriggerHighOrLow. Defaults to StimTriggerHighOrLow.HIGH.
        enabled: Whether stimulation is enabled for this channel. Defaults to True.
        post_trigger_delay: Delay after trigger before stimulation in μs. Defaults to 0.0.
        pulse_or_train: Pulse or train mode. References StimPulseMode. Defaults to StimPulseMode.SINGLE_PULSE.
        number_of_stim_pulses: Number of pulses in a train. Defaults to 1.
        pulse_train_period: Period of pulses in a train in μs. Defaults to 0.0.
        refractory_period: Refractory period in μs. Defaults to 0.0.
        pre_stim_amp_settle: Pre-stimulation amplifier settle time in μs. Defaults to 0.0.
        post_stim_amp_settle: Post-stimulation amplifier settle time in μs. Defaults to 0.0.
        maintain_amp_settle: Whether to maintain amplifier settle. Defaults to False.
        enable_amp_settle: Whether to enable amplifier settle. Defaults to False.
        headstage_global_amp_settle: Whether to enable global headstage amp settle. Defaults to False.
        post_stim_charge_recov_on: Post-stimulation charge recovery on time in μs. Defaults to 0.0.
        post_stim_charge_recov_off: Post-stimulation charge recovery off time in μs. Defaults to 0.0.
        enable_charge_recovery: Whether to enable charge recovery. Defaults to False.
        trigger_source_is_keypress: Whether the trigger source is a keypress. Defaults to False.
        trigger_source_idx: AIN/DIN system channel index, or keypress number. References StimKeypressIndex if keypress.
    """

    model_config = ConfigDict(frozen=True)

    stim_sys_chan_idx: int
    stim_shape: StimShape = StimShape.BIPHASIC
    stim_polarity: StimPolarity = StimPolarity.CATHODIC_FIRST
    first_phase_duration: float = 100.0
    second_phase_duration: float = 100.0
    interphase_delay: float = 0.0
    first_phase_amplitude: float = 0.1
    second_phase_amplitude: float = 0.0
    baseline_voltage: float = 0.0
    trigger_edge_or_level: StimTriggerEdgeOrLevel = StimTriggerEdgeOrLevel.EDGE
    trigger_high_or_low: StimTriggerHighOrLow = StimTriggerHighOrLow.HIGH
    enabled: bool = True
    post_trigger_delay: float = 0.0
    pulse_or_train: StimPulseMode = StimPulseMode.SINGLE_PULSE
    number_of_stim_pulses: int = 1
    pulse_train_period: float = 0.0
    refractory_period: float = 0.0
    pre_stim_amp_settle: float = 0.0
    post_stim_amp_settle: float = 0.0
    maintain_amp_settle: bool = False
    enable_amp_settle: bool = False
    headstage_global_amp_settle: bool = False
    post_stim_charge_recov_on: float = 0.0
    post_stim_charge_recov_off: float = 0.0
    enable_charge_recovery: bool = False
    trigger_source_is_keypress: bool = False
    trigger_source_idx: int | StimKeypressIndex

    @classmethod
    def biphasic(
        cls: type[StimParams],
        *,
        amplitude: float,
        pulse_duration: float,
        polarity: StimPolarity,
        stim_sys_chan_idx: int,
        enabled: bool = True,
        **overrides: object,
    ) -> StimParams:
        """Create biphasic stimulation parameters.

        Args:
            amplitude: First phase amplitude (μA for amp channels)
            pulse_duration: Phase duration in microseconds
            polarity: Stimulation polarity (cathodic or anodic first)
            stim_sys_chan_idx: System channel index
            **overrides: Override any other StimParams field

        Returns:
            StimParams configured for biphasic stimulation
        """
        return cls(
            stim_shape=StimShape.BIPHASIC,
            stim_polarity=polarity,
            first_phase_duration=pulse_duration,
            second_phase_duration=pulse_duration,
            first_phase_amplitude=amplitude,
            second_phase_amplitude=amplitude,
            stim_sys_chan_idx=stim_sys_chan_idx,
            enabled=enabled,
            **overrides,  # type: ignore[arg-type]
        )


class StimStepRequest(BaseModel):
    """Request to set stimulation step size.

    Attributes:
        stim_step: Target stimulation step size. References StimStep.
    """

    model_config = ConfigDict(frozen=True)

    stim_step: StimStep


class StimStepState(BaseModel):
    """Current stimulation step size setting.

    Attributes:
        stim_step: Current stimulation step size. References StimStep.
    """

    model_config = ConfigDict(frozen=True)

    stim_step: StimStep
