#!/usr/bin/env python3
"""Participant-held-out motor imagery benchmark for PhysioNet EEGMMIDB v1.0.0."""

import argparse
import hashlib
import json
from pathlib import Path

import numpy as np


DATA = 'https://physionet.org/files/eegmmidb/1.0.0/'
TRAIN = range(1, 17)
TEST = range(17, 25)
CHANNELS = ('C3', 'C4')
BANDS = ((8, 12), (13, 26))
SAMPLE_RATE = 160
WINDOW_SECONDS = 2
ONSET_DELAY = 0.5


def read_run(path):
    raw = path.read_bytes()
    count = int(raw[252:256])
    header = int(raw[184:192])
    record_count = int(raw[236:244])
    if raw[192:197] != b'EDF+C' or count != 65 or header != 256 * (count + 1) or float(raw[244:252]) != 1:
        raise ValueError(f'Unexpected EDF+ format in {path.name}')

    labels = [raw[256 + i * 16:256 + (i + 1) * 16].decode().strip().strip('.').upper() for i in range(count)]
    samples_at = 256 + 216 * count
    sizes = [int(raw[samples_at + i * 8:samples_at + (i + 1) * 8]) for i in range(count)]
    if sizes[:-1] != [SAMPLE_RATE] * 64 or labels[-1] != 'EDF ANNOTATIONS':
        raise ValueError(f'Unexpected channels in {path.name}')
    if header + record_count * sum(sizes) * 2 != len(raw):
        raise ValueError(f'Unexpected sample count in {path.name}')

    records = np.frombuffer(raw, dtype='<i2', offset=header).reshape(record_count, sum(sizes))
    signals = {}
    for channel in CHANNELS:
        index = labels.index(channel)
        start = SAMPLE_RATE * index
        def header_number(position):
            offset = 256 + position * count + index * 8
            return float(raw[offset:offset + 8])

        physical_min, physical_max = header_number(104), header_number(112)
        digital_min, digital_max = header_number(120), header_number(128)
        if physical_max <= physical_min or digital_max <= digital_min:
            raise ValueError(f'Invalid calibration for {channel} in {path.name}')
        gain = (physical_max - physical_min) / (digital_max - digital_min)
        # A constant offset disappears when each epoch is centered.
        signals[channel] = records[:, start:start + SAMPLE_RATE].ravel().astype(float) * gain

    events = []
    for annotation in records[:, SAMPLE_RATE * 64:]:
        for tal in annotation.astype('<i2', copy=False).tobytes().split(b'\x00'):
            parts = tal.split(b'\x14')
            timing = parts[0].split(b'\x15')
            for label in parts[1:]:
                if label in (b'T1', b'T2'):
                    if len(timing) != 2 or float(timing[1]) < ONSET_DELAY + WINDOW_SECONDS:
                        raise ValueError(f'Task interval too short in {path.name}')
                    events.append((float(timing[0]), int(label[1:]) - 1))
    if len({onset for onset, _ in events}) != len(events):
        raise ValueError(f'Duplicate event onset in {path.name}')
    if {label for _, label in events} != {0, 1}:
        raise ValueError(f'Missing a fist label in {path.name}')
    return signals, events, hashlib.sha256(raw).hexdigest(), len(raw)


def features_for(signals, events):
    samples = WINDOW_SECONDS * SAMPLE_RATE
    taper = np.hanning(samples)
    frequencies = np.fft.rfftfreq(samples, 1 / SAMPLE_RATE)
    features = []
    for onset, _ in events:
        start = round((onset + ONSET_DELAY) * SAMPLE_RATE)
        if start < 0:
            raise ValueError('Task interval before the recording')
        channel_powers = []
        for channel in CHANNELS:
            segment = signals[channel][start:start + samples]
            if len(segment) != samples:
                raise ValueError('Task interval after the recording')
            power = np.abs(np.fft.rfft((segment - segment.mean()) * taper)) ** 2
            channel_powers.append([power[(frequencies >= lo) & (frequencies <= hi)].sum() for lo, hi in BANDS])
        powers = np.asarray(channel_powers)
        if np.any(powers <= 0) or not np.all(np.isfinite(powers)):
            raise ValueError('Invalid EEG band power')
        features.append(np.log(powers[0]) - np.log(powers[1]))
    return np.asarray(features), np.asarray([label for _, label in events])


