"""EEG Annotation Tool adapter for SenuaLab EEG-DINO Medium."""

from __future__ import annotations

try:
    from braindecode.models import EEGDINO
except ImportError as exc:
    raise ImportError("This model requires braindecode[hub]==1.6.1") from exc

from .preprocessing import INTERNATIONAL_10_20_CHANNELS, preprocess_batch as _preprocess


MODEL_CHANNELS = INTERNATIONAL_10_20_CHANNELS
MODEL_INPUT_SAMPLES = 800
MODEL_SAMPLING_RATE_HZ = 200.0
MODEL_WINDOW_SECONDS = 4.0
MODEL_ENTRY_CLASS = "SenuaEEGDINOMedium"
MODEL_NUM_CLASSES = 2
MODEL_CLASS_LABELS = ["Non-IED", "IED"]
MODEL_NON_IED_CLASS_INDEX = 0
MODEL_IED_CLASS_INDICES = [1]
MODEL_DESCRIPTION = "SenuaLab partially fine-tuned EEG-DINO Medium binary IED detector"
MODEL_BATCH_PREPROCESSOR = "preprocess_batch"
MODEL_REQUIRED_REFERENCE = "common average (applied by model adapter)"
MODEL_REQUIRED_FILTERS = ["1-45 Hz zero-phase Butterworth (applied by model adapter)"]
MODEL_REQUIRED_NORMALIZATION = "global four-second window z-score, clipped to [-8,8]"
MODEL_INPUT_UNIT = "scale-invariant after window z-score"
MODEL_SOURCE_SIGNAL_POLICY = "raw"
MODEL_REQUIRES_FULL_WINDOW = True
MODEL_REQUIRES_ALL_CHANNELS = True
MODEL_DEFAULT_THRESHOLD = 0.619140625
MODEL_DEFAULT_STEP_MS = 500.0
MODEL_DEFAULT_PAD_POLICY = "skip"
MODEL_DEFAULT_BATCH_SIZE = 8
MODEL_DEFAULT_BATCH_MEMORY_MB = 384.0
MODEL_VALIDATION_NOTE = "Threshold selected on vEpiSet validation subjects; requires Braindecode 1.6.1 and is not externally validated."
MODEL_PREPROCESSING_NOTE = "The adapter applies the released 1-45 Hz, common-average, 200 Hz resampling, and global-window z-score pipeline."


def preprocess_batch(batch, source_sfreq=None, channel_names=None):
    return _preprocess(
        batch,
        source_sfreq=source_sfreq,
        target_sfreq=200,
        target_samples=800,
        channel_names=channel_names,
    )


class SenuaEEGDINOMedium(EEGDINO):
    def __init__(self, num_channels: int = 19, num_classes: int = 2, input_length: int = 800):
        if num_channels != 19 or input_length != 800:
            raise ValueError("SenuaEEGDINOMedium requires 19 channels and 800 samples")
        super().__init__(
            n_outputs=num_classes,
            n_chans=19,
            n_times=800,
            sfreq=200,
            patch_size=200,
            n_layer=16,
            nhead=8,
            dim_feedforward=1024,
            channels_kernel_stride_padding_norm=(
                (64, 49, 25, 24, (8, 64)),
                (128, 3, 1, 1, (8, 128)),
                (64, 3, 1, 1, (8, 64)),
            ),
            n_channel_embeddings=19,
            n_global_tokens=1,
            global_token_layer=1,
            drop_prob=0.1,
        )