data_loader.py

import io
import json
import os
import warnings
from functools import lru_cache
from typing import Optional

import numpy as np
import pandas as pd
import s3fs
from dotenv import load_dotenv

load_dotenv(os.path.join(os.path.dirname(__file__), ".env"))

warnings.filterwarnings("ignore", category=UserWarning, message=".*asynchronous.*")

MANIFEST_DATASET = "tankalapavankalyan/eeg-corpus-manifest"
PARQUET_BASE = (
    "https://huggingface.co/datasets/tankalapavankalyan/eeg-corpus-manifest"
    "/resolve/refs%2Fconvert%2Fparquet"
)

TUAR_DATASET_IDS = ["tuh_eeg_artifact"]

ARTIFACT_LABEL_MAP = {
    "eyem": "Eye Movement",
    "eyeb": "Eye Blink",
    "musc": "Muscle",
    "elec": "Electrode Pop",
    "chew": "Chewing",
    "shiv": "Shiver",
    "null": "Clean",
    "elpp": "Electrode Pop",
    "artf": "Artifact (Generic)",
    "bckg": "Background",
    "eyem_musc": "Eye Movement + Muscle",
    "musc_elec": "Muscle + Electrode Pop",
    "eyem_elec": "Eye Movement + Electrode Pop",
    "eyem_chew": "Eye Movement + Chewing",
    "chew_elec": "Chewing + Electrode Pop",
    "chew_musc": "Chewing + Muscle",
    "eyem_shiv": "Eye Movement + Shiver",
    "shiv_elec": "Shiver + Electrode Pop",
}

ARTIFACT_COLORS = {
    "Eye Movement": "rgba(30, 144, 255, 0.25)",
    "Eye Blink": "rgba(0, 100, 255, 0.25)",
    "Muscle": "rgba(220, 20, 60, 0.25)",
    "Electrode Pop": "rgba(255, 165, 0, 0.25)",
    "Chewing": "rgba(50, 205, 50, 0.25)",
    "Shiver": "rgba(148, 103, 189, 0.25)",
    "Artifact (Generic)": "rgba(128, 128, 128, 0.25)",
    "Background": "rgba(200, 200, 200, 0.08)",
    "Clean": "rgba(200, 200, 200, 0.08)",
    "Eye Movement + Muscle": "rgba(125, 82, 158, 0.25)",
    "Muscle + Electrode Pop": "rgba(238, 93, 30, 0.25)",
    "Eye Movement + Electrode Pop": "rgba(143, 155, 128, 0.25)",
    "Eye Movement + Chewing": "rgba(40, 175, 153, 0.25)",
    "Chewing + Electrode Pop": "rgba(153, 185, 30, 0.25)",
    "Chewing + Muscle": "rgba(135, 113, 56, 0.25)",
    "Eye Movement + Shiver": "rgba(89, 124, 222, 0.25)",
    "Shiver + Electrode Pop": "rgba(202, 134, 95, 0.25)",
}

_fs: Optional[s3fs.S3FileSystem] = None


def get_s3fs() -> s3fs.S3FileSystem:
    global _fs
    if _fs is None:
        key = os.environ.get("AWS_ACCESS_KEY_ID")
        secret = os.environ.get("AWS_SECRET_ACCESS_KEY")
        region = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
        if key and secret:
            _fs = s3fs.S3FileSystem(key=key, secret=secret, client_kwargs={"region_name": region})
        else:
            _fs = s3fs.S3FileSystem(anon=True, client_kwargs={"region_name": region})
    return _fs


def reset_s3fs():
    global _fs
    _fs = None
    _zarr_cache.clear()
    _scale_cache.clear()
    _annotation_cache.clear()


MANIFEST_COLUMNS = [
    "recording_id", "dataset_id", "subject_id_in_dataset", "session_id",
    "run_id", "task", "archival_uri", "archival_format", "duration_s",
    "n_channels", "n_eeg_channels", "sampling_rate_hz", "reference",
    "montage_name", "recording_type", "channel_names", "canonical_uri",
    "conversion_status", "roundtrip_class",
]


def get_tuar_recordings() -> pd.DataFrame:
    url = f"{PARQUET_BASE}/recordings/train/0000.parquet"
    df = pd.read_parquet(url, columns=MANIFEST_COLUMNS)
    mask = df["dataset_id"].isin(TUAR_DATASET_IDS) & (df["conversion_status"] == "ok")
    result = df[mask].copy()
    del df
    result = result.sort_values("subject_id_in_dataset").reset_index(drop=True)
    return result


def get_recording_display_list(df: pd.DataFrame) -> list[str]:
    entries = []
    for _, row in df.iterrows():
        label = (
            f"{row['recording_id'][:8]}… | "
            f"{row['dataset_id']} | "
            f"subj={row['subject_id_in_dataset']} | "
            f"ses={row.get('session_id', 'N/A')} | "
            f"dur={row['duration_s']:.0f}s | "
            f"{row['n_channels']:.0f}ch @ {row['sampling_rate_hz']:.0f}Hz"
        )
        entries.append(label)
    return entries


