ADABase and REVELIO: Loading the Public Sets¶

import os
import pathlib as ptl

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import tables
from matplotlib.colors import ListedColormap
from matplotlib.patches import Patch

# folder with one subfolder per Zenodo record, see the setup notes above
DATA_PATH = ptl.Path(os.environ.get("COGNITIVE_LOAD_DATA", "data")).expanduser()

# the two public Zenodo records, by short name
SUBSETS = {
    "drive": {
        "title": "ADABase: Cognitive Driving Public Set",
        "folder": "ADABase-driving-public",
        "doi": "10.5281/zenodo.24051552",
    },
    "game": {
        "title": "REVELIO: Cognitive Gaming Public Set",
        "folder": "REVELIO-gaming-public",
        "doi": "10.5281/zenodo.24060698",
    },
}

# one file per session, named <participant>-<code>.h5
FILES = {
    experiment: sorted(
        p for p in (DATA_PATH / subset["folder"]).glob("*.h5") if p.stem.count("-") == 1
    )
    for experiment, subset in SUBSETS.items()
}
for experiment, files in FILES.items():
    if not files:
        raise FileNotFoundError(f"No .h5 files in {DATA_PATH / SUBSETS[experiment]['folder']}")

EXPERIMENT = {f.stem: experiment for experiment, files in FILES.items() for f in files}
# the first session of each set serves as example
EXAMPLE = {experiment: files[0] for experiment, files in FILES.items()}

for experiment, files in FILES.items():
    size_gb = sum(f.stat().st_size for f in files) / 1e9
    print(f"{experiment}: {SUBSETS[experiment]['title']}, {len(files)} sessions, {size_gb:.1f} GB")

plt.rcParams.update(
    {
        "figure.dpi": 120,
        "font.size": 8,
        "axes.titlesize": 9,
        "axes.titleweight": "bold",
        "axes.titlelocation": "left",
        "axes.labelcolor": "#52514e",
        "axes.edgecolor": "#8a8984",
        "axes.spines.top": False,
        "axes.spines.right": False,
        "axes.grid": True,
        "grid.color": "#e8e7e3",
        "grid.linewidth": 0.6,
        "xtick.color": "#52514e",
        "ytick.color": "#52514e",
        "legend.frameon": False,
    }
)

NEUTRAL, LIGHT = "#b4b3ad", "#e3e2dd"
STUDY_COLORS = {
    "n-back": "#2a78d6",
    "k-drive": "#eb6834",
    "u-harry": "#1baf7a",
    "s-cook": "#6250d6",
}
STUDY_LIGHT = {
    "n-back": "#bcd6f5",
    "k-drive": "#f8c9b3",
    "u-harry": "#bfe8d8",
    "s-cook": "#d3cef5",
}
# light to dark, encodes the difficulty level
RAMPS = {
    "n-back": ["#a9cbf2", "#5f9ee6", "#1f5fae"],
    "k-drive": ["#f7c1a8", "#ef8a5c", "#c0491b"],
    "u-harry": ["#9fdcc4", "#4cc49a", "#12845b"],
    "s-cook": ["#c3bcf2", "#9085e9", "#4a3bb0"],
}
SPLIT_COLORS = {"train": "#2a78d6", "validation": "#eb6834", "test": "#1baf7a"}

TABLES = ["SIGNALS", "SUBJECTIVE", "PERFORMANCE", "META"]

rows = []
for experiment, path in EXAMPLE.items():
    with pd.HDFStore(path, "r") as store:
        for key in TABLES:
            storer = store.get_storer(key)
            rows.append(
                {
                    "table": key,
                    "experiment": experiment,
                    "rows": storer.nrows,
                    "columns": len(storer.non_index_axes[0][1]),
                }
            )

pd.DataFrame(rows).pivot(index="table", columns="experiment").loc[TABLES]

EXAMPLE_CHANNELS = [
    "STUDY",
    "PHASE",
    "LEVEL",
    "RAW_ECG_I",
    "RAW_EDA",
    "ECG_I_PEAKS",
    "EYE_LEFT_PUPIL_DIAMETER",
    "EYE_LEFT_GAZE_POINT_ON_DISPLAY_AREA_X",
    "FACE_AU12",
    "MOVE_X_NOSE",
]
pd.read_hdf(EXAMPLE["drive"], "SIGNALS", columns=EXAMPLE_CHANNELS, start=3_000_000, stop=3_000_010)

with pd.HDFStore(EXAMPLE["drive"], "r") as store:
    all_channels = list(store.get_storer("SIGNALS").non_index_axes[0][1])
