_preprocessing.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EEG preprocessing with robust handling of non-10–20 aux channels:

- Standardise names (incl. P7→T5, P8→T6; T7→T3, T8→T4)
- Classify aux channels (ROC/LOC→eog, EMG→emg, PHOTIC→stim, IBI/BURSTS/SUPPR→misc, unknown non-10–20→misc)
- Apply montage without dropping channels
- Run PyPREP on EEG-only with a NaN-free trimmed montage
- RANSAC interpolation + reset bads
- Optional ASR; ICA + ICLabel pruning (optionally keep/strip cardiac)
- ECG R-peak events (with gap interpolation)
- Save and TSV log
"""

import os
import csv
import logging
import re
from datetime import datetime
from pathlib import Path
from contextlib import redirect_stdout, redirect_stderr

import mne
import numpy as np
import neurokit2 as nk
import pandas as pd
from mne.preprocessing import ICA

try:
    import asrpy
except Exception:  # optional; only required when ASR is enabled
    asrpy = None

try:
    import pyprep
except Exception:  # optional; only required when PyPREP is enabled
    pyprep = None

try:
    from mne_icalabel import label_components
except Exception:  # optional; only required when ICA+ICLabel is enabled
    label_components = None

# ------------------------- Logging helpers ------------------------- #

def _now_utc_iso() -> str:
    return datetime.utcnow().replace(microsecond=0).isoformat() + "Z"

# Fixed TSV schema so all runs (success/fail) align.
_PREPROC_TSV_FIELDS = [
    "utc", "status", "stage",
    "input", "output",
    "n_channels", "n_events", "bad_channels",
    "remove_cfa", "remove_cfa_mode", "flip_ecg",
    "asr_threshold", "n_comp", "stim_keep",
    "montage_name", "ica_path", "ica_bads", "ica_bad_labels",
    "spirometry_channel", "spirometry_status", "spirometry_peaks",
    "spirometry_troughs", "spirometry_qc_png", "spirometry_signals_csv",
    "spirometry_summary_csv",
    "error_type", "error",
]

def _append_preproc_tsv(logging_path: str | None, rec: dict) -> None:
    """Append a preprocessing record (success or failure) to a TSV with stable columns."""
    if not logging_path:
        return
    try:
        logp = Path(logging_path)
        if str(logp.parent) not in ("", "."):
            logp.parent.mkdir(parents=True, exist_ok=True)
        file_exists = logp.exists() and logp.stat().st_size > 0
        row = {k: "" for k in _PREPROC_TSV_FIELDS}
        for k, v in (rec or {}).items():
            if k in row:
                row[k] = "" if v is None else str(v)
        with open(str(logp), "a", newline="") as f:
            w = csv.DictWriter(f, fieldnames=_PREPROC_TSV_FIELDS, delimiter="\t")
            if not file_exists:
                w.writeheader()
            w.writerow(row)
    except Exception:
        logging.exception("Failed to append preprocessing TSV record to %s", logging_path)


# ------------------------- Naming & typing helpers ------------------------- #

TEN_TWENTY_SET = {
    'Fp1','Fp2','F7','F3','Fz','F4','F8',
    'T3','C3','Cz','C4','T4',
    'T5','P3','Pz','P4','T6',
    'O1','O2','A1','A2'
}


from dataclasses import dataclass
from pathlib import Path

# Light wrapper so we can keep names below unchanged
@dataclass
class _RuntimeCfg:
    output_root: Path
    target_sfreq: float | None
    use_asr: float | None
    use_pyprep: bool
    remove_cfa: bool
    remove_cfa_mode: str
    log_file: str | None
    ref_chs: str | list
    reref_chs: str | list
    high_pass: float
    low_pass: float
    prep_ransac: bool
    line_freqs: tuple | list
    montage_name: str | None = None
    rename_to_1020: bool = True
    use_ica: bool = True
    spirometry_channel: str | None = None

_CFG_RT: _RuntimeCfg | None = None
# fallback defaults in case set_runtime_config is not called
prep_params = {
    "line_freqs": (50.0, 100.0),
    "ref_chs": "eeg",
    "reref_chs": "eeg",
    "l_freq": 1.0,
    "h_freq": 100.0,
    "ransac": True,
}

def set_runtime_config(cfg) -> None:
    """Call this once from the CLI or notebook to bind the active config."""
    global _CFG_RT, output_dir, target_sfreq, do_asr, remove_cfa, remove_cfa_mode, log_file, prep_params
    # derive ICA cardiac strategy
    mode = getattr(cfg, "remove_cfa_mode", None)
    legacy_flag = getattr(cfg, "remove_cfa", None)
    if mode is None:
        # backward compatibility: infer mode from legacy bool
        if legacy_flag is None:
            mode = "remove"
            legacy_flag = True
        else:
            mode = "remove" if bool(legacy_flag) else "keep"
    if legacy_flag is None:
        legacy_flag = True if str(mode).lower() != "keep" else False
    _CFG_RT = _RuntimeCfg(
        output_root=Path(cfg.output_root),
        target_sfreq=getattr(cfg, "target_sfreq", None),
        use_asr=getattr(cfg, "use_asr", None),
        use_pyprep=bool(getattr(cfg, "use_pyprep", getattr(cfg, "run_pyprep", True))),
        remove_cfa=bool(legacy_flag),
        remove_cfa_mode=str(mode).lower(),
        log_file=getattr(cfg, "log_file", None),
        ref_chs=getattr(cfg, "ref_chs", "eeg"),
        reref_chs=getattr(cfg, "reref_chs", "eeg"),
        high_pass=getattr(cfg, "high_pass", 1.0),
        low_pass=getattr(cfg, "low_pass", 100.0),
        prep_ransac=getattr(cfg, "prep_ransac", True),
        line_freqs=getattr(cfg, "line_freqs", (50.0, 100.0)),
        montage_name=getattr(cfg, "montage_name", None),
        rename_to_1020=getattr(cfg, "rename_to_1020", True),
        use_ica=bool(getattr(cfg, "use_ica", getattr(cfg, "run_ica", True))),
        spirometry_channel=getattr(cfg, "spirometry_channel", None),
    )

    # keep legacy names used in the rest of this module
    output_dir   = str(_CFG_RT.output_root)
    target_sfreq = float(_CFG_RT.target_sfreq) if _CFG_RT.target_sfreq else None
    do_asr       = _CFG_RT.use_asr
    do_pyprep    = _CFG_RT.use_pyprep
    remove_cfa   = _CFG_RT.remove_cfa
    remove_cfa_mode = _CFG_RT.remove_cfa_mode
    log_file     = _CFG_RT.log_file or ""

    prep_params = {
        "line_freqs": _CFG_RT.line_freqs,
        "ref_chs":   _CFG_RT.ref_chs,
        "reref_chs": _CFG_RT.reref_chs,
        "l_freq":    _CFG_RT.high_pass,
        "h_freq":    _CFG_RT.low_pass,
        "ransac":    _CFG_RT.prep_ransac,
    }

def _canonicalise_name(ch: str) -> str:
    nm = ch.strip()
    nm = re.sub(r'^\s*EEG\s+', '', nm, flags=re.IGNORECASE)
    parts = [p.strip() for p in nm.split('-') if p.strip()]
    if len(parts) > 1:
        nm = parts[1] if parts[0].isdigit() else parts[0]
    nm = nm.upper()
    nm = nm.replace('FP', 'Fp').replace('Z', 'z')
    repl = {'T7': 'T3', 'T8': 'T4', 'P7': 'T5', 'P8': 'T6'}
    return repl.get(nm, nm)

AUX_EOG = {'ROC', 'LOC'}
AUX_STIM = {'PHOTIC'}
AUX_MISC = {'IBI', 'BURSTS', 'SUPPR'}
RSP_PEAK_DESC = "RSP_Peak"
RSP_TROUGH_DESC = "RSP_Trough"
ECG_TOKENS = ("ECG", "EKG", "CARD", "EXG")

def _optional_name(value: str | None) -> str | None:
    text = str(value or "").strip()
    return text or None

def _find_channel(raw: mne.io.BaseRaw, preferred: str | None) -> str | None:
    preferred = _optional_name(preferred)
    if not preferred:
        return None
    if preferred in raw.ch_names:
        return preferred
    wanted = preferred.casefold()
    for ch in raw.ch_names:
        if ch.casefold() == wanted:
            return ch
    canonical = _canonicalise_name(preferred)
    if canonical in raw.ch_names:
        return canonical
    wanted = canonical.casefold()
    for ch in raw.ch_names:
        if ch.casefold() == wanted:
            return ch
        if _canonicalise_name(ch).casefold() == wanted:
            return ch
    return None

def _active_spirometry_channel(spirometry_channel: str | None = None) -> str | None:
    return _optional_name(spirometry_channel) or (_optional_name(_CFG_RT.spirometry_channel) if _CFG_RT else None)

def _mark_spirometry_channel(raw: mne.io.BaseRaw, spirometry_channel: str | None) -> str | None:
    ch = _find_channel(raw, _active_spirometry_channel(spirometry_channel))
    if ch:
        raw.set_channel_types({ch: "resp"})
    return ch

def _mark_ecg_channels(raw: mne.io.BaseRaw, ecg_channel: str | None = None) -> list[str]:
    names = []
    preferred = _find_channel(raw, ecg_channel)
    if preferred:
        names.append(preferred)
    for ch in raw.ch_names:
        if any(tok in ch.upper() for tok in ECG_TOKENS) and ch not in names:
            names.append(ch)
    for ch in names:
        try:
            raw.set_channel_types({ch: "ecg"})
        except Exception:
            pass
    return names

def _flip_ecg_channels(raw: mne.io.BaseRaw, ecg_channel: str | None = None) -> list[str]:
    names = _mark_ecg_channels(raw, ecg_channel)
    if not names:
        names = [raw.ch_names[i] for i in mne.pick_types(raw.info, ecg=True)]
    idx = [raw.ch_names.index(ch) for ch in names if ch in raw.ch_names]
    if idx:
        raw._data[idx] *= -1
    return [raw.ch_names[i] for i in idx]

def _preproc_base_from_output(path: str | Path) -> str:
    stem = Path(path).stem
    for suffix in ("_pp_raw_keepcfa", "_pp_raw", "_raw_keepcfa", "_raw"):
        if stem.endswith(suffix):
            return stem[:-len(suffix)]
    return stem

def _has_spirometry_annotations(raw: mne.io.BaseRaw) -> bool:
    if raw.annotations is None or len(raw.annotations) == 0:
        return False
    descriptions = set(np.asarray(raw.annotations.description).astype(str))
    return bool({RSP_PEAK_DESC, RSP_TROUGH_DESC} & descriptions)

def _add_spirometry_annotations(raw: mne.io.BaseRaw, peaks: np.ndarray, troughs: np.ndarray) -> None:
    if raw.annotations is not None and len(raw.annotations):
        old = [i for i, desc in enumerate(raw.annotations.description) if desc in (RSP_PEAK_DESC, RSP_TROUGH_DESC)]
        if old:
            annotations = raw.annotations.copy()
            annotations.delete(old)
            raw.set_annotations(annotations)
    samples = np.concatenate([peaks, troughs]).astype(int)
    if samples.size == 0:
        return
    descriptions = [RSP_PEAK_DESC] * len(peaks) + [RSP_TROUGH_DESC] * len(troughs)
    order = np.argsort(samples)
    annotations = mne.Annotations(
        onset=samples[order] / float(raw.info["sfreq"]),
        duration=np.zeros(samples.size),
        description=[descriptions[i] for i in order],
        orig_time=raw.annotations.orig_time,
    )
    raw.set_annotations(raw.annotations + annotations)

def save_spirometry_analysis(
    raw: mne.io.BaseRaw,
    spirometry_channel: str | None,
    output_dir: str | Path | None,
    base: str,
    redo: bool = False,
) -> dict:
    requested = _active_spirometry_channel(spirometry_channel)
    if not requested:
        return {"spirometry_status": "off"}
    ch = _find_channel(raw, requested)
    if not ch:
        return {"spirometry_channel": requested, "spirometry_status": "missing"}

    sfreq = float(raw.info["sfreq"])
    signals, info = nk.rsp_process(raw.get_data(picks=ch)[0], sampling_rate=sfreq)
    peaks = np.asarray(info.get("RSP_Peaks", []), dtype=int)
    troughs = np.asarray(info.get("RSP_Troughs", []), dtype=int)
    _add_spirometry_annotations(raw, peaks, troughs)

    rec = {
        "spirometry_channel": ch,
        "spirometry_status": "ok",
        "spirometry_peaks": str(len(peaks)),
        "spirometry_troughs": str(len(troughs)),
    }
    if output_dir:
        out_dir = Path(output_dir)
        out_dir.mkdir(parents=True, exist_ok=True)
        signals_csv = out_dir / f"{base}_rsp_signals.csv"
        summary_csv = out_dir / f"{base}_rsp_summary.csv"
        events_csv = out_dir / f"{base}_rsp_events.csv"
        qc_png = out_dir / f"{base}_rsp_qc.png"

        if redo or not signals_csv.exists():
            signals.to_csv(signals_csv, index=False)
        try:
            summary = nk.rsp_analyze(signals, sampling_rate=sfreq, method="interval-related")
        except Exception as exc:
            summary = pd.DataFrame([{"error": str(exc)}])
        if redo or not summary_csv.exists():
            summary.to_csv(summary_csv, index=False)
        if redo or not events_csv.exists():
            event_rows = (
                [{"sample": int(s), "time_s": float(s / sfreq), "kind": "peak"} for s in peaks]
                + [{"sample": int(s), "time_s": float(s / sfreq), "kind": "trough"} for s in troughs]
            )
            pd.DataFrame(event_rows, columns=["sample", "time_s", "kind"]).sort_values("sample").to_csv(events_csv, index=False)
        if redo or not qc_png.exists():
            import matplotlib.pyplot as plt
            nk.rsp_plot(signals, info, static=True)
            fig = plt.gcf()
            fig.savefig(qc_png, dpi=150, bbox_inches="tight")
            plt.close(fig)

        rec.update({
            "spirometry_qc_png": str(qc_png),
            "spirometry_signals_csv": str(signals_csv),
            "spirometry_summary_csv": str(summary_csv),
        })
    return rec

def spirometry_phase_at_samples(raw: mne.io.BaseRaw, samples: np.ndarray) -> pd.DataFrame | None:
    if raw.annotations is None or len(raw.annotations) == 0:
        return None
    sfreq = float(raw.info["sfreq"])
    desc = np.asarray(raw.annotations.description).astype(str)
    ann_samples = np.rint(np.asarray(raw.annotations.onset) * sfreq).astype(int)
    peaks = ann_samples[desc == RSP_PEAK_DESC]
    troughs = ann_samples[desc == RSP_TROUGH_DESC]
    if len(peaks) == 0 or len(troughs) == 0:
        return None
    phase = nk.rsp_phase(peaks=peaks, troughs=troughs, desired_length=raw.n_times)
    idx = np.clip(np.asarray(samples, dtype=int), 0, raw.n_times - 1)
    rows = phase.iloc[idx].reset_index(drop=True).rename(columns={
        "RSP_Phase": "resp_phase",
        "RSP_Phase_Completion": "resp_phase_completion",
    })
    vals = pd.to_numeric(rows["resp_phase"], errors="coerce").to_numpy()
    rows["resp_phase_label"] = np.where(
        np.isfinite(vals),
        np.where(vals == 1, "inspiratory", "expiratory"),
        "unknown",
    )
    return rows

def _classify_aux_channels(raw: mne.io.BaseRaw) -> None:
    """Classify obvious auxiliary channels but preserve all EEG leads.

    Previous logic demoted any EEG channel not in the 10-20 set to 'misc', causing
    large arrays (e.g., dense caps, BioSemi, EGI) to be dropped from EEG processing.
    We now only re-type clearly non-EEG auxiliary channels (EOG/stim/misc/ECG/EMG) and
    leave all remaining channels as EEG.
    """
    for ch in list(raw.ch_names):
        up = ch.upper()
        if up in AUX_EOG:
            raw.set_channel_types({ch: 'eog'})
        elif up in AUX_STIM:
            raw.set_channel_types({ch: 'stim'})
        elif up in AUX_MISC:
            raw.set_channel_types({ch: 'misc'})
        elif any(tok in up for tok in ECG_TOKENS):
            raw.set_channel_types({ch: 'ecg'})
        elif up == 'EMG':
            raw.set_channel_types({ch: 'emg'})
    # Do not demote remaining EEG channels; keep full set intact.

_DEF_MONTAGES = [
    "standard_1020", "standard_1005", "biosemi32", "biosemi64", "biosemi128",
    "GSN-HydroCel-32", "GSN-HydroCel-64", "GSN-HydroCel-128", "GSN-HydroCel-256",
    "easycap-M1", "easycap-M10", "egi256"
]

def _auto_detect_montage(raw: mne.io.BaseRaw, candidates: list[str] | None = None) -> str | None:
    """Optional montage auto-detection: choose built-in montage with highest channel name coverage."""
    if candidates is None:
        candidates = _DEF_MONTAGES
    ch_set = {
        _canonicalise_name(ch)
        for ch, kind in zip(raw.ch_names, raw.get_channel_types())
        if kind == "eeg" and not any(tok in _canonicalise_name(ch).upper() for tok in ECG_TOKENS)
    }
    best_name, best_cov = None, 0.0
    for name in candidates:
        try:
            mont = mne.channels.make_standard_montage(name)
        except Exception:
            continue
        mont_chs = set(mont.ch_names)
        inter = ch_set & mont_chs
        if not ch_set:
            continue
        cov = len(inter) / max(1, len(ch_set))
        if cov > best_cov:
            best_cov, best_name = cov, name
    # require minimal coverage threshold
    if best_cov >= 0.5:
        return best_name
    return None

def standardise_and_montage(raw: mne.io.BaseRaw) -> mne.io.BaseRaw:
    """Rename channels, apply montage early, then classify auxiliaries.

    Applying the montage before auxiliary classification avoids losing EEG leads
    when dense arrays are used. All EEG channels are preserved.
    """
    # Canonicalise channel names first to maximize montage matching
    raw.rename_channels({ch: _canonicalise_name(ch) for ch in raw.ch_names})
    # Decide montage: explicit, auto-detect, or fallback
    m_name = _CFG_RT.montage_name if _CFG_RT else None
    if m_name is None or str(m_name).lower() in ("", "auto", "none"):
        detected = _auto_detect_montage(raw)
        m_name = detected or "standard_1020"
    mont = mne.channels.make_standard_montage(m_name)
    raw.set_montage(mont, match_case=False, on_missing='ignore')
    logging.info(f"Applied montage: {m_name}")
    # Classify auxiliaries AFTER montage so EEG leads remain EEG
    eeg_before = sum(1 for t in raw.get_channel_types() if t == 'eeg')
    _classify_aux_channels(raw)
    eeg_after = sum(1 for t in raw.get_channel_types() if t == 'eeg')
    if eeg_after < eeg_before:
        logging.warning(f"EEG channel count decreased from {eeg_before} to {eeg_after}. Check classification rules.")

    # Drop EEG channels not present in montage (post-application)
    mont_chs = set(mont.ch_names)
    drop_eeg = [ch for ch, t in zip(raw.ch_names, raw.get_channel_types()) if t == "eeg" and ch not in mont_chs]
    if drop_eeg:
        raw.drop_channels(drop_eeg)

    # force the lead name types to be str (to avoid issues with numpy.str_)
    raw.rename_channels({ch: str(ch) for ch in raw.ch_names})
    
    return raw

def _finite_ch_pos(ch_pos: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
    ok = {}
    for k, v in (ch_pos or {}).items():
        arr = np.asarray(v, float)
        if arr.shape and np.all(np.isfinite(arr)):
            ok[k] = arr
    return ok

def _trim_eeg_montage_no_nan(raw_eeg: mne.io.BaseRaw) -> mne.channels.DigMontage | None:
    mont = raw_eeg.get_montage()
    if mont is None:
        return None
    pos = mont.get_positions()
    ch_pos = _finite_ch_pos(pos.get('ch_pos', {}))
    present = [ch for ch in raw_eeg.ch_names if ch in ch_pos]
    if not present:
        return None
    return mne.channels.make_dig_montage(
        ch_pos={k: ch_pos[k] for k in present},
        nasion=pos.get('nasion'), lpa=pos.get('lpa'), rpa=pos.get('rpa'),
        hpi=pos.get('hpi'), coord_frame=pos.get('coord_frame', 'head'),
    )

def _prune_prep_params_for_raw(params: dict, raw: mne.io.BaseRaw) -> dict:
    import copy
    pruned = copy.deepcopy(params)
    present = set(raw.ch_names)
    keys = {
        "ref_chs","reref_chs","eog_chs","corr_chs",
        "ransac_channel_picks","interpolation_channel_picks",
        "exclude","include","target_channels","bad_channel_prior"
    }
    def _walk(d):
        for k, v in list(d.items()):
            if isinstance(v, dict):
                _walk(v)
            elif isinstance(v, (list, tuple)) and any(isinstance(x, str) for x in v):
                d[k] = [x for x in v if isinstance(x, str) and x in present]
            if k in keys:
                vv = d.get(k, None)
                if isinstance(vv, (list, tuple)):
                    d[k] = [x for x in vv if isinstance(x, str) and x in present]
        return d
    return _walk(pruned)

# ------------------------------- Core steps -------------------------------- #

def run_pyprep(
    raw: mne.io.BaseRaw,
    param_dict: dict | None = None,
    random_seed: int = 42,
    ransac: bool = True,
    channel_wise: bool = True,
) -> mne.io.BaseRaw:
    if pyprep is None:
        raise RuntimeError("PyPREP is enabled but the 'pyprep' package is not installed. Disable 'Run PyPREP' or install pyprep.")
    if param_dict is None:
        param_dict = prep_params
    raw_eeg = raw.copy().pick('eeg')
    safe_params = _prune_prep_params_for_raw(param_dict, raw_eeg)
    trimmed_mont = _trim_eeg_montage_no_nan(raw_eeg)

    # Ensure keys PyPREP expects are present to avoid KeyError
    safe_params.setdefault("line_freqs", ())
    safe_params.setdefault("l_freq", None)
    safe_params.setdefault("h_freq", None)

    prep = pyprep.PrepPipeline(
        raw=raw_eeg, montage=trimmed_mont, prep_params=safe_params,
        random_state=random_seed, ransac=ransac, channel_wise=channel_wise,
    )
    try:
        with open(os.devnull, 'w') as fnull, redirect_stdout(fnull), redirect_stderr(fnull):
            result = prep.fit()
    except Exception as e:
        print("→ EEG picks:", sorted(raw_eeg.ch_names))
        if trimmed_mont is not None:
            print("→ trimmed EEG montage ch_pos:",
                  sorted((_finite_ch_pos(trimmed_mont.get_positions().get('ch_pos', {}))).keys()))
        print("→ params (ref/reref/eog/ransac/interp/corr):",
              {k: safe_params.get(k) for k in ("ref_chs","reref_chs","eog_chs",
                                               "ransac_channel_picks","interpolation_channel_picks","corr_chs")})
        raise RuntimeError(f"PyPREP failed: {e}") from e

    cleaned_eeg = result.raw_eeg
    try:
        cleaned_eeg.interpolate_bads(reset_bads=False)
    except Exception:
        raise RuntimeError("PyPREP bad channel interpolation failed.")
    
    cleaned_eeg.info['bads'] = []

    l_freq = safe_params.get('l_freq', None)
    h_freq = safe_params.get('h_freq', None)

    # if the raw is already filtered with h_freq = h_freq, skip filtering
    if h_freq is not None and raw_eeg.info.get('highpass', None) is not None:
        if np.isclose(raw_eeg.info['highpass'], float(h_freq)):
            h_freq = None
    # if the raw is already filtered with l_freq = l_freq, skip filtering
    if l_freq is not None and raw_eeg.info.get('lowpass', None) is not None:
        if np.isclose(raw_eeg.info['lowpass'], float(l_freq)):
            l_freq = None

    if l_freq is not None or h_freq is not None:
        cleaned_eeg.filter(l_freq=l_freq, h_freq=h_freq, fir_design='firwin')

    non_eeg = raw.copy().pick([ch for ch, t in zip(raw.ch_names, raw.get_channel_types()) if t != 'eeg'])
    if len(non_eeg.ch_names):
        cleaned_eeg.add_channels([non_eeg], force_update_info=True)

    return cleaned_eeg

def _max_eeg_time_std(inst: mne.io.BaseRaw) -> float:
    data = inst.get_data(picks="eeg")
    return float(np.nanmax(np.nanstd(data, axis=1))) if data.size else 0.0


def _apply_ica_excludes(inst: mne.io.BaseRaw, fitted_ica: ICA, exclude: list[int]) -> mne.io.BaseRaw:
    fitted_ica.exclude = list(exclude)
    if not fitted_ica.exclude:
        return inst.copy()
    before = _max_eeg_time_std(inst)
    cleaned = fitted_ica.apply(inst.copy())
    if before and _max_eeg_time_std(cleaned) <= before * 1e-6:
        logging.warning("ICA application flattened EEG; keeping pre-ICA data.")
        fitted_ica.exclude = []
        return inst.copy()
    return cleaned


def run_asr_ica(
    raw: mne.io.BaseRaw,
    asr_thresh: float | int | bool = 20,
    random_seed: int = 420,
    n_comp: int | None = None,
    remove_cfa_flag: bool = False,
) -> tuple[mne.io.BaseRaw, ICA]:
    raw_eeg = raw.copy().pick('eeg')
    raw_eeg, _ = mne.set_eeg_reference(raw_eeg, ref_channels='average')

    if asr_thresh:
        if asrpy is None:
            raise RuntimeError("ASR is enabled but the 'asrpy' package is not installed. Disable 'Run ASR' or install asrpy.")
        asr = asrpy.ASR(sfreq=raw_eeg.info['sfreq'], cutoff=asr_thresh)
        asr.fit(raw_eeg.copy().resample(100))
        raw_eeg = asr.transform(raw_eeg)

    n_channels = len(raw_eeg.ch_names) - len(raw_eeg.info.get('bads', []))
    if n_comp is None:
        n_comp = max(2, min(n_channels - 1, 48))
    ica = ICA(n_components=n_comp, method='infomax', random_state=random_seed, fit_params={"extended": True})
    ica.fit(raw_eeg, decim=3)

    if label_components is None:
        raise RuntimeError("ICA+ICLabel is enabled but the 'mne_icalabel' package is not installed. Disable 'Run ICA + ICLabel' or install mne-icalabel.")
    iclabel_result = label_components(raw_eeg, ica, method="iclabel")
    labels = iclabel_result.get('labels', [])
    if remove_cfa_flag:
        bads = [i for i, lab in enumerate(labels) if lab not in ('brain', 'other')]
        bad_label_dict = {i: label for i, label in enumerate(labels) if i in bads}
        logging.info(f"ICLabel-based ICA pruning (removing cardiac): {bad_label_dict}")
    else:
        bads = [i for i, lab in enumerate(labels) if lab not in ('brain', 'heart', 'other')]
        bad_label_dict = {i: label for i, label in enumerate(labels) if i in bads}
        logging.info(f"ICLabel-based ICA pruning (not removing cardiac): {bad_label_dict}")

    # see if there are any EOG leads to help identify EOG components
    if any(t == 'eog' for t in raw.get_channel_types()):
        try:
            eog_inds, _ = ica.find_bads_eog(raw_eeg)
            bads = sorted(set(bads) | set(eog_inds))
        except Exception:
            pass

    eeg_clean = _apply_ica_excludes(raw_eeg, ica, bads)
    if not ica.exclude:
        bad_label_dict = {}
    eeg_clean.interpolate_bads(reset_bads=True)

    non_eeg = raw.copy().pick([ch for ch, t in zip(raw.ch_names, raw.get_channel_types()) if t != 'eeg'])
    if len(non_eeg.ch_names):
        eeg_clean.add_channels([non_eeg], force_update_info=True)

    return eeg_clean, ica, bad_label_dict


def run_asr_only(
    raw: mne.io.BaseRaw,
    asr_thresh: float | int | bool = 20,
) -> mne.io.BaseRaw:
    """Apply ASR without ICA/ICLabel, preserving non-EEG channels."""
    if not asr_thresh:
        return raw
    if asrpy is None:
        raise RuntimeError("ASR is enabled but the 'asrpy' package is not installed. Disable 'Run ASR' or install asrpy.")
    raw_eeg = raw.copy().pick('eeg')
    raw_eeg, _ = mne.set_eeg_reference(raw_eeg, ref_channels='average')
    asr = asrpy.ASR(sfreq=raw_eeg.info['sfreq'], cutoff=asr_thresh)
    asr.fit(raw_eeg.copy().resample(100))
    eeg_clean = asr.transform(raw_eeg)
    non_eeg = raw.copy().pick([ch for ch, t in zip(raw.ch_names, raw.get_channel_types()) if t != 'eeg'])
    if len(non_eeg.ch_names):
        eeg_clean.add_channels([non_eeg], force_update_info=True)
    return eeg_clean

def detect_r_peaks(raw: mne.io.BaseRaw, stim_channel: str = 'STI 014', gap_threshold_factor: float = 2.0):
    sfreq = raw.info['sfreq']
    ecg_chs = [ch for ch, t in zip(raw.ch_names, raw.get_channel_types()) if t == 'ecg']
    if not ecg_chs:
        return raw, []
    ecg_channel = ecg_chs[0]
    ecg = raw.get_data(picks=ecg_channel)[0]
    ecg_cleaned = nk.ecg_clean(ecg, sampling_rate=sfreq)
    signals, info = nk.ecg_peaks(ecg_cleaned, sampling_rate=sfreq)
    peaks = info.get('ECG_R_Peaks', None)
    if peaks is None or len(peaks) < 2:
        return raw, []
    rr_samples = np.diff(peaks)
    rr_sec = rr_samples / sfreq
    median_rr = float(np.median(rr_sec))
    max_gap = gap_threshold_factor * median_rr
    all_peaks = list(peaks)
    for idx, interval in enumerate(rr_sec):
        if interval > max_gap:
            n_missing = int(np.round(interval / median_rr)) - 1
            for j in range(1, n_missing + 1):
                new_sample = peaks[idx] + int(j * median_rr * sfreq)
                all_peaks.append(new_sample)
    all_peaks = np.unique(all_peaks).astype(int)
    events = [(int(s), 0, 1) for s in all_peaks if 0 <= int(s) < raw.n_times]
    stim_data = np.zeros((1, raw.n_times), dtype=int)
    for sample, _, _ in events:
        stim_data[0, sample] = 1
    info_stim = mne.create_info([stim_channel], sfreq, ch_types=['stim'])
    stim_raw = mne.io.RawArray(stim_data, info_stim)
    raw.add_channels([stim_raw], force_update_info=True)
    return raw, events

# ------------------------------- Orchestration ---------------------------- #

def _read_raw_any(input_path: str, *, preload=True, verbose=False) -> mne.io.BaseRaw:
    suffix = Path(input_path).suffix.lower()
    if suffix == ".bdf":
        return mne.io.read_raw_bdf(input_path, preload=preload, verbose=verbose)
    if suffix == ".edf":
        return mne.io.read_raw_edf(input_path, preload=preload, verbose=verbose)
    if suffix == ".fif":
        return mne.io.read_raw_fif(input_path, preload=preload, verbose=verbose)
    if suffix == ".vhdr":
        return mne.io.read_raw_brainvision(input_path, preload=preload, verbose=verbose)
    raise ValueError(f"Unsupported input format: {input_path}")

def preprocess_edf(
    input_path: str,
    output_path: str | None = None,
    redo: bool = False,
    pyprep_dict: dict | None = None,
    asr_threshold=None,
    random_seed: int = 42,
    n_comp: int | None = None,
    logging_path: str | None = None,
    remove_cfa_override: bool | None = None,
    flip_ecg: bool = False,
    ecg_channel: str | None = None,
    stim_keep: list | None = None,
    spirometry_channel: str | None = None,
    spirometry_dir: str | Path | None = None,
):
    if pyprep_dict is None:
        try:
            pyprep_dict = prep_params
        except Exception:
            if _CFG_RT is not None:
                pyprep_dict = {
                    "line_freqs": _CFG_RT.line_freqs,
                    "ref_chs": _CFG_RT.ref_chs,
                    "reref_chs": _CFG_RT.reref_chs,
                    "l_freq": _CFG_RT.high_pass,
                    "h_freq": _CFG_RT.low_pass,
                    "ransac": _CFG_RT.prep_ransac,
                }
            else:
                pyprep_dict = {
                    "line_freqs": (50.0, 100.0),
                    "ref_chs": "eeg",
                    "reref_chs": "eeg",
                    "l_freq": 1.0,
                    "h_freq": 100.0,
                    "ransac": True,
                }
    if asr_threshold is None:
        asr_threshold = do_asr

    # resolve output path
    if output_path is None:
        base = Path(input_path).stem + "_pp_raw.fif"
        output_path = str(Path(output_dir) / base)

    # resolve log path (safe default if config.log_file is empty)
    if logging_path is None or str(logging_path).strip() == "":
        logging_path = str(Path(output_dir) / "logs" / "preprocessing.tsv")

    # idempotency
    if not redo and Path(output_path).exists():
        logging.info(f"Output exists, skipping: {output_path}")
        raw_cached = mne.io.read_raw_fif(output_path, preload=True)
        had_rsp = _has_spirometry_annotations(raw_cached)
        spirometry_rec = save_spirometry_analysis(
            raw_cached,
            spirometry_channel,
            spirometry_dir or (Path(output_dir) / "spirometry"),
            _preproc_base_from_output(output_path),
            redo=False,
        )
        if spirometry_rec.get("spirometry_status") == "ok" and not had_rsp:
            raw_cached.save(str(output_path), overwrite=True)
        rec = {
            "utc": _now_utc_iso(),
            "status": "OK",
            "stage": "cached",
            "input": input_path,
            "output": str(output_path),
            "n_channels": len(raw_cached.ch_names),
            "remove_cfa": str(remove_cfa_override),
            "remove_cfa_mode": str(getattr(_CFG_RT, "remove_cfa_mode", "")),
            "flip_ecg": str(bool(flip_ecg)),
        }
        rec.update(spirometry_rec)
        _append_preproc_tsv(logging_path, rec)
        return raw_cached

    # 1) Read raw input (preload)
    raw = _read_raw_any(input_path, preload=True, verbose=False)

    # 2) Type auxiliaries early (no n_times change)
    aux_like = ["IBI", "BURSTS", "SUPPR", "T1", "T2", "26", "27", "28", "29", "30"]
    present_aux = [ch for ch in aux_like if ch in raw.ch_names]
    if present_aux:
        raw.set_channel_types({ch: "misc" for ch in present_aux})
    _mark_spirometry_channel(raw, spirometry_channel)
    _mark_ecg_channels(raw, ecg_channel)

    # Optional ECG polarity flip
    if flip_ecg:
        _flip_ecg_channels(raw, ecg_channel)

    # 3) Global resample once (for identical n_times across all channels)
    if target_sfreq and not np.isclose(raw.info["sfreq"], float(target_sfreq)):
        raw.resample(float(target_sfreq), npad="auto")

    # 4) Rename / classify / montage
    
    # set stim leads to stim type before montage
    if stim_keep:
        for ch in stim_keep:
            if ch in raw.ch_names:
                raw.set_channel_types({ch: "stim"})

    _mark_ecg_channels(raw, ecg_channel)

    raw = standardise_and_montage(raw)

    # Sanity: ensure EEG remains
    if sum(1 for t in raw.get_channel_types() if t == "eeg") == 0:
        raise ValueError("No EEG channels remain after montage/drop; check montage name and channel labels.")

    # 5) PyPREP on EEG-only; recombine with non-EEG
    try:
        do_pyprep_flag = do_pyprep  # global from set_runtime_config
    except NameError:
        do_pyprep_flag = getattr(_CFG_RT, "use_pyprep", True)
    if do_pyprep_flag:
        raw = run_pyprep(raw, pyprep_dict, random_seed=random_seed, ransac=bool(_CFG_RT.prep_ransac))

    # 6) ASR + ICA (+/- remove cardiac via ICLabel), or ASR-only when ICA is disabled
    if remove_cfa_override is None:
        remove_flag = bool(remove_cfa)
    else:
        remove_flag = bool(remove_cfa_override)
    use_ica_flag = bool(getattr(_CFG_RT, "use_ica", True)) if _CFG_RT is not None else True
    if use_ica_flag:
        raw, ica_model, bad_label_dict = run_asr_ica(
            raw,
            asr_thresh=asr_threshold,
            random_seed=random_seed,
            n_comp=n_comp,
            remove_cfa_flag=remove_flag,
        )
    else:
        raw = run_asr_only(raw, asr_threshold)
        ica_model = None
        bad_label_dict = {}

    # 7) ECG → R-peaks → stim events
    raw, events = detect_r_peaks(raw)
    outp = Path(output_path)
    spirometry_rec = save_spirometry_analysis(
        raw,
        spirometry_channel,
        spirometry_dir or (Path(output_dir) / "spirometry"),
        _preproc_base_from_output(outp),
        redo=redo,
    )

    raw.info["bads"] = []  # clear any stragglers

    # 8) Save (guard empty parents)
    if str(outp.parent) not in ("", "."):
        outp.parent.mkdir(parents=True, exist_ok=True)
    raw.set_meas_date(1)
    raw.save(str(outp), overwrite=True)

    # save ICA separately (fail fast if it cannot be written)
    ica_path = ""
    if ica_model is not None:
        stem_no_pp = outp.stem.replace("_pp_raw", "")
        ica_path = outp.with_name(stem_no_pp + "_ica.fif")
        ica_path.parent.mkdir(parents=True, exist_ok=True)
        ica_model.save(str(ica_path), overwrite=True)

    # 9) Log (guard empty parents)
    rec = {
        'input': input_path,
        'output': str(outp),
        'n_channels': len(raw.ch_names),
        'n_events': len(events),
        'bad_channels': ';'.join(raw.info.get('bads', [])) if raw.info.get('bads') else '',
        'remove_cfa': str(remove_flag),
        'remove_cfa_mode': str(getattr(_CFG_RT, "remove_cfa_mode", "")),
        'flip_ecg': str(bool(flip_ecg)),
        'asr_threshold': str(asr_threshold),
        'n_comp': str(n_comp),
        'stim_keep': ";".join(stim_keep) if stim_keep else "",
        'ica_path': ica_path,
        'channels_dropped': str(len(drop_candidates)) if 'drop_candidates' in locals() else "0",
        'montage_name': str(getattr(_CFG_RT, "montage_name", "auto")),
        'ica_bads': ';'.join(str(key) for key in bad_label_dict.keys()),
        'ica_bad_labels': ';'.join(str(val) for val in bad_label_dict.values()),
    }
    rec.update(spirometry_rec)
    logp = Path(logging_path)
    if str(logp.parent) not in ("", "."):
        logp.parent.mkdir(parents=True, exist_ok=True)
    file_exists = logp.exists() and logp.stat().st_size > 0
    with open(str(logp), 'a', newline='') as f:
        w = csv.DictWriter(f, fieldnames=rec.keys(), delimiter='\t')
        if not file_exists:
            w.writeheader()
        w.writerow(rec)

    return raw




def _fit_asr_ica_iclabel_once(
    raw: mne.io.BaseRaw,
    asr_thresh: float | int | bool = 20,
    random_seed: int = 420,
    n_comp: int | None = None,
):
    """Run (optional) ASR, fit ICA once, and run ICLabel.

    Returns
    -------
    raw_eeg_proc : mne.io.BaseRaw
        EEG-only data after referencing and optional ASR, used for ICA fitting.
    ica : mne.preprocessing.ICA
        Fitted ICA model.
    labels : list[str]
        ICLabel class label per component.
    bad_label_dict : dict[int, str]
        Mapping from excluded component index to ICLabel label for the *remove-CFA* rule
        (i.e. excluding anything not in ('brain','other')).
    bad_label_dict_keep : dict[int, str]
        Mapping for the *keep-CFA* rule (i.e. excluding anything not in ('brain','heart','other')).
    eog_inds : list[int]
        Components additionally flagged by find_bads_eog (may be empty).
    """
    raw_eeg = raw.copy().pick('eeg')
    raw_eeg, _ = mne.set_eeg_reference(raw_eeg, ref_channels='average')

    # Optional ASR (fit at 100 Hz to reduce cost, consistent with existing code)
    if asr_thresh:
        if asrpy is None:
            raise RuntimeError("ASR is enabled but the 'asrpy' package is not installed. Disable 'Run ASR' or install asrpy.")
        asr = asrpy.ASR(sfreq=raw_eeg.info['sfreq'], cutoff=asr_thresh)
        asr.fit(raw_eeg.copy().resample(100))
        raw_eeg = asr.transform(raw_eeg)

    n_channels = len(raw_eeg.ch_names) - len(raw_eeg.info.get('bads', []))
    if n_comp is None:
        n_comp = max(2, min(n_channels - 1, 48))

    ica = ICA(
        n_components=n_comp,
        method='infomax',
        random_state=random_seed,
        fit_params={"extended": True},
    )
    ica.fit(raw_eeg, decim=3)

    # ICLabel expects 1–100 Hz bandpass; use a filtered *copy* for feature extraction.
    inst_for_iclabel = raw_eeg.copy()
    try:
        inst_for_iclabel.filter(l_freq=1.0, h_freq=100.0, fir_design='firwin', verbose=False)
    except Exception:
        # If filtering fails for some reason, proceed (mne_icalabel will likely warn / fail).
        pass

    if label_components is None:
        raise RuntimeError("ICA+ICLabel is enabled but the 'mne_icalabel' package is not installed. Disable 'Run ICA + ICLabel' or install mne-icalabel.")
    iclabel_result = label_components(inst_for_iclabel, ica, method="iclabel")
    labels = list(iclabel_result.get('labels', []))

    bads_remove = [i for i, lab in enumerate(labels) if lab not in ('brain', 'other')]
    bads_keep = [i for i, lab in enumerate(labels) if lab not in ('brain', 'heart', 'other')]

    bad_label_dict_remove = {i: labels[i] for i in bads_remove if i < len(labels)}
    bad_label_dict_keep = {i: labels[i] for i in bads_keep if i < len(labels)}

    # EOG-based IC detection (optional)
    eog_inds = []
    if any(t == 'eog' for t in raw.get_channel_types()):
        try:
            eog_inds, _ = ica.find_bads_eog(raw_eeg)
        except Exception:
            eog_inds = []

    if eog_inds:
        # EOG inds should be excluded in both variants
        bads_remove = sorted(set(bads_remove) | set(eog_inds))
        bads_keep = sorted(set(bads_keep) | set(eog_inds))
        for i in eog_inds:
            if i < len(labels):
                bad_label_dict_remove.setdefault(i, "eog")
                bad_label_dict_keep.setdefault(i, "eog")

    return raw_eeg, ica, labels, bads_remove, bads_keep, bad_label_dict_remove, bad_label_dict_keep, eog_inds


def preprocess_edf_both(
    input_path: str,
    output_path_remove: str,
    output_path_keep: str,
    redo: bool = False,
    pyprep_dict: dict | None = None,
    asr_threshold=None,
    random_seed: int = 42,
    n_comp: int | None = None,
    logging_path: str | None = None,
    flip_ecg: bool = False,
    ecg_channel: str | None = None,
    stim_keep: list | None = None,
    spirometry_channel: str | None = None,
    spirometry_dir: str | Path | None = None,
):
    """Preprocess once up to ICA+ICLabel, then write *two* outputs.

    - **remove**: excludes ICLabel components not in ('brain','other')
    - **keep**: excludes ICLabel components not in ('brain','heart','other')

    This avoids running read/resample/montage/PyPREP/ASR/ICA/ICLabel twice.
    """
    # Resolve log path
    if logging_path is None or str(logging_path).strip() == "":
        logging_path = str(Path(output_dir) / "logs" / "preprocessing.tsv")

    # Idempotency
    if (not redo) and Path(output_path_remove).exists() and Path(output_path_keep).exists():
        logging.info("Both outputs exist, skipping: %s AND %s", output_path_remove, output_path_keep)
        raw_remove = mne.io.read_raw_fif(output_path_remove, preload=True)
        raw_keep = mne.io.read_raw_fif(output_path_keep, preload=True)
        for raw_cached, out_path in ((raw_remove, output_path_remove), (raw_keep, output_path_keep)):
            had_rsp = _has_spirometry_annotations(raw_cached)
            spirometry_rec = save_spirometry_analysis(
                raw_cached,
                spirometry_channel,
                spirometry_dir or (Path(output_dir) / "spirometry"),
                _preproc_base_from_output(out_path),
                redo=False,
            )
            if spirometry_rec.get("spirometry_status") == "ok" and not had_rsp:
                raw_cached.save(str(out_path), overwrite=True)
            rec = {
                "utc": _now_utc_iso(),
                "status": "OK",
                "stage": "cached",
                "input": input_path,
                "output": str(out_path),
                "n_channels": len(raw_cached.ch_names),
                "remove_cfa_mode": "both",
                "flip_ecg": str(bool(flip_ecg)),
            }
            rec.update(spirometry_rec)
            _append_preproc_tsv(logging_path, rec)
        return raw_remove, raw_keep

    # Shared steps: read, resample, montage, PyPREP
    try:
        # (1) Read
        raw = _read_raw_any(input_path, preload=True, verbose=False)

        # (2) Re-type obvious aux channels (keep consistent with preprocess_edf)
        aux_like = ["IBI", "BURSTS", "SUPPR", "T1", "T2", "26", "27", "28", "29", "30"]
        present_aux = [ch for ch in aux_like if ch in raw.ch_names]
        if present_aux:
            raw.set_channel_types({ch: "misc" for ch in present_aux})
        _mark_spirometry_channel(raw, spirometry_channel)
        _mark_ecg_channels(raw, ecg_channel)

        # Flip ECG if requested
        if flip_ecg:
            _flip_ecg_channels(raw, ecg_channel)

        # (3) Resample
        if target_sfreq and not np.isclose(raw.info["sfreq"], float(target_sfreq)):
            raw.resample(float(target_sfreq), npad="auto")

        # (4) Mark stim + ecg by name before montage
        if stim_keep:
            for ch in stim_keep:
                if ch in raw.ch_names:
                    raw.set_channel_types({ch: "stim"})
        _mark_ecg_channels(raw, ecg_channel)

        # (5) Standardise and apply montage (drops EEG channels not in montage)
        raw = standardise_and_montage(raw)

        if sum(1 for t in raw.get_channel_types() if t == "eeg") == 0:
            raise ValueError("No EEG channels remain after montage/drop; check montage name and channel labels.")

        # (6) PyPREP
        if pyprep_dict is None:
            pyprep_dict = prep_params
        do_pyprep_flag = True
        try:
            do_pyprep_flag = bool(do_pyprep)
        except Exception:
            if _CFG_RT is not None:
                do_pyprep_flag = bool(_CFG_RT.use_pyprep)
        if do_pyprep_flag:
            raw = run_pyprep(raw, pyprep_dict, random_seed=random_seed, ransac=bool(_CFG_RT.prep_ransac))

        # (7) ASR + ICA + ICLabel once
        if asr_threshold is None:
            asr_threshold = do_asr

        raw_eeg_proc, ica, labels, bads_remove, bads_keep, bad_map_remove, bad_map_keep, eog_inds = _fit_asr_ica_iclabel_once(
            raw,
            asr_thresh=asr_threshold,
            random_seed=420 + int(random_seed),
            n_comp=n_comp,
        )

        non_eeg = raw.copy().pick([ch for ch, t in zip(raw.ch_names, raw.get_channel_types()) if t != 'eeg'])

        def _apply_and_recombine(exclude, variant_label_dict):
            ica_local = ica.copy()
            eeg_clean = _apply_ica_excludes(raw_eeg_proc, ica_local, list(exclude))
            if not ica_local.exclude:
                variant_label_dict = {}
            try:
                eeg_clean.interpolate_bads(reset_bads=True)
            except Exception:
                pass
            if len(non_eeg.ch_names):
                eeg_clean.add_channels([non_eeg], force_update_info=True)
            # R-peak events
            eeg_clean, events = detect_r_peaks(eeg_clean)
            spirometry_rec = save_spirometry_analysis(
                eeg_clean,
                spirometry_channel,
                spirometry_dir or (Path(output_dir) / "spirometry"),
                _preproc_base_from_output(output_path_remove),
                redo=redo,
            )
            eeg_clean.info["bads"] = []
            return eeg_clean, events, variant_label_dict, ica_local, spirometry_rec

        raw_remove, events_remove, bad_labels_remove, ica_remove, spirometry_remove = _apply_and_recombine(bads_remove, bad_map_remove)
        raw_keep, events_keep, bad_labels_keep, ica_keep, spirometry_keep = _apply_and_recombine(bads_keep, bad_map_keep)

        # (8) Save both raws + ICA models
        for outp_str, raw_out, ica_out, bad_labels, events, mode in [
            (output_path_remove, raw_remove, ica_remove, bad_labels_remove, events_remove, "remove"),
            (output_path_keep, raw_keep, ica_keep, bad_labels_keep, events_keep, "keep"),
        ]:
            outp = Path(outp_str)
            outp.parent.mkdir(parents=True, exist_ok=True)
            raw_out.set_meas_date(1)
            raw_out.save(str(outp), overwrite=True)

            ica_path = ""
            try:
                stem_no_pp = outp.stem.replace("_pp_raw", "")
                ica_path = str(outp.with_name(stem_no_pp + "_ica.fif"))
                Path(ica_path).parent.mkdir(parents=True, exist_ok=True)
                ica_out.save(ica_path, overwrite=True)
            except Exception:
                ica_path = ""

            rec = {
                "utc": _now_utc_iso(),
                "status": "OK",
                "stage": "save",
                "input": input_path,
                "output": str(outp),
                "n_channels": len(raw_out.ch_names),
                "n_events": len(events),
                "bad_channels": ';'.join(raw_out.info.get('bads', [])) if raw_out.info.get('bads') else '',
                "remove_cfa": str(mode == "remove"),
                "remove_cfa_mode": "both",
                "flip_ecg": str(bool(flip_ecg)),
                "asr_threshold": str(asr_threshold),
                "n_comp": str(n_comp),
                "stim_keep": ";".join(stim_keep) if stim_keep else "",
                "montage_name": str(getattr(_CFG_RT, "montage_name", "auto")),
                "ica_path": ica_path,
                "ica_bads": ';'.join(str(k) for k in sorted(bad_labels.keys())),
                "ica_bad_labels": ';'.join(str(bad_labels[k]) for k in sorted(bad_labels.keys())),
            }
            rec.update(spirometry_remove if mode == "remove" else spirometry_keep)
            _append_preproc_tsv(logging_path, rec)

        return raw_remove, raw_keep

    except Exception as e:
        logging.exception("preprocess_edf_both failed for %s", input_path)
        _append_preproc_tsv(logging_path, {
            "utc": _now_utc_iso(),
            "status": "FAIL",
            "stage": "preprocess_edf_both",
            "input": input_path,
            "output": f"{output_path_remove} | {output_path_keep}",
            "error_type": type(e).__name__,
            "error": str(e),
        })
        raise

# ----------------------------------- CLI (optional) ----------------------- #
if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="EDF→FIF EEG preprocessing")
    parser.add_argument("edf", help="Path to input .edf file")
    parser.add_argument("--out", help="Output .fif path (default: output_dir/<name>_pp_raw.fif)")
    parser.add_argument("--redo", action="store_true", help="Overwrite existing outputs")
    parser.add_argument("--n_comp", type=int, default=None, help="ICA components (default: auto)")
    parser.add_argument("--asr", type=float, default=None, help="ASR cutoff (default: config.do_asr)")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)
    preprocess_edf(
        args.edf,
        output_path=args.out,
        redo=args.redo,
        n_comp=args.n_comp,
        asr_threshold=args.asr,
    )