#!/usr/bin/env python3
"""Regenerate statistical and performance summaries from retained raw result files."""

from __future__ import annotations

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

MODELS = ["qwen38", "qwen36", "gemma4"]
MODEL_NAMES = {"qwen38": "Qwen3.8-27B", "qwen36": "Qwen3.6-27B", "gemma4": "Gemma 4 31B-it"}
GPU_TOTAL_MIB = 24564


def load(path: Path):
    return json.loads(path.read_text())


def maybe(path: Path):
    return load(path) if path.exists() else None


def wilson(successes: int, total: int, z: float = 1.959963984540054) -> dict:
    if total <= 0:
        return {"successes": successes, "total": total, "rate": None, "low": None, "high": None}
    p = successes / total
    den = 1 + z * z / total
    center = (p + z * z / (2 * total)) / den
    margin = z * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total)) / den
    return {"successes": successes, "total": total, "rate": p, "low": max(0, center - margin), "high": min(1, center + margin)}


def mcnemar_exact(left_only: int, right_only: int) -> float:
    n = left_only + right_only
    if n == 0:
        return 1.0
    k = min(left_only, right_only)
    tail = sum(math.comb(n, i) for i in range(k + 1)) / (2 ** n)
    return min(1.0, 2 * tail)


def rate_from(summary: dict, success_key: str, total_key: str) -> dict:
    return wilson(int(summary[success_key]), int(summary[total_key]))


def gpu_telemetry(path: Path) -> dict:
    power_values = []
    memory_values = []
    if path.exists():
        for line in path.read_text(errors="replace").splitlines():
            fields = [field.strip() for field in line.split(",")]
            if len(fields) < 3:
                continue
            try:
                memory_values.append(float(fields[1]))
                power_values.append(float(fields[2]))
            except ValueError:
                continue
    return {
        "samples": len(power_values),
        "median_power_watts": statistics.median(power_values) if power_values else None,
        "peak_power_watts": max(power_values) if power_values else None,
        "peak_memory_mib": max(memory_values) if memory_values else None,
    }


def quality_rates(root: Path, model: str) -> dict:
    base = root / model
    out = {}
    specs = {
        "instruction": ("instruction-summary.json", "passed", "total"),
        "writing_mechanical": ("writing-summary.json", "passed", "total"),
        "multilingual_mechanical": ("multilingual-summary.json", "passed", "total"),
        "document_qa": ("documents-summary.json", "passed", "total"),
        "vision": ("vision-summary.json", "passed", "total"),
        "reasoning_runs": ("reasoning-summary.json", "passed_runs", "total_runs"),
        "reasoning_pass_at_3": ("reasoning-summary.json", "problems_any_pass", "problems"),
        "tools_single_runs": ("tools-single-summary.json", "passed_runs", "total_runs"),
        "tools_single_pass_at_3": ("tools-single-summary.json", "cases_any_pass", "cases"),
        "tools_multi_runs": ("tools-multi-summary.json", "passed_runs", "total_runs"),
        "tools_multi_pass_at_3": ("tools-multi-summary.json", "cases_any_pass", "cases"),
        "coding_pass_at_1": ("coding-summary.json", "pass_at_1", "tasks"),
        "coding_pass_at_3": ("coding-summary.json", "pass_at_3", "tasks"),
        "coding_runs": ("coding-summary.json", "passed_runs", "total_runs"),
    }
    for label, (filename, success, total) in specs.items():
        data = maybe(base / filename)
        if data is not None:
            out[label] = rate_from(data, success, total)
    reasoning_files = sorted((base / "reasoning").glob("*/*.json"))
    if reasoning_files:
        rows = [load(path) for path in reasoning_files]
        out["reasoning_answer_correct"] = wilson(sum(bool(row.get("answer_correct")) for row in rows), len(rows))
        out["reasoning_format_compliant"] = wilson(sum(bool(row.get("format_compliant")) for row in rows), len(rows))
    return out


