scripts/phase3_inventory_and_sample.py

"""
Phase 3, Steps 1-3: inventory the real uploaded Parquet shard
(pipecat-ai/smart-turn-data-v3.2-train, shard 0 of 10, uploaded by the user
to /mnt/user-data/uploads/train-00000-of-00010.parquet) and materialize a
reproducible stratified sample of real audio for pipeline validation.

Uses the hand-written pure-Python Parquet reader in
src/turn_detector/parquet_reader.py, built specifically because this
sandbox has no network access to install pyarrow/fastparquet/datasets (see
docs/PHASE3_REAL_AUDIO_VALIDATION.md for the full story, including the
real bugs caught and fixed while building that reader against this file).
"""

from __future__ import annotations

import json
import subprocess
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from turn_detector.parquet_reader import (
    read_footer_metadata, read_column_values, parse_flac_streaminfo,
)
from turn_detector.data import stratified_reservoir_sample, duration_bucket

SHARD_PATH = "/mnt/user-data/uploads/train-00000-of-00010.parquet"
OUT_DIR = Path("data/raw/phase3_sample")
SCALAR_COLUMNS = ["id", "language", "endpoint_bool", "midfiller", "endfiller", "synthetic", "dataset"]


def decode_full_shard_metadata(shard_path: str) -> list[dict]:
    """Decode every scalar metadata column + audio duration (via FLAC
    STREAMINFO header only, no full audio decode) for EVERY row in the
    shard. This is cheap (~5s total, verified) since it doesn't touch the
    compressed audio payload beyond decompressing it to read a 34-byte
    header. Returns one dict per row, in file order.
    """
    pf = read_footer_metadata(shard_path)
    print(f"Shard: {pf.num_rows} total rows, {len(pf.row_groups)} row groups", file=sys.stderr)

    all_records = []
    corrupted = []

    for rg_idx, rg in enumerate(pf.row_groups):
        col_values = {}
        for colname in SCALAR_COLUMNS:
            col = rg.columns[colname]
            mdl = pf.max_def_levels[colname]
            col_values[colname] = read_column_values(shard_path, col, max_rows=rg.num_rows, max_def_level=mdl)

        audio_col = rg.columns["audio.bytes"]
        audio_mdl = pf.max_def_levels["audio.bytes"]
        audio_bytes_list = read_column_values(shard_path, audio_col, max_rows=rg.num_rows, max_def_level=audio_mdl)

        for i in range(rg.num_rows):
            rec = {c: col_values[c][i] for c in SCALAR_COLUMNS}
            for k in ("id", "language", "dataset"):
                if rec[k] is not None:
                    rec[k] = rec[k].decode("utf-8", errors="replace")
            audio_bytes = audio_bytes_list[i]
            if audio_bytes is None:
                corrupted.append({"row_group": rg_idx, "index_in_group": i, "id": rec["id"], "reason": "audio.bytes is null"})
                rec["duration_sec"] = None
                rec["sample_rate"] = None
                rec["channels"] = None
                rec["_corrupted"] = True
            else:
                try:
                    info = parse_flac_streaminfo(audio_bytes)
                    rec["duration_sec"] = info["duration_sec"]
                    rec["sample_rate"] = info["sample_rate"]
                    rec["channels"] = info["channels"]
                    rec["_corrupted"] = False
                except Exception as e:
                    corrupted.append({"row_group": rg_idx, "index_in_group": i, "id": rec["id"], "reason": str(e)})
                    rec["duration_sec"] = None
                    rec["sample_rate"] = None
                    rec["channels"] = None
                    rec["_corrupted"] = True
            rec["_row_group"] = rg_idx
            rec["_index_in_group"] = i
            all_records.append(rec)

        if (rg_idx + 1) % 8 == 0:
            print(f"  processed {rg_idx + 1}/{len(pf.row_groups)} row groups...", file=sys.stderr)

    print(f"Decoded {len(all_records)} rows. Corrupted/unreadable: {len(corrupted)}", file=sys.stderr)
    return all_records, corrupted, pf


