# SPDX-License-Identifier: Apache-2.0
"""Aggregate a run_matrix.py results directory into latency.csv, benchmark.json and README.md.

Pure stdlib. Safe to run at any time, including on a partial or interrupted run:

    python3 aggregate.py --out ./results-<host>-<date>

``run_matrix.py`` imports this module and calls it after every cell, so the
summary files always reflect what has been measured so far.

Column names deliberately keep the Spark spelling (``nvfp4_36l_480p_mean_s``)
so the CSV joins with the DGX Spark one. On the Mac the NVFP4 VAE decoders run
as W4A16 (weights NVFP4-packed, activations bf16), not as a native FP4 GEMM.
"""
from __future__ import annotations

import argparse
import csv
import json
import statistics
import time
from pathlib import Path

DEFAULT_COMPOSITIONS = ("image", "image-voice", "images3", "clip", "video")
EXTRA_COMPOSITIONS = ("mixed", "clips3", "images9")
LABELS = {
    "image": "1 image",
    "image-voice": "1 image + voice",
    "images3": "3 images",
    "images9": "9 images",
    "clip": "1 short clip",
    "video": "1 long clip",
    "clips3": "3 clips",
    "mixed": "mixed references",
    "full": "9 images + 3 clips",
}
STANDARD_CANVASES = {"480p": (832, 480), "768p": (1344, 768)}
DECODERS = ("36l", "8l")
LANE_PREFIX = "fastvideo_omni_mlx"
WARMUP_SEED = 42
DEFAULT_SEEDS = (0, 1, 2)
SPARK_COLUMN_ORDER = (("36l", "480p"), ("8l", "480p"), ("36l", "768p"), ("8l", "768p"))
RECORD_SCHEMA = "omniref-mlx-request-v1"


def cell_id_for(composition: str, seconds: int, decoder: str, canvas: str) -> str:
    return f"{composition}-{seconds}s-{decoder}-{canvas}"


def lane_key(decoder: str, canvas: str) -> str:
    return f"nvfp4_{decoder}_{canvas}"


def read_jsonl(path: Path) -> list[dict]:
    if not path.is_file():
        return []
    rows = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line:
            continue
        try:
            rows.append(json.loads(line))
        except json.JSONDecodeError:
            continue  # a torn final line from a hard kill must not break aggregation
    return rows


def read_json(path: Path) -> dict:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return {}


def sample_std(values: list[float]) -> float | None:
    return statistics.stdev(values) if len(values) >= 2 else None


def _ok(record: dict) -> bool:
    return record.get("status") == "ok" and bool(record.get("output", {}).get("ok"))


def summarize_cell(cell: dict, records: list[dict], status: dict) -> dict:
    """Merge one cell's records and status file into its measured/failed classification."""
    smoke = bool(status.get("smoke") or cell.get("smoke"))
    requested = tuple(status.get("requested_seeds") or ((WARMUP_SEED,) if smoke else DEFAULT_SEEDS))
    partition = "smoke" if smoke else "measured"
    own = [r for r in records if r.get("cell_id") == cell["cell_id"]]
    by_seed: dict[int, dict] = {}
    for record in own:  # last successful record per seed wins (resume re-runs)
        if record.get("partition") == partition and _ok(record):
            by_seed[int(record["seed"])] = record
    failures = [r for r in own if r.get("status") == "failed"]
    state = status.get("state", "not_started")
    if requested and all(seed in by_seed for seed in requested):
        classification = "measured"
    elif state in ("starting", "loading_model", "request_running"):
        classification = "in_progress"
    elif state == "interrupted":
        classification = "interrupted"
    elif failures or state == "failed":
        classification = "failed"
    else:
        classification = "not_measured"
    seconds = [float(by_seed[s]["wall_s"]) for s in requested if s in by_seed]
    return {
        **cell,
        "smoke": smoke,
        "requested_seeds": list(requested),
        "records": own,
        "by_seed": by_seed,
        "failures": failures,
        "status_file": status,
        "classification": classification,
        "mean_s": statistics.mean(seconds) if seconds else None,
        "std_s": sample_std(seconds),
        "seconds": seconds,
    }


