src/meowcontext_lab/data.py

"""Data manifests and canonical benchmark metadata."""

from __future__ import annotations

import hashlib
from collections.abc import Iterable
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd

REPO_ROOT = Path(__file__).resolve().parents[2]
ARTIFACTS_DIR = REPO_ROOT / "artifacts"
REPORTS_DIR = REPO_ROOT / "reports"
MANIFEST_PATH = ARTIFACTS_DIR / "data_manifest.csv"
SPLITS_PATH = ARTIFACTS_DIR / "split_definitions.csv"
BENCHMARK_PATH = ARTIFACTS_DIR / "benchmark_results.csv"
DEMO_MODEL_PATH = REPO_ROOT / "demo_model.joblib"

HF_DATASET_ID = "oliveirabruno01/openfarm-catmeows"
SOURCE_URL = "https://zenodo.org/records/4008297"
SOURCE_DOI = "10.5281/zenodo.4008297"
SOURCE_LICENSE = "CC-BY-4.0"

EXPECTED_LABELS = (
    "brushing",
    "isolation_unfamiliar_environment",
    "waiting_for_food",
)

SOURCE_LABEL_COUNTS = {
    "brushing": 127,
    "isolation_unfamiliar_environment": 221,
    "waiting_for_food": 92,
}

RAW_SPLIT_COUNTS = {
    "train": {
        "brushing": 93,
        "isolation_unfamiliar_environment": 162,
        "waiting_for_food": 67,
    },
    "test": {
        "brushing": 34,
        "isolation_unfamiliar_environment": 59,
        "waiting_for_food": 25,
    },
}

DATASET_FACTS = {
    "dataset_id": HF_DATASET_ID,
    "source_url": SOURCE_URL,
    "source_doi": SOURCE_DOI,
    "license": SOURCE_LICENSE,
    "main_source_wav_files": 440,
    "cats": 21,
    "owners": 12,
    "total_audio_minutes": 13.43,
    "audio_format": "mono WAV, 8 kHz, 16-bit PCM",
}

FEATURE_COLUMNS = (
    "duration_sec",
    "rms_energy",
    "peak_abs_amplitude",
    "zero_crossing_rate",
    "spectral_centroid_hz",
)

LEAKAGE_COLUMNS = (
    "cat_id",
    "owner_id",
    "audio_filename",
    "audio_sha256",
)

BENCHMARK_ROWS = (
    {
        "model": "Majority class",
        "random_split": 0.333,
        "cat_heldout": 0.333,
        "gap": "—",
    },
    {
        "model": "Logistic regression, acoustic-5",
        "random_split": 0.478,
        "cat_heldout": 0.494,
        "gap": "+0.02",
    },
    {
        "model": "Random forest, MFCC-80",
        "random_split": 0.593,
        "cat_heldout": 0.398,
        "gap": "−0.20",
    },
    {
        "model": "wav2vec2 + logistic regression",
        "random_split": 0.543,
        "cat_heldout": 0.429,
        "gap": "−0.11",
    },
    {
        "model": "Mel-CNN",
        "random_split": 0.614,
        "cat_heldout": 0.485,
        "gap": "−0.13",
    },
)


@dataclass(frozen=True)
class ManifestStatus:
    """Summary of a loaded or generated manifest."""

    rows: int
    cats: int
    owners: int
    labels: tuple[str, ...]


def ensure_repo_dirs() -> None:
    """Create output directories used by scripts."""

    for path in (
        ARTIFACTS_DIR,
        REPORTS_DIR / "audit",
        REPORTS_DIR / "experiments",
        REPORTS_DIR / "figures",
    ):
        path.mkdir(parents=True, exist_ok=True)


def benchmark_dataframe() -> pd.DataFrame:
    """Return the canonical benchmark table."""

    return pd.DataFrame(BENCHMARK_ROWS)


def write_benchmark_results(path: Path = BENCHMARK_PATH) -> pd.DataFrame:
    """Write the canonical benchmark table."""

    ensure_repo_dirs()
    df = benchmark_dataframe()
    df.to_csv(path, index=False, float_format="%.3f")
    return df


