aggregate_stats.py

#!/usr/bin/env python3
"""Aggregate per-dataset stats across DMA, NOAA, Piraeus, Norway for REPORT.md.

Walks each dataset root for:
  - data_interim/{stage}/summary.json   (per-stage in/out row counts)
  - benchmark/*.export_summary.json     (per-partition window counts)
  - benchmark/core_manifest.csv         (split list, optional)
and prints a markdown comparison table to stdout.
"""
from __future__ import annotations

import os
import argparse
import glob
import json
import sys
from pathlib import Path


def agg_export_summary(root: Path) -> dict:
    total: dict[str, int] = {}
    n_files = 0
    bench = root / "benchmark"
    # Pattern 1: per-partition *.export_summary.json (DMA / NOAA layout)
    for s in sorted(bench.glob("*.export_summary.json")):
        try:
            d = json.loads(s.read_text())
            for k, v in d.get("counts", {}).items():
                total[k] = total.get(k, 0) + v
            n_files += 1
        except Exception:
            continue
    # Pattern 2: single root export_summary.json (smaller datasets like Norway)
    if not total and (bench / "export_summary.json").exists():
        try:
            d = json.loads((bench / "export_summary.json").read_text())
            for k, v in d.get("counts", {}).items():
                total[k] = total.get(k, 0) + v
            n_files = 1
        except Exception:
            pass
    total["_partition_files"] = n_files
    return total


def read_stage_summary(root: Path, stage_name: str) -> dict:
    """Find data_interim/<stage_name>/summary.json (or first .json found)."""
    p = root / "data_interim" / stage_name / "summary.json"
    if p.exists():
        try:
            return json.loads(p.read_text())
        except Exception:
            return {}
    return {}


def dataset_stats(root: Path, label: str) -> dict:
    info = {"label": label, "root": str(root)}
    if not root.exists():
        info["status"] = "MISSING"
        return info
    # Stage summaries (best-effort)
    for s in ["02_filtered", "03_deduped", "04_shiptype_speed_filtered",
              "05_segmented", "06_underway_only", "07_gap_imputed",
              "08_resampled_20s", "10_minlen_filtered", "11_windows", "12_quality_labeled"]:
        info[f"stage_{s.split('_',1)[0]}"] = read_stage_summary(root, s)
    info["benchmark"] = agg_export_summary(root)
    return info


def fmt_int(x) -> str:
    if x is None:
        return "—"
    if isinstance(x, int):
        return f"{x:,}"
    return str(x)


def get(d: dict, *path):
    cur = d
    for k in path:
        if cur is None:
            return None
        cur = cur.get(k) if isinstance(cur, dict) else None
    return cur


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument("--root", default=".", help="parent dir containing all dataset subdirs")
    p.add_argument("--dma-root", default=os.environ.get("DMA_ROOT", "data/dma"))
    p.add_argument("--noaa-root", default=os.environ.get("NOAA_ROOT", "data/noaa"))
    p.add_argument("--piraeus-name", default="Piraeus_ship_trajectory_datasets")
    p.add_argument("--norway-name", default="norway_ship_trajectory_datasets")
    args = p.parse_args()

    root = Path(args.root).resolve()
    datasets = [
        dataset_stats(Path(args.dma_root),   "DMA (2025-09)"),
        dataset_stats(Path(args.noaa_root),  "NOAA (2025-03)"),
        dataset_stats(root / args.piraeus_name, "Piraeus (2019-12)"),
        dataset_stats(root / args.norway_name,  "Norway (2025-09 wk1)"),
    ]

    # Markdown table
    metrics = [
        ("Total benchmark windows", lambda d: sum(v for k, v in d.get("benchmark", {}).items()
                                                   if k.startswith(("core_", "full_")))),
        ("Core train windows",     lambda d: get(d, "benchmark", "core_train")),
        ("Core val windows",       lambda d: get(d, "benchmark", "core_val")),
        ("Core test windows",      lambda d: get(d, "benchmark", "core_test")),
        ("Full train windows",     lambda d: get(d, "benchmark", "full_train")),
        ("Partition files",        lambda d: get(d, "benchmark", "_partition_files")),
        ("rows after stage-03 dedup", lambda d: get(d, "stage_03", "rows_after")),
        ("rows after stage-04 filter", lambda d: get(d, "stage_04", "rows_after")),
        ("segments after stage-10",   lambda d: get(d, "stage_10", "rows_after")),
        ("sliding windows (stage 11)",lambda d: get(d, "stage_11", "rows_after")),
    ]

    print()
    print("## Cross-dataset comparison table")
    print()
    headers = ["Metric"] + [d["label"] for d in datasets]
    print("| " + " | ".join(headers) + " |")
    print("|" + "|".join(["---"] * len(headers)) + "|")
    for name, fn in metrics:
        row = [name]
        for d in datasets:
            try:
                row.append(fmt_int(fn(d)))
            except Exception:
                row.append("—")
        print("| " + " | ".join(row) + " |")
    print()
    print("(DMA / NOAA stats from `benchmark/*.export_summary.json`; new datasets from this build.)")
    return 0


if __name__ == "__main__":
    sys.exit(main())