"""MIPHEI-ViT vs GigaTIME virtual mIF on ORION, against the real mIF of the same tiles (tile means).

MIPHEI was trained on the ORION official train split -> 'clean' = official TEST tiles (2 patients, never seen);
'seen' = official TRAIN tiles (training data of MIPHEI, optimistic). GigaTIME never saw ORION (clean everywhere).
  (a) per-channel Spearman of tile means (pooled and within-slide)
  (b) direct real-mean <-> virtual-mean retrieval (rank-gaussianised means, cosine; no learned head), 100-galleries
Output results/miphei_diag.csv, results/miphei_diag_retrieval.csv
"""
import json
import os

import numpy as np
import pandas as pd
from scipy.special import ndtri
from scipy.stats import rankdata, spearmanr

from common import CACHE, ROOT, SPLITS, load_all_tiles

OUT = os.path.join(ROOT, "results")
ORION_CH = ["Hoechst", "CD31", "CD45", "CD68", "CD4", "FOXP3", "CD8a", "CD45RO", "CD20", "PD-L1", "CD3e",
            "CD163", "E-cadherin", "PD-1", "Ki67", "Pan-CK", "SMA"]  # order of orion_real_if_means.npy (25_)
GT = {"Hoechst": ["DAPI"], "CD68": ["CD68"], "CD4": ["CD4"], "CD8a": ["CD8"], "CD20": ["CD20"], "PD-L1": ["PD-L1"],
      "CD3e": ["CD3"], "Ki67": ["Ki67"], "Pan-CK": ["CK"], "SMA": ["Actin-D", "Transgelin"]}
df = load_all_tiles()
slide = df.in_slide_name.values
real = np.load(os.path.join(CACHE, "gigatime", "orion_real_if_means.npy"))
mch = json.load(open(os.path.join(CACHE, "miphei", "channels.json")))["channels"]
mip = np.load(os.path.join(CACHE, "miphei", "orion_means.npy"))
gch = json.load(open(os.path.join(CACHE, "gigatime", "channels.json")))["channels"]
gm = np.load(os.path.join(CACHE, "gigatime", "orion_means.npy"))
off = pd.read_csv(os.path.join(SPLITS, "orion_official", "fold0.csv")).split.values
sets = {"clean (official test)": off == "test", "seen (official train)": off == "train"}


def within(a, b, m):
    r, w = [], []
    for s in np.unique(slide[m]):
        i = m & (slide == s)
        if i.sum() >= 15 and a[i].std() > 0 and b[i].std() > 0:
            r.append(spearmanr(a[i], b[i])[0]); w.append(i.sum())
    return float(np.average(r, weights=w)) if r else np.nan


rows = []
for ch in mch:
    r = real[:, ORION_CH.index(ch)]
    for sname, m in sets.items():
        row = {"channel": ch, "set": sname,
               "MIPHEI pooled": spearmanr(mip[m, mch.index(ch)], r[m])[0],
               "MIPHEI within-slide": within(mip[:, mch.index(ch)], r, m)}
        if ch in GT:
            g = gm[:, [gch.index(x) for x in GT[ch]]].mean(1)
            row.update({"GigaTIME pooled": spearmanr(g[m], r[m])[0], "GigaTIME within-slide": within(g, r, m)})
        rows.append(row)
C = pd.DataFrame(rows)
rg = lambda x: ndtri((np.apply_along_axis(rankdata, 0, x) - 0.5) / len(x))


def g100(a, b, draws=50, seed=0):
    a = a / np.linalg.norm(a, axis=1, keepdims=True); b = b / np.linalg.norm(b, axis=1, keepdims=True)
    rng = np.random.default_rng(seed); r = []
    for _ in range(draws):
        s = rng.choice(len(a), min(100, len(a)), replace=False)
        sim = a[s] @ b[s].T; d = np.diag(sim)[:, None]
        r.append(((sim > d).sum(1) == 0).mean() / 2 + ((sim.T > d).sum(1) == 0).mean() / 2)
    return float(np.mean(r))


ten = list(GT)
Rt = []
for sname, m in sets.items():
    R16 = real[m][:, [ORION_CH.index(c) for c in mch]]
    R10 = real[m][:, [ORION_CH.index(c) for c in ten]]
    G10 = np.stack([gm[m][:, [gch.index(x) for x in GT[c]]].mean(1) for c in ten], 1)
    Rt += [{"set": sname, "n": int(m.sum()), "pair": "real <-> MIPHEI (16 ch)", "R@1_g100": g100(rg(R16), rg(mip[m]))},
           {"set": sname, "n": int(m.sum()), "pair": "real <-> MIPHEI (10 shared ch)", "R@1_g100": g100(rg(R10), rg(mip[m][:, [mch.index(c) for c in ten]]))},
           {"set": sname, "n": int(m.sum()), "pair": "real <-> GigaTIME (10 shared ch)", "R@1_g100": g100(rg(R10), rg(G10))}]
Rt = pd.DataFrame(Rt)
pd.set_option("display.width", 200)
print(C.round(3).to_string(index=False))
print("\nmean over channels\n", C.groupby("set")[[c for c in C.columns if c not in ("channel", "set")]].mean().round(3).to_string())
print("\nmean over the 10 GigaTIME-shared channels\n", C[C.channel.isin(ten)].groupby("set")[[c for c in C.columns if c not in ("channel", "set")]].mean().round(3).to_string())
print("\n", Rt.round(3).to_string(index=False))
C.round(4).to_csv(os.path.join(OUT, "miphei_diag.csv"), index=False)
Rt.round(4).to_csv(os.path.join(OUT, "miphei_diag_retrieval.csv"), index=False)