src/analysis/load.py

#!/usr/bin/env python
"""Load a cap recording and cut it into labelled MI epochs.

This is the "录制 → 分析加载, 按 trigger 切 epoch" link. A recording saved by cap_gui is
either a `.npz` (raw µV + a per-sample label track) or an MNE `.fif` (with the labels as
annotations). Here we:

    recording ──► MNE Raw (µV→V, montage) ──► filter (notch + band-pass)
              ──► bad-channel detect + interpolate     (simple artifact handling)
              ──► events from the label track's rising edges
              ──► Epochs, one per imagery onset, baseline-corrected

The label track is the software `marker` (small MI codes 1–5, see common.mi_events) when
present, else the hardware `trigger` column. Heavier artifact handling (ICA / autoreject /
EOG-EMG regression) is deliberately NOT here yet — see analysis/artifacts.py / the plan.

CLI:
  python src/analysis/load.py recordings/cap32_*.npz              # summary + event counts
  python src/analysis/load.py recordings/cap32_*.npz --epochs     # build epochs, print shape
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))                          # src/
from common.montage import CAP32_CHANNELS                      # noqa: E402
from common.mi_events import CODE_TO_LABEL                      # noqa: E402

DEFAULT_TMIN, DEFAULT_TMAX = -2.0, 4.0
DEFAULT_BASELINE = (-1.5, -0.5)
FULLSCALE_UV = (2 ** 23 - 1) * 0.02235175   # ADS1299 ±full scale ≈ 187500 µV


# --------------------------------------------------------------------- raw npz
def load_npz(path):
    z = np.load(path, allow_pickle=True)
    data = z["data"].astype(np.float64)                        # (n_ch, N) µV
    fs = float(z["fs"])
    ch = [str(c) for c in z["ch_names"]] if "ch_names" in z else list(CAP32_CHANNELS)
    trigger = z["trigger"].astype(np.int64) if "trigger" in z else np.zeros(data.shape[1], int)
    marker = z["marker"].astype(np.int64) if "marker" in z else None
    gap = z["gap"].astype(np.int8) if "gap" in z else None      # 1 = reconstructed sample
    rec = dict(data=data, fs=fs, ch_names=ch, trigger=trigger, marker=marker, gap=gap)
    # format v2: explicit trial table + metadata written by the paradigm
    if "trial_onset" in z:
        rec["trials"] = dict(onset=z["trial_onset"].astype(np.int64),
                             code=z["trial_code"].astype(np.int64),
                             name=[str(s) for s in z["trial_name"]],
                             cue=z["trial_cue_onset"].astype(np.int64)
                             if "trial_cue_onset" in z else None)
    if "meta_json" in z:
        import json
        try:
            rec["meta"] = json.loads(str(z["meta_json"]))
        except Exception:
            pass
    return rec


def events_from_trials(trials):
    """Exact events from the paradigm's trial table (preferred over edge detection)."""
    on = np.asarray(trials["onset"]); code = np.asarray(trials["code"])
    keep = on >= 0
    on, code = on[keep], code[keep]
    events = np.column_stack([on, np.zeros_like(on), code]).astype(int)
    event_id = {}
    for c, n in zip(code.tolist(), np.asarray(trials["name"])[keep].tolist()):
        event_id.setdefault(str(n), int(c))
    return events, event_id


def label_track(rec):
    """Pick the label track: software `marker` if it carries any event, else hardware."""
    m = rec.get("marker")
    if m is not None and np.any(m != 0):
        return m, "marker"
    return rec["trigger"], "trigger"


# --------------------------------------------------------------- events / raw
def events_from_track(track, fs=None):
    """Rising edges of a per-sample integer label -> MNE events (n,3) + event_id.
    An event starts where the label changes to a nonzero value."""
    track = np.asarray(track)
    onsets = np.where((track[1:] != track[:-1]) & (track[1:] != 0))[0] + 1
    codes = track[onsets].astype(int)
    events = np.column_stack([onsets, np.zeros_like(onsets), codes]).astype(int)
    present = sorted(set(codes.tolist()))
    event_id = {CODE_TO_LABEL.get(c, f"code{c}"): int(c) for c in present}
    return events, event_id


def to_raw(rec, montage=True):
    import mne
    info = mne.create_info(list(rec["ch_names"]), rec["fs"], "eeg")
    raw = mne.io.RawArray(rec["data"] * 1e-6, info, verbose="ERROR")   # µV -> V
    if montage:
        raw.set_montage(mne.channels.make_standard_montage("standard_1020"),
                        match_case=False, on_missing="ignore", verbose="ERROR")
    return raw