def discover_cells(out: Path, planned: list[dict] | None) -> list[dict]:
    cells: dict[str, dict] = {}
    for item in planned or []:
        cells[item["cell_id"]] = {k: item[k] for k in ("cell_id", "composition", "target_seconds", "decoder", "canvas",
                                                       "width", "height", "dit", "quant_gemm") if k in item}
    for path in sorted((out / "status").glob("*.json")) if (out / "status").is_dir() else []:
        if path.name.endswith(".spec.json"):
            continue
        status = read_json(path)
        if "cell_id" in status and status["cell_id"] not in cells:
            cells[status["cell_id"]] = {k: status.get(k) for k in ("cell_id", "composition", "target_seconds",
                                                                  "decoder", "canvas", "width", "height")}
            cells[status["cell_id"]].update({k: status[k] for k in ("dit", "quant_gemm") if k in status})
    return list(cells.values())


def _sort_key(cell: dict):
    comps = DEFAULT_COMPOSITIONS + EXTRA_COMPOSITIONS
    comp_rank = comps.index(cell["composition"]) if cell["composition"] in comps else len(comps)
    return (cell["target_seconds"], comp_rank, cell["composition"], cell["canvas"], cell["decoder"])


def build_summaries(out: Path, planned: list[dict] | None = None) -> list[dict]:
    records = read_jsonl(out / "per-request-records.jsonl")
    summaries = []
    for cell in discover_cells(out, planned):
        status = read_json(out / "status" / f"{cell['cell_id']}.json")
        summaries.append(summarize_cell(cell, records, status))
    return sorted(summaries, key=_sort_key)


def _lane_columns(summaries: list[dict]) -> list[tuple[str, str]]:
    present = {(c["decoder"], c["canvas"]) for c in summaries}
    standard = [lane for lane in SPARK_COLUMN_ORDER]  # always emit the Spark columns, even when empty
    extras = sorted(lane for lane in present if lane not in standard)
    return standard + extras


