"""Readout-matched comparison and the campaign's final result record.

Everything above used the frozen probe on the MAE side, which is what EAT's own
proxy ranking used and so is the like-for-like comparison against +0.630. This
adds the arm the user asked for: the MAE side scored with the same AS-20K LoRA
readout at 3 pinned seeds that produced the EAT base ground truth, so proxy and
target differ in objective and scale but not in how they are read out.
"""
import csv
import glob
import json
import statistics as st
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())
lora = json.loads(Path("/workspace/runs/mae-64/lora_summary.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")
delta = np.bincount(ev, minlength=20) / len(ev) - train_prop

dids = sorted(d for d in (int(v["dist_id"]) for v in spec.values())
              if d in eat_proxy and f"d{d}-s0" in lora)
base = np.array([spec[str(d)]["lora_mean"] for d in dids])
eatp = np.array([eat_proxy[d] for d in dids])
mae_probe = 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])
mae_lora = np.array([lora[f"d{d}-s0"]["mean"] for d in dids])
lora_sd = np.array([lora[f"d{d}-s0"]["sd"] for d in dids])
em = np.array([float(np.array(spec[str(d)]["weights"]) @ delta) 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)


sigma_ft = float(np.sqrt(np.mean(lora_sd**2)))
between = float(mae_lora.std(ddof=1))
rel_lora = (between**2 - (sigma_ft / np.sqrt(3))**2) / between**2

print(f"n = {len(dids)} arms\n")
print("MAE-side LoRA readout (4 epochs, lr 3e-3, seeds 0/1/2 pinned):")
print(f"  mean {mae_lora.mean():.5f}   between-arm sd {between:.5f}")
print(f"  sigma_ft (pooled within arm)  {sigma_ft:.5f}   "
      f"SE of 3-seed mean {sigma_ft/np.sqrt(3):.5f}")
print(f"  reliability of the 3-seed mean {rel_lora:.3f}")
print(f"  max within-arm sd {lora_sd.max():.5f}  (collapse screen threshold 0.02)")

print(f"\n{'predictor -> EAT base (LoRA x3)':<38}{'rho':>9}")
rows = {
    "EAT proxy, frozen probe": eatp,
    "MAE proxy, frozen probe": mae_probe,
    "MAE proxy, LoRA x3 (readout-matched)": mae_lora,
    "eval-match (no training)": em,
    "EAT probe + MAE LoRA": z(eatp) + z(mae_lora),
    "EAT probe + MAE LoRA + eval-match": z(eatp) + z(mae_lora) + z(em),
}
for name, v in rows.items():
    print(f"{name:<38}{rho(v, base):>+9.3f}")

print(f"\nMAE probe vs MAE LoRA agreement: {rho(mae_probe, mae_lora):+.3f}")
print(f"EAT proxy vs MAE LoRA agreement : {rho(eatp, mae_lora):+.3f}")

rng = np.random.default_rng(0)
d = np.array([rho(mae_lora[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 (MAE LoRA - EAT probe): {rho(mae_lora, 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}")

out = {
    "n_arms": len(dids),
    "rho": {k: float(rho(v, base)) for k, v in rows.items()},
    "mae_lora": {"mean": float(mae_lora.mean()), "between_sd": between,
                 "sigma_ft": sigma_ft, "reliability_3seed": float(rel_lora)},
    "agreement": {"mae_probe_vs_mae_lora": float(rho(mae_probe, mae_lora)),
                  "eat_vs_mae_lora": float(rho(eatp, mae_lora))},
    "recipe": {"objective": "mae", "scale": "15m", "patch": "16x16", "lr": 5e-4,
               "budget_clips": 1912024, "mask_prob": 0.75, "clone_batch": 1,
               "compile": False, "readout_lora_epochs": 4, "readout_lora_lr": 3e-3,
               "node": "4x RTX 3090", "floor_p16": 0.03861},
}
Path("/workspace/analysis/FINAL_RESULT.json").write_text(json.dumps(out, indent=1))
print("\nwrote /workspace/analysis/FINAL_RESULT.json")