# --- Signal access via direct Zarr + S3 ---

_zarr_cache: dict[str, object] = {}
_scale_cache: dict[str, tuple[np.ndarray, np.ndarray]] = {}


def _register_flac_codec():
    """Register the FLAC codec with zarr v3 if not already done."""
    try:
        from zarr.registry import get_codec_class
        get_codec_class("numcodecs.flac")
    except KeyError:
        import numcodecs
        from flac_numcodecs import Flac
        numcodecs.register_codec(Flac)
        from zarr.codecs.numcodecs._codecs import _NumcodecsBytesBytesCodec
        from zarr.codecs import register_codec

        class FlacCodec(_NumcodecsBytesBytesCodec):
            codec_name = "numcodecs.flac"
            def __init__(self, **kwargs):
                super().__init__(codec_id="flac", codec_config=kwargs)

        register_codec("numcodecs.flac", FlacCodec)


def _open_zarr(canonical_uri: str):
    if canonical_uri in _zarr_cache:
        return _zarr_cache[canonical_uri]

    _register_flac_codec()
    fs = get_s3fs()
    s3_path = canonical_uri.replace("s3://", "")
    if not s3_path.endswith("/"):
        s3_path += "/"

    import zarr
    fsspec_store = zarr.storage.FsspecStore(fs=fs, path=s3_path, read_only=True)
    root = zarr.open_group(fsspec_store, mode="r")
    _zarr_cache[canonical_uri] = root

    if len(_zarr_cache) > 50:
        oldest = next(iter(_zarr_cache))
        del _zarr_cache[oldest]
        _scale_cache.pop(oldest, None)

    return root


def _get_scale_offset(canonical_uri: str):
    if canonical_uri in _scale_cache:
        return _scale_cache[canonical_uri]

    root = _open_zarr(canonical_uri)
    ch_grp = root["channels"]
    phys_min = np.array(ch_grp["physical_min"][:], dtype=np.float64)
    phys_max = np.array(ch_grp["physical_max"][:], dtype=np.float64)
    dig_min = np.array(ch_grp["digital_min"][:], dtype=np.float64)
    dig_max = np.array(ch_grp["digital_max"][:], dtype=np.float64)

    scale = (phys_max - phys_min) / (dig_max - dig_min + 1e-12)
    offset = phys_min - dig_min * scale
    _scale_cache[canonical_uri] = (scale.astype(np.float32), offset.astype(np.float32))
    return _scale_cache[canonical_uri]


def read_signal_window(
    canonical_uri: str,
    start_sample: int,
    end_sample: int,
    channel_indices: Optional[list[int]] = None,
) -> np.ndarray:
    root = _open_zarr(canonical_uri)
    sig_arr = root["signal"]

    if channel_indices is not None:
        raw = np.array(sig_arr[channel_indices, start_sample:end_sample], dtype=np.float32)
    else:
        raw = np.array(sig_arr[:, start_sample:end_sample], dtype=np.float32)

    scale, offset = _get_scale_offset(canonical_uri)
    if channel_indices is not None:
        scale = scale[channel_indices]
        offset = offset[channel_indices]

    data = raw * scale[:, None] + offset[:, None]

    kernel_size = 5
    if data.shape[1] > kernel_size:
        kernel = np.ones(kernel_size) / kernel_size
        for i in range(data.shape[0]):
            data[i] = np.convolve(data[i], kernel, mode="same")

    return data


def get_channel_names(canonical_uri: str) -> list[str]:
    root = _open_zarr(canonical_uri)
    return list(root["channels"]["name"][:])


def get_store_metadata(canonical_uri: str) -> dict:
    root = _open_zarr(canonical_uri)
    attrs = dict(root.attrs)
    return {
        "n_channels": root["signal"].shape[0],
        "n_samples": root["signal"].shape[1],
        "sampling_rate_hz": attrs.get("sampling_rate_hz"),
        "duration_s": attrs.get("duration_s"),
        "channel_names": list(root["channels"]["name"][:]),
        "reference": attrs.get("reference"),
        "montage_name": attrs.get("montage_name"),
        "recording_type": attrs.get("recording_type"),
        "manufacturer": attrs.get("manufacturer"),
        "source_uri": attrs.get("source_uri"),
    }


# --- TUAR artifact annotations from CSV companion files ---

_annotation_cache: dict[str, list[dict]] = {}


def _get_csv_path_from_source_uri(source_uri: str) -> Optional[str]:
    if not source_uri or not source_uri.endswith(".edf"):
        return None
    return source_uri.replace("s3://", "").rsplit(".", 1)[0] + ".csv"