def write_latency_csv(out: Path, summaries: list[dict]) -> list[dict]:
    lanes = _lane_columns(summaries)
    header = ["composition", "target_seconds", "gpu_count"]
    for decoder, canvas in lanes:
        key = lane_key(decoder, canvas)
        header += [f"{key}_mean_s", f"{key}_sample_std_s", f"{key}_status"]
    grouped: dict[tuple[str, int], dict[tuple[str, str], dict]] = {}
    for cell in summaries:
        grouped.setdefault((cell["composition"], cell["target_seconds"]), {})[(cell["decoder"], cell["canvas"])] = cell
    rows = []
    for (composition, seconds), by_lane in grouped.items():
        row: dict = {"composition": composition, "target_seconds": seconds, "gpu_count": 1}
        for decoder, canvas in lanes:
            key = lane_key(decoder, canvas)
            cell = by_lane.get((decoder, canvas))
            measured = cell is not None and cell["classification"] == "measured"
            row[f"{key}_mean_s"] = cell["mean_s"] if measured else ""
            row[f"{key}_sample_std_s"] = cell["std_s"] if measured and cell["std_s"] is not None else ""
            row[f"{key}_status"] = cell["classification"] if cell is not None else "not_measured"
        rows.append(row)
    with open(out / "latency.csv", "w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=header)
        writer.writeheader()
        writer.writerows(rows)
    return rows


def _output_entry(record: dict) -> dict:
    output = record.get("output", {})
    video = (output.get("video_streams") or [{}])[0]
    audio = (output.get("audio_streams") or [{}])[0]
    return {
        "seed": record.get("seed"),
        "file": output.get("relative_path"),
        "sha256": output.get("sha256"),
        "bytes": output.get("bytes"),
        "video_codec": video.get("codec_name"),
        "width": video.get("width"),
        "height": video.get("height"),
        "nb_frames": video.get("nb_frames"),
        "r_frame_rate": video.get("r_frame_rate"),
        "audio_codec": audio.get("codec_name"),
        "audio_sample_rate": audio.get("sample_rate"),
        "audio_channels": audio.get("channels"),
        "duration_s": float(output.get("format", {}).get("duration", 0.0) or 0.0),
    }


def _profile(cell: dict, index_by_key: dict) -> dict:
    lane = f"{LANE_PREFIX}_{cell['decoder']}_{cell['canvas']}"
    profile = {
        "id": cell["cell_id"],
        "composition": cell["composition"],
        "label": LABELS.get(cell["composition"], cell["composition"]),
        "target_seconds": cell["target_seconds"],
        "decoder": cell["decoder"],
        "canvas": cell["canvas"],
        "width": cell.get("width"),
        "height": cell.get("height"),
        "fps": 24,
        "status": cell["classification"],
        "requested_seeds": cell["requested_seeds"],
        "conditions": None,
        "outputs": {},
        "performance": {},
    }
    first = next(iter(cell["by_seed"].values()), None) or (cell["records"][0] if cell["records"] else {})
    profile["prompt_sha256"] = first.get("prompt_sha256")
    profile["conditions"] = first.get("condition_media")
    profile["frames"] = first.get("num_frames")
    if cell["classification"] == "measured":
        seed0 = cell["by_seed"].get(min(cell["by_seed"]))
        profile["outputs"][lane] = {**_output_entry(seed0), "decoder": cell["decoder"]}
        profile["performance"][lane] = {
            "mean_s": cell["mean_s"],
            "sample_std_s": cell["std_s"],
            "measured_seconds": cell["seconds"],
            "measured_seeds": cell["requested_seeds"],
        }
        profile["validation"] = {"outputs": [_output_entry(cell["by_seed"][s]) for s in sorted(cell["by_seed"])]}
    if cell["failures"]:
        failure = cell["failures"][-1]
        profile["failure"] = {k: failure.get(k) for k in ("error_type", "error_message", "exit_code", "signal")}
    return profile


def write_benchmark_json(out: Path, summaries: list[dict]) -> None:
    receipt = read_json(out / "environment_receipt.json")
    settings = receipt.get("run_settings", {})
    benchmark = {
        "schema_version": 1,
        "protocol": {
            "snapshot_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
            "hardware": receipt.get("hardware", {}).get("chip"),
            "memory_gib": receipt.get("hardware", {}).get("memory_gib"),
            "gpu_count": 1,
            "concurrency": 1,
            "fps": 24,
            "warmup_seed": WARMUP_SEED,
            "measured_seeds": list(DEFAULT_SEEDS),
            "warmup_requests_per_profile": 1,
            "latency": "client end-to-end to completed MP4 (pipeline.generate); model load excluded",
            "aggregation": "arithmetic mean and sample standard deviation across measured seeds",
            "fast_denoiser_evaluations": settings.get("num_steps"),
            "lane": "FastVideo OmniRef MLX (Apple Silicon). Column names keep the Spark 'nvfp4' spelling; the "
                    "Mac runs the NVFP4 VAE decoders as W4A16 and the DiT as the converted quantization named in "
                    "'dit'. Not directly an FP4-tensor-core comparison.",
            "source": receipt.get("software", {}),
            "host": receipt.get("hostname"),
            "settings": settings,
            "input_identity": receipt.get("data", {}),
            "model_identity": receipt.get("weights", {}),
        },
        "profiles": [_profile(c, {}) for c in summaries],
    }
    (out / "benchmark.json").write_text(json.dumps(benchmark, indent=2, default=str) + "\n", encoding="utf-8")


def _cell_text(cell: dict | None) -> str:
    if cell is None:
        return "**not_measured**"
    if cell["classification"] != "measured":
        return f"**{cell['classification']}**"
    if cell["smoke"] or cell["std_s"] is None:
        return f"{cell['mean_s']:.2f} s (n={len(cell['seconds'])})"
    return f"{cell['mean_s']:.2f} s ± {cell['std_s']:.2f}"


def _peak(cell: dict) -> str:
    peaks = [r.get("peak_memory_gib", {}) for r in cell["by_seed"].values()]
    values = [max(p.values()) for p in peaks if p]
    return f"{max(values):.1f}" if values else "-"


def _swap(cell: dict) -> str:
    values = [r.get("memory", {}).get("max_swap_used_mb") for r in cell["by_seed"].values()]
    values = [v for v in values if v is not None]
    return f"{max(values) / 1024:.1f}" if values else "-"


def write_readme(out: Path, summaries: list[dict]) -> None:
    receipt = read_json(out / "environment_receipt.json")
    hw = receipt.get("hardware", {})
    sw = receipt.get("software", {})
    lanes = [lane for lane in _lane_columns(summaries) if any((c["decoder"], c["canvas"]) == lane for c in summaries)]
    by_key = {(c["composition"], c["target_seconds"], c["decoder"], c["canvas"]): c for c in summaries}
    rows = sorted({(c["composition"], c["target_seconds"]) for c in summaries},
                  key=lambda k: _sort_key({"composition": k[0], "target_seconds": k[1], "canvas": "", "decoder": ""}))
    measured = [c for c in summaries if c["classification"] == "measured"]
    title = "# FastH3 OmniRef on Apple Silicon (MLX): benchmark results"
    lines = [
        title,
        "",
        f"Host `{receipt.get('hostname', '?')}` ({hw.get('chip', '?')}, {hw.get('memory_gib', '?')} GiB unified, "
        f"macOS {hw.get('macos', '?')}); one GPU; concurrency 1. FastVideo OmniRef MLX pipeline "
        f"(DiT `{receipt.get('run_settings', {}).get('dit', '?')}`, text encoder NVFP4, NVFP4 VAE decoders run as "
        "W4A16). Column names in `latency.csv` keep the Spark `nvfp4_*` spelling for comparability.",
        "",
        "## Results (mean ± sample SD over seeds 0/1/2)",
        "",
        "| Composition | " + " | ".join(f"{canvas} {decoder.upper()}" for decoder, canvas in lanes) + " |",
        "| --- | " + " | ".join("---" for _ in lanes) + " |",
    ]
    for composition, seconds in rows:
        cells = [_cell_text(by_key.get((composition, seconds, d, c))) for d, c in lanes]
        label = composition if seconds == 5 else f"{composition} ({seconds} s)"
        lines.append(f"| {label} | " + " | ".join(cells) + " |")
    lines += [
        "",
        f"Measured cells: {len(measured)}/{len(summaries)}. Each full cell has an excluded seed-42 warmup, then "
        "seeds 0, 1, 2 end-to-end to MP4 (model load excluded). A smoke run reports a single seed-42 request (n=1).",
        "",
        "## Peak memory",
        "",
        "Largest MLX peak across seeds (GiB, from the pipeline's `peak_memory_gib`) and largest swap used while "
        "generating (GiB, `sysctl vm.swapusage`).",
        "",
        "| Cell | MLX peak GiB | Max swap GiB | Status |",
        "| --- | --- | --- | --- |",
    ]
    for cell in summaries:
        lines.append(f"| {cell['cell_id']} | {_peak(cell)} | {_swap(cell)} | {cell['classification']} |")
    failed = [c for c in summaries if c["classification"] in ("failed", "interrupted") and c["failures"]]
    if failed:
        lines += ["", "## Failures", ""]
        for cell in failed:
            failure = cell["failures"][-1]
            message = (failure.get("error_message") or failure.get("error_type") or "no detail").strip()
            lines.append(f"- `{cell['cell_id']}`: {failure.get('error_type')}: {message[:300]}")
    warned = [(c["cell_id"], r["seed"], r["validation_warnings"]) for c in summaries for r in c["by_seed"].values()
              if r.get("validation_warnings")]
    if warned:
        lines += ["", "## Output validation warnings", ""]
        lines += [f"- `{cid}` seed {seed}: {'; '.join(w)}" for cid, seed, w in warned]
    lines += [
        "",
        "## Method and identity",
        "",
        f"- FastVideo `{sw.get('fastvideo', {}).get('git_sha', '?')}` (dirty={sw.get('fastvideo', {}).get('dirty')}), "
        f"MLX {sw.get('mlx', '?')}, Python {sw.get('python', '?')}.",
        f"- Data bundle sha256sums: `{receipt.get('data', {}).get('sha256sums_sha256', '?')}` "
        f"({receipt.get('data', {}).get('verification', '?')}).",
        f"- Settings: `{json.dumps(receipt.get('run_settings', {}), sort_keys=True)}`",
        "",
        "## Files",
        "",
        "- `latency.csv`: one row per composition, Spark column schema.",
        "- `benchmark.json`: profiles and performance. `per-request-records.jsonl`: every warmup/measured request "
        "(pipeline timings, peaks, memory-sampler summary, ffprobe, hashes).",
        "- `environment_receipt.json`: host, software and weight identity. `media/`: MP4s plus a `.mem.jsonl` "
        "memory trace per request. `logs/`: one log per cell. `status/`: per-cell state.",
    ]
    (out / "README.md").write_text("\n".join(lines) + "\n", encoding="utf-8")


def aggregate(out: Path, planned: list[dict] | None = None) -> list[dict]:
    summaries = build_summaries(out, planned)
    write_latency_csv(out, summaries)
    write_benchmark_json(out, summaries)
    write_readme(out, summaries)
    return summaries


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--out", type=Path, required=True, help="results directory written by run_matrix.py")
    args = parser.parse_args()
    summaries = aggregate(args.out)
    done = sum(1 for c in summaries if c["classification"] == "measured")
    print(f"aggregated {len(summaries)} cells ({done} measured) into {args.out}")


if __name__ == "__main__":
    main()