print(len(all_channels), "channels, for example", all_channels[:8])

ten_seconds = pd.read_hdf(
    EXAMPLE["drive"],
    "SIGNALS",
    columns=["RAW_ECG_I", "EYE_LEFT_PUPIL_VALIDITY", "FACE_AU12"],
    start=3_000_000,
    stop=3_010_000,
)
for channel in ten_seconds:
    interval_ms = np.median(np.diff(ten_seconds[channel].dropna().index))
    print(f"{channel:<26} one sample every {interval_ms:.0f} ms = {1000 / interval_ms:.0f} Hz")

pd.read_hdf(EXAMPLE["drive"], "SUBJECTIVE").head()

pd.read_hdf(EXAMPLE["game"], "PERFORMANCE").dropna(axis=1, how="all").head()

columns = {e: pd.read_hdf(path, "META").columns for e, path in EXAMPLE.items()}
pd.Series(
    {e: ", ".join(dict.fromkeys(c.get_level_values(0))) for e, c in columns.items()},
    name="column groups",
).to_frame()

def demographics(files):
    person = pd.concat([pd.read_hdf(f, "META") for f in files])["Subject"]
    age = person["Age"].astype(float)
    sexes = person["Sex"].value_counts()
    return {
        "Sessions": len(files),
        "Female / male / not stated": f"{sexes.get('female', 0)} / {sexes.get('male', 0)} / {sexes.get('unknown', 0)}",
        "Age [years]": f"{age.mean():.1f} ± {age.std():.1f} (range {age.min():.0f} to {age.max():.0f})",
    }


pd.DataFrame({experiment: demographics(files) for experiment, files in FILES.items()})

person_of = {stem: stem.split("-")[0] for stem in EXPERIMENT}
people = {e: {person_of[f.stem] for f in files} for e, files in FILES.items()}
print(f"{len(person_of)} sessions from {len(set(person_of.values()))} people")
print(f"{len(people['drive'] & people['game'])} people took part in both experiments")
print(
    f"{len(FILES['game']) - len(people['game'])} gaming sessions come from a person with an earlier gaming session"
)

def read_annotation(path, step=100):
    # STUDY, PHASE and LEVEL every `step` ms, indexed by time in seconds
    columns = {}
    with tables.open_file(path, "r") as h5:
        table = h5.root.SIGNALS.table
        rows = table.read(step=step)
        for block in table.colnames[1:]:
            names = list(getattr(table.attrs, f"{block}_kind"))
            for i, name in enumerate(names):
                if name in ("STUDY", "PHASE", "LEVEL"):
                    columns[name] = (block, rows[block][:, i])
        categorical = {
            b for b, _ in columns.values() if getattr(table.attrs, f"{b}_meta", None) == "category"
        }
    frame = pd.DataFrame(
        {n: v for n, (_, v) in columns.items()},
        index=pd.Index(rows["index"] / 1000, name="time [s]"),
    )
    for name, (block, _) in columns.items():
        if block in categorical:
            frame[name] = np.asarray(pd.read_hdf(path, f"SIGNALS/meta/{block}/meta"))[frame[name]]
        elif name != "LEVEL":
            frame[name] = frame[name].str.decode("utf-8")
    frame["LEVEL"] = frame["LEVEL"].astype(int)
    return frame[["STUDY", "PHASE", "LEVEL"]]


annotations = {f.stem: read_annotation(f) for files in FILES.values() for f in files}

rows = []
for stem, annotation in annotations.items():
    active = annotation[annotation["PHASE"] != "n/a"]
    for key, samples in active.value_counts().items():
        rows.append(
            {
                "experiment": EXPERIMENT[stem],
                "STUDY": key[0],
                "PHASE": key[1],
                "LEVEL": key[2],
                "minutes": samples / 600,
            }
        )

(
    pd.DataFrame(rows)
    .groupby(["experiment", "STUDY", "PHASE", "LEVEL"])["minutes"]
    .agg(sessions="count", mean="mean", min="min", max="max")
    .round(1)
)

# baseline and test levels per experiment and study
LEVELS = {
    "drive": {"n-back": ([1, 2], [1, 2, 3, 4, 5, 6]), "k-drive": ([1, 2, 3], [1, 2, 3])},
    "game": {
        "n-back": ([1, 2], [1, 2, 3]),
        "u-harry": ([1, 2], [1, 2]),
        "s-cook": ([1], [1, 2, 3]),
    },
}