def read_recording(path):
    """Unified reader → (raw, events, event_id). Accepts .npz or .fif."""
    import mne
    path = str(path)
    if path.endswith(".fif"):
        raw = mne.io.read_raw_fif(path, preload=True, verbose="ERROR")
        try:
            events, ann_id = mne.events_from_annotations(raw, verbose="ERROR")
            # map annotation descriptions ("left"/"right"/"T2"/…) back to MI codes
            from common.mi_events import MI_TASK_CODES
            event_id = {k: v for k, v in ann_id.items()}
            return raw, events, event_id
        except Exception:
            return raw, np.empty((0, 3), int), {}
    rec = load_npz(path)
    raw = to_raw(rec)
    if rec.get("trials") is not None and np.any(np.asarray(rec["trials"]["onset"]) >= 0):
        events, event_id = events_from_trials(rec["trials"])   # v2: exact, from the paradigm
    else:
        track, _src = label_track(rec)
        events, event_id = events_from_track(track, rec["fs"])
    return raw, events, event_id


# ------------------------------------------------------- simple artifact steps
def detect_bad_channels(raw, flat_uv=0.5, noisy_uv=150.0, rail_uv=1.5e5, hf_z=6.0,
                        min_bad_frac=0.5):
    """Conservative bad-channel flags (no ICA) using ABSOLUTE thresholds so real signal is
    never touched — a channel with more mu/beta is NOT bad. Flags only:
      • flat / dead  (std < flat_uv),
      • railed       (peak-to-peak > rail_uv ≈ 150 mV, ADS1299 saturated),
      • grossly noisy (std > noisy_uv µV — non-physiological for scalp EEG),
      • high-frequency junk (sample-to-sample jitter a strong outlier vs the array).
    Returns a list of channel names. Deliberately misses the "slightly noisy posterior"
    case — that is left to ICA/autoreject, which we add separately."""
    # An INTERMITTENT channel must not be condemned for the whole recording: only call it
    # globally bad when it is bad for at least `min_bad_frac` of the time. Segments where an
    # otherwise-good channel rails are handled per-epoch instead (see channel_uptime).
    frac, _ = channel_uptime(raw)
    # Judge amplitude on a HIGH-PASSED copy: dry electrodes carry huge DC drift, so on
    # unfiltered data every channel's std exceeds `noisy_uv` and the whole array would be
    # condemned. This makes the function safe to call on raw or filtered input alike.
    hp = raw.copy().filter(1.0, None, verbose="ERROR") if raw.info["highpass"] < 0.5 else raw
    x = hp.get_data() * 1e6                                    # V -> µV
    sd = x.std(axis=1)
    ptp = np.ptp(x, axis=1)
    hf = np.std(np.diff(x, axis=1), axis=1)                    # high-freq content proxy
    del ptp
    bad = set()
    for i, name in enumerate(raw.ch_names):
        if frac[i] >= min_bad_frac or sd[i] > noisy_uv:
            bad.add(name)
    # high-freq outlier: broken electrodes jitter far more than the median channel
    keep = np.array([raw.ch_names[i] not in bad for i in range(len(hf))])
    if keep.sum() >= 6:
        med = np.median(hf[keep]); mad = np.median(np.abs(hf[keep] - med)) + 1e-9
        for i, name in enumerate(raw.ch_names):
            if name not in bad and (hf[i] - med) / (1.4826 * mad) > hf_z and hf[i] > 40.0:
                bad.add(name)
    return sorted(bad)


def channel_uptime(raw, win_s=1.0, rail_frac=0.5):
    """Time-RESOLVED channel health: (frac_bad per channel, mask (n_ch, n_win)).

    The cap has at least one INTERMITTENT electrode (F7 sat at +full scale for four whole
    sessions, yet reads normally in others — a loose connector, not a permanent break).
    A single whole-recording verdict handles that badly: a channel that is fine for half a
    session gets interpolated away entirely, and one that rails only briefly may not be
    flagged at all while its rail artifact still poisons CAR and any epoch it touches.
    So score every channel per window and let the caller act per segment."""
    x = raw.get_data() * 1e6
    fs = raw.info["sfreq"]
    w = max(1, int(win_s * fs))
    n = x.shape[1] // w
    if n == 0:
        return np.zeros(x.shape[0]), np.zeros((x.shape[0], 0), bool)
    seg = x[:, :n * w].reshape(x.shape[0], n, w)
    railed = (np.abs(seg) > 0.97 * FULLSCALE_UV).mean(-1) > rail_frac
    flat = seg.std(-1) < 0.5
    bad = railed | flat
    return bad.mean(1), bad


