shared_data.py

"""Shared data, windowing, checkpoint, and metric helpers for SEST."""

import glob
import os
from pathlib import Path

import numpy as np
import torch
import yaml


def _expand(value):
    if isinstance(value, dict):
        return {key: _expand(item) for key, item in value.items()}
    if isinstance(value, list):
        return [_expand(item) for item in value]
    if isinstance(value, str):
        return os.path.expandvars(value)
    return value


def load_yaml(path):
    path = Path(path).resolve()
    with path.open("r", encoding="utf-8") as handle:
        config = yaml.safe_load(handle)
    if not isinstance(config, dict):
        raise ValueError(f"Expected a YAML mapping: {path}")
    config = _expand(config)
    config["_path"] = str(path)
    return config


def numeric_session_dirs(root):
    root = Path(root)
    if not root.is_dir():
        raise FileNotFoundError(f"Dataset directory does not exist: {root}")
    sessions = sorted(path for path in root.iterdir() if path.is_dir() and path.name.isdigit())
    if not sessions:
        raise FileNotFoundError(f"No numeric session folders found under: {root}")
    return sessions


def discover_roles(session_dir, modality, prefix):
    session_dir = Path(session_dir)
    suffix = modality if modality.startswith(".") else f".{modality}"
    roles = []
    for path in session_dir.glob(f"{prefix}*{suffix}"):
        role = path.name[: -len(suffix)]
        if role.startswith(prefix):
            roles.append(role)
    return sorted(set(roles))


def read_regression_labels(path):
    raw = np.atleast_1d(np.genfromtxt(path, delimiter="\n", dtype=str))
    values = []
    for value in raw:
        try:
            parsed = float(value)
        except (TypeError, ValueError):
            parsed = 0.0
        values.append(parsed if np.isfinite(parsed) else 0.0)
    return np.asarray(values, dtype=np.float32)


def read_classification_labels(path, classes, ignore_index=-100):
    raw = np.atleast_1d(np.genfromtxt(path, delimiter="\n", dtype=str))
    class_to_id = {name: index for index, name in enumerate(classes)}
    labels = []
    for value in raw:
        text = str(value).strip()
        if text in class_to_id:
            labels.append(class_to_id[text])
            continue
        try:
            numeric = int(float(text))
        except (TypeError, ValueError):
            numeric = ignore_index
        labels.append(numeric if 0 <= numeric < len(classes) else ignore_index)
    return np.asarray(labels, dtype=np.int64)


def core_starts(num_frames, core_length=96, stride=288):
    if num_frames < 1:
        return []
    if stride < 1:
        raise ValueError("stride must be >= 1.")
    if num_frames <= core_length:
        return [0]
    starts = list(range(0, num_frames - core_length + 1, stride))
    tail = num_frames - core_length
    if starts[-1] != tail:
        starts.append(tail)
    return starts


def edge_slice(values, start, end, pad_value=None):
    if values.ndim == 1:
        values = values[:, None]
        squeeze = True
    else:
        squeeze = False
    if values.shape[0] == 0:
        raise ValueError("Cannot slice an empty sequence.")
    pieces = []
    if start < 0:
        if pad_value is None:
            pieces.append(np.repeat(values[:1], -start, axis=0))
        else:
            pieces.append(np.full((-start, values.shape[1]), pad_value, dtype=values.dtype))
    valid_start = max(0, start)
    valid_end = min(values.shape[0], end)
    if valid_end > valid_start:
        pieces.append(values[valid_start:valid_end])
    if end > values.shape[0]:
        count = end - values.shape[0]
        if pad_value is None:
            pieces.append(np.repeat(values[-1:], count, axis=0))
        else:
            pieces.append(np.full((count, values.shape[1]), pad_value, dtype=values.dtype))
    result = np.concatenate(pieces, axis=0)
    expected = end - start
    if result.shape[0] != expected:
        raise ValueError(f"Window length mismatch: got {result.shape[0]}, expected {expected}.")
    return result[:, 0] if squeeze else result


