"""Signal selection and signal data array models."""

import numpy as np
import numpy.typing as npt
from pydantic import BaseModel, ConfigDict

from radiens_core.models.enums import KeyIndex


class SigSelect(BaseModel):
    """Signal selection specifying channels per signal type.

    Each field is a numpy array of native channel indices for that signal type.
    """

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)

    amp: npt.NDArray[np.int64]
    gpio_ain: npt.NDArray[np.int64]
    gpio_din: npt.NDArray[np.int64]
    gpio_dout: npt.NDArray[np.int64]
    key_idx: KeyIndex


class SignalArrays(BaseModel):
    """Container for signal data arrays, one per signal type.

    Each field is either a 2D numpy array ``(channels, samples)``
    or ``None`` if that signal type was not requested.
    """

    model_config = ConfigDict(frozen=True, arbitrary_types_allowed=True)

    amp: npt.NDArray[np.float64] | None = None
    gpio_ain: npt.NDArray[np.float64] | None = None
    gpio_din: npt.NDArray[np.float64] | None = None
    gpio_dout: npt.NDArray[np.float64] | None = None
