#!/usr/bin/env python3
"""Audit and analyze the frozen E07 live-agent experiment."""

from __future__ import annotations

import argparse
from collections import Counter
import csv
from hashlib import sha256
import itertools
import json
import math
from pathlib import Path
import random
import statistics
from typing import Any, Sequence

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt


HARNESSES = ("H000", "H003", "H007", "H008", "H011", "H016", "H018")
CONTRASTS = (
    ("C1_full_vs_exact", "H007", "H000", "primary"),
    ("C2_no_search_vs_exact", "H016", "H000", "primary"),
    ("C3_oracle_file_vs_full", "H018", "H007", "primary"),
    ("C4_dense_vs_exact", "H003", "H000", "secondary"),
    ("C5_graph_vs_full", "H008", "H007", "secondary"),
    ("C6_specialized_vs_graph_unified", "H011", "H008", "secondary"),
)
TASKS = (
    "TASK_CR_001", "TASK_CR_002", "TASK_CR_003", "TASK_CR_005", "TASK_CR_006",
    "TASK_CR_007", "TASK_CR_008", "TASK_CR_009", "TASK_CR_012", "TASK_CR_013",
)
BOOTSTRAPS = 20_000
SEED = 20260718


class AnalysisError(RuntimeError):
    pass


def percentile(values: Sequence[float], probability: float) -> float:
    ordered = sorted(values)
    position = (len(ordered) - 1) * probability
    low, high = math.floor(position), math.ceil(position)
    if low == high:
        return ordered[low]
    return ordered[low] * (high - position) + ordered[high] * (position - low)


def bootstrap_ci(values: Sequence[float], rng: random.Random) -> tuple[float, float]:
    samples = [statistics.fmean(rng.choice(values) for _ in values) for _ in range(BOOTSTRAPS)]
    return percentile(samples, 0.025), percentile(samples, 0.975)


def exact_mcnemar(left: Sequence[int], right: Sequence[int]) -> tuple[int, int, float]:
    n10 = sum(a == 1 and b == 0 for a, b in zip(left, right))
    n01 = sum(a == 0 and b == 1 for a, b in zip(left, right))
    discordant = n10 + n01
    if discordant == 0:
        return n10, n01, 1.0
    tail = sum(math.comb(discordant, k) for k in range(min(n10, n01) + 1)) / 2**discordant
    return n10, n01, min(1.0, 2 * tail)


def exact_sign_flip(differences: Sequence[float]) -> float:
    values = [float(value) for value in differences if value != 0]
    if not values:
        return 1.0
    observed = abs(statistics.fmean(values))
    extreme = sum(
        abs(statistics.fmean(sign * value for sign, value in zip(signs, values))) >= observed - 1e-12
        for signs in itertools.product((-1.0, 1.0), repeat=len(values))
    )
    return extreme / 2 ** len(values)


def holm(p_values: Sequence[float]) -> list[float]:
    result = [1.0] * len(p_values)
    running = 0.0
    for rank, index in enumerate(sorted(range(len(p_values)), key=p_values.__getitem__)):
        running = max(running, min(1.0, (len(p_values) - rank) * p_values[index]))
        result[index] = running
    return result