def make_window(values, start, context=96, core=96, sparse_num=3):
    dense = edge_slice(values, start - context, start + core + context)
    return dense[::sparse_num] if sparse_num > 1 else dense


def make_dense_labels(labels, start, context=96, core=96, ignore_index=None):
    return edge_slice(
        labels,
        start - context,
        start + core + context,
        pad_value=ignore_index,
    )


class GeminiTable:
    def __init__(self, path):
        self.path = Path(path)
        if not self.path.is_file():
            raise FileNotFoundError(f"Missing Gemini feature file: {self.path}")
        with np.load(self.path, allow_pickle=False) as data:
            self.embedding = np.asarray(data["embedding"], dtype=np.float32)
            self.starts = np.asarray(data["core_starts"], dtype=np.int64).reshape(-1)
        if self.embedding.ndim != 2 or self.embedding.shape[0] != self.starts.shape[0]:
            raise ValueError(f"Invalid Gemini arrays in {self.path}")
        self.index = {int(start): index for index, start in enumerate(self.starts)}

    @property
    def dim(self):
        return int(self.embedding.shape[1])

    def get(self, start, missing="nearest"):
        index = self.index.get(int(start))
        if index is not None:
            return self.embedding[index]
        if missing == "zero":
            return np.zeros(self.dim, dtype=np.float32)
        if missing == "nearest":
            index = int(np.argmin(np.abs(self.starts - int(start))))
            return self.embedding[index]
        raise KeyError(f"Missing Gemini core_start={start} in {self.path}")


def reconstruct_window_outputs(outputs, starts, num_frames, context=96):
    outputs = np.asarray(outputs)
    trailing_shape = outputs.shape[2:]
    sums = np.zeros((num_frames, *trailing_shape), dtype=np.float64)
    counts = np.zeros(num_frames, dtype=np.float64)
    for output, start in zip(outputs, starts):
        dense_start = int(start) - context
        dense_end = dense_start + output.shape[0]
        valid_start = max(0, dense_start)
        valid_end = min(num_frames, dense_end)
        if valid_end <= valid_start:
            continue
        local_start = valid_start - dense_start
        local_end = local_start + valid_end - valid_start
        sums[valid_start:valid_end] += output[local_start:local_end]
        counts[valid_start:valid_end] += 1
    if np.any(counts == 0):
        raise ValueError(f"Window reconstruction left {np.count_nonzero(counts == 0)} frames uncovered.")
    divisor = counts.reshape((num_frames,) + (1,) * len(trailing_shape))
    return (sums / divisor).astype(np.float32)


def reconstruct_core_outputs(outputs, starts, num_frames, context=96, core=96):
    outputs = np.asarray(outputs)
    trailing_shape = outputs.shape[2:]
    sums = np.zeros((num_frames, *trailing_shape), dtype=np.float64)
    counts = np.zeros(num_frames, dtype=np.float64)
    for output, start in zip(outputs, starts):
        valid_start = int(start)
        valid_end = min(num_frames, valid_start + core)
        if valid_end <= valid_start:
            continue
        length = valid_end - valid_start
        sums[valid_start:valid_end] += output[context : context + length]
        counts[valid_start:valid_end] += 1
    if np.any(counts == 0):
        raise ValueError(f"Core reconstruction left {np.count_nonzero(counts == 0)} frames uncovered.")
    divisor = counts.reshape((num_frames,) + (1,) * len(trailing_shape))
    return (sums / divisor).astype(np.float32)


def ccc_score(labels, predictions):
    labels = np.asarray(labels, dtype=np.float64).reshape(-1)
    predictions = np.asarray(predictions, dtype=np.float64).reshape(-1)
    label_mean = labels.mean()
    prediction_mean = predictions.mean()
    covariance = np.mean((labels - label_mean) * (predictions - prediction_mean))
    denominator = labels.var() + predictions.var() + (label_mean - prediction_mean) ** 2
    return float(2 * covariance / (denominator + 1e-12))


