"""Shared paths / helpers for the ORION ImageBind-style 4-modality experiments.

Modalities (all spatially paired per 333x333 tile at 20x):
  he      : H&E RGB tile                      -> anchor, frozen ImageBind-Huge vision encoder
  if      : 17-channel multiplex IF (uint8)   -> trainable encoder
  nuclei  : nuclei instance mask (int32)      -> trainable encoder (binary mask + boundary channels)
  markers : per-tile marker vector (17 *_prop + 16 *_count from the dataframes) -> trainable MLP
"""
import os
import sys
import types

import numpy as np
import pandas as pd

ROOT = os.path.dirname(os.path.abspath(__file__))
ORION = os.path.join(ROOT, "..", "ORION", "ORIONCRC_dataset_tile_20x")
CACHE = os.path.join(ROOT, "cache")
SPLITS = os.path.join(ROOT, "splits")
RUNS = os.path.join(ROOT, "runs")
IB_WEIGHTS = os.path.join(ROOT, "weights", "imagebind_huge.pth")
IMG = 224

MODALITIES = ["he", "if", "nuclei", "markers"]


def load_all_tiles():
    """All tiles across the official train/val/test csvs, with a stable global index."""
    dfs = []
    for s in ["train", "val", "test"]:
        d = pd.read_csv(os.path.join(ORION, f"{s}_dataframe.csv"))
        d["orion_split"] = s
        dfs.append(d)
    df = pd.concat(dfs, ignore_index=True)
    slides = pd.read_csv(os.path.join(ORION, "slide_dataframe.csv"))
    df = df.merge(slides[["in_slide_name", "orion_slide_id", "targ_slide_path"]], on="in_slide_name", how="left")
    df["patient"] = df.orion_slide_id.str.split("_").str[0]          # CRC33_01 / CRC33_02 -> CRC33
    df["if_batch"] = df.targ_slide_path.str.extract(r"_(C59[a-z]X)_")[0]  # IF acquisition batch
    df["tile_id"] = np.arange(len(df))
    return df


def marker_columns(df):
    prop = [c for c in df.columns if c.endswith("_prop")]
    cnt = [c for c in df.columns if c.endswith("_count")]
    return prop, cnt


def marker_matrix(df):
    prop, cnt = marker_columns(df)
    return np.concatenate([df[prop].values, np.log1p(df[cnt].values)], 1).astype(np.float32)


def dihedral(x, d):
    """Apply one of 8 dihedral transforms to the last two (spatial) dims of a tensor/array."""
    import torch
    k, flip = d % 4, d // 4
    if flip:
        x = torch.flip(x, dims=(-1,))
    return torch.rot90(x, k, dims=(-2, -1))


def load_imagebind_vision(device="cuda"):
    """Build ImageBind-Huge and load the official checkpoint without importing imagebind.data
    (which needs torchaudio/pytorchvideo)."""
    src = os.path.join(ROOT, "ImageBind_src", "imagebind")
    if "imagebind" not in sys.modules:
        pkg = types.ModuleType("imagebind")
        pkg.__path__ = [src]
        sys.modules["imagebind"] = pkg
    import torch
    from imagebind.models import imagebind_model
    from imagebind.models.imagebind_model import ModalityType

    model = imagebind_model.imagebind_huge(pretrained=False)
    sd = torch.load(IB_WEIGHTS, map_location="cpu")
    missing, unexpected = model.load_state_dict(sd, strict=True), None
    model.eval().to(device)
    return model, ModalityType


def metrics(zq, zg, g100=True, eps=1e-6):
    """Tie-aware retrieval metrics (expected rank under random tie-breaking; template texts can tie)."""
    import torch
    from scipy.optimize import linear_sum_assignment
    sim = (zq @ zg.T).numpy()
    n = len(sim)
    d = sim[np.arange(n), np.arange(n)][:, None]
    gt = (sim > d + eps).sum(1)
    tie = (np.abs(sim - d) <= eps).sum(1) - 1  # other candidates tied with the true match

    def rk(k):  # P(true match within top-k) with random tie-breaking
        return np.clip((k - gt) / (tie + 1), 0, 1).mean()

    r, c = linear_sum_assignment(-sim)
    res = {"R@1": rk(1), "R@5": rk(5), "R@10": rk(10), "MRR": float(np.mean(1 / (gt + tie / 2 + 1))),
           "pair_acc": float((c[np.argsort(r)] == np.arange(n)).mean()), "N": n}
    if g100:
        rng = np.random.default_rng(0)
        if n <= 100:
            res["R@1_g100"] = res["R@1"]
        else:
            res["R@1_g100"] = float(np.mean([metrics(zq[s], zg[s], g100=False)["R@1"]
                                             for s in (torch.from_numpy(rng.choice(n, 100, replace=False)) for _ in range(50))]))
    return res