def performance(root: Path, model: str) -> dict:
    base = root / model / "performance" / "f16-f16"
    bench = maybe(base / "llama-bench.stdout.json") or []
    conditions = {}
    for row in bench:
        if row.get("n_prompt"):
            key = f"prompt_{row['n_prompt']}"
        else:
            key = f"decode_{row.get('n_gen')}"
        samples = [float(value) for value in row.get("samples_ts") or []]
        conditions[key] = {
            "median_tokens_per_second": statistics.median(samples) if samples else row.get("avg_ts"),
            "mean_tokens_per_second": row.get("avg_ts"),
            "min_tokens_per_second": min(samples) if samples else None,
            "max_tokens_per_second": max(samples) if samples else None,
            "samples_tokens_per_second": samples,
            "repetitions": len(samples),
            "n_prompt": row.get("n_prompt"),
            "n_gen": row.get("n_gen"),
            "model_size_bytes_runtime": row.get("model_size"),
            "model_n_params_runtime": row.get("model_n_params"),
            "build_commit": row.get("build_commit"),
            "build_number": row.get("build_number"),
        }
    contexts = maybe(base / "context-summary.json") or {"rows": []}
    context_rows = []
    for row in contexts.get("rows", []):
        item = {
            "ctx_allocated": row.get("ctx_allocated"), "cache_k": row.get("cache_k"), "cache_v": row.get("cache_v"),
            "passed": bool(row.get("passed")), "input_tokens": row.get("input_tokens_tokenize"),
            "peak_gpu_mib": row.get("peak_gpu_mib"), "error": row.get("error"),
            "characters_per_second_first_turn": row.get("first_turn_characters_per_second"),
            "server_ready_seconds": row.get("ready_seconds"),
        }
        if row.get("peak_gpu_mib") is not None:
            item["headroom_mib"] = GPU_TOTAL_MIB - int(row["peak_gpu_mib"])
        if row.get("turns"):
            item["first_turn"] = row["turns"][0]
        context_rows.append(item)
    best = maybe(root / model / "performance" / "q8_0-q8_0" / "context-summary.json")
    best_rows = []
    if best:
        for row in best.get("rows", []):
            best_rows.append({
                "ctx_allocated": row.get("ctx_allocated"), "cache_k": row.get("cache_k"), "cache_v": row.get("cache_v"),
                "passed": bool(row.get("passed")), "input_tokens": row.get("input_tokens_tokenize"), "peak_gpu_mib": row.get("peak_gpu_mib"),
                "headroom_mib": GPU_TOTAL_MIB - int(row["peak_gpu_mib"]) if row.get("peak_gpu_mib") is not None else None,
                "characters_per_second_first_turn": row.get("first_turn_characters_per_second"), "turns": row.get("turns"), "error": row.get("error"),
                "server_ready_seconds": row.get("ready_seconds"),
            })
    stream = maybe(root / model / "stream-timing" / "summary.json")
    telemetry = gpu_telemetry(base / "llama-bench-gpu.csv")
    return {
        "synthetic": conditions,
        "equal_q4_f16_context": context_rows,
        "best_fit_context": best_rows,
        "stream_timing": stream,
        "synthetic_bench_telemetry": telemetry,
    }


def mtp(root: Path, model: str) -> dict:
    micro = maybe(root / model / "mtp" / "summary.json")
    practical = maybe(root / model / "mtp-practical-cpu-draft-sampling" / "summary.json")
    if practical is None:
        practical = maybe(root / model / "mtp-practical" / "summary.json")
    return {"microbenchmark": micro, "practical": practical}


def collect_outcomes(root: Path, model: str) -> dict[str, dict[str, bool]]:
    base = root / model
    patterns = {
        "instruction": "instruction/*.json",
        "writing": "writing/*.json",
        "multilingual": "multilingual/*.json",
        "documents": "documents/*.json",
        "vision": "vision/*.json",
        "reasoning": "reasoning/*/*.json",
        "tools_single": "tools-single/*/*.json",
        "tools_multi": "tools-multi/*/*.json",
        "coding": "coding/*/*.json",
    }
    out = {}
    for category, pattern in patterns.items():
        rows = {}
        for path in sorted(base.glob(pattern)):
            data = load(path)
            rows[str(path.relative_to(base))] = bool(data.get("passed"))
        if rows:
            out[category] = rows
    return out


def paired(root: Path, left: str, right: str) -> dict:
    lout = collect_outcomes(root, left)
    rout = collect_outcomes(root, right)
    result = {}
    for category in sorted(set(lout) & set(rout)):
        keys = sorted(set(lout[category]) & set(rout[category]))
        left_only = sum(lout[category][key] and not rout[category][key] for key in keys)
        right_only = sum(rout[category][key] and not lout[category][key] for key in keys)
        both_pass = sum(lout[category][key] and rout[category][key] for key in keys)
        both_fail = len(keys) - left_only - right_only - both_pass
        result[category] = {
            "paired_items": len(keys), "left_only_pass": left_only, "right_only_pass": right_only,
            "both_pass": both_pass, "both_fail": both_fail,
            "mcnemar_exact_two_sided_p": mcnemar_exact(left_only, right_only),
        }
    return result


def coding_efficiency(root: Path, model: str) -> dict:
    files = sorted((root / model / "coding").glob("*/*.json"))
    rows = []
    for path in files:
        data = load(path)
        transcript = data.get("transcript") or []
        tool_calls = 0
        for run in transcript:
            try:
                tool_calls += len(run["response"]["choices"][0]["message"].get("tool_calls") or [])
            except Exception:
                pass
        rows.append({
            "task": path.parent.name, "seed": path.stem, "passed": bool(data.get("passed")),
            "model_elapsed_seconds": sum(float(run.get("elapsed_seconds") or 0) for run in transcript),
            "assistant_turns": len(transcript), "tool_calls": tool_calls, "invalid_calls": int(data.get("invalid_calls") or 0),
        })
    successful = [row for row in rows if row["passed"]]
    return {
        "runs": rows, "successful_runs": len(successful), "total_runs": len(rows),
        "median_model_seconds_success": statistics.median([row["model_elapsed_seconds"] for row in successful]) if successful else None,
        "median_tool_calls_success": statistics.median([row["tool_calls"] for row in successful]) if successful else None,
        "invalid_calls_total": sum(row["invalid_calls"] for row in rows),
    }


