tetrad_results_aggregate.py

"""Aggregate multiple tetrad scaling benchmark JSONL outputs.

Read-only consolidation of Structural Field Tetrad timing/value records.
Preserves TNFR invariants (no mutation, pure telemetry post-processing).

Input: one or more JSONL files generated by `tetrad_scaling_benchmark.py`.
Each line example:
{
  "topology": "ring",
  "n_nodes": 20,
  "seed": 42,
  "precision_mode": "standard",
  "telemetry_density": "low",
  "timings": {"phi_s": ..., "phase_grad": ...},
  "tetrad_values": {"phi_s_mean": ..., "phase_grad_mean": ..., "xi_c": null},
  "snapshot_size": 449,
  "total_time": 0.0027
}

Output: aggregated CSV (always) and Parquet (if pyarrow available).
Schema columns (flat):
  topology, n_nodes, seed, precision_mode, telemetry_density,
    timing_phi_s, timing_phase_grad,
    timing_phase_curv, timing_xi_c,
    timing_tetrad_snapshot,
  phi_s_mean, phi_s_std, phase_grad_mean, phase_grad_std,
  phase_curv_mean, phase_curv_std, xi_c, snapshot_size, total_time

Usage examples:
    python benchmarks/tetrad_results_aggregate.py --glob \
        results/tetrad_scaling_*.jsonl
    python benchmarks/tetrad_results_aggregate.py --files \
        results/tetrad_scaling_20251115_*.jsonl --output-dir results

The aggregator is deterministic; ordering of input lines is preserved.
"""

from __future__ import annotations

import argparse
import datetime as _dt
import glob
import json
import os
import sys
from typing import Any, Dict, Iterable, List

try:  # Optional dependency for Parquet
    import pyarrow as pa  # type: ignore
    import pyarrow.parquet as pq  # type: ignore

    _PARQUET_AVAILABLE = True
except Exception:  # pragma: no cover - availability branch
    _PARQUET_AVAILABLE = False


REQUIRED_TOP_LEVEL = {
    "topology",
    "n_nodes",
    "seed",
    "precision_mode",
    "telemetry_density",
    "timings",
    "tetrad_values",
    "snapshot_size",
    "total_time",
}

REQUIRED_TIMINGS = {
    "phi_s",
    "phase_grad",
    "phase_curv",
    "xi_c",
    "tetrad_snapshot",
}
REQUIRED_TETRAD_VALUES = {
    "phi_s_mean",
    "phi_s_std",
    "phase_grad_mean",
    "phase_grad_std",
    "phase_curv_mean",
    "phase_curv_std",
    "xi_c",  # may be null
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Aggregate tetrad scaling benchmark JSONL outputs into CSV/Parquet"
        ),
    )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument(
        "--glob",
        help=(
            "Glob pattern for input JSONL files "
            "(e.g. results/tetrad_scaling_*.jsonl)"
        ),
    )
    group.add_argument(
        "--files",
        nargs="+",
        help="Explicit list of JSONL files to aggregate",
    )
    parser.add_argument(
        "--output-dir",
        default="results",
        help="Directory for aggregated outputs (default: results)",
    )
    parser.add_argument(
        "--prefix",
        default="tetrad_scaling_aggregate",
        help="Output filename prefix (default: tetrad_scaling_aggregate)",
    )
    parser.add_argument(
        "--fail-on-missing",
        action="store_true",
        help=("Fail if any required field missing (default: skip invalid " "records)"),
    )
    parser.add_argument(
        "--no-parquet",
        action="store_true",
        help="Force skip Parquet even if pyarrow available",
    )
    return parser.parse_args()