def lanes(experiment, study):
    # (phase, levels, label, color) for every lane of a study
    baselines, tests = LEVELS[experiment][study]
    ramp = RAMPS[study]
    colors = [
        ramp[i % 3] if len(tests) > 3 else ramp[i + 3 - len(tests)] for i in range(len(tests))
    ]
    return (
        [("baseline", [b], f"baseline {b}", NEUTRAL) for b in baselines]
        + [("train", tests, "train", LIGHT)]
        + [("test", [t], f"test {t}", c) for t, c in zip(tests, colors)]
    )


def runs(mask):
    # start and stop index of every contiguous True run
    edges = np.flatnonzero(np.diff(np.r_[0, mask.astype(int), 0]))
    return list(zip(edges[::2], edges[1::2]))


def plot_lanes(ax, annotation, experiment, study):
    minutes = annotation.index.to_numpy() / 60
    study_lanes = lanes(experiment, study)
    for row, (phase, levels, label, color) in enumerate(study_lanes):
        mask = (
            (annotation["STUDY"] == study)
            & (annotation["PHASE"] == phase)
            & annotation["LEVEL"].isin(levels)
        ).to_numpy()
        for start, stop in runs(mask):
            ax.broken_barh(
                [(minutes[start], minutes[stop - 1] - minutes[start])],
                (row - 0.38, 0.76),
                color=color,
            )
    ax.set_yticks(range(len(study_lanes)), [lane[2] for lane in study_lanes])
    ax.set_ylim(len(study_lanes) - 0.5, -0.5)
    ax.grid(axis="y", visible=False)
    ax.set_title(study, color=STUDY_COLORS[study])


def plot_protocol(path, experiment):
    annotation = annotations[path.stem]
    studies = list(LEVELS[experiment])
    heights = [len(lanes(experiment, s)) for s in studies]
    fig, axes = plt.subplots(
        len(studies),
        1,
        figsize=(7.5, 0.28 * sum(heights) + 0.6 * len(studies)),
        sharex=True,
        gridspec_kw={"height_ratios": heights},
    )
    for ax, study in zip(axes, studies):
        plot_lanes(ax, annotation, experiment, study)
    axes[-1].set_xlabel("Time since start of recording [min]")
    axes[-1].set_xlim(0, annotation.index[-1] / 60)
    fig.suptitle(f"Session {path.stem} ({experiment})", x=0.01, ha="left", fontweight="bold")
    fig.tight_layout()
    plt.show()


plot_protocol(EXAMPLE["drive"], "drive")
plot_protocol(EXAMPLE["game"], "game")

STUDIES = list(STUDY_COLORS)
CODE_COLORS = ["#f4f3ef"] + [c for s in STUDIES for c in (STUDY_LIGHT[s], STUDY_COLORS[s])]


def encode(annotation):
    # one code per second, 0 outside the tasks, then a light and a dark code per study
    frame = annotation.iloc[::10]
    code = np.zeros(len(frame))
    for i, study in enumerate(STUDIES):
        active = (frame["STUDY"] == study) & (frame["PHASE"] != "n/a")
        code[active.to_numpy()] = 2 * i + 1
        code[(active & (frame["PHASE"] == "test")).to_numpy()] = 2 * i + 2
    return code


def task_order(code):
    onsets = {
        s: np.argmax(code == 2 * i + 2) for i, s in enumerate(STUDIES) if (code == 2 * i + 2).any()
    }
    return tuple(sorted(onsets, key=onsets.get))


fig, axes = plt.subplots(2, 1, figsize=(7.5, 8.2))
for ax, (experiment, files) in zip(axes, FILES.items()):
    encoded = {f.stem: encode(annotations[f.stem]) for f in files}
    order = sorted(encoded, key=lambda s: task_order(encoded[s]))
    grid = np.full((len(order), max(len(c) for c in encoded.values())), np.nan)
    for row, stem in enumerate(order):
        grid[row, : len(encoded[stem])] = encoded[stem]
    cmap = ListedColormap(CODE_COLORS)
    cmap.set_bad("white")
    ax.imshow(
        np.ma.masked_invalid(grid),
        aspect="auto",
        cmap=cmap,
        vmin=-0.5,
        vmax=len(CODE_COLORS) - 0.5,
        interpolation="nearest",
        extent=(0, grid.shape[1] / 60, len(order) - 0.5, -0.5),
    )
    ax.set_yticks(range(len(order)), order, fontsize=6)
    ax.grid(False)
    ax.set_title(experiment)
    ax.set_xlim(0, 210)
axes[-1].set_xlabel("Time since start of recording [min]")
axes[0].legend(
    handles=[Patch(color=STUDY_COLORS[s], label=s) for s in STUDIES],
    ncol=4,
    loc="lower right",
    bbox_to_anchor=(1, 1.0),
)
fig.tight_layout()
plt.show()

