icassp2027-experiment/prepare_easycom_oracle_features.py

#!/usr/bin/env python3
"""Build explicit, causal EasyCom oracle-person sidecars for information tests.

This is not a deployable feature extractor. It combines the already extracted
recent mixed-audio vector with ground-truth, participant-resolved speech
activity from EasyCom transcripts. Only fully observed 200 ms bins ending no
later than each state cutoff are exposed. The output format is deliberately
rejected by ordinary local-feature loaders.
"""

from __future__ import annotations

import argparse
import os
import tempfile
from pathlib import Path
from typing import Any, Sequence

import numpy as np
import torch

from derive_easycom_interaction_events import (
    CAMERA_WEARER_ID,
    FRAMES_PER_TASK_STEP,
    Rules,
    build_clips,
    load_utterances,
)
from easycom_head import (
    NUM_CLASSES,
    STEP_SECONDS,
    normalize_sample_name,
    resolve_feature_path,
)
from easycom_manifest import ALL_SESSIONS
from local_turn_features import (
    ORACLE_FORMAT,
    load_frame_set_features,
    reference_alignment,
)


ORACLE_CONDITIONS = (
    "visibility-only",
    "aggregate-visible",
    "per-person",
    "shuffled-id",
    "delayed",
)
UNKNOWN_PERSON_IDS = {"", "-1", "none", "null", "unknown"}


def exact_steps(seconds: float, name: str) -> int:
    if not np.isfinite(seconds) or seconds <= 0:
        raise ValueError(f"{name} must be finite and positive")
    steps = round(seconds / STEP_SECONDS)
    if not np.isclose(seconds, steps * STEP_SECONDS, atol=1e-9, rtol=0):
        raise ValueError(f"{name} must be a multiple of {STEP_SECONDS:g}s")
    return int(steps)


def participant_activity(
    utterances: Sequence[Any],
    *,
    num_steps: int,
    wearer_id: int,
) -> tuple[dict[int, torch.BoolTensor], torch.BoolTensor, dict[str, int]]:
    """Rasterize inclusive 20 Hz transcript spans into observed 5 Hz bins."""

    activity: dict[int, torch.BoolTensor] = {}
    wearer = torch.zeros(num_steps, dtype=torch.bool)
    discarded_terminal_frames = 0
    grid_frames = num_steps * FRAMES_PER_TASK_STEP
    for row in utterances:
        dense_start = int(row.dense_start_frame)
        dense_end = int(row.dense_end_frame_exclusive)
        if dense_end > grid_frames:
            overrun = dense_end - grid_frames
            if overrun > FRAMES_PER_TASK_STEP - 1:
                raise ValueError(
                    f"Transcript exceeds the discarded terminal sub-step by "
                    f"{overrun} frames: {row.uid}"
                )
            discarded_terminal_frames += max(
                0, dense_end - max(dense_start, grid_frames)
            )
            dense_end = grid_frames
        if dense_start >= dense_end:
            continue
        first = dense_start // FRAMES_PER_TASK_STEP
        last = (dense_end - 1) // FRAMES_PER_TASK_STEP
        if first < 0 or last < first or last >= num_steps:
            raise ValueError(
                f"Transcript interval leaves target grid: {row.uid}, "
                f"bins=[{first},{last}], target_steps={num_steps}"
            )
        destination = (
            wearer
            if int(row.participant_id) == wearer_id
            else activity.setdefault(
                int(row.participant_id),
                torch.zeros(num_steps, dtype=torch.bool),
            )
        )
        destination[first : last + 1] = True
    return activity, wearer, {
        "discarded_terminal_transcript_frames": discarded_terminal_frames
    }


