"""Digital Signal Processing (DSP) domain models and enums."""

from __future__ import annotations

from enum import IntEnum

from pydantic import BaseModel, ConfigDict

from radiens_core.models.ports import Port


class DSPType(IntEnum):
    """Digital Signal Processing filter type. Values match proto ``common.DSPType``."""

    LOWPASS = 0
    HIGHPASS = 1
    BANDPASS = 2
    NOTCH = 3
    BANDSTOP = 4
    CAR = 5
    VIRTUAL_REFERENCE = 6
    PAIRED_REFERENCE = 7


class FilterStage(IntEnum):
    """Digital Signal Processing filter stage. Values match proto ``common.FilterStage``."""

    STAGE_HARDWARE = 0
    STAGE1 = 1
    STAGE2 = 2
    SPIKE_SORTER = 3


class FreqSpecBand(BaseModel):
    """Frequency band specification for filters.

    Attributes:
        low_freq: Low frequency in Hertz.
        high_freq: High frequency in Hertz.
    """

    model_config = ConfigDict(frozen=True)

    low_freq: float
    high_freq: float


class DSPParams(BaseModel):
    """Configuration parameters for a single DSP filter stage.

    Attributes:
        type: DSP filter type. References DSPType.
        stage: DSP filter stage. References FilterStage.
        freq: Filter frequency. Defaults to None.
        freq_spec_band: Frequency band specification. Defaults to None.
        ref_ntv_chan_idx: Reference native channel index. Defaults to None.
        notch_bandwidth: Notch filter bandwidth. Defaults to None.
        user_label: User-defined label for the filter. Defaults to "".
        port: Port configuration. References Port. Defaults to None.
        target_ntv_chan_idx: Target native channel index. Defaults to None.
        is_aux: Whether this is an auxiliary filter. Defaults to False.
        aux_chan_idx: Auxiliary channel index. Defaults to None.
        filter_order: Filter order. Defaults to None.
        zero_phase: Whether to use zero-phase filtering. Defaults to False.
        force_sos: Whether to force Second-Order Sections. Defaults to False.
    """

    model_config = ConfigDict(frozen=True)

    type: DSPType
    stage: FilterStage
    freq: float | None = None
    freq_spec_band: FreqSpecBand | None = None
    ref_ntv_chan_idx: int | None = None
    notch_bandwidth: float | None = None
    user_label: str = ""
    port: Port | None = None
    target_ntv_chan_idx: float | None = None
    is_aux: bool = False
    aux_chan_idx: int | None = None
    filter_order: int | None = None
    zero_phase: bool = False
    force_sos: bool = False


class DSPGroup(BaseModel):
    """Group of DSP filter stages, organized by hardware/software layers.

    Attributes:
        hardware: List of hardware filter stages.
        software_stage1: List of software stage 1 filters.
        software_stage2: List of software stage 2 filters.
        spike_sorter: List of spike sorter filters.
    """

    model_config = ConfigDict(frozen=True)

    hardware: list[DSPParams]
    software_stage1: list[DSPParams]
    software_stage2: list[DSPParams]
    spike_sorter: list[DSPParams]