def iter_records(path: str) -> Iterable[Dict[str, Any]]:
    with open(path, "r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            try:
                yield json.loads(line)
            except json.JSONDecodeError as e:  # pragma: no cover - rare
                print(f"[WARN] Skipping malformed JSON line in {path}: {e}")
                continue


def validate_record(record: Dict[str, Any]) -> bool:
    # Top level
    if not REQUIRED_TOP_LEVEL.issubset(record.keys()):
        return False
    timings = record.get("timings", {})
    tetrad_vals = record.get("tetrad_values", {})
    if not REQUIRED_TIMINGS.issubset(timings.keys()):
        return False
    if not REQUIRED_TETRAD_VALUES.issubset(tetrad_vals.keys()):
        return False
    return True


def flatten_record(record: Dict[str, Any]) -> Dict[str, Any]:
    timings = record["timings"]
    tv = record["tetrad_values"]
    return {
        "topology": record["topology"],
        "n_nodes": record["n_nodes"],
        "seed": record["seed"],
        "precision_mode": record["precision_mode"],
        "telemetry_density": record["telemetry_density"],
        "timing_phi_s": timings["phi_s"],
        "timing_phase_grad": timings["phase_grad"],
        "timing_phase_curv": timings["phase_curv"],
        "timing_xi_c": timings["xi_c"],
        "timing_tetrad_snapshot": timings["tetrad_snapshot"],
        "phi_s_mean": tv["phi_s_mean"],
        "phi_s_std": tv["phi_s_std"],
        "phase_grad_mean": tv["phase_grad_mean"],
        "phase_grad_std": tv["phase_grad_std"],
        "phase_curv_mean": tv["phase_curv_mean"],
        "phase_curv_std": tv["phase_curv_std"],
        "xi_c": tv.get("xi_c"),  # may be None
        "snapshot_size": record["snapshot_size"],
        "total_time": record["total_time"],
    }


def write_csv(rows: List[Dict[str, Any]], path: str) -> None:
    if not rows:
        raise SystemExit("No rows to write.")
    headers = list(rows[0].keys())
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(",".join(headers) + "\n")
        for r in rows:
            row = ",".join(_serialize_csv_value(r[h]) for h in headers)
            fh.write(row + "\n")


def _serialize_csv_value(v: Any) -> str:
    if v is None:
        return ""
    if isinstance(v, float):
        # Use up to 8 significant digits to avoid overly long tails while
        # remaining stable for scaling analyses. This trims pathological
        # sequences like '123456789' verified in tests.
        return f"{v:.8g}"
    return str(v)


def write_parquet(rows: List[Dict[str, Any]], path: str) -> None:
    if not _PARQUET_AVAILABLE:
        return
    table = pa.Table.from_pylist(rows)
    pq.write_table(table, path)


def main() -> None:
    args = parse_args()
    if args.glob:
        input_files = sorted(glob.glob(args.glob))
    else:
        input_files = [f for f in args.files if os.path.isfile(f)]
    if not input_files:
        print("[ERROR] No input files matched.")
        sys.exit(1)

    rows: List[Dict[str, Any]] = []
    invalid_count = 0
    for fpath in input_files:
        for rec in iter_records(fpath):
            if not validate_record(rec):
                invalid_count += 1
                if args.fail_on_missing:
                    print(
                        f"[ERROR] Missing required fields in record from " f"{fpath}."
                    )
                    sys.exit(2)
                continue
            rows.append(flatten_record(rec))

    os.makedirs(args.output_dir, exist_ok=True)
    timestamp = _dt.datetime.utcnow().strftime("%Y%m%d_%H%M%S")
    base = f"{args.prefix}_{timestamp}"
    csv_path = os.path.join(args.output_dir, base + ".csv")
    write_csv(rows, csv_path)

    parquet_path = os.path.join(args.output_dir, base + ".parquet")
    if _PARQUET_AVAILABLE and not args.no_parquet:
        write_parquet(rows, parquet_path)
        parquet_msg = f"Parquet written: {parquet_path}"
    else:
        parquet_msg = "Parquet skipped (pyarrow not available or disabled)."

    print(f"Aggregated rows: {len(rows)} (invalid skipped: {invalid_count})")
    print(f"CSV written: {csv_path}")
    print(parquet_msg)
    print("Invariants preserved (read-only aggregation).")


if __name__ == "__main__":  # pragma: no cover
    main()