def classification_metrics(labels, predictions, num_classes, ignore_index=-100):
    labels = np.asarray(labels).reshape(-1)
    predictions = np.asarray(predictions).reshape(-1)
    valid = (labels != ignore_index) & (labels >= 0) & (labels < num_classes)
    if not np.any(valid):
        raise ValueError("Cannot compute classification metrics without valid labels.")
    labels = labels[valid].astype(np.int64)
    predictions = predictions[valid].astype(np.int64)
    confusion = np.zeros((num_classes, num_classes), dtype=np.float64)
    np.add.at(confusion, (labels, predictions), 1)
    total = confusion.sum()
    observed = np.trace(confusion) / total
    expected = np.dot(confusion.sum(axis=1), confusion.sum(axis=0)) / (total * total)
    kappa = (observed - expected) / (1 - expected + 1e-12)
    return {"accuracy": float(observed), "kappa": float(kappa), "valid": int(total)}


def resolve_checkpoints(checkpoint_prefix, patterns):
    paths = []
    root = Path(checkpoint_prefix).resolve() if checkpoint_prefix else None
    for pattern in patterns:
        candidate = str(root / pattern) if root is not None and not os.path.isabs(pattern) else pattern
        matches = sorted(glob.glob(candidate))
        if not matches and Path(candidate).is_file():
            matches = [candidate]
        paths.extend(matches)
    paths = list(dict.fromkeys(str(Path(path).resolve()) for path in paths))
    if not paths:
        raise FileNotFoundError("No checkpoints matched the requested paths.")
    return paths


def resolve_normalization_stats(checkpoints, explicit_path=None, required=False):
    if explicit_path:
        path = Path(explicit_path).resolve()
        if not path.is_file():
            raise FileNotFoundError(f"Normalization statistics do not exist: {path}")
        return str(path)
    candidates = sorted(
        {
            str(Path(checkpoint).parent / "normalization_stats.npz")
            for checkpoint in checkpoints
            if (Path(checkpoint).parent / "normalization_stats.npz").is_file()
        }
    )
    if not candidates:
        if required:
            raise FileNotFoundError(
                "Normalization is enabled but normalization_stats.npz was not found next to the checkpoints. "
                "Pass --normalization-stats explicitly."
            )
        return None
    if len(candidates) == 1:
        return candidates[0]
    reference = None
    for candidate in candidates:
        with np.load(candidate) as data:
            arrays = {key: np.asarray(data[key]) for key in data.files}
        if reference is None:
            reference = arrays
            continue
        if arrays.keys() != reference.keys() or any(
            not np.array_equal(arrays[key], reference[key]) for key in reference
        ):
            raise ValueError(
                "Ensemble checkpoints use different normalization statistics; "
                "pass one compatible --normalization-stats file explicitly."
            )
    return candidates[0]


def clean_state_dict(payload):
    if isinstance(payload, dict) and "state_dict" in payload:
        payload = payload["state_dict"]
    if not isinstance(payload, dict):
        raise ValueError("Checkpoint must contain a state dictionary.")
    result = {}
    for key, value in payload.items():
        for prefix in ("module.", "_orig_mod."):
            if key.startswith(prefix):
                key = key[len(prefix) :]
        result[key] = value
    return result


def load_model_checkpoint(model, path, device):
    payload = torch.load(path, map_location=device)
    model.load_state_dict(clean_state_dict(payload), strict=True)
    model.to(device).eval()
    return model


def write_prediction_csv(path, values, overwrite=False):
    path = Path(path)
    if path.exists() and not overwrite:
        raise FileExistsError(f"Prediction file already exists: {path}")
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for value in values:
            handle.write(f"{value}\n")