def causal_history(
    sequence: torch.BoolTensor,
    *,
    row: int,
    history_steps: int,
    delay_steps: int = 0,
    observed_sequence: torch.BoolTensor | None = None,
) -> torch.Tensor:
    """Return [activity, observed-mask] ending at row-delay, never after row."""

    if sequence.ndim != 1 or sequence.dtype != torch.bool:
        raise ValueError("Activity sequence must be bool [T]")
    if observed_sequence is not None and (
        observed_sequence.dtype != torch.bool
        or observed_sequence.shape != sequence.shape
    ):
        raise ValueError("Observed sequence must be bool [T] aligned to activity")
    if not 0 <= row < len(sequence):
        raise IndexError(row)
    if history_steps <= 0 or delay_steps < 0:
        raise ValueError("Invalid history/delay")
    end = row - delay_steps
    values = torch.zeros(history_steps, dtype=torch.float32)
    observed = torch.zeros(history_steps, dtype=torch.float32)
    if end >= 0:
        start = max(0, end - history_steps + 1)
        count = end - start + 1
        available = (
            torch.ones(count, dtype=torch.bool)
            if observed_sequence is None
            else observed_sequence[start : end + 1]
        )
        values[-count:] = (sequence[start : end + 1] & available).float()
        observed[-count:] = available.float()
    return torch.cat((values, observed), dim=0)


def visible_people_by_row(
    source_payload: dict[str, Any],
    cutoffs: torch.Tensor,
    *,
    wearer_id: int,
) -> tuple[list[list[int]], list[list[dict[str, Any]]]]:
    """Use the latest provided-box frame strictly before every state cutoff."""

    persons = source_payload["persons"]
    frame_times = torch.as_tensor(persons["frame_times"], dtype=torch.float64)
    offsets = torch.as_tensor(persons["frame_offsets"], dtype=torch.long)
    records = persons["records"]
    row_to_frame = torch.searchsorted(frame_times, cutoffs, right=False) - 1
    ids_by_row: list[list[int]] = []
    records_by_row: list[list[dict[str, Any]]] = []
    for frame_index in row_to_frame.tolist():
        if frame_index < 0:
            ids_by_row.append([])
            records_by_row.append([])
            continue
        start = int(offsets[frame_index])
        stop = int(offsets[frame_index + 1])
        seen: set[int] = set()
        row_ids: list[int] = []
        row_records: list[dict[str, Any]] = []
        for source_record in records[start:stop]:
            raw_id = source_record.get("id")
            text_id = "" if raw_id is None else str(raw_id).strip()
            if text_id.lower() in UNKNOWN_PERSON_IDS:
                continue
            try:
                participant_id = int(text_id)
            except ValueError as error:
                raise ValueError(f"Non-numeric EasyCom participant ID {text_id!r}") from error
            if participant_id == wearer_id or participant_id in seen:
                continue
            seen.add(participant_id)
            row_ids.append(participant_id)
            row_records.append(source_record)
        ids_by_row.append(row_ids)
        records_by_row.append(row_records)
    return ids_by_row, records_by_row