ORDER = [
    (e, s, level)
    for e, studies in LEVELS.items()
    for s, (_, tests) in studies.items()
    for level in tests
]


def read_table(key):
    # one table for all sessions, restricted to the test levels listed in LEVELS
    frames = []
    for experiment, files in FILES.items():
        for f in files:
            frame = pd.read_hdf(f, key).reset_index(drop=True)
            frame.insert(0, "session", f.stem)
            frame.insert(0, "experiment", experiment)
            frames.append(frame)
    frame = pd.concat(frames, ignore_index=True)
    frame[["STUDY", "PHASE"]] = frame[["STUDY", "PHASE"]].astype(str)
    frame["LEVEL"] = frame["LEVEL"].astype(int)
    known = [
        phase == "test" and (e, study, level) in ORDER
        for e, study, phase, level in zip(
            frame["experiment"], frame["STUDY"], frame["PHASE"], frame["LEVEL"]
        )
    ]
    return frame[known].copy()


DIMENSIONS = ["MENTAL", "PHYSICAL", "TEMPORAL", "PERFORMANCE", "EFFORT", "FRUSTRATION"]

subjective = read_table("SUBJECTIVE")
subjective["RTLX"] = subjective[DIMENSIONS].mean(axis=1)
weights = subjective[[f"WEIGHT {d}" for d in DIMENSIONS]].to_numpy()
subjective["TLX"] = (subjective[DIMENSIONS].to_numpy() * weights).sum(axis=1) / weights.sum(axis=1)

table = (
    subjective.groupby(["experiment", "STUDY", "LEVEL"])[DIMENSIONS + ["RTLX", "TLX"]]
    .mean()
    .loc[ORDER]
    .rename(columns=str.title)
    .rename(columns={"Rtlx": "RTLX", "Tlx": "TLX"})
)
(
    table.style.format("{:.1f}")
    .background_gradient(cmap="Blues", vmin=0, vmax=20)
    .set_caption("Mean NASA-TLX rating per test level (0 to 20)")
)

def plot_by_level(ax, experiment, study, column="RTLX"):
    data = subjective[(subjective["experiment"] == experiment) & (subjective["STUDY"] == study)]
    for _, person in data.groupby("session"):
        person = person.sort_values("LEVEL")
        ax.plot(person["LEVEL"], person[column], color=NEUTRAL, lw=0.5, alpha=0.5)
    stats = data.groupby("LEVEL")[column].agg(["mean", "std"])
    ax.errorbar(
        stats.index,
        stats["mean"],
        yerr=stats["std"],
        color=STUDY_COLORS[study],
        lw=2,
        marker="o",
        ms=4,
        capsize=2,
    )
    levels = LEVELS[experiment][study][1]
    ax.set_xticks(levels)
    ax.set_xlim(levels[0] - 0.4, levels[-1] + 0.4)
    ax.set_xlabel("LEVEL")
    ax.set_title(f"{study} ({experiment})", color=STUDY_COLORS[study])


fig = plt.figure(figsize=(7.5, 5.4))
for subfig, (experiment, studies) in zip(fig.subfigures(2, 1), LEVELS.items()):
    widths = [len(tests) + 0.8 for _, tests in studies.values()]
    axes = subfig.subplots(1, len(studies), sharey=True, gridspec_kw={"width_ratios": widths})
    for ax, study in zip(axes, studies):
        plot_by_level(ax, experiment, study)
    axes[0].set_ylabel("RTLX (0 to 20)")
    axes[0].set_ylim(-1, 20)
plt.show()

def race_seconds(text):
    # race times are stored as text, for example "02:41,56"
    if pd.isna(text):
        return np.nan
    minutes, seconds = str(text).split(":")
    return 60 * int(minutes) + float(seconds.replace(",", "."))


performance = read_table("PERFORMANCE")
performance["F1 score"] = performance[["VISUAL F1", "AUDITIVE F1", "F1"]].mean(axis=1)
performance["Reaction time [s]"] = performance[
    ["VISUAL MEAN REACTION TIME", "AUDITIVE MEAN REACTION TIME", "REACTION TIME"]
].mean(axis=1)
performance["RACE TIME [s]"] = performance["RACE TIME"].map(race_seconds)

