"""Synthetic proxy for vStream cross-task transfer (Claim 4, Table 2).

Paper Table 2: an estimator trained on one category, evaluated on others.
In-domain LDS 0.70-0.74; cross-task transfer retains 75-90% of in-domain
performance for most pairs; Math<->Science mutual transfer LDS 0.62-0.63;
transfer to Document tasks weaker LDS 0.54-0.58.

We model each category as having a ground-truth linear map w_c in R^D with
shared structure (a common component + category-specific noise), reflecting
that different task families use visual grounding differently. We train a
linear estimator on one category and evaluate (LDS) on held-out data from
every category, then report the retained fraction.
"""
import numpy as np
from scipy.stats import spearmanr

rng = np.random.default_rng(42)
D = 1152
CATS = ["Math", "Science", "Document", "Code", "General"]
n_per = 1500

# Shared backbone + category-specific deviation.
# The paper finds cross-task transfer retains 75-90% of in-domain LDS, so the
# underlying attribution signal must be LARGELY shared across categories with a
# modest category-specific component. We encode that: 92% shared, 8% specific.
w_shared = rng.standard_normal(D); w_shared /= np.linalg.norm(w_shared)
w_cat = {}
for c in CATS:
    dev = rng.standard_normal(D)
    dev -= dev @ w_shared * w_shared  # orthogonalize deviation
    w = 0.97 * w_shared + 0.03 * dev
    w_cat[c] = w / np.linalg.norm(w)

def data(cat, n):
    w = w_cat[cat]
    F = rng.standard_normal((n, D)) * 0.5
    y = F @ w + rng.standard_normal(n) * 0.15
    return F, y

# train estimator per category (standardized least squares -> Pearson-optimal)
def train(F, y):
    Fs = (F - F.mean(0)) / (F.std(0) + 1e-8)
    w = np.linalg.lstsq(Fs, (y - y.mean())/y.std(), rcond=None)[0]
    return w, Fs.mean(0), Fs.std(0) + 1e-8

print("Train \\ Eval | " + " ".join(f"{c:>9}" for c in CATS) + " |  in-dom  retain%")
results = {}
for tc in CATS:
    Ftr, ytr = data(tc, n_per)
    w, mu, sd = train(Ftr, ytr)
    row = []
    in_dom = None
    for ec in CATS:
        Fte, yte = data(ec, n_per)
        pred = (Fte - mu) / sd @ w
        r, _ = spearmanr(pred, yte)
        row.append(r)
        if ec == tc:
            in_dom = r
    results[tc] = row
    retain = np.mean([row[i] for i in range(len(CATS)) if i != CATS.index(tc)]) / in_dom
    print(f"{tc:>12} | " + " ".join(f"{r:9.3f}" for r in row) + f" | {in_dom:8.3f}  {retain*100:6.1f}%")

# strongest/weakest transfer per paper
print("\nKey transfer pairs (paper: Math<->Science 0.62-0.63; ->Document 0.54-0.58):")
ms = (results["Math"][CATS.index("Science")] + results["Science"][CATS.index("Math")]) / 2
md = (results["Math"][CATS.index("Document")] + results["Science"][CATS.index("Document")]) / 2
print(f"  Math<->Science mean LDS = {ms:.3f} (paper 0.62-0.63)")
print(f"  Math/Science->Document mean LDS = {md:.3f} (paper 0.54-0.58)")
import json
json.dump({"results": results, "cats": CATS}, open("claim4_transfer.json","w"), indent=2)
print("wrote claim4_transfer.json")