def write_csv(path: Path, rows: Sequence[dict[str, Any]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)


def discover(root: Path) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    paths = sorted((root / "results/raw/E07").glob("H*/TASK_*/*/final_metrics.json"))
    rows: list[dict[str, Any]] = []
    revisions: set[str] = set()
    run_ids: set[str] = set()
    response_models: set[str] = set()
    fingerprints: set[str] = set()
    unparsed_intent_cells: set[str] = set()
    unparsed_intent_responses = 0
    empty_no_tool_responses = 0
    response_count = 0
    for path in paths:
        directory = path.parent
        required = ("run_manifest.json", "messages.json", "model.patch", "validation.json", "trajectory.jsonl")
        if any(not (directory / name).exists() for name in required):
            raise AnalysisError(f"missing immutable artifact in {directory}")
        row = json.loads(path.read_text(encoding="utf-8"))
        manifest = json.loads((directory / "run_manifest.json").read_text(encoding="utf-8"))
        revisions.add(manifest["identity"]["code_revision"])
        if row["run_id"] in run_ids:
            raise AnalysisError(f"duplicate run ID {row['run_id']}")
        run_ids.add(row["run_id"])
        if any(len(item.get("after_instances", [])) != 1 for item in row["residency_transitions"]):
            raise AnalysisError(f"non-exclusive residency in {directory}")
        responses = sorted(directory.glob("model_response_*.json"))
        if len(responses) != int(row["model_calls"]):
            raise AnalysisError(f"model-call artifact mismatch in {directory}")
        for response_path in responses:
            response = json.loads(response_path.read_text(encoding="utf-8"))
            response_count += 1
            response_models.add(str(response.get("model")))
            if response.get("system_fingerprint"):
                fingerprints.add(str(response["system_fingerprint"]))
            message = response.get("choices", [{}])[0].get("message", {})
            calls = message.get("tool_calls") or []
            reasoning = str(message.get("reasoning_content") or "")
            content = str(message.get("content") or "")
            if not calls and ("<tool_call>" in reasoning or "<function=" in reasoning):
                unparsed_intent_responses += 1
                unparsed_intent_cells.add(row["run_id"])
            if not calls and not content:
                empty_no_tool_responses += 1
        rows.append(row)
    expected = {(task, harness) for task in TASKS for harness in HARNESSES}
    observed = {(row["task_id"], row["harness_id"]) for row in rows}
    if len(rows) != 70 or observed != expected:
        raise AnalysisError(f"E07 grid mismatch: {len(rows)} rows, missing={sorted(expected-observed)}")
    if len(revisions) != 1:
        raise AnalysisError(f"mixed E07 code revisions: {sorted(revisions)}")
    if response_models != {"qwen/qwen3.6-35b-a3b"}:
        raise AnalysisError(f"unexpected response models: {sorted(response_models)}")
    audit = {
        "cells": len(rows), "unique_run_ids": len(run_ids), "code_revision": next(iter(revisions)),
        "model_responses": response_count, "response_models": sorted(response_models),
        "system_fingerprints": sorted(fingerprints),
        "all_cells_have_tool_calls": all(int(row["tool_calls"]) > 0 for row in rows),
        "total_prompt_tokens": sum(int(row["usage"]["prompt_tokens"]) for row in rows),
        "total_completion_tokens": sum(int(row["usage"]["completion_tokens"]) for row in rows),
        "total_tokens": sum(int(row["usage"]["total_tokens"]) for row in rows),
        "exploratory_unparsed_tool_intent_responses": unparsed_intent_responses,
        "exploratory_unparsed_tool_intent_cells": len(unparsed_intent_cells),
        "empty_no_tool_responses": empty_no_tool_responses,
    }
    return rows, audit


def cell_table(raw: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    for row in sorted(raw, key=lambda item: (item["task_id"], item["harness_id"])):
        rows.append({
            "task_id": row["task_id"], "harness_id": row["harness_id"],
            "resolved_at_1": int(bool(row["resolved_at_1"])),
            "patch_applied": int(bool(row["patch_applied"])),
            "fail_to_pass": int(bool(row["fail_to_pass"])),
            "pass_to_pass": int(bool(row["pass_to_pass"])),
            "all_gold_modified": int(bool(row["localization_metrics"]["all_gold_in_top_10"])),
            "search_all_gold": int(bool(row["search_localization_metrics"]["all_gold_in_top_10"])),
            "read_all_gold": int(bool(row["read_localization_metrics"]["all_gold_in_top_10"])),
            "model_calls": int(row["model_calls"]), "tool_calls": int(row["tool_calls"]),
            "test_runs": int(row["test_runs"]), "prompt_tokens": int(row["usage"]["prompt_tokens"]),
            "completion_tokens": int(row["usage"]["completion_tokens"]),
            "total_tokens": int(row["usage"]["total_tokens"]),
            "elapsed_seconds": float(row["elapsed_seconds"]),
            "model_elapsed_seconds": float(row["model_elapsed_seconds"]),
            "model_switch_count": int(row["model_switch_count"]),
            "model_switch_seconds": float(row["model_switch_seconds"]),
            "finished_reason": row["finished_reason"], "failure_stage": row["failure_stage"],
            "protocol_violation_count": len(row["protocol_violations"]),
        })
    return rows


def summaries(cells: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    rng = random.Random(SEED)
    rows: list[dict[str, Any]] = []
    for harness in HARNESSES:
        group = [row for row in cells if row["harness_id"] == harness]
        resolved = [float(row["resolved_at_1"]) for row in group]
        low, high = bootstrap_ci(resolved, rng)
        rows.append({
            "harness_id": harness, "tasks": len(group), "resolved_count": int(sum(resolved)),
            "resolved_rate": statistics.fmean(resolved), "resolved_ci_low": low,
            "resolved_ci_high": high,
            "patch_apply_rate": statistics.fmean(float(row["patch_applied"]) for row in group),
            "all_gold_modified_rate": statistics.fmean(float(row["all_gold_modified"]) for row in group),
            "mean_tool_calls": statistics.fmean(float(row["tool_calls"]) for row in group),
            "mean_model_calls": statistics.fmean(float(row["model_calls"]) for row in group),
            "mean_total_tokens": statistics.fmean(float(row["total_tokens"]) for row in group),
            "mean_elapsed_seconds": statistics.fmean(float(row["elapsed_seconds"]) for row in group),
            "mean_switch_seconds": statistics.fmean(float(row["model_switch_seconds"]) for row in group),
        })
    return rows


def contrasts(cells: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
    lookup = {(row["task_id"], row["harness_id"]): row for row in cells}
    rng = random.Random(SEED + 1)
    rows: list[dict[str, Any]] = []
    for name, left_id, right_id, family in CONTRASTS:
        left = [lookup[(task, left_id)] for task in TASKS]
        right = [lookup[(task, right_id)] for task in TASKS]
        a = [int(row["resolved_at_1"]) for row in left]
        b = [int(row["resolved_at_1"]) for row in right]
        binary_diff = [float(x-y) for x, y in zip(a, b)]
        low, high = bootstrap_ci(binary_diff, rng)
        n10, n01, p_value = exact_mcnemar(a, b)
        token_diff = [float(x["total_tokens"]-y["total_tokens"]) for x, y in zip(left, right)]
        time_diff = [float(x["elapsed_seconds"]-y["elapsed_seconds"]) for x, y in zip(left, right)]
        rows.append({
            "contrast": name, "family": family, "left_harness": left_id, "right_harness": right_id,
            "left_resolved": sum(a), "right_resolved": sum(b),
            "paired_risk_difference": statistics.fmean(binary_diff),
            "risk_difference_ci_low": low, "risk_difference_ci_high": high,
            "discordant_left_only": n10, "discordant_right_only": n01,
            "mcnemar_p": p_value, "mcnemar_p_holm": math.nan,
            "mean_token_difference": statistics.fmean(token_diff),
            "token_sign_flip_p": exact_sign_flip(token_diff),
            "mean_elapsed_difference": statistics.fmean(time_diff),
            "elapsed_sign_flip_p": exact_sign_flip(time_diff),
        })
    for row, adjusted in zip(rows, holm([float(row["mcnemar_p"]) for row in rows])):
        row["mcnemar_p_holm"] = adjusted
    return rows


def failure_analysis(raw: Sequence[dict[str, Any]]) -> dict[str, Any]:
    finish = Counter(str(row["finished_reason"]) for row in raw)
    failure = Counter(str(row["failure_stage"]) for row in raw)
    tool_counts: dict[str, dict[str, int]] = {}
    for harness in HARNESSES:
        combined: Counter[str] = Counter()
        for row in raw:
            if row["harness_id"] == harness:
                combined.update({str(key): int(value) for key, value in row["tool_counts"].items()})
        tool_counts[harness] = dict(sorted(combined.items()))
    read_complete = [row for row in raw if row["read_localization_metrics"]["all_gold_in_top_10"]]
    search_complete = [row for row in raw if row["search_localization_metrics"]["all_gold_in_top_10"]]
    return {
        "status": "exploratory_descriptive",
        "finish_reasons": dict(sorted(finish.items())),
        "failure_stages": dict(sorted(failure.items())),
        "cells_with_accepted_patch": sum(bool(row["patch_applied"]) for row in raw),
        "cells_with_any_edit": sum(bool(row["modified_files"]) for row in raw),
        "cells_with_protocol_violation": sum(bool(row["protocol_violations"]) for row in raw),
        "protocol_violation_events": sum(len(row["protocol_violations"]) for row in raw),
        "search_complete_cells": len(search_complete),
        "resolutions_given_search_complete": sum(bool(row["resolved_at_1"]) for row in search_complete),
        "read_complete_cells": len(read_complete),
        "resolutions_given_read_complete": sum(bool(row["resolved_at_1"]) for row in read_complete),
        "model_calls": sum(int(row["model_calls"]) for row in raw),
        "tool_calls": sum(int(row["tool_calls"]) for row in raw),
        "public_test_tool_calls": sum(int(row["test_runs"]) for row in raw),
        "non_reused_residency_transitions": sum(int(row["model_switch_count"]) for row in raw),
        "model_switch_seconds": sum(float(row["model_switch_seconds"]) for row in raw),
        "cell_elapsed_seconds": sum(float(row["elapsed_seconds"]) for row in raw),
        "model_elapsed_seconds": sum(float(row["model_elapsed_seconds"]) for row in raw),
        "tool_counts_by_harness": tool_counts,
    }


def plots(directory: Path, rows: Sequence[dict[str, Any]]) -> list[Path]:
    labels = [row["harness_id"] for row in rows]
    rates = [float(row["resolved_rate"]) for row in rows]
    errors = [[rate-float(row["resolved_ci_low"]) for rate, row in zip(rates, rows)],
              [float(row["resolved_ci_high"])-rate for rate, row in zip(rates, rows)]]
    colors = ["#355070", "#6d597a", "#2a9d8f", "#e9c46a", "#f4a261", "#b56576", "#457b9d"]
    fig, axis = plt.subplots(figsize=(7.2, 3.8))
    axis.bar(labels, rates, color=colors, edgecolor="black", linewidth=.5)
    axis.errorbar(labels, rates, yerr=errors, fmt="none", ecolor="black", capsize=3)
    axis.set(ylim=(0, 1.05), ylabel="Hidden-test resolved@1", xlabel="Live-agent harness")
    axis.grid(axis="y", alpha=.25); fig.tight_layout()
    path = directory / "e07_resolved_rate.pdf"; fig.savefig(path); fig.savefig(directory / "e07_resolved_rate.png", dpi=220); plt.close(fig)
    fig, axis = plt.subplots(figsize=(6.4, 4.2))
    for row, color in zip(rows, colors):
        axis.scatter(row["mean_total_tokens"], row["resolved_rate"], label=row["harness_id"], color=color, s=55, edgecolor="black", linewidth=.5)
    axis.set(xlabel="Mean LM Studio tokens per task", ylabel="Hidden-test resolved@1", ylim=(-.03, 1.03))
    axis.grid(alpha=.25); axis.legend(ncol=2, frameon=False); fig.tight_layout()
    path2 = directory / "e07_quality_efficiency.pdf"; fig.savefig(path2); fig.savefig(directory / "e07_quality_efficiency.png", dpi=220); plt.close(fig)
    return [path, path2]


def analyze(root: Path) -> dict[str, Any]:
    raw, audit = discover(root); cells = cell_table(raw); summary = summaries(cells); contrast = contrasts(cells)
    failures = failure_analysis(raw)
    output = root / "results/derived/e07"; output.mkdir(parents=True, exist_ok=True)
    tables = [output/"e07_cell_metrics.csv", output/"e07_harness_summary.csv", output/"e07_contrasts.csv"]
    for path, rows in zip(tables, (cells, summary, contrast)): write_csv(path, rows)
    figures = plots(output, summary)
    report = {"schema_version": 1, "experiment_id": "E07",
              "analysis_plan": {"task_is_statistical_unit": True, "bootstrap_samples": BOOTSTRAPS,
              "bootstrap_seed": SEED, "binary_test": "exact two-sided McNemar",
              "multiplicity": "Holm across six declared contrasts", "continuous_test": "exact paired sign-flip"},
              "audit": audit, "harness_summaries": summary, "contrasts": contrast,
              "failure_analysis": failures}
    analysis_path = output/"e07_analysis.json"; analysis_path.write_text(json.dumps(report, indent=2, sort_keys=True)+"\n")
    files = [*tables, analysis_path, *figures]
    checksums = {path.name: sha256(path.read_bytes()).hexdigest() for path in files}
    (output/"SHA256SUMS.json").write_text(json.dumps(checksums, indent=2, sort_keys=True)+"\n")
    return {**report, "output_directory": str(output), "checksums": checksums}


def main() -> int:
    parser = argparse.ArgumentParser(); parser.add_argument("--root", type=Path, default=Path(__file__).resolve().parents[1]); args = parser.parse_args()
    try: result = analyze(args.root.resolve())
    except (AnalysisError, OSError, ValueError, KeyError) as exc: print(f"E07 ANALYSIS FAILED: {exc}"); return 1
    print(json.dumps(result, indent=2, sort_keys=True)); return 0


if __name__ == "__main__": raise SystemExit(main())