metrics = [
    "F1 score",
    "Reaction time [s]",
    "SONGS RECALL",
    "RACE TIME [s]",
    "MISSED RINGS",
    "SCORE",
]
summary = (
    performance.groupby(["experiment", "STUDY", "LEVEL"])[metrics]
    .agg(["mean", "std"])
    .reindex(ORDER)
)
pd.DataFrame(
    {
        m: [f"{a:.2f} ± {b:.2f}" if not np.isnan(a) else "" for a, b in summary[m].to_numpy()]
        for m in metrics
    },
    index=summary.index,
)

def read_signals(path, columns, start_s, stop_s):
    # selected SIGNALS columns between two time points in seconds, index in seconds
    chunks = pd.read_hdf(
        path,
        "SIGNALS",
        columns=columns,
        start=int(start_s * 1000),
        stop=int(stop_s * 1000),
        chunksize=200_000,
    )
    frame = pd.concat(chunks)
    frame.index = frame.index / 1000
    return frame


annotation = annotations[EXAMPLE["drive"].stem]
three_back = annotation[
    (annotation["STUDY"] == "n-back")
    & (annotation["PHASE"] == "test")
    & (annotation["LEVEL"] == 3)
]
t0 = three_back.index[0] + 40
segment = read_signals(
    EXAMPLE["drive"],
    [
        "RAW_ECG_I",
        "RAW_ECG_II",
        "RAW_PPG",
        "RAW_EDA",
        "RAW_RSP",
        "RAW_EMG",
        "RAW_SKT",
        "ECG_I_PEAKS",
        "PPG_PEAKS",
        "RSP_PEAKS",
        "EYE_LEFT_PUPIL_DIAMETER",
        "EYE_RIGHT_PUPIL_DIAMETER",
    ],
    t0,
    t0 + 15,
)
segment.head()

PANELS = [
    ("RAW_ECG_I", "ECG I", "ECG_I_PEAKS"),
    ("RAW_ECG_II", "ECG II", None),
    ("RAW_PPG", "PPG", "PPG_PEAKS"),
    ("RAW_RSP", "Respiration", "RSP_PEAKS"),
    ("RAW_EDA", "EDA", None),
    ("RAW_EMG", "EMG", None),
    ("RAW_SKT", "Skin temp.", None),
]

fig, axes = plt.subplots(len(PANELS) + 1, 1, figsize=(7.5, 7.2), sharex=True)
time = segment.index - t0
for ax, (column, label, events) in zip(axes, PANELS):
    ax.plot(time, segment[column], color="#0b0b0b", lw=0.6)
    if events is not None:
        hits = segment[events].fillna(False).astype(bool).to_numpy()
        ax.plot(time[hits], segment.loc[hits, column], "o", ms=3, color=STUDY_COLORS["k-drive"])
    ax.set_ylabel(label, rotation=0, ha="right", va="center")
    ax.set_yticks([])

ax = axes[-1]
for column, side, color in [
    ("EYE_LEFT_PUPIL_DIAMETER", "left", STUDY_COLORS["n-back"]),
    ("EYE_RIGHT_PUPIL_DIAMETER", "right", STUDY_COLORS["u-harry"]),
]:
    pupil = segment[column].dropna()
    ax.plot(pupil.index - t0, pupil, lw=0.8, color=color, label=side)
ax.set_ylabel("Pupil [mm]", rotation=0, ha="right", va="center")
ax.legend(ncol=2, loc="lower right", bbox_to_anchor=(1, 0.95))
ax.set_xlabel("Time [s]")
ax.set_xlim(0, 15)
fig.suptitle(
    f"15 s of n-back test 3, session {EXAMPLE['drive'].stem}", x=0.01, ha="left", fontweight="bold"
)
fig.align_ylabels()
fig.tight_layout()
plt.show()

FS = 100  # target rate in Hz


