"""End-to-end: EEG classifier output flows into fusion as the `eeg` modality."""
from __future__ import annotations

from pathlib import Path

import numpy as np

from src.fusion import engine
from src.fusion.types import (
    FusionInput,
    ModalityClassProb,
    ModalityPrediction,
)
from src.models import eeg_model
from tests.fixtures.build_dummy_eeg_clf import build as build_dummy_eeg


def _eeg_pred_from_features(model, features: np.ndarray) -> ModalityPrediction:
    raw = eeg_model.predict_features(model, features)
    return ModalityPrediction(
        label_text=raw["label_text"],
        label=raw["label"],
        confidence=raw["confidence"],
        probabilities=[
            ModalityClassProb(label_text=p["label_text"], probability=p["probability"])
            for p in raw["probabilities"]
        ],
    )


class TestEEGFusionFlow:
    def test_alzheimers_eeg_lifts_alzheimers_disease_score(self, tmp_path: Path) -> None:
        ckpt = build_dummy_eeg(tmp_path / "eeg.joblib", n_features=16)
        model = eeg_model.load(ckpt)
        eeg_pred = _eeg_pred_from_features(model, np.full((16,), 2.0, dtype=np.float32))

        out = engine.fuse(FusionInput(eeg=eeg_pred))

        alz = next(d for d in out.diseases if d.disease == "alzheimers")
        assert alz.probability > 0.5
        assert any(c.modality == "eeg" for c in alz.contributions)
        assert "mri" in out.missing_inputs

    def test_control_eeg_does_not_inflate_alzheimers(self, tmp_path: Path) -> None:
        ckpt = build_dummy_eeg(tmp_path / "eeg.joblib", n_features=16)
        model = eeg_model.load(ckpt)
        eeg_pred = _eeg_pred_from_features(model, np.zeros((16,), dtype=np.float32))

        out = engine.fuse(FusionInput(eeg=eeg_pred))

        alz = next(d for d in out.diseases if d.disease == "alzheimers")
        assert alz.probability < 0.5