def oracle_person_rows(
    activity: dict[int, torch.BoolTensor],
    visibility: dict[int, torch.BoolTensor],
    visible_ids: list[list[int]],
    source_records: list[list[dict[str, Any]]],
    *,
    condition: str,
    history_steps: int,
    delay_steps: int,
    seed: int,
) -> tuple[torch.Tensor, torch.LongTensor, list[dict[str, Any]], dict[str, int]]:
    """Create a 5 Hz ragged set of causal oracle histories."""

    if condition not in ORACLE_CONDITIONS:
        raise ValueError(f"Unknown oracle condition: {condition}")
    if len(visible_ids) != len(source_records):
        raise ValueError("Visible ID/record row counts differ")
    num_steps = len(visible_ids)
    zero_sequence = torch.zeros(num_steps, dtype=torch.bool)
    rng = np.random.default_rng(seed)
    feature_rows: list[torch.Tensor] = []
    records: list[dict[str, Any]] = []
    offsets = [0]
    shuffled_values_changed = 0
    shuffled_values_total = 0

    for row, (person_ids, row_records) in enumerate(zip(visible_ids, source_records)):
        if len(person_ids) != len(row_records):
            raise ValueError("Visible person IDs/records differ within a row")
        base = torch.stack(
            [
                causal_history(
                    activity.get(person_id, zero_sequence),
                    row=row,
                    history_steps=history_steps,
                    delay_steps=delay_steps if condition == "delayed" else 0,
                    observed_sequence=visibility.get(person_id, zero_sequence),
                )
                for person_id in person_ids
            ],
            dim=0,
        ) if person_ids else torch.empty((0, 2 * history_steps))

        if condition == "visibility-only":
            output = base.clone()
            output[:, :history_steps] = 0
        elif condition == "aggregate-visible":
            if len(base):
                output = torch.cat(
                    (
                        base[:, :history_steps].amax(dim=0, keepdim=True),
                        base[:, history_steps:].amax(dim=0, keepdim=True),
                    ),
                    dim=-1,
                )
            else:
                output = base
        elif condition == "shuffled-id":
            output = base.clone()
            # Independently reassign each historical time bin across the
            # currently visible slots. Per-bin activity counts remain exact,
            # while cross-time person continuity is destroyed.
            for column in range(history_steps):
                if len(base) > 1:
                    permutation = torch.from_numpy(rng.permutation(len(base)))
                    output[:, column] = base[permutation, column]
                    output[:, history_steps + column] = base[
                        permutation, history_steps + column
                    ]
            shuffled_values_changed += int(
                (output[:, :history_steps] != base[:, :history_steps]).sum()
            )
            shuffled_values_total += int(base[:, :history_steps].numel())
        else:
            output = base

        feature_rows.extend(output.unbind(dim=0))
        if condition == "aggregate-visible" and person_ids:
            records.append(
                {
                    "id": "aggregate-visible",
                    "source_person_ids": person_ids,
                    "source_visual_times": [
                        float(item.get("source_time", float("nan")))
                        for item in row_records
                    ],
                }
            )
        else:
            for person_id, source_record in zip(person_ids, row_records):
                records.append(
                    {
                        "id": str(person_id),
                        "source_clip_index": int(
                            source_record.get("source_clip_index", 0)
                        ),
                        "source_visual_time": float(
                            source_record.get("source_time", float("nan"))
                        ),
                    }
                )
        offsets.append(offsets[-1] + len(output))

    feature_dim = 2 * history_steps
    features = (
        torch.stack(feature_rows)
        if feature_rows
        else torch.empty((0, feature_dim), dtype=torch.float32)
    )
    return (
        features,
        torch.tensor(offsets, dtype=torch.long),
        records,
        {
            "shuffled_values_changed": shuffled_values_changed,
            "shuffled_values_total": shuffled_values_total,
        },
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-root", type=Path, required=True)
    parser.add_argument("--feature-root", type=Path, required=True)
    parser.add_argument("--target-root", type=Path, required=True)
    parser.add_argument("--local-feature-root", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--sessions", nargs="+", default=list(ALL_SESSIONS))
    parser.add_argument("--conditions", nargs="+", choices=ORACLE_CONDITIONS,
                        default=list(ORACLE_CONDITIONS))
    parser.add_argument("--history-seconds", type=float, default=2.0)
    parser.add_argument("--delay-seconds", type=float, default=1.0)
    parser.add_argument("--camera-wearer-id", type=int, default=CAMERA_WEARER_ID)
    parser.add_argument("--shuffle-seed", type=int, default=2027)
    parser.add_argument("--max-label-mismatch-rows", type=int, default=0)
    parser.add_argument("--expected-mode", default="open-stream-replay")
    parser.add_argument("--reuse-existing", action="store_true")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    history_steps = exact_steps(args.history_seconds, "--history-seconds")
    delay_steps = exact_steps(args.delay_seconds, "--delay-seconds")
    if args.max_label_mismatch_rows < 0:
        raise ValueError("--max-label-mismatch-rows cannot be negative")
    sessions = [normalize_sample_name(value, "easycom") for value in args.sessions]
    unknown = sorted(set(sessions) - set(ALL_SESSIONS))
    if unknown:
        raise ValueError(f"Unknown EasyCom sessions: {unknown}")
    if len(sessions) != len(set(sessions)):
        raise ValueError("Duplicate sessions")
    conditions = list(dict.fromkeys(args.conditions))
    rules = Rules()

    print(
        f"Oracle analysis only: conditions={','.join(conditions)}, "
        f"history={history_steps}x{STEP_SECONDS:g}s, "
        f"delay={delay_steps}x{STEP_SECONDS:g}s"
    )
    for session_position, session in enumerate(sessions, 1):
        reference_path = resolve_feature_path(args.feature_root, session)
        target_path = args.target_root / f"{session}.npy"
        source_local_path = resolve_feature_path(args.local_feature_root, session)
        reference = torch.load(reference_path, map_location="cpu", weights_only=False)
        times, label_indices = reference_alignment(reference)
        if args.expected_mode != "any" and reference.get("mode") != args.expected_mode:
            raise ValueError(
                f"{reference_path} mode is {reference.get('mode')!r}, "
                f"expected {args.expected_mode!r}"
            )
        target = np.load(target_path, allow_pickle=False)
        if target.shape != (len(times), NUM_CLASSES):
            raise ValueError(
                f"Target/reference mismatch: {target_path}, "
                f"{target.shape} vs {(len(times), NUM_CLASSES)}"
            )

        direct_audio, _source_people, source_contract = load_frame_set_features(
            source_local_path,
            reference,
            sample_id=session,
            dataset="easycom",
        )
        source_payload = torch.load(
            source_local_path, map_location="cpu", weights_only=False
        )
        audio_dim = int(source_contract["config"]["audio_dim"])
        audio = direct_audio[:, :audio_dim]
        audio_valid = direct_audio[:, audio_dim].bool()

        clips = build_clips(args.dataset_root.resolve(), session, rules)
        utterances = load_utterances(session, clips)
        activity, wearer_activity, transcript_audit = participant_activity(
            utterances,
            num_steps=len(times),
            wearer_id=args.camera_wearer_id,
        )
        all_other_activity = torch.zeros(len(times), dtype=torch.bool)
        for sequence in activity.values():
            all_other_activity |= sequence
        target_tensor = torch.from_numpy(np.asarray(target))
        wearer_mismatch = int(
            (wearer_activity != target_tensor[:, 1].bool()).sum()
        )
        other_mismatch = int(
            (all_other_activity != target_tensor[:, 2].bool()).sum()
        )
        mismatch_union = int(
            (
                (wearer_activity != target_tensor[:, 1].bool())
                | (all_other_activity != target_tensor[:, 2].bool())
            ).sum()
        )
        if mismatch_union > args.max_label_mismatch_rows:
            raise ValueError(
                f"{session}: transcript/target activity differs on "
                f"{mismatch_union} rows (wearer={wearer_mismatch}, "
                f"others={other_mismatch}); inspect alignment or explicitly "
                "raise --max-label-mismatch-rows"
            )

        visible_ids, source_records = visible_people_by_row(
            source_payload,
            times,
            wearer_id=args.camera_wearer_id,
        )
        visibility: dict[int, torch.BoolTensor] = {
            participant_id: torch.zeros(len(times), dtype=torch.bool)
            for participant_id in activity
        }
        for row, participant_ids in enumerate(visible_ids):
            for participant_id in participant_ids:
                visibility.setdefault(
                    participant_id, torch.zeros(len(times), dtype=torch.bool)
                )[row] = True
        visible_counts = [len(row) for row in visible_ids]
        print(
            f"[{session_position}/{len(sessions)}] {session}: "
            f"states={len(times)}, participants={len(activity)}, "
            f"visible_rows={sum(count > 0 for count in visible_counts)}, "
            f"max_visible={max(visible_counts, default=0)}, "
            f"target_mismatch={mismatch_union}"
        )

        for condition in conditions:
            destination = args.output_root / condition / f"{session}.pt"
            if destination.exists() and args.reuse_existing and not args.overwrite:
                _, existing_people, existing_contract = load_frame_set_features(
                    destination,
                    reference,
                    sample_id=session,
                    dataset="easycom",
                    oracle_analysis=True,
                )
                existing_config = existing_contract["config"]
                expected_existing = {
                    "oracle_condition": condition,
                    "history_steps": history_steps,
                    "delay_steps": delay_steps if condition == "delayed" else 0,
                    "camera_wearer_id": args.camera_wearer_id,
                    "audio_dim": audio_dim,
                    "source_local_contract": source_contract,
                }
                mismatches = {
                    key: (existing_config.get(key), value)
                    for key, value in expected_existing.items()
                    if existing_config.get(key) != value
                }
                if mismatches:
                    raise ValueError(
                        f"Cannot reuse incompatible oracle sidecar {destination}: "
                        f"{mismatches}"
                    )
                print(
                    f"  {condition}: reuse {destination}; "
                    f"observations={len(existing_people.features)}"
                )
                continue
            if destination.exists() and not args.overwrite:
                raise FileExistsError(
                    f"Refusing to overwrite {destination}; pass --overwrite "
                    "or --reuse-existing"
                )
            condition_seed = (
                args.shuffle_seed
                + 1009 * session_position
                + ORACLE_CONDITIONS.index(condition)
            )
            features, frame_offsets, records, corruption = oracle_person_rows(
                activity,
                visibility,
                visible_ids,
                source_records,
                condition=condition,
                history_steps=history_steps,
                delay_steps=delay_steps,
                seed=condition_seed,
            )
            visual = torch.zeros((len(times), 2 * history_steps))
            visual_valid = torch.zeros(len(times), dtype=torch.bool)
            for row in range(len(times)):
                start = int(frame_offsets[row])
                stop = int(frame_offsets[row + 1])
                if stop > start:
                    visual[row] = features[start:stop].mean(dim=0)
                    visual_valid[row] = True
            config = {
                "oracle_analysis": True,
                "uses_speech_labels": True,
                "causal_policy": "fully-observed-half-open-5hz-bins",
                "oracle_condition": condition,
                "oracle_scope": "currently-visible-nonwearer-participants",
                "history_steps": history_steps,
                "history_seconds": history_steps * STEP_SECONDS,
                "delay_steps": delay_steps if condition == "delayed" else 0,
                "delay_seconds": (
                    delay_steps * STEP_SECONDS if condition == "delayed" else 0.0
                ),
                "shuffle_seed": args.shuffle_seed if condition == "shuffled-id" else None,
                "speaker_activity_source": "EasyCom Speech_Transcriptions inclusive 20Hz intervals",
                "visibility_source": "provided face boxes in source local sidecar",
                "camera_wearer_id": args.camera_wearer_id,
                "visual_dim": 2 * history_steps,
                "audio_dim": audio_dim,
                "source_local_contract": source_contract,
            }
            payload = {
                "format": ORACLE_FORMAT,
                "dataset": "easycom",
                "sample_id": session,
                "times": times,
                "label_indices": label_indices,
                "partial": False,
                "visual": visual,
                "audio": audio,
                "visual_valid": visual_valid,
                "audio_valid": audio_valid,
                "persons": {
                    "features": features,
                    # Bin k covers [k/5,(k+1)/5); its feature is available at
                    # the matching state cutoff (k+1)/5.
                    "frame_times": torch.arange(len(times), dtype=torch.float64)
                    * STEP_SECONDS,
                    "frame_offsets": frame_offsets,
                    "records": records,
                },
                "config": config,
                "reference_features": str(reference_path.resolve()),
                "source_local_features": str(source_local_path.resolve()),
                "audit": {
                    "strictly_past_or_current_bins_only": True,
                    "transcript_target_mismatch_rows": mismatch_union,
                    "wearer_mismatch_rows": wearer_mismatch,
                    "others_mismatch_rows": other_mismatch,
                    "rows_with_visible_people": sum(
                        count > 0 for count in visible_counts
                    ),
                    "max_visible_people": max(visible_counts, default=0),
                    "person_observations": len(features),
                    "generation_seed": condition_seed,
                    **transcript_audit,
                    **corruption,
                },
            }
            if args.dry_run:
                print(
                    f"  {condition}: would write {destination}; "
                    f"observations={len(features)}"
                )
                continue
            destination.parent.mkdir(parents=True, exist_ok=True)
            with tempfile.NamedTemporaryFile(
                dir=destination.parent,
                prefix=f".{session}.",
                suffix=".tmp",
                delete=False,
            ) as handle:
                temporary = Path(handle.name)
            try:
                torch.save(payload, temporary)
                os.replace(temporary, destination)
            finally:
                temporary.unlink(missing_ok=True)
            print(
                f"  {condition}: saved {destination}; "
                f"observations={len(features)}"
            )


if __name__ == "__main__":
    main()