def reasoning_efficiency(root: Path, model: str) -> dict:
    files = sorted((root / model / "reasoning").glob("*/*.json"))
    rows = []
    for path in files:
        data = load(path)
        raw = data.get("raw") or {}
        response = raw.get("response") or {}
        usage = response.get("usage") or {}
        completion_tokens = int(usage.get("completion_tokens") or 0)
        rows.append({
            "problem": path.parent.name,
            "seed": path.stem,
            "passed": bool(data.get("passed")),
            "visible_final": bool((data.get("output") or "").strip()),
            "completion_tokens": completion_tokens,
            "elapsed_seconds": float(raw.get("elapsed_seconds") or 0),
            "finish_reason": ((response.get("choices") or [{}])[0].get("finish_reason")),
        })
    passed = [row for row in rows if row["passed"]]
    failed = [row for row in rows if not row["passed"]]
    return {
        "runs": len(rows),
        "no_visible_final_runs": sum(not row["visible_final"] for row in rows),
        "length_finish_runs": sum(row["finish_reason"] == "length" for row in rows),
        "median_completion_tokens_all": statistics.median(row["completion_tokens"] for row in rows) if rows else None,
        "median_completion_tokens_pass": statistics.median(row["completion_tokens"] for row in passed) if passed else None,
        "median_completion_tokens_fail": statistics.median(row["completion_tokens"] for row in failed) if failed else None,
        "median_elapsed_seconds_all": statistics.median(row["elapsed_seconds"] for row in rows) if rows else None,
        "median_elapsed_seconds_pass": statistics.median(row["elapsed_seconds"] for row in passed) if passed else None,
        "median_elapsed_seconds_fail": statistics.median(row["elapsed_seconds"] for row in failed) if failed else None,
    }


def markdown(summary: dict) -> str:
    lines = ["# Kingy.ai 24GB comparison — regenerated result summary", "", f"Generated: {summary['generated_utc']}", ""]
    lines += ["## Quality and reliability", ""]
    headers = ["Measure"] + [MODEL_NAMES[m] for m in MODELS]
    lines.append("| " + " | ".join(headers) + " |")
    lines.append("|" + "|".join(["---"] + ["---:" for _ in MODELS]) + "|")
    measures = sorted(set().union(*(summary["models"][m]["quality"].keys() for m in MODELS)))
    for measure in measures:
        cells = []
        for model in MODELS:
            row = summary["models"][model]["quality"].get(measure)
            if not row:
                cells.append("—")
            else:
                cells.append(f"{row['successes']}/{row['total']} ({row['rate']*100:.1f}%; 95% CI {row['low']*100:.1f}–{row['high']*100:.1f})")
        lines.append("| " + measure.replace("_", " ") + " | " + " | ".join(cells) + " |")
    lines += ["", "## Synthetic speed", ""]
    lines.append("| Condition | " + " | ".join(MODEL_NAMES[m] for m in MODELS) + " |")
    lines.append("|---|" + "|".join("---:" for _ in MODELS) + "|")
    conditions = ["prompt_512", "prompt_8192", "prompt_32768", "decode_256"]
    for condition in conditions:
        cells = []
        for model in MODELS:
            row = summary["models"][model]["performance"]["synthetic"].get(condition)
            cells.append(f"{row['median_tokens_per_second']:.2f} tok/s" if row else "—")
        lines.append("| " + condition.replace("_", " ") + " | " + " | ".join(cells) + " |")
    lines += ["", "All percentages are accompanied by retained counts; binary intervals are Wilson 95% intervals. See summary.json for paired outcomes and exact McNemar p-values.", ""]
    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--results-root", type=Path, required=True)
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()
    summary = {
        "generated_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "gpu_total_mib": GPU_TOTAL_MIB,
        "models": {},
        "paired": {},
    }
    for model in MODELS:
        summary["models"][model] = {
            "display_name": MODEL_NAMES[model],
            "quality": quality_rates(args.results_root, model),
            "performance": performance(args.results_root, model),
            "mtp": mtp(args.results_root, model),
            "coding_efficiency": coding_efficiency(args.results_root, model),
            "reasoning_efficiency": reasoning_efficiency(args.results_root, model),
        }
    for left, right in [("qwen38", "qwen36"), ("qwen38", "gemma4"), ("qwen36", "gemma4")]:
        summary["paired"][f"{left}_vs_{right}"] = paired(args.results_root, left, right)
    args.out.mkdir(parents=True, exist_ok=True)
    (args.out / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2) + "\n")
    (args.out / "summary.md").write_text(markdown(summary))
    print(json.dumps({"out": str(args.out), "models": list(summary["models"]), "paired": list(summary["paired"])}))


if __name__ == "__main__":
    main()