def clean_raw(raw, l_freq=1.0, h_freq=40.0, notch=50.0, car=False,
              interpolate=True, verbose=False):
    """Filter + (optional) notch + bad-channel interpolation. Returns a *copy*.
    This is the light, dependency-free cleaning; ICA/autoreject live elsewhere."""
    raw = raw.copy()
    if notch:
        raw.notch_filter(notch, verbose="ERROR")
    raw.filter(l_freq, h_freq, verbose="ERROR")
    bad = detect_bad_channels(raw)
    raw.info["bads"] = bad
    if verbose and bad:
        print(f"  bad channels: {bad}")
    if interpolate and bad and raw.get_montage() is not None:
        raw.interpolate_bads(reset_bads=True, verbose="ERROR")
    if car:
        raw.set_eeg_reference("average", verbose="ERROR")
    return raw, bad


# --------------------------------------------------------------------- epochs
def make_epochs(path, tmin=DEFAULT_TMIN, tmax=DEFAULT_TMAX, baseline=DEFAULT_BASELINE,
                picks=None, l_freq=1.0, h_freq=40.0, notch=50.0, car=False,
                interpolate=True, reject_uv=None, verbose=True, drop_filled=False):
    """Recording path → cleaned, baseline-corrected MI Epochs (labelled by task name)."""
    import mne
    raw, events, event_id = read_recording(path)
    if len(events) == 0:
        raise SystemExit(f"no events/markers found in {path} — was a paradigm run recorded?")
    if drop_filled and str(path).endswith(".npz"):
        gap = load_npz(path).get("gap")
        if gap is not None and gap.any():        # drop trials overlapping reconstructed data
            fs0 = raw.info["sfreq"]
            lo, hi = int(tmin * fs0), int(tmax * fs0)
            keep = [not gap[max(0, o + lo): o + hi].any() for o in events[:, 0]]
            n_drop = len(keep) - sum(keep)
            events = events[np.array(keep, bool)]
            if verbose and n_drop:
                print(f"  dropped {n_drop} epoch(s) containing UDP-gap-filled samples")
    raw, bad = clean_raw(raw, l_freq, h_freq, notch, car, interpolate, verbose)
    reject = dict(eeg=reject_uv * 1e-6) if reject_uv else None
    ep = mne.Epochs(raw, events, event_id=event_id, tmin=tmin, tmax=tmax,
                    baseline=baseline, picks=picks, preload=True, reject=reject,
                    reject_by_annotation=False, verbose="ERROR")
    if verbose:
        counts = {k: int((ep.events[:, 2] == v).sum()) for k, v in ep.event_id.items()}
        print(f"epochs: {len(ep)} × {len(ep.ch_names)}ch × {ep.times.size} samp "
              f"[{tmin}, {tmax}]s @ {raw.info['sfreq']:.0f}Hz  ·  {counts}"
              + (f"  ·  interpolated {bad}" if bad else ""))
    return ep


# --------------------------------------------------------------------- CLI
def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("path")
    ap.add_argument("--epochs", action="store_true", help="build epochs (not just summarise)")
    ap.add_argument("--tmin", type=float, default=DEFAULT_TMIN)
    ap.add_argument("--tmax", type=float, default=DEFAULT_TMAX)
    ap.add_argument("--no-interp", action="store_true")
    args = ap.parse_args()

    raw, events, event_id = read_recording(args.path)
    print(f"{args.path}\n  {len(raw.ch_names)} ch @ {raw.info['sfreq']:.0f} Hz, "
          f"{raw.n_times} samples ({raw.times[-1]:.1f}s)")
    print(f"  events: {len(events)}  event_id: {event_id}")
    if events.size:
        for name, code in event_id.items():
            print(f"    {name:8s} (code {code}): {(events[:,2]==code).sum()} trials")
    if args.epochs:
        make_epochs(args.path, tmin=args.tmin, tmax=args.tmax,
                    interpolate=not args.no_interp)


if __name__ == "__main__":
    main()