"""有限分布上的逻辑回归、内部干预和可靠性算例；只用 Python 标准库。"""
from argparse import ArgumentParser
from collections import Counter
from itertools import product
from pathlib import Path
import hashlib
import json
import math
import platform
import time

HERE = Path(__file__).resolve().parent
PLAN = HERE / "data/interpretability-plan.json"
RECORD = HERE / "data/interpretability-record.json"
MODES = ("original", "balanced", "shape_only")
NAMES = ("原分布训练", "背景平衡训练", "只用形状")


def digest(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def save_new(path, value):
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("x", encoding="utf-8") as stream:
        json.dump(value, stream, ensure_ascii=False, indent=2, allow_nan=False)
        stream.write("\n")


def make_plan():
    return {
        "schema": 1, "shape_agreement": 0.8,
        "train_background_agreement": 0.95,
        "test_background_agreement": {"original": 0.95, "independent": 0.5, "reversed": 0.05},
        "initial_weights": [0.0, 0.0], "steps": 2000, "learning_rate": 0.5,
        "modes": list(MODES), "thresholds": [0.5, 0.8, 0.9, 0.95, 0.99],
        "data": "exact enumeration of eight (label, shape, background) states; not samples",
        "selection": "fixed final weights, no validation or hyperparameter selection",
        "superposition_activation": [0.05, 0.2, 0.5, 1.0],
    }


def sigmoid(z):
    return 1 / (1 + math.exp(-z))


def population(shape_agreement, background_agreement):
    rows = []
    for y, s, b in product((-1, 1), repeat=3):
        ps = shape_agreement if s == y else 1 - shape_agreement
        pb = background_agreement if b == y else 1 - background_agreement
        rows.append({"y": y, "x": [s, b], "mass": 0.5 * ps * pb})
    return rows


def objective(weights, rows):
    loss, gradient = 0.0, [0.0, 0.0]
    for row in rows:
        x, y, mass = row["x"], row["y"], row["mass"]
        margin = y * sum(w * v for w, v in zip(weights, x))
        loss += mass * math.log1p(math.exp(-margin))
        for j in range(2):
            gradient[j] -= mass * y * x[j] * sigmoid(-margin)
    return loss, gradient


def fit(rows, plan, shape_only=False):
    weights = plan["initial_weights"].copy()
    losses = []
    for _ in range(plan["steps"]):
        loss, gradient = objective(weights, rows)
        losses.append(loss)
        for j in range(1 if shape_only else 2):
            weights[j] -= plan["learning_rate"] * gradient[j]
    loss, gradient = objective(weights, rows)
    losses.append(loss)
    return {"weights": weights, "losses": losses, "final_gradient": gradient}


def evaluate(weights, rows, thresholds):
    predictions = []
    for row in rows:
        z = sum(w * v for w, v in zip(weights, row["x"]))
        prediction = 1 if z >= 0 else -1
        predictions.append({**row, "logit": z, "positive_probability": sigmoid(z),
                            "confidence": sigmoid(abs(z)), "prediction": prediction,
                            "correct": prediction == row["y"]})
    accuracy = sum(r["mass"] for r in predictions if r["correct"])
    groups = {}
    for y, agreement in product((-1, 1), (False, True)):
        group = [r for r in predictions if r["y"] == y and (r["x"][1] == y) == agreement]
        mass = sum(r["mass"] for r in group)
        groups[f"y={y},background_agrees={agreement}"] = (
            sum(r["mass"] for r in group if r["correct"]) / mass
        )
    selection = []
    for threshold in thresholds:
        chosen = [r for r in predictions if r["confidence"] >= threshold]
        coverage = sum(r["mass"] for r in chosen)
        risk = sum(r["mass"] for r in chosen if not r["correct"]) / coverage if coverage else None
        selection.append({"threshold": threshold, "coverage": coverage, "risk": risk})
    flips = 0.0
    for row in predictions:
        s, b = row["x"]
        new = 1 if weights[0] * s - weights[1] * b >= 0 else -1
        flips += row["mass"] * (new != row["prediction"])
    return {"predictions": predictions, "accuracy": accuracy, "groups": groups,
            "worst_group_accuracy": min(groups.values()), "selection": selection,
            "background_flip_rate": flips, "nll": objective(weights, rows)[0]}


def interventions(weights):
    # 人为等价展开，不声称这些通道由神经网络自动学出。
    readout = [weights[0], weights[1] / 2, weights[1] / 2, 0.0]
    clean, changed = [1, 1, 1, -1], [1, -1, -1, -1]
    score = lambda h: sum(w * v for w, v in zip(readout, h))
    baseline, target = score(changed), score(clean)
    patches = []
    for indices in ([], [0], [1], [2], [1, 2], [3]):
        h = changed.copy()
        for j in indices:
            h[j] = clean[j]
        patches.append({"indices": indices, "hidden": h, "logit": score(h),
                        "recovery": (score(h) - baseline) / (target - baseline)})
    ablations = []
    for indices in ([0], [1], [2], [1, 2], [3]):
        h = clean.copy()
        for j in indices:
            h[j] = 0
        ablations.append({"indices": indices, "hidden": h, "logit": score(h)})
    # 第四维完全可以读出，改变它却不影响输出。
    probe = [{"hidden": [s, b, b, t], "label": t, "prediction": t,
              "logit": score([s, b, b, t])}
             for s, b, t in product((-1, 1), repeat=3)]
    return {"readout": readout, "clean": clean, "changed": changed,
            "clean_logit": target, "changed_logit": baseline,
            "patches": patches, "ablations": ablations, "unused_probe": probe}


def superposition(probabilities):
    directions = [[1.0, 0.0], [-0.5, math.sqrt(3) / 2], [-0.5, -math.sqrt(3) / 2]]
    patterns = []
    for x in product((0, 1), repeat=3):
        h = [sum(x[j] * directions[j][k] for j in range(3)) for k in range(2)]
        decoded = [max(0.0, sum(a * b for a, b in zip(d, h))) for d in directions]
        patterns.append({"input": list(x), "hidden": h, "decoded": decoded,
                         "squared_error": sum((a - b) ** 2 for a, b in zip(x, decoded))})
    expectations = []
    for p in probabilities:
        loss = sum(p**sum(r["input"]) * (1 - p)**(3 - sum(r["input"])) * r["squared_error"]
                   for r in patterns)
        expectations.append({"activation_probability": p, "triangle_loss": loss,
                             "keep_two_loss": p})
    return {"directions": directions, "patterns": patterns, "expectations": expectations}


def entropy(labels):
    counts = Counter(labels)
    return max(0.0, -sum((n / len(labels)) * math.log(n / len(labels)) for n in counts.values()))


def evidence_toy():
    # 回答和语义分组均手工给定，没有调用语言模型或自动蕴含判断器。
    cases = [
        {"name": "同义正确", "answers": ["二楼", "在二楼", "2 楼", "第二层"],
         "meanings": ["floor2"] * 4, "correct": [True] * 4},
        {"name": "一致错误", "answers": ["三楼", "在三楼", "3 楼", "第三层"],
         "meanings": ["floor3"] * 4, "correct": [False] * 4},
        {"name": "答案摇摆", "answers": ["二楼", "三楼", "二楼", "三楼"],
         "meanings": ["floor2", "floor3", "floor2", "floor3"],
         "correct": [True, False, True, False]},
    ]
    for case in cases:
        case["string_entropy"] = entropy(case["answers"])
        case["meaning_entropy"] = entropy(case["meanings"])
    sources = {"floor": {"value": 2, "source": "notice-A"},
               "closing_hour": {"value": 20, "source": "notice-A"}}
    claims = [{"key": "floor", "value": 2}, {"key": "closing_hour", "value": 21},
              {"key": "has_color_printer", "value": True}]
    for claim in claims:
        evidence = sources.get(claim["key"])
        claim["status"] = ("unknown" if evidence is None else
                           "supported" if claim["value"] == evidence["value"] else "contradicted")
        claim["source"] = evidence["source"] if evidence else None
    return {"cases": cases, "sources": sources, "claims": claims,
            "supported_fraction": sum(c["status"] == "supported" for c in claims) / len(claims)}


def run(plan_path, out):
    plan = json.loads(plan_path.read_text())
    assert plan == make_plan(), "方案与代码不一致，请重新冻结到新文件。"
    if out.exists():
        raise FileExistsError(out)
    start = time.perf_counter()
    original = population(plan["shape_agreement"], plan["train_background_agreement"])
    balanced = population(plan["shape_agreement"], 0.5)
    models = {}
    for mode in MODES:
        rows = balanced if mode == "balanced" else original
        model = fit(rows, plan, shape_only=mode == "shape_only")
        model["training_population"] = rows
        model["evaluation"] = {
            name: evaluate(model["weights"], population(plan["shape_agreement"], q),
                           plan["thresholds"])
            for name, q in plan["test_background_agreement"].items()
        }
        models[mode] = model
    record = {"complete": True, "plan": plan, "plan_sha256": digest(plan_path),
              "source_sha256": digest(Path(__file__)), "python": platform.python_version(),
              "models": models, "interventions": interventions(models["original"]["weights"]),
              "superposition": superposition(plan["superposition_activation"]),
              "evidence": evidence_toy(), "elapsed_seconds": time.perf_counter() - start}
    save_new(out, record)
    return record


def show(record):
    print("三个环境的正确率：原分布 / 背景独立 / 背景反转")
    for mode, name in zip(MODES, NAMES):
        values = [record["models"][mode]["evaluation"][e]["accuracy"]
                  for e in ("original", "independent", "reversed")]
        print(f"{name}：" + " / ".join(f"{v:.1%}" for v in values))
    for env in ("original", "reversed"):
        selected = record["models"]["original"]["evaluation"][env]["selection"][2]
        print(f"{env}，置信度至少 0.9：覆盖 {selected['coverage']:.1%}，错误 {selected['risk']:.1%}")
    patches = record["interventions"]["patches"]
    print("恢复一个 / 两个背景通道：" + " / ".join(f"{patches[i]['recovery']:.1%}" for i in (2, 4)))
    print("各组回答的语义熵：" + " / ".join(f"{r['meaning_entropy']:.3f}" for r in record["evidence"]["cases"]))
    print("三条陈述的证据状态：" + " / ".join(r["status"] for r in record["evidence"]["claims"]))


def main():
    parser = ArgumentParser(description=__doc__)
    parser.add_argument("--freeze", type=Path)
    parser.add_argument("--run", type=Path, metavar="NEW_RECORD")
    parser.add_argument("--plan", type=Path, default=PLAN)
    args = parser.parse_args()
    if args.freeze:
        save_new(args.freeze, make_plan())
    elif args.run:
        show(run(args.plan, args.run))
    else:
        show(json.loads(RECORD.read_text()))


if __name__ == "__main__":
    main()