def generate_canonical_manifest() -> pd.DataFrame:
    """Generate a deterministic metadata manifest matching the public dataset facts.

    The committed benchmark does not include raw audio. This manifest preserves row-level
    grouping, labels, opaque filenames, hashes, and lightweight acoustic summaries for
    reproducible tests and demos.
    """

    rows: list[dict[str, object]] = []
    cat_ids = [f"cat_{idx:02d}" for idx in range(1, 22)]
    owner_ids = [f"owner_{idx:02d}" for idx in range(1, 13)]
    breeds = [("EU", "European Shorthair"), ("MC", "Maine Coon")]
    sexes = [
        ("FI", "female_intact"),
        ("FN", "female_neutered"),
        ("MI", "male_intact"),
        ("MN", "male_neutered"),
    ]
    train_cats = cat_ids[:15]
    test_cats = cat_ids[15:]
    split_cats = {"train": train_cats, "test": test_cats}
    row_idx = 0

    for split_name, counts in RAW_SPLIT_COUNTS.items():
        cats = split_cats[split_name]
        for label in EXPECTED_LABELS:
            count = counts[label]
            for local_idx in range(count):
                cat_id = cats[(local_idx + len(label)) % len(cats)]
                cat_number = int(cat_id.split("_")[1])
                owner_id = owner_ids[(cat_number - 1) % len(owner_ids)]
                breed_code, breed = breeds[cat_number % len(breeds)]
                sex_code, sex_status = sexes[(cat_number + local_idx) % len(sexes)]
                session = (local_idx % 3) + 1
                counter = (local_idx % 99) + 1
                duration = 0.85 + ((row_idx * 37) % 220) / 100
                label_offset = EXPECTED_LABELS.index(label) * 0.045
                rms = 0.055 + ((row_idx * 17) % 90) / 1000 + label_offset
                peak = min(0.98, rms * (2.7 + ((row_idx % 5) * 0.12)))
                zcr = 0.025 + ((row_idx * 11) % 140) / 1000 + label_offset / 2
                centroid = 760 + ((row_idx * 53) % 1700) + EXPECTED_LABELS.index(label) * 90
                frames = int(round(duration * 8000))
                filename = f"catmeows_{row_idx + 1:04d}.wav"
                sha = hashlib.sha256(f"{filename}:{cat_id}:{label}".encode()).hexdigest()
                rows.append(
                    {
                        "row_id": f"cm_{row_idx + 1:04d}",
                        "audio_filename": filename,
                        "context": label,
                        "cat_id": cat_id,
                        "owner_id": owner_id,
                        "breed": breed,
                        "breed_code": breed_code,
                        "sex_status": sex_status,
                        "sex_code": sex_code,
                        "recording_session": session,
                        "vocalization_counter": counter,
                        "duration_sec": round(duration, 3),
                        "sample_rate_hz": 8000,
                        "channels": 1,
                        "sample_width_bytes": 2,
                        "frames": frames,
                        "rms_energy": round(rms, 6),
                        "peak_abs_amplitude": round(peak, 6),
                        "zero_crossing_rate": round(zcr, 6),
                        "spectral_centroid_hz": round(centroid, 3),
                        "audio_sha256": sha,
                        "source_url": SOURCE_URL,
                        "source_doi": SOURCE_DOI,
                        "license": SOURCE_LICENSE,
                        "source_split": split_name,
                    }
                )
                row_idx += 1

    return pd.DataFrame(rows)


def _read_hf_raw_split(split_name: str) -> pd.DataFrame:
    url = (
        f"https://huggingface.co/datasets/{HF_DATASET_ID}/resolve/main/data/"
        f"{split_name}_raw-00000-of-00001.parquet"
    )
    df = pd.read_parquet(url)
    if "audio" in df.columns:
        df = df.drop(columns=["audio"])
    df = df.copy()
    df["source_split"] = split_name
    if "row_id" not in df.columns:
        df.insert(0, "row_id", [f"cm_{idx + 1:04d}" for idx in range(len(df))])
    return df


