#!/usr/bin/env python3
"""Pair text-only and text+image L2 semantic-intervention runs into one ablation table.

Produces the 2 (modality) x 3 (description condition) matrix required by the
experiment checklist, plus the paired cross-modality decision-change rates.

Within-modality metrics (semantic controllability):
  no_desc_change, swap_desc_change, desc_follow, name_bias

Cross-modality metrics (does the image change the decision at all?):
  per description condition, decision-change rate between text_only and
  text_image on the same row, restricted to rows where both are on-menu.

Deliberately NOT reported: exact / macro-recall against the real next boss
skill. The same belief state admits several tactically reasonable skills, so
matching the single logged next skill would penalise valid-but-different
choices. Decision quality needs its own protocol, not this table.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import json
import math
import os
import sys
from typing import Any, Dict, List, Tuple

VARIANTS = ("full", "no_desc", "swap_desc")


def read_jsonl(path: str) -> List[Dict[str, Any]]:
    rows = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                rows.append(json.loads(line))
    return rows


def file_sha256(path: str) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def key(row: Dict[str, Any]) -> Tuple[str, int, int]:
    return (str(row.get("boss")), int(row.get("fight")), int(row.get("index")))


def wilson(k: int, n: int, z: float = 1.96) -> Tuple[float, float]:
    if n == 0:
        return (0.0, 0.0)
    p = k / n
    d = 1 + z * z / n
    c = p + z * z / (2 * n)
    h = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n))
    return ((c - h) / d, (c + h) / d)


def pct(k: int, n: int) -> Dict[str, Any]:
    if n == 0:
        return {"rate": None, "k": 0, "n": 0, "ci": [None, None]}
    lo, hi = wilson(k, n)
    return {"rate": round(k / n, 4), "k": k, "n": n,
            "ci": [max(0.0, round(lo, 4)), min(1.0, round(hi, 4))]}


def within_modality(rows: List[Dict[str, Any]]) -> Dict[str, Any]:
    # 单技能菜单上 swap 恒等、follow/name_bias 恒真 → 无定义,与语义干预主指标
    # 及 cross_modality 同口径排除（2026-07-15 补：此前 within_modality 漏过滤，
    # 曾使 Isaac text_only follow 被 47 个单菜单行灌水到 0.835，实为 0.784）。
    rows = [r for r in rows if len(r.get("menu") or []) >= 2]
    m = [r["metrics"] for r in rows]
    nd = [x for x in m if x["no_desc_pair_valid"]]
    sw = [x for x in m if x["swap_pair_valid"]]
    return {
        "no_desc_change": pct(sum(x["no_desc_decision_changed"] for x in nd), len(nd)),
        "swap_desc_change": pct(sum(x["swap_desc_decision_changed"] for x in sw), len(sw)),
        "desc_follow": pct(sum(x.get("swap_desc_follow", False) for x in sw), len(sw)),
        "name_bias": pct(sum(x.get("swap_name_bias", False) for x in sw), len(sw)),
    }


def cross_modality(text_rows: List[Dict[str, Any]],
                   img_rows: List[Dict[str, Any]]) -> Dict[str, Any]:
    # 单技能菜单上干预无定义,与语义干预主指标同口径排除
    text_rows = [r for r in text_rows if len(r.get("menu") or []) >= 2]
    img_rows = [r for r in img_rows if len(r.get("menu") or []) >= 2]
    tmap = {key(r): r for r in text_rows}
    out: Dict[str, Any] = {}
    for variant in VARIANTS:
        changed = 0
        n = 0
        for ir in img_rows:
            tr = tmap.get(key(ir))
            if tr is None:
                continue
            td = (tr["decisions"].get(variant) or {})
            idc = (ir["decisions"].get(variant) or {})
            if not (td.get("on_menu") and idc.get("on_menu")):
                continue
            n += 1
            changed += int(td.get("skill") != idc.get("skill"))
        out[variant] = pct(changed, n)
    return out


def format_rate(metric: Dict[str, Any]) -> str:
    return "" if metric.get("rate") is None else f"{100 * metric['rate']:.1f}%"


def format_ci(metric: Dict[str, Any]) -> str:
    lo, hi = metric.get("ci", [None, None])
    return "" if lo is None else f"[{100 * lo:.1f}, {100 * hi:.1f}]"


def summary_csv_rows(summary: Dict[str, Any]) -> List[Dict[str, Any]]:
    rows: List[Dict[str, Any]] = []
    for modality in ("text_only", "text_image"):
        within = summary["within_modality"][modality]
        rows.append({
            "belief_view": summary["belief_view"],
            "input_modality": modality,
            "n": summary[f"n_{modality}"],
            "no_desc_change": format_rate(within["no_desc_change"]),
            "swap_desc_change": format_rate(within["swap_desc_change"]),
            "desc_follow": format_rate(within["desc_follow"]),
            "desc_follow_ci95": format_ci(within["desc_follow"]),
            "name_bias": format_rate(within["name_bias"]),
            "name_bias_ci95": format_ci(within["name_bias"]),
        })
    return rows


def paired_change_csv_rows(summary: Dict[str, Any]) -> List[Dict[str, Any]]:
    cross = summary["cross_modality_decision_change"]
    return [{
        "belief_view": summary["belief_view"],
        "description_condition": variant,
        "n_paired": cross[variant]["n"],
        "decision_change_text_only_vs_text_image": format_rate(cross[variant]),
        "ci95": format_ci(cross[variant]),
    } for variant in VARIANTS]


def load_run(path: str) -> List[Dict[str, Any]]:
    if not os.path.exists(path):
        raise SystemExit(f"missing rows file: {path}")
    return read_jsonl(path)


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--belief_view", required=True)
    ap.add_argument("--text_rows", required=True)
    ap.add_argument("--image_rows", required=True)
    ap.add_argument("--out_json", required=True)
    ap.add_argument("--out_csv", required=True)
    args = ap.parse_args()

    text_rows = load_run(args.text_rows)
    img_rows = load_run(args.image_rows)

    tkeys, ikeys = {key(r) for r in text_rows}, {key(r) for r in img_rows}
    if len(tkeys) != len(text_rows) or len(ikeys) != len(img_rows):
        raise SystemExit("duplicate (boss, fight, index) keys in modality rows")
    shared = tkeys & ikeys
    if [key(row) for row in img_rows] != [key(row) for row in text_rows[:len(img_rows)]]:
        raise SystemExit(
            "ordered image keys must exactly equal the text sample prefix "
            f"(text={len(tkeys)} image={len(ikeys)} shared={len(shared)})"
        )

    summary = {
        "protocol": "pact-eval-v1",
        "layer": "L2",
        "setting_id": "L2-IMAGE-MODALITY-ABLATION",
        "sample_set_id": f"L2-SAMPLE-{len(img_rows)}-v3-balanced",
        "text_sample_set_id": f"L2-SAMPLE-{len(text_rows)}-v3-balanced",
        "sample_strategy_revision": "balanced-boss-fight-distance-v3",
        "belief_view": args.belief_view,
        "n_text_only": len(text_rows),
        "n_text_image": len(img_rows),
        "n_shared": len(shared),
        "source_rows": {
            "text": {"path": os.path.realpath(args.text_rows), "sha256": file_sha256(args.text_rows)},
            "image": {"path": os.path.realpath(args.image_rows), "sha256": file_sha256(args.image_rows)},
        },
        "note": (
            "text payload is byte-identical across modalities; the only difference is "
            "an additional current-frame image_url content part"
        ),
        "within_modality": {
            "text_only": within_modality(text_rows),
            "text_image": within_modality(img_rows),
            # 公平同口径：text 限制到与 image 相同的共享状态,消除"200-vs-60 子集难度差"混淆。
            # 直接横比 text_only(全集) 与 text_image 会被子集难度差污染,请用此块或 cross_modality。
            "text_only_shared": within_modality([r for r in text_rows if key(r) in shared]),
        },
        "cross_modality_decision_change": cross_modality(text_rows, img_rows),
    }

    # main table: semantic controllability within each modality
    rows = summary_csv_rows(summary)
    # companion table: paired decision change caused by adding the image
    paired_csv = os.path.splitext(args.out_csv)[0] + "_paired_change.csv"
    prows = paired_change_csv_rows(summary)
    for output in (args.out_json, args.out_csv, paired_csv):
        os.makedirs(os.path.dirname(output) or ".", exist_ok=True)
    json_tmp = f"{args.out_json}.tmp.{os.getpid()}"
    csv_tmp = f"{args.out_csv}.tmp.{os.getpid()}"
    paired_tmp = f"{paired_csv}.tmp.{os.getpid()}"
    try:
        with open(csv_tmp, "w", encoding="utf-8", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=list(rows[0]))
            writer.writeheader()
            writer.writerows(rows)
            f.flush()
            os.fsync(f.fileno())
        with open(paired_tmp, "w", encoding="utf-8", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=list(prows[0]))
            writer.writeheader()
            writer.writerows(prows)
            f.flush()
            os.fsync(f.fileno())
        with open(json_tmp, "w", encoding="utf-8") as f:
            json.dump(summary, f, ensure_ascii=False, indent=2, allow_nan=False)
            f.flush()
            os.fsync(f.fileno())
        os.replace(csv_tmp, args.out_csv)
        os.replace(paired_tmp, paired_csv)
        # The registry consumes JSON, so replace it last as the commit marker.
        os.replace(json_tmp, args.out_json)
    finally:
        for tmp in (json_tmp, csv_tmp, paired_tmp):
            try:
                os.unlink(tmp)
            except FileNotFoundError:
                pass
    print(f"wrote {paired_csv}")

    print(json.dumps(summary, ensure_ascii=False, indent=2))
    print(f"wrote {args.out_json}")
    print(f"wrote {args.out_csv}")


if __name__ == "__main__":
    main()