def get_annotations(canonical_uri: str, source_uri: Optional[str] = None) -> list[dict]:
    cache_key = canonical_uri
    if cache_key in _annotation_cache:
        return _annotation_cache[cache_key]

    if source_uri is None:
        try:
            rec = _open_recording(canonical_uri)
            source_uri = rec.metadata.source_uri
        except Exception:
            _annotation_cache[cache_key] = []
            return []

    csv_path = _get_csv_path_from_source_uri(source_uri)
    if csv_path is None:
        _annotation_cache[cache_key] = []
        return []

    try:
        fs = get_s3fs()
        raw = fs.cat(csv_path).decode("utf-8")
    except Exception:
        _annotation_cache[cache_key] = []
        return []

    annotations = _parse_tuar_csv(raw)
    _annotation_cache[cache_key] = annotations
    return annotations


def preload_all_annotations(df: pd.DataFrame) -> None:
    """Bulk-fetch all CSV annotation files in one batch S3 call."""
    paths_map: dict[str, str] = {}

    for _, row in df.iterrows():
        canonical_uri = row.get("canonical_uri", "")
        archival_uri = row.get("archival_uri", "")
        if not canonical_uri or canonical_uri in _annotation_cache:
            continue
        csv_path = _get_csv_path_from_source_uri(archival_uri)
        if csv_path:
            paths_map[csv_path] = canonical_uri

    if not paths_map:
        return

    fs = get_s3fs()
    csv_paths = list(paths_map.keys())

    BATCH = 50
    for i in range(0, len(csv_paths), BATCH):
        batch = csv_paths[i:i + BATCH]
        try:
            results = fs.cat(batch, on_error="return")
        except Exception:
            continue

        if isinstance(results, dict):
            for csv_path, content in results.items():
                canonical_uri = paths_map[csv_path]
                if isinstance(content, bytes):
                    try:
                        _annotation_cache[canonical_uri] = _parse_tuar_csv(content.decode("utf-8"))
                    except Exception:
                        _annotation_cache[canonical_uri] = []
                else:
                    _annotation_cache[canonical_uri] = []


def _parse_tuar_csv(raw: str) -> list[dict]:
    lines = [l for l in raw.strip().split("\n") if not l.startswith("#") and l.strip()]
    if not lines:
        return []

    header_line = lines[0]
    if "channel" in header_line and "start_time" in header_line:
        lines = lines[1:]

    seen = set()
    annotations = []

    for line in lines:
        parts = line.strip().split(",")
        if len(parts) < 4:
            continue

        channel = parts[0].strip()
        try:
            start = float(parts[1].strip())
            stop = float(parts[2].strip())
        except ValueError:
            continue
        raw_label = parts[3].strip().lower()
        confidence = float(parts[4].strip()) if len(parts) > 4 else 1.0

        dedup_key = (round(start, 3), round(stop, 3), raw_label)
        if dedup_key in seen:
            continue
        seen.add(dedup_key)

        label = ARTIFACT_LABEL_MAP.get(raw_label, raw_label.title())
        color = ARTIFACT_COLORS.get(label, "rgba(128, 128, 128, 0.2)")

        annotations.append({
            "onset_s": start,
            "duration_s": stop - start,
            "end_s": stop,
            "raw_label": raw_label,
            "label": label,
            "color": color,
            "channel": channel,
            "confidence": confidence,
        })

    annotations.sort(key=lambda a: a["onset_s"])
    return annotations


def get_annotations_in_window(
    canonical_uri: str,
    start_s: float,
    end_s: float,
    source_uri: Optional[str] = None,
) -> list[dict]:
    all_ann = get_annotations(canonical_uri, source_uri)
    return [a for a in all_ann if a["end_s"] > start_s and a["onset_s"] < end_s]


def get_recording_info(row: pd.Series) -> dict:
    channel_names = row.get("channel_names", [])
    if isinstance(channel_names, str):
        try:
            channel_names = json.loads(channel_names)
        except (json.JSONDecodeError, TypeError):
            channel_names = []
    if not isinstance(channel_names, list):
        channel_names = list(channel_names)

    return {
        "recording_id": row["recording_id"],
        "dataset_id": row["dataset_id"],
        "subject": row.get("subject_id_in_dataset", "N/A"),
        "session": row.get("session_id", "N/A"),
        "task": row.get("task", "N/A"),
        "duration_s": row.get("duration_s", 0),
        "n_channels": row.get("n_channels", 0),
        "n_eeg_channels": row.get("n_eeg_channels", 0),
        "sampling_rate_hz": row.get("sampling_rate_hz", 0),
        "reference": row.get("reference", "N/A"),
        "montage_name": row.get("montage_name", "N/A"),
        "recording_type": row.get("recording_type", "N/A"),
        "archival_format": row.get("archival_format", "N/A"),
        "canonical_uri": row.get("canonical_uri", ""),
        "archival_uri": row.get("archival_uri", ""),
        "channel_names": channel_names,
        "roundtrip_class": row.get("roundtrip_class", "N/A"),
    }