"""Matched image and performance trials. Run from the repository root."""

from __future__ import annotations

import argparse
import gc
import json
import logging
import statistics
from pathlib import Path

from h3_mlx.runtime import Pipeline

PROFILES = {
    "native": {},
    "fused": {"mlp_chunk": 0, "kernel": "metal-qk", "threads": 256, "vae_batch": 8},
    "chunked": {"mlp_chunk": 1024},
    "wide-mlp": {"mlp_chunk": 0},
    "tuned": {"mlp_chunk": 0, "kernel": "metal", "threads": 256, "vae_batch": 8},
    "q8": {"bits": 8, "mlp_chunk": 0, "vae_batch": 8},
    "q4": {"bits": 4, "mlp_chunk": 0, "vae_batch": 8},
    "q8-g32": {"bits": 8, "group_size": 32, "mlp_chunk": 0, "vae_batch": 8},
    "preview": {"mlp_chunk": 0, "kernel": "metal", "threads": 256, "vae_batch": 8},
    "stream": {"dit_cache_gb": 0},
}


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument("--text-encoder", required=True)
    p.add_argument(
        "--profiles",
        nargs="+",
        choices=PROFILES,
        default=["native", "tuned", "q8", "q4", "preview"],
    )
    p.add_argument("--cases", nargs="+", default=["product", "portrait"])
    p.add_argument("--case-file", default="experiments/cases.json")
    p.add_argument("--output", default="outputs/suite")
    p.add_argument(
        "--repeat",
        type=int,
        default=2,
        help="First trial includes model loading; later trials reuse resident models",
    )
    p.add_argument("--save-trace", action="store_true")
    a = p.parse_args()
    if a.repeat < 1:
        p.error("repeat must be positive")
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S"
    )
    cases = [c for c in json.loads(Path(a.case_file).read_text()) if c["id"] in a.cases]
    if {c["id"] for c in cases} != set(a.cases):
        p.error("Requested case not found in case file")
    root = Path(a.output)
    root.mkdir(parents=True, exist_ok=True)
    ledger = root / "trials.jsonl"
    for name in a.profiles:
        logging.info("PROFILE %s", name)
        pipeline = Pipeline(a.model, a.text_encoder, **PROFILES[name])
        try:
            for c in cases:
                for trial in range(a.repeat):
                    output = root / c["id"] / f"{name}-{trial}.png"
                    _, row = pipeline.generate(
                        c["prompt"],
                        width=c["width"],
                        height=c["height"],
                        seed=c["seed"],
                        output=output,
                        steps=[0, 2, 4] if name == "preview" else None,
                        save_trace=a.save_trace,
                        images=c.get("images", ()),
                    )
                    row.update(
                        profile=name,
                        case=c["id"],
                        trial=trial,
                        sample_type="resident"
                        if row["dit_was_resident"]
                        else "load_included",
                    )
                    with ledger.open("a") as h:
                        h.write(json.dumps(row) + "\n")
        finally:
            pipeline.clear()
            del pipeline
            gc.collect()
    rows = [json.loads(line) for line in ledger.read_text().splitlines()]
    summary = []
    for profile in a.profiles:
        for case in a.cases:
            group = [r for r in rows if r["profile"] == profile and r["case"] == case]
            resident = [
                r["timings"]["total_s"] for r in group if r["sample_type"] == "resident"
            ]
            summary.append(
                {
                    "profile": profile,
                    "case": case,
                    "samples": len(group),
                    "resident_samples": len(resident),
                    "resident_median_s": statistics.median(resident)
                    if resident
                    else None,
                    "load_included_s": [
                        r["timings"]["total_s"]
                        for r in group
                        if r["sample_type"] == "load_included"
                    ],
                }
            )
    (root / "summary.json").write_text(json.dumps(summary, indent=2) + "\n")


if __name__ == "__main__":
    main()