def downsample(frame, fs=FS):
    # mean over bins of 1000 / fs ms, index is the sample number at the new rate
    ms = np.round(frame.index.to_numpy() * 1000).astype(int)
    return frame.groupby(ms // (1000 // fs)).mean()


def fill_gaps(series, max_gap_s=10, fs=FS):
    missing = series.isna()
    gap_length = missing.groupby((~missing).cumsum()).transform("sum")
    filled = series.interpolate(method="pchip", limit_area="inside")
    return filled.mask(missing & (gap_length > max_gap_s * fs))


raw = read_signals(EXAMPLE["drive"], ["RAW_ECG_I", "EYE_LEFT_PUPIL_DIAMETER"], t0, t0 + 10)
processed = downsample(raw).apply(fill_gaps)
print(f"{len(raw)} rows at 1 kHz -> {len(processed)} rows at {FS} Hz")

fig, axes = plt.subplots(2, 1, figsize=(7.5, 3.8))

ax = axes[0]
ecg = raw["RAW_ECG_I"].loc[t0 + 2 : t0 + 4]
ecg_100 = processed["RAW_ECG_I"]
ecg_100.index = ecg_100.index / FS - t0
ecg_100 = ecg_100.loc[2:4]
ax.plot(ecg.index - t0, ecg, color=NEUTRAL, lw=1.6, label="1000 Hz")
ax.plot(ecg_100.index, ecg_100, ".-", color=STUDY_COLORS["n-back"], ms=3, lw=0.8, label="100 Hz")
ax.set_title("ECG I, downsampled")
ax.legend(ncol=2, loc="lower right", bbox_to_anchor=(1, 0.98))

ax = axes[1]
pupil = raw["EYE_LEFT_PUPIL_DIAMETER"].dropna()
pupil_100 = processed["EYE_LEFT_PUPIL_DIAMETER"]
ax.plot(
    pupil_100.index / FS - t0,
    pupil_100,
    color=STUDY_COLORS["n-back"],
    lw=1.2,
    label="100 Hz, gaps filled",
)
ax.plot(pupil.index - t0, pupil, ".", color="#0b0b0b", ms=1.5, label="250 Hz, raw")
ax.set_title("Left pupil diameter [mm], blinks interpolated")
ax.set_xlabel("Time [s]")
ax.legend(ncol=2, loc="lower right", bbox_to_anchor=(1, 0.98))

fig.tight_layout()
plt.show()

WINDOW_S, STEP_S, MIN_SHARE = 40, 20, 0.9

LOW_LOAD = [
    ("drive", "n-back", "baseline", 1),
    ("drive", "n-back", "baseline", 2),
    ("drive", "n-back", "test", 1),
    ("drive", "n-back", "test", 4),
    ("drive", "k-drive", "baseline", 1),
    ("drive", "k-drive", "baseline", 2),
    ("drive", "k-drive", "baseline", 3),
    ("game", "n-back", "baseline", 1),
    ("game", "n-back", "baseline", 2),
    ("game", "n-back", "test", 1),
    ("game", "u-harry", "baseline", 1),
    ("game", "u-harry", "baseline", 2),
    ("game", "s-cook", "baseline", 1),
    ("game", "s-cook", "test", 1),
]
HIGH_LOAD = [
    ("drive", "n-back", "test", 2),
    ("drive", "n-back", "test", 3),
    ("drive", "n-back", "test", 5),
    ("drive", "n-back", "test", 6),
    ("drive", "k-drive", "test", 1),
    ("drive", "k-drive", "test", 2),
    ("drive", "k-drive", "test", 3),
    ("game", "n-back", "test", 2),
    ("game", "n-back", "test", 3),
    ("game", "u-harry", "test", 1),
    ("game", "u-harry", "test", 2),
    ("game", "s-cook", "test", 2),
    ("game", "s-cook", "test", 3),
]
CONDITIONS = LOW_LOAD + HIGH_LOAD


pd.DataFrame(
    {"load": ["low"] * len(LOW_LOAD) + ["high"] * len(HIGH_LOAD)},
    index=pd.MultiIndex.from_tuples(CONDITIONS, names=["experiment", "study", "phase", "level"]),
).unstack("experiment").fillna("")

def select_windows(annotation, experiment, min_share=MIN_SHARE, fs=10):
    # keep a window if one condition covers at least min_share of it, annotation at fs Hz
    keys = list(
        zip(
            [experiment] * len(annotation),
            annotation["STUDY"],
            annotation["PHASE"],
            annotation["LEVEL"],
        )
    )
    size, step = WINDOW_S * fs, STEP_S * fs
    rows = []
    for first in range(0, len(keys) - size, step):
        shares = pd.Series(keys[first : first + size]).value_counts(normalize=True)
        condition, share = shares.index[0], shares.iloc[0]
        keep = condition in CONDITIONS and share >= min_share
        rows.append(
            {
                "start [s]": annotation.index[first],
                "condition": condition,
                "share": share,
                "kept": keep,
            }
        )
    return pd.DataFrame(rows)


windows = select_windows(annotation, "drive")
windows[windows["condition"].str[1] == "n-back"].head(8)

nback = annotation[annotation["STUDY"] == "n-back"]

fig, axes = plt.subplots(
    2, 1, figsize=(7.5, 4.6), sharex=True, gridspec_kw={"height_ratios": [9, 2.2]}
)
plot_lanes(axes[0], annotation, "drive", "n-back")
ax = axes[1]
for i, w in windows.iterrows():
    color = STUDY_COLORS["n-back"] if w["kept"] else LIGHT
    ax.broken_barh([(w["start [s]"] / 60, WINDOW_S / 60 - 0.04)], (i % 2 - 0.35, 0.7), color=color)
ax.set_yticks([])
ax.set_ylim(1.6, -0.6)
ax.grid(False)
ax.set_title("40 s windows with 20 s step, blue kept, gray dropped", color="#52514e")
axes[0].set_xlim(nback.index[0] / 60 - 1, nback.index[-1] / 60 + 1)
axes[-1].set_xlabel("Time since start of recording [min]")
fig.tight_layout()
plt.show()

counts = {}
for experiment, path in EXAMPLE.items():
    for share in [0.5, 0.9, 1.0]:
        selected = select_windows(annotations[path.stem], experiment, min_share=share)
        counts[(experiment, f"share {share}")] = selected.loc[
            selected["kept"], "condition"
        ].value_counts()
counts = pd.DataFrame(counts).reindex(CONDITIONS).fillna(0).astype(int)
counts.index = [" ".join(map(str, c)) for c in counts.index]
counts.loc["total"] = counts.sum()
counts

CHANNELS = [
    "RAW_ECG_I",
    "RAW_ECG_II",
    "RAW_PPG",
    "RAW_EDA",
    "RAW_RSP",
    "RAW_SKT",
    "EYE_LEFT_PUPIL_DIAMETER",
    "EYE_RIGHT_PUPIL_DIAMETER",
]


def build_session(path, experiment):
    # preprocessed windows of one session
    annotation = annotations[path.stem]
    blocks = []
    for study in annotation["STUDY"].unique():
        if study == "n/a":
            continue
        times = annotation.index[annotation["STUDY"] == study]
        blocks.append(downsample(read_signals(path, CHANNELS, times[0], times[-1] + 0.1)))
    signals = pd.concat(blocks).sort_index().apply(fill_gaps)
    signals = ((signals - signals.mean()) / signals.std()).fillna(0.0)

    selected = select_windows(annotation, experiment)
    selected = selected[selected["kept"]]
    X = []
    for start in selected["start [s]"]:
        first = int(round(start * FS))
        X.append(signals.reindex(range(first, first + WINDOW_S * FS)).fillna(0.0).to_numpy().T)
    n = len(selected)
    return {
        "X": np.stack(X).astype(np.float32),
        "y": np.array([c in HIGH_LOAD for c in selected["condition"]], dtype=np.int64),
        "condition": np.array([" ".join(map(str, c)) for c in selected["condition"]]),
        "experiment": np.array([experiment] * n),
        "session": np.array([path.stem] * n),
    }


CACHE = DATA_PATH / f"public-sets-windows-{FS}hz.npz"
if CACHE.exists():
    dataset = dict(np.load(CACHE))
else:
    parts = [build_session(f, experiment) for experiment, files in FILES.items() for f in files]
    dataset = {key: np.concatenate([p[key] for p in parts]) for key in parts[0]}
    np.savez(CACHE, **dataset)

X, y = dataset["X"], dataset["y"]
print(f"X {X.shape}, y {y.shape}, {y.mean():.0%} high load")

per_condition = pd.crosstab(dataset["condition"], dataset["session"])
per_condition = per_condition.reindex([" ".join(map(str, c)) for c in CONDITIONS])
pd.DataFrame(
    {
        "load": ["low"] * len(LOW_LOAD) + ["high"] * len(HIGH_LOAD),
        "windows": per_condition.sum(axis=1),
        "sessions": (per_condition > 0).sum(axis=1),
        "per session": per_condition.replace(0, np.nan).mean(axis=1).round(1),
    }
).rename_axis("condition")

index = np.flatnonzero(dataset["condition"] == "drive n-back test 3")[0]
time = np.arange(X.shape[-1]) / FS

fig, axes = plt.subplots(len(CHANNELS), 1, figsize=(7.5, 5.2), sharex=True)
for ax, channel, values in zip(axes, CHANNELS, X[index]):
    color = "#0b0b0b" if channel.startswith("RAW_") else STUDY_COLORS["n-back"]
    ax.plot(time, values, color=color, lw=0.6)
    label = (
        channel.replace("RAW_", "").replace("EYE_", "").replace("_DIAMETER", "").replace("_", " ")
    )
    ax.set_ylabel(label.lower(), rotation=0, ha="right", va="center")
    ax.set_yticks([])
axes[-1].set_xlabel("Time [s]")
axes[-1].set_xlim(0, time[-1])
fig.suptitle(
    f"Window {index}, {dataset['session'][index]}, label y = {y[index]}",
    x=0.01,
    ha="left",
    fontweight="bold",
)
fig.align_ylabels()
fig.tight_layout()
plt.show()

def stratified_grouped_split(X, y, g, n_splits=5):
    # assign each group greedily to the fold that keeps the label shares most even
    labels = sorted(set(y))
    counts = {}
    for label, group in zip(y, g):
        counts.setdefault(group, np.zeros(len(labels)))[labels.index(label)] += 1
    total = sum(counts.values())
    fold_counts = np.zeros((n_splits, len(labels)))
    fold_groups = [[] for _ in range(n_splits)]
    for group, count in sorted(counts.items(), key=lambda item: -np.std(item[1])):
        cost = []
        for fold in range(n_splits):
            fold_counts[fold] += count
            cost.append(np.mean(np.std(fold_counts / total, axis=0)))
            fold_counts[fold] -= count
        best = int(np.argmin(cost))
        fold_counts[best] += count
        fold_groups[best].append(group)

    members = [[x for x, group in zip(X, g) if group in groups] for groups in fold_groups]
    splits = []
    for test in range(n_splits):
        validation = 1 if test == 0 else 0
        train = [
            x for fold in range(n_splits) if fold not in (test, validation) for x in members[fold]
        ]
        splits.append((train, members[validation], members[test]))
    return splits


sessions = sorted(EXPERIMENT, key=lambda s: (person_of[s], EXPERIMENT[s]))
splits = stratified_grouped_split(
    sessions, [EXPERIMENT[s] for s in sessions], [person_of[s] for s in sessions]
)

# no person in two sets of the same fold
for train, validation, test in splits:
    groups = [{person_of[s] for s in part} for part in (train, validation, test)]
    assert not (groups[0] & groups[1] or groups[0] & groups[2] or groups[1] & groups[2])

matrix = np.zeros((len(splits), len(sessions)))
for fold, sets in enumerate(splits):
    for code, members in enumerate(sets):
        matrix[fold, [sessions.index(m) for m in members]] = code

fig, (top, ax) = plt.subplots(
    2, 1, figsize=(7.5, 2.6), sharex=True, gridspec_kw={"height_ratios": [1, 5]}
)
top.imshow(
    [[0 if EXPERIMENT[s] == "drive" else 1 for s in sessions]],
    aspect="auto",
    cmap=ListedColormap(["#52514e", LIGHT]),
    vmin=-0.5,
    vmax=1.5,
)
top.set_yticks([0], ["experiment"])
top.grid(False)
top.legend(
    handles=[Patch(color="#52514e", label="drive"), Patch(color=LIGHT, label="game")]
    + [Patch(color=color, label=name) for name, color in SPLIT_COLORS.items()],
    ncol=5,
    loc="lower left",
    bbox_to_anchor=(0, 1.0),
)
ax.imshow(
    matrix, aspect="auto", cmap=ListedColormap(list(SPLIT_COLORS.values())), vmin=-0.5, vmax=2.5
)
ax.set_xticks(np.arange(-0.5, len(sessions)), minor=True)
ax.set_yticks(np.arange(-0.5, len(splits)), minor=True)
ax.grid(which="minor", color="white", lw=1)
ax.grid(which="major", visible=False)
ax.tick_params(which="minor", length=0)
ax.set_xticks(range(len(sessions)), sessions, rotation=90, fontsize=5)
ax.set_yticks(range(len(splits)), [f"fold {i}" for i in range(len(splits))])
plt.show()

rows = []
for fold, sets in enumerate(splits):
    for name, members in zip(SPLIT_COLORS, sets):
        mask = np.isin(dataset["session"], members)
        rows.append(
            {
                "fold": fold,
                "set": name,
                "people": len({person_of[s] for s in members}),
                "drive sessions": sum(EXPERIMENT[s] == "drive" for s in members),
                "game sessions": sum(EXPERIMENT[s] == "game" for s in members),
                "windows": mask.sum(),
                "high load": f"{y[mask].mean():.0%}",
            }
        )
pd.DataFrame(rows).pivot(index="fold", columns="set").reindex(
    columns=list(SPLIT_COLORS), level="set"
)

masks = {
    name: np.isin(dataset["session"], members) for name, members in zip(SPLIT_COLORS, splits[0])
}
X_train, y_train = X[masks["train"]], y[masks["train"]]
X_val, y_val = X[masks["validation"]], y[masks["validation"]]
X_test, y_test = X[masks["test"]], y[masks["test"]]
print("train", X_train.shape, "validation", X_val.shape, "test", X_test.shape)