Research-Stack/4-Infrastructure/shim/pist_train_ground_truth.py

#!/usr/bin/env python3
"""Train on ground-truth labels (proof_method, domain, RRCShape).

Usage:
    python3 pist_train_ground_truth.py
"""

import json
import os
import sys
from collections import Counter, defaultdict
from math import sqrt

FEATURE_VECTORS_PATH = os.path.join(os.path.dirname(__file__), "../..",
                                    "shared-data/pist_canary_feature_vectors.jsonl")
GROUND_TRUTH_PATH = os.path.join(os.path.dirname(__file__), "../..",
                                 "shared-data/pist_canary_ground_truth_labels.jsonl")
REPORT_PATH = os.path.join(os.path.dirname(__file__), "../..",
                           "shared-data/pist_canary_ground_truth_report.json")


def load_vectors(path: str) -> list[dict]:
    rows = []
    with open(path) as f:
        for line in f:
            rows.append(json.loads(line))
    return rows


def load_labels(path: str) -> dict[str, dict]:
    labels = {}
    with open(path) as f:
        for line in f:
            row = json.loads(line)
            labels[row["theorem_name"]] = row
    return labels


def normalize(vectors):
    n = len(vectors)
    if n == 0:
        return vectors, [], []
    dim = len(vectors[0])
    means = [sum(v[i] for v in vectors) / n for i in range(dim)]
    stds = [sqrt(sum((v[i] - means[i]) ** 2 for v in vectors) / max(n - 1, 1)) for i in range(dim)]
    stds = [s if s > 1e-9 else 1.0 for s in stds]
    return [[(v[i] - means[i]) / stds[i] for i in range(dim)] for v in vectors], means, stds


def centroid(vecs):
    if not vecs:
        return []
    return [sum(v[i] for v in vecs) / len(vecs) for i in range(len(vecs[0]))]


def euclidean(a, b):
    return sqrt(sum((a[i] - b[i]) ** 2 for i in range(len(a))))


def knn(train_v, train_l, test_v, k):
    dists = sorted(
        [(euclidean(test_v, tv), tl) for tv, tl in zip(train_v, train_l)],
        key=lambda x: x[0],
    )
    votes = Counter(tl for _, tl in dists[:k])
    return votes.most_common(1)[0][0]


def eval_loocv(vectors, labels, method="centroid", k=3):
    n = len(vectors)
    correct = 0
    top2_correct = 0
    n_classes = len(set(labels))
    confusion = defaultdict(lambda: defaultdict(int))

    for i in range(n):
        tv = vectors[:i] + vectors[i + 1 :]
        tl = labels[:i] + labels[i + 1 :]
        test_v = vectors[i]
        test_l = labels[i]

        if method == "centroid":
            cls_vecs = defaultdict(list)
            for v, l in zip(tv, tl):
                cls_vecs[l].append(v)
            cents = {l: centroid(vecs) for l, vecs in cls_vecs.items()}
            preds = sorted(
                cents.keys(), key=lambda l: euclidean(test_v, cents[l])
            )
            pred = preds[0]
            top2 = preds[:2]
        else:
            dists = sorted(
                [(euclidean(test_v, tv[j]), tl[j]) for j in range(len(tv))],
                key=lambda x: x[0],
            )
            votes = Counter(tl for _, tl in dists[:k])
            pred = votes.most_common(1)[0][0]
            top2 = [t[1] for t in dists[:2]]

        confusion[test_l][pred] += 1
        if pred == test_l:
            correct += 1
        if test_l in top2:
            top2_correct += 1

    baseline = max(Counter(labels).values()) / n
    accuracy = correct / n
    top2_acc = top2_correct / n

    return accuracy, top2_acc, dict(confusion), baseline


def main():
    vectors = load_vectors(FEATURE_VECTORS_PATH)
    gt_labels = load_labels(GROUND_TRUTH_PATH)

    # Align by theorem name
    aligned = []
    for v in vectors:
        name = v.get("name", "")
        if name in gt_labels:
            aligned.append((v["vector"], gt_labels[name]))

    n = len(aligned)
    print(f"Aligned vectors: {n}", flush=True)

    raw_vecs = [a[0] for a in aligned]
    normed, _, _ = normalize(raw_vecs)

    TARGETS = [
        ("proof_method_label", "Proof method"),
        ("domain_label", "Domain"),
        ("manual_rrc_shape", "Manual RRCShape"),
    ]

    results = {}

    for key, label in TARGETS:
        print(f"\n{'=' * 60}", flush=True)
        print(f"TARGET: {label} ({key})", flush=True)
        print(f"{'=' * 60}", flush=True)

        targets = []
        for _, lbl in aligned:
            val = lbl.get(key, "unknown")
            if val is None:
                val = "none"
            targets.append(val)

        classes = sorted(set(targets))
        dist = Counter(targets)
        print(f"Classes ({len(classes)}): {dict(dist)}", flush=True)
        print(f"Majority baseline: {max(dist.values()) / n:.1%}", flush=True)

        for method, k in [
            ("centroid", None),
            ("knn_1", 1),
            ("knn_3", 3),
            ("knn_5", 5),
        ]:
            acc, top2, conf, baseline = eval_loocv(
                normed, targets, "centroid" if method == "centroid" else "knn", k or 3
            )
            print(f"  {method:15s} LOOCV: {acc:.1%} ({int(acc * n)}/{n})  "
                  f"top-2: {top2:.1%}  baseline: {baseline:.1%}",
                  flush=True)

            results[f"{key}_{method}"] = {
                "accuracy": round(acc, 4),
                "top2_accuracy": round(top2, 4),
                "baseline": round(baseline, 4),
                "n_classes": len(classes),
                "improvement_over_baseline": round(
                    (acc - baseline) / max(baseline, 0.01), 3
                ),
            }

        # Per-class accuracy for centroid
        acc, top2, conf, baseline = eval_loocv(normed, targets, "centroid", 3)
        print(f"\n  Per-class (centroid):")
        for cls in sorted(classes):
            tp = conf.get(cls, {}).get(cls, 0)
            total = dist[cls]
            print(f"    {cls:30s} {tp:3d}/{total:3d} = {tp / max(total, 1):.0%}")

    # Summary
    print(f"\n{'=' * 60}", flush=True)
    print("SUMMARY", flush=True)
    print(f"{'=' * 60}", flush=True)
    print(f"{'Target':35s} {'Method':12s} {'Acc':>6} {'Top-2':>6} {'Base':>6} {'Impr':>6}",
          flush=True)
    print(f"{'-' * 35} {'-' * 12} {'-' * 6} {'-' * 6} {'-' * 6} {'-' * 6}", flush=True)
    for key, label in TARGETS:
        for method in ["centroid", "knn_3"]:
            r = results.get(f"{key}_{method}", {})
            impr = r.get("improvement_over_baseline", 0)
            print(
                f"{label:35s} {method:12s} "
                f"{r.get('accuracy', 0):6.1%} "
                f"{r.get('top2_accuracy', 0):6.1%} "
                f"{r.get('baseline', 0):6.1%} "
                f"{impr:>+5.1f}x",
                flush=True,
            )

    # Save report
    report = {
        "n_samples": n,
        "targets": {key: label for key, label in TARGETS},
        "results": results,
        "class_distributions": {
            key: dict(Counter([lbl.get(key, "unknown") for _, lbl in aligned]))
            for key, _ in TARGETS
        },
    }
    with open(REPORT_PATH, "w") as f:
        json.dump(report, f, indent=2)
    print(f"\nReport: {REPORT_PATH}", flush=True)
    return 0


if __name__ == "__main__":
    main()