def summarize(records: list[dict]) -> dict:
    n = len(records)
    def dist(key):
        c = Counter(str(r.get(key)) for r in records)
        return dict(sorted(c.items(), key=lambda kv: -kv[1]))

    durations = [r["duration_sec"] for r in records if r.get("duration_sec") is not None]
    sample_rates = Counter(r["sample_rate"] for r in records if r.get("sample_rate") is not None)
    channels = Counter(r["channels"] for r in records if r.get("channels") is not None)

    return {
        "n_rows": n,
        "endpoint_bool_distribution": dist("endpoint_bool"),
        "language_distribution": dist("language"),
        "dataset_source_distribution": dist("dataset"),
        "synthetic_distribution": dist("synthetic"),
        "midfiller_distribution_including_null": dist("midfiller"),
        "endfiller_distribution_including_null": dist("endfiller"),
        "sample_rates_found": dict(sample_rates),
        "channel_counts_found": dict(channels),
        "n_corrupted": sum(1 for r in records if r.get("_corrupted")),
        "duration_stats_sec": {
            "n": len(durations),
            "min": min(durations) if durations else None,
            "max": max(durations) if durations else None,
            "mean": sum(durations) / len(durations) if durations else None,
            "median": sorted(durations)[len(durations) // 2] if durations else None,
        },
    }


def main():
    print("=" * 70)
    print("STEP 1: Full-shard inventory (real data, actually decoded)")
    print("=" * 70)
    records, corrupted, pf = decode_full_shard_metadata(SHARD_PATH)
    summary = summarize(records)
    print(json.dumps(summary, indent=2))

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    with open(OUT_DIR / "full_shard_inventory_summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    with open(OUT_DIR / "corrupted_rows.json", "w") as f:
        json.dump(corrupted, f, indent=2)

    print()
    print("=" * 70)
    print("STEP 2: Stratified development sample (target 300, seed=42)")
    print("=" * 70)
    for r in records:
        r["duration_sec_for_strat"] = r.get("duration_sec")
    strat_input = [
        {**r, "duration_sec": r.get("duration_sec") or 0.0}
        for r in records
    ]
    sample = stratified_reservoir_sample(strat_input, target_n=300, seed=42)
    print(f"Sampled {len(sample)} rows (target 300) from {len(records)} available.")

    sample_summary = summarize(sample)
    print(json.dumps(sample_summary, indent=2))
    with open(OUT_DIR / "dev_sample_summary.json", "w") as f:
        json.dump(sample_summary, f, indent=2)

    print()
    print("=" * 70)
    print("STEP 3: Materializing sampled audio as real WAV files")
    print("=" * 70)

    pf2 = read_footer_metadata(SHARD_PATH)  # re-read for column offsets
    audio_dir = OUT_DIR / "audio"
    audio_dir.mkdir(parents=True, exist_ok=True)

    # Group sample by row_group to minimize re-reading column chunks
    by_rg: dict[int, list[dict]] = {}
    for r in sample:
        by_rg.setdefault(r["_row_group"], []).append(r)

    metadata_rows = []
    n_written = 0
    n_ffmpeg_failed = 0
    for rg_idx, recs in sorted(by_rg.items()):
        rg = pf2.row_groups[rg_idx]
        audio_col = rg.columns["audio.bytes"]
        audio_mdl = pf2.max_def_levels["audio.bytes"]
        audio_vals = read_column_values(SHARD_PATH, audio_col, max_rows=rg.num_rows, max_def_level=audio_mdl)

        for r in recs:
            idx = r["_index_in_group"]
            audio_bytes = audio_vals[idx]
            if audio_bytes is None:
                continue
            out_path = audio_dir / f"{r['id']}.wav"
            proc = subprocess.run(
                ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
                 "-i", "pipe:0", "-ar", "16000", "-ac", "1", str(out_path)],
                input=audio_bytes, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            )
            if proc.returncode != 0 or not out_path.exists():
                n_ffmpeg_failed += 1
                continue
            n_written += 1
            metadata_rows.append({
                "id": r["id"],
                "audio_path": str(out_path),
                "language": r["language"],
                "endpoint_bool": r["endpoint_bool"],
                "midfiller": r["midfiller"],
                "endfiller": r["endfiller"],
                "synthetic": r["synthetic"],
                "dataset": r["dataset"],
                "duration_sec": r.get("duration_sec"),
            })

    print(f"Materialized {n_written} real WAV files ({n_ffmpeg_failed} ffmpeg failures).")

    import csv
    with open(OUT_DIR / "metadata.csv", "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["id", "audio_path", "language", "endpoint_bool",
                                                  "midfiller", "endfiller", "synthetic", "dataset", "duration_sec"])
        writer.writeheader()
        for row in metadata_rows:
            writer.writerow(row)

    print(f"Wrote metadata for {len(metadata_rows)} rows to {OUT_DIR / 'metadata.csv'}")


if __name__ == "__main__":
    main()