"""Do the two objectives carry independent signal about the target ranking?

rho(EAT proxy, MAE proxy) = +0.482 is lower than either proxy's correlation with
the target (+0.630, +0.668). Two noisy measurements of one underlying quantity
cannot do that beyond what attenuation explains, so the objectives may be seeing
partly different things -- in which case combining them should beat either, the
same structure cluster-weight-channels.md found for the eval-match channel
(+0.630 -> +0.802 from a term that was near-orthogonal to the proxy score).
"""
import csv
import glob
import json
from pathlib import Path

import numpy as np
from scipy import stats

spec = json.loads(Path("/workspace/analysis/canonical64.json").read_text())
curve = json.loads(Path("/workspace/analysis/transfer_curve.json").read_text())
final = max(r["step"] for r in curve)

eat_proxy = {}
for r in csv.DictReader(open("/workspace/analysis/eat-map-regmix/tables/runs.csv")):
    if r.get("campaign") != "regmix256":
        continue
    try:
        if float(r["pretrain_seed"]) == 0.0:
            eat_proxy[int(float(r["dist_id"]))] = float(r["as20k_map"])
    except (ValueError, KeyError):
        pass

ART = "/workspace/artifacts/audio-mixture-scaling/artifacts/audioset-dasheng-0.6b-k20-r31d6389"
train_prop = np.array(json.loads(
    Path("/workspace/data/asmel_flat/train.cluster_index.json").read_text())["proportions"])
ev = np.load(f"{ART}/test_cluster_index.npy")
eval_prop = np.bincount(ev, minlength=20).astype(float)
delta = eval_prop / eval_prop.sum() - train_prop

dids = sorted(int(v["dist_id"]) for v in spec.values() if int(v["dist_id"]) in eat_proxy)
base = np.array([spec[str(d)]["lora_mean"] for d in dids])
eatp = np.array([eat_proxy[d] for d in dids])
maep = np.array([json.loads(Path(
    glob.glob(f"/workspace/runs/mae-64/d{d}-s0/exports/step_{final:08d}/probe.json")[0]
).read_text())["probe/map"] for d in dids])
em = np.array([float(np.array(spec[str(d)]["weights"]) @ delta) for d in dids])
batch = np.array([spec[str(d)]["batch"] for d in dids])


def rho(a, b):
    return stats.spearmanr(a, b).statistic


def z(x):
    return (x - x.mean()) / x.std(ddof=1)


combos = {
    "EAT proxy alone": eatp,
    "MAE proxy alone": maep,
    "eval-match alone": em,
    "EAT + MAE": z(eatp) + z(maep),
    "EAT + eval-match": z(eatp) + z(em),
    "MAE + eval-match": z(maep) + z(em),
    "EAT + MAE + eval-match": z(eatp) + z(maep) + z(em),
}
print(f"{'predictor':<26}{'rho vs base':>12}{'  test A':>10}{'  test B':>10}")
A, B = batch == "A", batch == "B"
for name, v in combos.items():
    print(f"{name:<26}{rho(v, base):>+12.3f}{rho(v[A], base[A]):>+10.3f}"
          f"{rho(v[B], base[B]):>+10.3f}")

rng = np.random.default_rng(0)
best = z(eatp) + z(maep)
d = np.array([rho(best[i], base[i]) - rho(eatp[i], base[i])
              for i in (rng.integers(0, len(dids), len(dids)) for _ in range(4000))])
print(f"\nDelta rho, (EAT+MAE) vs EAT alone: {rho(best, base) - rho(eatp, base):+.3f}   "
      f"95% CI [{np.percentile(d, 2.5):+.3f}, {np.percentile(d, 97.5):+.3f}]   "
      f"P(>0) = {(d > 0).mean():.3f}")

# variance decomposition on ranks, the framing cluster-weight-channels.md used
r = {k: stats.rankdata(v) for k, v in
     dict(base=base, eat=eatp, mae=maep, em=em).items()}
for k in ("eat", "mae", "em"):
    print(f"R2(base ~ {k:>3}) = {np.corrcoef(r['base'], r[k])[0,1]**2:.3f}", end="   ")
X = np.column_stack([z(r["eat"]), z(r["mae"])])
beta = np.linalg.lstsq(X, z(r["base"]), rcond=None)[0]
pred = X @ beta
print(f"\nR2(base ~ eat + mae) = {np.corrcoef(pred, z(r['base']))[0,1]**2:.3f}")

print("\nselection value -- mean base LoRA mAP of the top-k picked:")
print(f"{'rule':<26}{'top1':>9}{'top3':>9}{'top5':>9}{'top10':>9}")
for name, v in list(combos.items()) + [("oracle", base)]:
    order = np.argsort(-v)
    print(f"{name:<26}" + "".join(f"{base[order[:k]].mean():>9.5f}" for k in (1, 3, 5, 10)))
print(f"{'no selection':<26}" + f"{base.mean():>9.5f}" * 4)