"""
Feature probe
=============
Question: do the vectors produced by extract_features.py actually CONTAIN the
class information?

Why this matters: when the LLM side underperforms there are two very different
causes —
  (a) the information is not in the vector      → encoder / feature problem
  (b) it is there but the LLM cannot read it    → channel problem
These call for opposite fixes. Fitting a plain LOGISTIC REGRESSION on the vector
separates them: if a linear layer can recover the class, the information is
present and the fault lies in the channel.

Measurement on the released checkpoint (subject-level split, 123 test cases):
    fused (512)   bal_acc 0.880   ← the head itself scores 0.895
    tab   (256)   bal_acc 0.855
    mri   (512)   bal_acc 0.410   (chance is 0.333)
Conclusion: the information is present and linearly decodable. Early failures of
projector + LoRA training were NOT a data or feature problem — they were a
channel problem.

Needs no GPU; runs in seconds.

Run:
    python probe_features.py --features features.pt
"""
from __future__ import annotations
import argparse

import numpy as np
import torch
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import balanced_accuracy_score, f1_score, confusion_matrix
from sklearn.preprocessing import StandardScaler


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--features", required=True)
    ap.add_argument("--C", type=float, default=1.0,
                    help="logistic regression regularisation")
    args = ap.parse_args()

    d = torch.load(args.features, map_location="cpu", weights_only=False)
    y = d["label"].numpy()
    sp = np.array(d["split"])
    cls = list(d["class_names"])
    tr, te = sp == "train", sp == "test"

    print(f"[data] {len(y)} records | "
          f"{ {cls[i]: int((y == i).sum()) for i in range(len(cls))} }")
    print(f"[split] train {int(tr.sum())} / test {int(te.sum())}\n")

    sets = {
        "fused (512)": d["fused_features"],
        "mri (512)": d["mri_features"],
        "tab (256)": d["tab_features"],
        "all (1280)": torch.cat(
            [d["fused_features"], d["mri_features"], d["tab_features"]], dim=-1),
    }

    best_name, best_ba = None, -1.0
    for name, X in sets.items():
        X = X.numpy()
        sc = StandardScaler().fit(X[tr])
        clf = LogisticRegression(max_iter=3000, C=args.C,
                                 class_weight="balanced").fit(sc.transform(X[tr]), y[tr])
        p = clf.predict(sc.transform(X[te]))
        ba = balanced_accuracy_score(y[te], p)
        print(f"  {name:14s} bal_acc={ba:.3f}  macroF1={f1_score(y[te], p, average='macro'):.3f}")
        if ba > best_ba:
            best_name, best_ba, best_pred = name, ba, p

    hp = d["class_probs"].argmax(-1).numpy()
    hb = balanced_accuracy_score(y[te], hp[te])
    print(f"\n  {'head (reference)':14s} bal_acc={hb:.3f}  "
          f"macroF1={f1_score(y[te], hp[te], average='macro'):.3f}")

    print(f"\n  Best vector: {best_name}")
    print(f"  Confusion (rows = true, columns = predicted {cls}):")
    print("  " + str(confusion_matrix(y[te], best_pred)).replace("\n", "\n  "))

    print("\n" + "-" * 62)
    if best_ba >= 0.8 * hb:
        print("  → The information is present and linearly decodable.")
        print("    If the LLM side misbehaves, the fault is in the CHANNEL")
        print("    (projector / LoRA), not the features. No need to retrain")
        print("    the encoder.")
    else:
        print("  → The vector does not carry what the head knows. Check feature")
        print("    extraction first, and confirm the checkpoint matches the")
        print("    input modality.")
    print("-" * 62)


if __name__ == "__main__":
    main()