def classify(training, labels, testing):
    mean = training.mean(axis=0)
    scale = training.std(axis=0)
    if np.any(scale <= 0):
        raise ValueError('Constant training feature')
    training = (training - mean) / scale
    centers = np.stack([training[labels == label].mean(axis=0) for label in (0, 1)])
    distances = ((testing[:, None, :] - mean) / scale - centers) ** 2
    return distances.sum(axis=2).argmin(axis=1)


def confusion(actual, predicted):
    return np.bincount(actual * 2 + predicted, minlength=4).reshape(2, 2)


def run(folder, verified_on):
    listing = (folder / 'RECORDS').read_bytes()
    items = listing.decode().splitlines()
    if len(items) != len(set(items)):
        raise ValueError('Duplicate path in the published file list')

    train_x, train_y, test = [], [], []
    inputs = []
    for cohort, members, run_number in (('train', TRAIN, 4), ('test', TEST, 8)):
        for participant in members:
            subject = f'S{participant:03}'
            filename = f'{subject}R{run_number:02}.edf'
            relative = f'{subject}/{filename}'
            if relative.removesuffix('.edf') not in items and relative not in items:
                raise ValueError(f'{relative} is missing from the published file list')
            signals, events, digest, size = read_run(folder / filename)
            x, y = features_for(signals, events)
            inputs.append({'file': relative, 'group': cohort, 'sha256': digest, 'bytes': size, 'epochs': len(y)})
            if cohort == 'train':
                train_x.append(x)
                train_y.append(y)
            else:
                test.append((subject, x, y))

    training = np.concatenate(train_x)
    labels = np.concatenate(train_y)
    result = []
    for subject, x, y in test:
        counts = confusion(y, classify(training, labels, x))
        result.append({'participant': subject, 'confusion': counts.tolist()})
    matrices = np.asarray([row['confusion'] for row in result])
    total = matrices.sum(axis=0)
    # Resample participants, since epochs within the same recording are dependent.
    rng = np.random.default_rng(20261114)
    sampled = matrices[rng.integers(len(matrices), size=(20000, len(matrices)))].sum(axis=1)
    balanced = (sampled.diagonal(axis1=1, axis2=2) / sampled.sum(axis=2)).mean(axis=1)
    interval = np.quantile(balanced, (0.025, 0.975))
    return {
        'checkedOn': verified_on,
        'dataset': 'EEG Motor Movement/Imagery Dataset v1.0.0',
        'sourceUrl': 'https://physionet.org/content/eegmmidb/1.0.0/',
        'dataBase': DATA,
        'fileListSha256': hashlib.sha256(listing).hexdigest(),
        'listedRuns': len(items),
        'numpy': np.__version__,
        'trainPeople': len(TRAIN),
        'testPeople': len(TEST),
        'trainEpochs': len(labels),
        'testEpochs': int(total.sum()),
        'accuracy': round(float(total.trace() / total.sum()), 6),
        'balancedAccuracy': round(float((total.diagonal() / total.sum(axis=1)).mean()), 6),
        'balancedBootstrap95': [round(float(bound), 6) for bound in interval],
        'confusion': total.tolist(),
        'participants': result,
        'inputs': inputs,
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('folder', type=Path, help='Folder with RECORDS and the EDF files from the input manifest')
    parser.add_argument('--checked-on', default='2026-11-14')
    args = parser.parse_args()
    print(json.dumps(run(args.folder, args.checked_on), indent=2, allow_nan=False))


if __name__ == '__main__':
    main()