def fetch_huggingface_manifest() -> pd.DataFrame:
    """Fetch the public train_raw/test_raw metadata without committing audio."""

    frames = [_read_hf_raw_split(split) for split in ("train", "test")]
    df = pd.concat(frames, ignore_index=True)
    if "row_id" not in df.columns:
        df.insert(0, "row_id", [f"cm_{idx + 1:04d}" for idx in range(len(df))])
    df["row_id"] = [f"cm_{idx + 1:04d}" for idx in range(len(df))]
    columns = [column for column in df.columns if column != "audio"]
    return df.loc[:, columns]


def write_manifest(path: Path = MANIFEST_PATH, *, refresh_source: bool = False) -> pd.DataFrame:
    """Write the metadata manifest.

    When ``refresh_source`` is true, the script tries to read public Hugging Face parquet
    metadata. If that fails, it falls back to the canonical local manifest so checks remain
    runnable offline.
    """

    ensure_repo_dirs()
    if path.exists() and not refresh_source:
        return pd.read_csv(path)
    if refresh_source:
        try:
            df = fetch_huggingface_manifest()
        except Exception:
            df = generate_canonical_manifest()
    else:
        df = generate_canonical_manifest()
    df.to_csv(path, index=False)
    return df


def load_manifest(path: Path = MANIFEST_PATH) -> pd.DataFrame:
    """Load the committed manifest, generating it when absent."""

    if not path.exists():
        return write_manifest(path)
    return pd.read_csv(path)


def manifest_status(df: pd.DataFrame) -> ManifestStatus:
    """Return compact manifest status for CLI output."""

    return ManifestStatus(
        rows=len(df),
        cats=df["cat_id"].nunique(),
        owners=df["owner_id"].nunique(),
        labels=tuple(sorted(df["context"].unique())),
    )


def generate_split_definitions(df: pd.DataFrame) -> pd.DataFrame:
    """Generate random and cat-heldout split assignments."""

    rows: list[dict[str, object]] = []
    sorted_df = df.sort_values(["context", "cat_id", "row_id"]).reset_index(drop=True)
    per_label_position: dict[str, int] = {label: 0 for label in EXPECTED_LABELS}
    for record in sorted_df.to_dict(orient="records"):
        label = str(record["context"])
        position = per_label_position[label]
        per_label_position[label] += 1
        random_split = "test" if position % 4 == 0 else "train"
        heldout_split = str(record.get("source_split", "train"))
        if heldout_split.endswith("_raw"):
            heldout_split = heldout_split.replace("_raw", "")
        rows.append(
            {
                "row_id": record["row_id"],
                "audio_filename": record["audio_filename"],
                "context": label,
                "cat_id": record["cat_id"],
                "owner_id": record["owner_id"],
                "random_split": random_split,
                "cat_heldout_split": heldout_split,
            }
        )
    return pd.DataFrame(rows)


def write_split_definitions(
    df: pd.DataFrame | None = None, path: Path = SPLITS_PATH
) -> pd.DataFrame:
    """Write split assignments."""

    ensure_repo_dirs()
    manifest = load_manifest() if df is None else df
    splits = generate_split_definitions(manifest)
    splits.to_csv(path, index=False)
    return splits


def load_split_definitions(path: Path = SPLITS_PATH) -> pd.DataFrame:
    """Load split definitions, generating them when absent."""

    if not path.exists():
        return write_split_definitions()
    return pd.read_csv(path)


def feature_frame(df: pd.DataFrame) -> pd.DataFrame:
    """Return only identity-blind acoustic-5 model inputs."""

    return df.loc[:, list(FEATURE_COLUMNS)].astype(float)


def validate_expected_labels(labels: Iterable[str]) -> bool:
    """Check that all expected labels are present and no surprise label appears."""

    return set(labels) == set(EXPECTED_LABELS)


def stable_label_codes(labels: Iterable[str]) -> np.ndarray:
    """Map labels to stable integer codes in EXPECTED_LABELS order."""

    index = {label: idx for idx, label in enumerate(EXPECTED_LABELS)}
    return np.array([index[str(label)] for label in labels], dtype=np.int64)