"""Sanity checks of the external results (38_external_eval.py; embeddings cache/external/emb_<set>_fold<k>.npz).

Galleries in 38_ mix spots of several samples (IDC 4 samples, DLPFC 12 sections). A model that only recognises the
SAMPLE (batch / stain / platform effects) and is random within it already beats chance there. Per pair and set:
  mixed G         : the 38_ metric (G = 100 random spots from all samples), both directions
  sample oracle   : R@1 of a model that knows the sample exactly and is random within it = mean 1/n_same_sample
  top-1 same samp : fraction of top-1 hits from the query's own sample (base rate = share of same-sample items)
  within-sample G : galleries drawn from ONE sample only (IDC G = 100; DLPFC G = 35 since sections have 35-48 spots,
                    with the mixed metric at G = 35 for comparison) -> spot-level signal beyond sample identity
  shuffled null   : the mixed metric after permuting one side (checks the metric code; ~1/G)
Output results/external_check.csv
"""
import os

import numpy as np
import pandas as pd

from common import CACHE, ROOT

E = os.path.join(CACHE, "external")
PAIRS = [("he", "rna"), ("rna", "mif"), ("he", "mif"), ("rna", "text"), ("he", "text"), ("mif", "text")]
NAME = {"he": "H&E", "rna": "RNA", "mif": "mIF(MIPHEI)", "text": "text(synthetic)"}
DRAWS = 200


def r1(a, b, s):
    sim = a[s] @ b[s].T; d = np.diag(sim)[:, None]
    return ((sim > d).sum(1) == 0).mean() / 2 + ((sim > d.T).sum(0) == 0).mean() / 2, sim


def run(a, b, smp, G, within, rng):
    out = {"R@1": [], "oracle": [], "top1_same": [], "base_same": []}
    uniq = [u for u in np.unique(smp) if (smp == u).sum() >= G]
    for _ in range(DRAWS):
        if within:
            idx = np.flatnonzero(smp == uniq[rng.integers(len(uniq))])
            s = rng.choice(idx, G, replace=False)
        else:
            s = rng.choice(len(a), G, replace=False)
        v, sim = r1(a, b, s)
        g = smp[s]; same = g[:, None] == g[None, :]
        out["R@1"].append(v)
        out["oracle"].append((1 / same.sum(1)).mean())
        out["top1_same"].append((same[np.arange(G), sim.argmax(1)].mean() + same[sim.argmax(0), np.arange(G)].mean()) / 2)
        out["base_same"].append(same.mean())
    return {k: float(np.mean(v)) for k, v in out.items()}


rows = []
for st, Gw in [("dlpfc", 35)]:  # IDC dropped from the report 2026-10-18
    smp = pd.read_csv(os.path.join(E, f"{st}_meta.csv"))["sample"].astype(str).values
    for k in range(5):
        Z = np.load(os.path.join(E, f"emb_{st}_fold{k}.npz"))
        for x, y in PAIRS:
            if x not in Z or y not in Z:
                continue
            a, b = Z[x], Z[y]
            rng = np.random.default_rng(k)
            base = {"set": st, "fold": k, "pair": f"{NAME[x]}<->{NAME[y]}"}
            m = run(a, b, smp, 100, False, rng)
            rows.append({**base, "gallery": "mixed G=100", "chance": 0.01, **m})
            rows.append({**base, "gallery": "shuffled null G=100", "chance": 0.01,
                         **run(a, b[rng.permutation(len(b))], smp, 100, False, rng)})
            if Gw != 100:
                rows.append({**base, "gallery": f"mixed G={Gw}", "chance": 1 / Gw, **run(a, b, smp, Gw, False, rng)})
            rows.append({**base, "gallery": f"within-sample G={Gw}", "chance": 1 / Gw, **run(a, b, smp, Gw, True, rng)})

R = pd.DataFrame(rows)
R.to_csv(os.path.join(ROOT, "results", "external_check.csv"), index=False)
S = R.groupby(["set", "pair", "gallery"], sort=False).agg(
    R1=("R@1", "mean"), R1_sd=("R@1", "std"), chance=("chance", "first"), oracle=("oracle", "mean"),
    top1_same=("top1_same", "mean"), base_same=("base_same", "mean"))
S["x_chance"] = S.R1 / S.chance
pd.set_option("display.width", 250)
print(S.round(3).to_string())