src/data/affective_text.py

"""Utilities for SemEval-2007 Affective Text emotion-composition benchmark."""
from __future__ import annotations

import json
import re
import xml.etree.ElementTree as ET
from pathlib import Path

import numpy as np

EMOTION_NAMES = ["anger", "disgust", "fear", "joy", "sadness", "surprise"]


def _candidate_xml_files(root: Path) -> list[Path]:
    candidates = []
    for pattern in [
        "**/*test*.xml",
        "**/*headline*.xml",
        "**/*.xml",
    ]:
        candidates.extend(sorted(root.glob(pattern)))
    seen = set()
    out = []
    for path in candidates:
        if path in seen:
            continue
        seen.add(path)
        out.append(path)
    return out


def _candidate_gold_files(root: Path) -> list[Path]:
    candidates = []
    for pattern in [
        "**/*test*emotion*.gold*",
        "**/*emotion*.gold*",
        "**/*test*gold*",
        "**/*emotion*txt",
        "**/*gold*",
    ]:
        candidates.extend(sorted(root.glob(pattern)))
    seen = set()
    out = []
    for path in candidates:
        if path in seen:
            continue
        seen.add(path)
        out.append(path)
    return out


def _parse_headline_xml(path: Path) -> dict[str, str]:
    text = path.read_text(encoding="utf-8", errors="replace")
    matches = re.findall(r'<instance\s+id="([^"]+)">(.*?)</instance>', text, flags=re.DOTALL)
    if matches:
        return {str(idx): body.strip() for idx, body in matches if body.strip()}

    tree = ET.parse(path)
    root = tree.getroot()
    headlines = {}
    for elem in root.iter():
        elem_id = elem.attrib.get("id")
        body = (elem.text or "").strip()
        if elem_id and body:
            headlines[str(elem_id)] = body
    return headlines


def _parse_gold_file(path: Path) -> dict[str, np.ndarray]:
    gold = {}
    for raw_line in path.read_text(encoding="utf-8").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        parts = re.split(r"[\s,;]+", line)
        if len(parts) < 7:
            continue
        idx = parts[0]
        values = [float(x) for x in parts[1:7]]
        gold[str(idx)] = np.asarray(values, dtype=float)
    return gold


def _normalize_rows(x: np.ndarray) -> np.ndarray:
    x = np.maximum(np.asarray(x, dtype=float), 1e-8)
    return x / x.sum(axis=1, keepdims=True)


def load_affective_text(data_dir: str | Path) -> dict[str, object]:
    """Load SemEval-2007 Affective Text headlines and gold emotion scores.

    Expects the downloaded archive to be extracted under `data_dir`.
    The loader is intentionally permissive about internal filenames.
    """
    root = Path(data_dir)
    if not root.exists():
        raise FileNotFoundError(f"Data directory not found: {root}")

    headline_map = {}
    best_n = -1
    for path in _candidate_xml_files(root):
        parsed = _parse_headline_xml(path)
        if len(parsed) > best_n:
            headline_map = parsed
            best_n = len(parsed)
    if not headline_map:
        raise FileNotFoundError(f"Could not find headline XML under {root}")

    gold_map = {}
    best_n = -1
    for path in _candidate_gold_files(root):
        parsed = _parse_gold_file(path)
        if len(parsed) > best_n:
            gold_map = parsed
            best_n = len(parsed)
    if not gold_map:
        raise FileNotFoundError(f"Could not find gold emotion scores under {root}")

    common_ids = sorted(set(headline_map) & set(gold_map), key=lambda x: int(re.sub(r"\D", "", x) or x))
    if not common_ids:
        raise ValueError("No shared instance ids between headlines and gold scores")

    headlines = [headline_map[i] for i in common_ids]
    raw_scores = np.vstack([gold_map[i] for i in common_ids])
    Y = _normalize_rows(raw_scores)
    return {
        "ids": common_ids,
        "headlines": headlines,
        "raw_scores": raw_scores,
        "Y": Y,
        "emotions": EMOTION_NAMES,
    }


def load_prediction_cache(cache_path: str | Path) -> dict[str, dict]:
    cache = {}
    with open(cache_path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            cache[str(row["id"])] = row
    return cache


def build_prediction_matrix(
    ids: list[str],
    cache_path: str | Path,
) -> tuple[np.ndarray, np.ndarray]:
    cache = load_prediction_cache(cache_path)
    missing = [idx for idx in ids if idx not in cache]
    if missing:
        raise ValueError(f"Missing predictions for {len(missing)} ids in {cache_path}")

    raw_scores = np.vstack([np.asarray(cache[idx]["scores"], dtype=float) for idx in ids])
    U = _normalize_rows(raw_scores)
    return raw_scores, U