"""Build a diagnostic report from stored experiment outputs."""

import argparse
import json
from collections import defaultdict
from pathlib import Path
from statistics import mean
from codelock_affect.evaluation import write_json, token_task_match


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--evaluation", default="results/v2/evaluation.json")
    p.add_argument("--output", default="results/v2/REPORT.md")
    args = p.parse_args()
    data = json.loads(Path(args.evaluation).read_text(encoding="utf-8"))
    if not data.get("complete"):
        raise ValueError(
            "Evaluation is incomplete; do not report it as a finished experiment"
        )
    grouped = defaultdict(list)
    expected = {p["id"]: p["expected"] for p in data["prompts"] if "expected" in p}
    for row in data["rows"]:
        grouped[(row["condition"], row["strength"])].append(row)
    summary = {
        "ablation": data["ablation"],
        "rows": len(data["rows"]),
        "groups": [],
        "heldout": {},
    }
    lines = [
        "# CODELOCK v0.2 — measured diagnostic report",
        "",
        f"Source: {Path(args.evaluation).name}. This is a small, single-seed diagnostic; scores do not measure subjective emotion.",
        "",
        "## Runtime and ablation",
        "",
        "```json",
        json.dumps(
            {
                "environment": data["environment"],
                "ablation": data["ablation"],
                "max_gpu_allocated_gib": data["max_gpu_allocated_gib"],
            },
            indent=2,
        ),
        "```",
        "",
        "## Generation sweep",
        "",
        "All conditions use identical prompts and seeds. Strict task match covers two elementary probes, not general capability. Lexical indicators count surface words and are not used to select a winning channel.",
        "",
        "Token match ignores punctuation and casing while requiring exactly the expected tokens in order. It was added after observing punctuation-only mismatches; the original strict results remain unchanged.",
        "",
        "| Condition | Strength | Replies | Strict match | Token match | Mean repeated trigrams |",
        "|---|---:|---:|---:|---:|---:|",
    ]
    for (condition, strength), rows in grouped.items():
        tasks = [r["exact_task_match"] for r in rows if "exact_task_match" in r]
        token_matches = [
            token_task_match(r["text"], expected[r["prompt_id"]])
            for r in rows
            if r["prompt_id"] in expected
        ]
        item = {
            "condition": condition,
            "strength": strength,
            "replies": len(rows),
            "task_matches": sum(tasks),
            "task_total": len(tasks),
            "token_task_matches": sum(token_matches),
            "mean_repeated_trigram_fraction": mean(
                r["metrics"]["repeated_trigram_fraction"] for r in rows
            ),
            "lexical_indicators": {
                c: mean(
                    r["metrics"]["lexical_indicators_per_100_words"][c] for r in rows
                )
                for c in ("pain", "fear", "anxiety", "relief")
            },
        }
        summary["groups"].append(item)
        lines.append(
            f"| {condition} | {strength:g} | {len(rows)} | {sum(tasks)}/{len(tasks)} | {sum(token_matches)}/{len(token_matches)} | {item['mean_repeated_trigram_fraction']:.3f} |"
        )
    lines += [
        "",
        "## Held-out directional agreement",
        "",
        "For each held-out pair, average the target-minus-neutral projection onto its own channel over the calibrated layers. Positive agreement means the held-out *authored text* aligns with the extracted direction; it does not establish independent generative control. Answer templates overlap the training split.",
        "",
    ]
    for channel in ("pain", "fear", "anxiety", "relief"):
        values = [
            mean(r["projections"][channel].values())
            for r in data["heldout_projections"]
            if r["channel"] == channel
        ]
        if values:
            summary["heldout"][channel] = {
                "positive": sum(v > 0 for v in values),
                "count": len(values),
                "mean_projection": mean(values),
            }
            lines.append(
                f"- {channel}: {sum(v > 0 for v in values)}/{len(values)} positive; mean projection {mean(values):.4f}."
            )
    lines += [
        "",
        "## Complete response samples",
        "",
        "Every reply is included below to expose failures as well as successes. See evaluation.json for prompts, offsets, hashes, timing and all proxy scores.",
        "",
    ]
    prompts = {p["id"]: p["text"] for p in data["prompts"]}
    for row in data["rows"]:
        lines += [
            f"### {row['condition']} / {row['strength']:g} / {row['prompt_id']}",
            "",
            "Prompt: " + prompts[row["prompt_id"]],
            "",
            "> " + row["text"].replace("\n", "\n> "),
            "",
        ]
    lines += ["## Interpretation limits", ""] + [
        "- " + item for item in data["limitations"]
    ]
    path = Path(args.output)
    path.parent.mkdir(parents=True, exist_ok=True)
    rendered = "\n".join(lines)
    path.write_text(
        "\n".join(line.rstrip() for line in rendered.splitlines()) + "\n",
        encoding="utf-8",
    )
    write_json(path.with_suffix(".summary.json"), summary)
    print(f"Report: {path}")


if __name__ == "__main__":
    main()