scripts/image_traversals_compare.py

#!/usr/bin/env python3
"""
MERU-style image traversals comparing RadIR/RadIR (hyperbolic) vs RadIR_compare (Euclidean).

For each query CXR:
  1. Encode the image with both models.
  2. Walk along the geodesic / ray from the image embedding toward the origin [ROOT].
  3. At each step, retrieve the nearest neighbor from a shared multi-granularity concept pool.
     Hyperbolic path optionally keeps only texts that entail the current point
     (MERU-style cone filter), then takes NN in that subset.
  4. Render a MERU Figure-5 style side-by-side panel.

Usage (from RadIR/RadIR):
  python scripts/image_traversals_compare.py \\
    --images path1.jpg path2.jpg \\
    --out outputs/traversals_compare.png

  # Auto-pick diverse disease cases from valid.jsonl + disease labels:
  python scripts/image_traversals_compare.py --auto_pick 4 --out outputs/traversals_compare.png

  # Disable entailment filter (pure hyperbolic NN):
  python scripts/image_traversals_compare.py --auto_pick 4 --no_entailment
"""

from __future__ import annotations

import argparse
import csv
import gc
import json
import os
import re
import sys
import textwrap
from pathlib import Path
from typing import Sequence

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from matplotlib.patches import FancyBboxPatch
from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
# Workspace layout: HMClip/RadIR/RadIR and HMClip/RadIR_compare
_hmclip = ROOT.parent.parent
if (_hmclip / "RadIR_compare").is_dir():
    COMPARE_ROOT = _hmclip / "RadIR_compare"
elif (ROOT.parent / "RadIR_compare").is_dir():
    COMPARE_ROOT = ROOT.parent / "RadIR_compare"
else:
    COMPARE_ROOT = _hmclip / "RadIR_compare"

ROOT_LABEL = "[ROOT]"


# ---------------------------------------------------------------------------
# Concept pool
# ---------------------------------------------------------------------------

def load_concept_pool(
    concepts_file: str | None,
    extra_captions: int = 0,
    caption_jsonl: str | None = None,
    seed: int = 42,
) -> list[str]:
    texts: list[str] = []
    seen: set[str] = set()

    def _add(s: str) -> None:
        s = " ".join(s.strip().split())
        if not s:
            return
        key = s.lower()
        if key in seen or key == ROOT_LABEL.lower():
            return
        seen.add(key)
        texts.append(s)

    path = concepts_file or str(ROOT / "scripts" / "traversal_concepts.txt")
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            _add(line)

    if extra_captions > 0 and caption_jsonl and os.path.isfile(caption_jsonl):
        rng = np.random.default_rng(seed)
        rows = []
        with open(caption_jsonl, encoding="utf-8") as f:
            for line in f:
                item = json.loads(line)
                text = str(item.get("text", "")).strip()
                if not text:
                    continue
                first = text.split(".")[0].strip()
                if 8 <= len(first) <= 80:
                    rows.append(first)
        if rows:
            idx = rng.choice(len(rows), size=min(extra_captions, len(rows)), replace=False)
            for i in idx:
                _add(rows[int(i)])

    texts.append(ROOT_LABEL)
    return texts


# ---------------------------------------------------------------------------
# Image I/O
# ---------------------------------------------------------------------------

def resolve_image_path(path: str, dataset_root: str) -> str:
    candidates = [
        path,
        os.path.join(dataset_root, path),
        os.path.join(dataset_root, path.replace("dataset/", "", 1)),
        os.path.join(str(ROOT), path),
        os.path.join(str(ROOT), path.replace("dataset/", "dataset/", 1)),
    ]
    for c in candidates:
        if os.path.isfile(c):
            return os.path.abspath(c)
    raise FileNotFoundError(f"Image not found: {path}")


def load_image_tensor(path: str, imsize: int, with_depth: bool) -> torch.Tensor:
    """
    Grayscale -> [-1, 1].
    with_depth=False -> [1, H, W]   (RadIR/RadIR hyp)
    with_depth=True  -> [1, 1, H, W] (RadIR_compare, C×D×H×W)
    """
    img = Image.open(path).convert("L")
    arr = np.array(img, dtype=np.float32) / 255.0
    t = torch.from_numpy(arr).unsqueeze(0)  # [1, H, W]
    t = F.interpolate(
        t.unsqueeze(0), size=(imsize, imsize), mode="bilinear", align_corners=False
    ).squeeze(0)
    t = t * 2.0 - 1.0
    if with_depth:
        t = t.unsqueeze(0)  # [1, 1, H, W]
    return t


def preprocess_hyp_batch(image: torch.Tensor) -> torch.Tensor:
    if image.ndim == 3:
        image = image.unsqueeze(0)
    if image.ndim == 4:
        if image.shape[1] != 1:
            image = image.mean(dim=1, keepdim=True)
    elif image.ndim == 5:
        if image.shape[1] != 1:
            image = image.mean(dim=1, keepdim=True)
    else:
        raise ValueError(f"Unexpected image shape {tuple(image.shape)}")
    return image


# ---------------------------------------------------------------------------
# Auto-pick demo images
# ---------------------------------------------------------------------------

def auto_pick_images(
    dataset_root: str,
    labels_csv: str,
    jsonl: str,
    n: int,
    seed: int = 42,
) -> list[str]:
    labels = {}
    with open(labels_csv, encoding="utf-8") as f:
        for row in csv.DictReader(f):
            labels[row["dicom_id"]] = row

    wanted = [
        "normal_only",
        "Pneumonia_pos",
        "Effusion_pos",
        "Cardiomegaly_pos",
        "Edema_pos",
        "Atelectasis_pos",
        "Pneumothorax_pos",
    ]
    found: dict[str, str] = {}
    with open(jsonl, encoding="utf-8") as f:
        for line in f:
            item = json.loads(line)
            did = os.path.splitext(os.path.basename(item["img_path"]))[0]
            lab = labels.get(did)
            if lab is None:
                continue
            for k in wanted:
                if k in found or lab.get(k) != "1":
                    continue
                if k != "normal_only":
                    others = sum(
                        int(lab[x])
                        for x in wanted
                        if x not in ("normal_only", k)
                    )
                    if others > 0:
                        continue
                try:
                    full = resolve_image_path(item["img_path"], dataset_root)
                except FileNotFoundError:
                    continue
                found[k] = full
                break
            if len(found) >= n:
                break

    paths = [found[k] for k in wanted if k in found][:n]
    if len(paths) < n:
        with open(jsonl, encoding="utf-8") as f:
            for line in f:
                item = json.loads(line)
                try:
                    full = resolve_image_path(item["img_path"], dataset_root)
                except FileNotFoundError:
                    continue
                if full not in paths:
                    paths.append(full)
                if len(paths) >= n:
                    break
    return paths


# ---------------------------------------------------------------------------
# Module path helpers (avoid radir / CTViT collisions between trees)
# ---------------------------------------------------------------------------

def _purge_modules(prefixes: Sequence[str]) -> None:
    doomed = [
        k
        for k in list(sys.modules)
        if any(k == p or k.startswith(p + ".") for p in prefixes)
    ]
    for k in doomed:
        del sys.modules[k]


def _push_paths(paths: Sequence[str]) -> list[str]:
    old = list(sys.path)
    for p in reversed(paths):
        if p in sys.path:
            sys.path.remove(p)
        sys.path.insert(0, p)
    return old


# ---------------------------------------------------------------------------
# Hyperbolic (RadIR/RadIR) encode + traverse
# ---------------------------------------------------------------------------

@torch.no_grad()
def run_hyperbolic(
    image_paths: list[str],
    concepts: list[str],
    ckpt: str,
    model_path: str,
    device: str,
    imsize: int = 224,
    steps: int = 50,
    text_batch: int = 64,
    use_entailment: bool = True,
    aperture_thresh: float = 1.2,
) -> tuple[list[list[str]], float]:
    _purge_modules(["radir", "transformer_maskgit", "Rad_IR"])
    rad_ir_root = str(ROOT / "Rad_IR")
    # Inner package lives at ROOT/transformer_maskgit/transformer_maskgit/
    old = _push_paths(
        [str(ROOT), rad_ir_root, str(ROOT / "transformer_maskgit")]
    )

    try:
        from transformers import AutoModel, AutoTokenizer

        from radir import lorentz as L
        from radir.RadIR import RADIR, l2norm
        from transformer_maskgit.ctvit import CTViT

        tokenizer = AutoTokenizer.from_pretrained(
            model_path, trust_remote_code=True, local_files_only=True
        )
        text_model = AutoModel.from_pretrained(
            model_path, trust_remote_code=True, local_files_only=True
        )
        image_encoder = CTViT(
            dim=768,
            codebook_size=8192,
            image_size=imsize,
            patch_size=16,
            temporal_patch_size=10,
            spatial_depth=8,
            temporal_depth=6,
            cls_depth=4,
            dim_head=32,
            heads=8,
            channels=1,
        )
        model = RADIR(
            tokenizer=tokenizer,
            image_encoder=image_encoder,
            text_encoder=text_model,
            dim_text=768,
            dim_image=512,
            dim_latent=512,
            use_mlm=False,
            use_all_token_embeds=False,
        )
        model.load(ckpt)
        model.to(device)
        model.eval()
        curv = model.get_curv().detach()

        txt_feats = []
        for i in range(0, len(concepts), text_batch):
            batch = concepts[i : i + text_batch]
            real = [t for t in batch if t != ROOT_LABEL]
            feats_map: dict[str, torch.Tensor] = {}
            if real:
                tok = tokenizer(
                    real,
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                    max_length=128,
                )
                tok = {k: v.to(device) for k, v in tok.items()}
                if "token_type_ids" not in tok:
                    tok["token_type_ids"] = torch.zeros_like(tok["input_ids"])
                enc = model.text_transformer(
                    input_ids=tok["input_ids"],
                    token_type_ids=tok["token_type_ids"],
                    attention_mask=tok["attention_mask"],
                )[0]
                cls = enc[:, 0, :]
                hyp = model.project_txt(cls)
                for t, h in zip(real, hyp):
                    feats_map[t] = h
            for t in batch:
                if t == ROOT_LABEL:
                    dim = next(iter(feats_map.values())).shape[-1] if feats_map else 768
                    txt_feats.append(torch.zeros(dim, device=device))
                else:
                    txt_feats.append(feats_map[t])
        text_pool = torch.stack(txt_feats, dim=0).float()
        root_idx = concepts.index(ROOT_LABEL) if ROOT_LABEL in concepts else -1
        text_norms = torch.norm(text_pool, dim=-1)
        text_aperture = None
        if use_entailment:
            text_aperture = L.half_aperture(text_pool, curv)
            print(
                f"Entailment filter ON (soft; aperture_thresh={aperture_thresh}, "
                f"no hard radius gate — looser than training for traversal viz)"
            )

        all_paths: list[list[str]] = []
        fallback_count = 0
        total_steps = 0
        for ip in image_paths:
            img = load_image_tensor(ip, imsize=imsize, with_depth=False)
            imgs = preprocess_hyp_batch(img.unsqueeze(0).to(device))
            enc = model.visual_transformer(
                imgs,
                return_encoded_tokens=True,
                modal_embedding=False,
                modal_indexs=None,
                is_condition=False,
            )
            if enc.dim() == 2:
                image_global = l2norm(enc)
            else:
                image_global = l2norm(enc[:, 0, :])
            z0 = model.project_img(image_global)

            u0 = L.log_map0(z0.float(), curv)
            labels: list[str] = []
            ts = torch.linspace(1.0, 0.0, steps, device=device)
            for t in ts:
                u_t = (t * u0).float()
                z_t = L.exp_map0(u_t, curv)
                dist = L.pairwise_dist(z_t, text_pool, curv)[0]
                total_steps += 1

                if use_entailment and text_aperture is not None:
                    # MERU-style cone check, but softened for visualization:
                    # 1) wider aperture_thresh than training (default 1.2 vs 0.7)
                    # 2) soft penalty instead of hard mask (outside-cone texts can
                    #    still win if much closer)
                    # 3) no hard text_norm <= z_norm gate (that collapsed mid-path
                    #    to only [ROOT] in practice)
                    z_exp = z_t.expand_as(text_pool)
                    angle = L.oxy_angle(text_pool, z_exp, curv)
                    energy = torch.clamp(
                        angle - aperture_thresh * text_aperture, min=0.0
                    )
                    if root_idx >= 0:
                        energy[root_idx] = 0.0
                    # soft_weight ~1: energy (radians) added to hyperbolic dist
                    soft_weight = 1.0
                    scores = dist + soft_weight * energy
                    nn_idx = int(scores.argmin().item())
                else:
                    nn_idx = int(dist.argmin().item())
                labels.append(concepts[nn_idx])
            all_paths.append(labels)

        if use_entailment and total_steps > 0:
            print(
                f"Entailment: soft penalty (weight=1.0), "
                f"aperture_thresh={aperture_thresh}"
            )

        curv_val = float(curv.item())
        del model, text_model, image_encoder, text_pool
        return all_paths, curv_val
    finally:
        sys.path[:] = old
        _purge_modules(["radir", "transformer_maskgit", "Rad_IR"])
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


# ---------------------------------------------------------------------------
# Euclidean (RadIR_compare) encode + traverse
# ---------------------------------------------------------------------------

@torch.no_grad()
def run_euclidean(
    image_paths: list[str],
    concepts: list[str],
    ckpt: str,
    model_path: str,
    device: str,
    imsize: int = 480,
    steps: int = 50,
    text_batch: int = 32,
) -> list[list[str]]:
    _purge_modules(["radir", "transformer_maskgit", "Rad_IR", "hf_model_paths"])
    compare_scripts = str(COMPARE_ROOT / "scripts")
    compare_rad = str(COMPARE_ROOT / "Rad_IR")
    old = _push_paths(
        [
            str(COMPARE_ROOT),
            compare_rad,
            compare_scripts,
            str(COMPARE_ROOT / "transformer_maskgit"),
        ]
    )

    try:
        from transformers import BertModel, BertTokenizer

        from radir.RadIR import RADIR, l2norm
        from transformer_maskgit.ctvit import CTViT

        tokenizer = BertTokenizer.from_pretrained(
            model_path, local_files_only=True, do_lower_case=True
        )
        text_encoder = BertModel.from_pretrained(model_path, local_files_only=True)
        image_encoder = CTViT(
            dim=512,
            codebook_size=8192,
            image_size=imsize,
            patch_size=20,
            temporal_patch_size=10,
            spatial_depth=8,
            temporal_depth=6,
            cls_depth=4,
            dim_head=32,
            heads=8,
        )
        model = RADIR(
            image_encoder=image_encoder,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            dim_text=768,
            dim_image=512,
            dim_latent=512,
            use_mlm=False,
            use_all_token_embeds=False,
        )
        model.load(ckpt)
        model.to(device)
        model.eval()

        # Encode non-ROOT concepts only. On the unit sphere, ambient Euclidean
        # distance to the origin always beats weakly-aligned texts, so ROOT must
        # NOT compete via cdist — select it only when the query is near 0.
        non_root = [t for t in concepts if t != ROOT_LABEL]
        txt_feats = []
        for i in range(0, len(non_root), text_batch):
            batch = non_root[i : i + text_batch]
            dummy = torch.zeros(len(batch), 1, 1, imsize, imsize, device=device)
            tok = tokenizer(
                batch,
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=128,
            ).to(device)
            text_latents, _, _, _ = model(
                tok,
                dummy,
                return_latents=True,
                device=device,
                is_condition=False,
                modal_embedding=False,
                modal_indexs=None,
            )
            for h in text_latents:
                txt_feats.append(h if h.ndim == 1 else h.squeeze(0))
        text_pool = l2norm(torch.stack(txt_feats, dim=0))
        # MERU-CLIP protocol: interpolate image → [ROOT] on the sphere.
        # With root = 0, lerp+normalize keeps direction fixed for t>0, so NN
        # is typically sticky (few unique concepts); last steps → [ROOT].
        root_feat = torch.zeros(text_pool.shape[-1], device=device)

        all_paths: list[list[str]] = []
        for ip in image_paths:
            img = load_image_tensor(ip, imsize=imsize, with_depth=True)
            video = img.unsqueeze(0).to(device)
            tok = tokenizer(
                ["chest radiograph"],
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=128,
            ).to(device)
            _, image_latents, _, _ = model(
                tok,
                video,
                return_latents=True,
                device=device,
                is_condition=False,
                modal_embedding=False,
                modal_indexs=None,
            )
            z0 = image_latents
            if z0.ndim == 1:
                z0 = z0.unsqueeze(0)
            z0 = l2norm(z0).squeeze(0)

            # MERU interpolate(): lerp(root, image, w), w: 0→1, then flip
            # so order is image-first → root-last.
            weights = torch.linspace(0.0, 1.0, steps, device=device)
            interp = torch.stack(
                [torch.lerp(root_feat, z0, float(w.item())) for w in weights],
                dim=0,
            )
            interp = F.normalize(interp, dim=-1).flip(0)

            labels: list[str] = []
            for i, z_t in enumerate(interp):
                # Near the root end of the walk, emit [ROOT] (zero vector is
                # unstable under normalize; matches CLIP figure ending at ROOT).
                if i >= int(steps * 0.88):
                    labels.append(ROOT_LABEL)
                    continue
                sims = (z_t.float() @ text_pool.float().T)
                nn_idx = int(sims.argmax().item())
                labels.append(non_root[nn_idx])
            all_paths.append(labels)

        del model, text_encoder, image_encoder, text_pool
        return all_paths
    finally:
        sys.path[:] = old
        _purge_modules(["radir", "transformer_maskgit", "Rad_IR", "hf_model_paths"])
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


# ---------------------------------------------------------------------------
# Path compression (MERU figure style) + min-depth helpers
# ---------------------------------------------------------------------------

def _norm_concept(s: str) -> str:
    """Normalize for near-duplicate detection (case / plural / spacing)."""
    s = re.sub(r"\s+", " ", s.strip().lower())
    toks = []
    for t in s.split():
        if len(t) > 4 and t.endswith("ies"):
            t = t[:-3] + "y"  # abnormalities -> abnormality
        elif len(t) > 3 and t.endswith("s") and not t.endswith(
            ("ss", "ous", "us", "is", "as", "ies")
        ):
            t = t[:-1]
        toks.append(t)
    return " ".join(toks)


def _pick_unused(
    scores: torch.Tensor,
    concepts: Sequence[str],
    used: set[str],
    root_idx: int = -1,
    higher_is_better: bool = False,
) -> int | None:
    """Argmin/argmax among unused non-ROOT concepts."""
    s = scores.clone()
    ban = float("-inf") if higher_is_better else float("inf")
    if root_idx >= 0:
        s[root_idx] = ban
    for i, c in enumerate(concepts):
        if c == ROOT_LABEL or _norm_concept(c) in used:
            s[i] = ban
    if higher_is_better:
        if not torch.isfinite(s).any():
            return None
        return int(s.argmax().item())
    if not torch.isfinite(s).any():
        return None
    return int(s.argmin().item())


def unique_body_from_score_rows(
    score_rows: torch.Tensor,
    concepts: Sequence[str],
    root_idx: int,
    min_body: int,
    higher_is_better: bool = False,
) -> list[str]:
    """
    Build an ordered unique concept list along a traversal.

    1) Greedy: at each time step take nearest unused concept.
    2) If fewer than min_body, re-sample evenly in time (stratified) forcing
       unused picks so the figure can show ≥ min_concepts levels.
    """
    t_steps = int(score_rows.shape[0])
    used: set[str] = set()
    ordered: list[str] = []

    for ti in range(t_steps):
        idx = _pick_unused(
            score_rows[ti],
            concepts,
            used,
            root_idx=root_idx,
            higher_is_better=higher_is_better,
        )
        if idx is None:
            continue
        lab = concepts[idx]
        key = _norm_concept(lab)
        if key in used:
            continue
        used.add(key)
        ordered.append(lab)

    if len(ordered) >= min_body:
        return ordered

    # Stratified fill: force one new concept per evenly spaced geodesic step.
    used = set()
    ordered = []
    # Slightly more candidate slots than needed in case of ties / empties
    n_slots = max(min_body, min(min_body + 2, t_steps))
    slot_ids = np.linspace(0, t_steps - 1, n_slots)
    slot_ids = sorted(set(int(round(i)) for i in slot_ids))
    for ti in slot_ids:
        idx = _pick_unused(
            score_rows[ti],
            concepts,
            used,
            root_idx=root_idx,
            higher_is_better=higher_is_better,
        )
        if idx is None:
            continue
        lab = concepts[idx]
        used.add(_norm_concept(lab))
        ordered.append(lab)
        if len(ordered) >= min_body:
            break

    # Last-resort: dump next-best at the first / mid / last rows
    if len(ordered) < min_body:
        for ti in (0, t_steps // 2, t_steps - 1):
            idx = _pick_unused(
                score_rows[ti],
                concepts,
                used,
                root_idx=root_idx,
                higher_is_better=higher_is_better,
            )
            if idx is None:
                continue
            lab = concepts[idx]
            used.add(_norm_concept(lab))
            ordered.append(lab)
            if len(ordered) >= min_body:
                break

    return ordered


def synthesize_dense_labels(body: list[str], steps: int) -> list[str]:
    """Repeat each body label so compress_traversal can subsample evenly."""
    if not body:
        return [ROOT_LABEL] * max(steps, 1)
    labels: list[str] = []
    per = max(1, steps // len(body))
    for lab in body:
        labels.extend([lab] * per)
    while len(labels) < steps:
        labels.append(body[-1])
    labels = labels[:steps]
    # Reserve last few steps for ROOT (MERU ends at origin)
    n_root = max(1, steps // 10)
    labels[-n_root:] = [ROOT_LABEL] * n_root
    return labels


def compress_traversal(
    labels: list[str],
    keep: int = 6,
) -> list[str]:
    """Dedup consecutive (near-)identical labels; keep up to `keep` steps + [ROOT]."""
    dedup: list[str] = []
    for lab in labels:
        if lab == ROOT_LABEL:
            continue
        if not dedup:
            dedup.append(lab)
            continue
        # exact or near-duplicate of previous (e.g. abnormality / abnormalities)
        if lab == dedup[-1] or _norm_concept(lab) == _norm_concept(dedup[-1]):
            continue
        dedup.append(lab)
    if not dedup:
        return [ROOT_LABEL]

    # keep includes [ROOT]; body slots = keep - 1
    body_keep = max(1, keep - 1)
    if len(dedup) > body_keep:
        idx = np.linspace(0, len(dedup) - 1, body_keep)
        idx = sorted(set(int(round(i)) for i in idx))
        dedup = [dedup[i] for i in idx]
    dedup.append(ROOT_LABEL)
    return dedup


def align_two_paths(
    left: list[str],
    right: list[str],
) -> tuple[list[str], list[str]]:
    def split(p: list[str]) -> tuple[list[str], str]:
        if p and p[-1] == ROOT_LABEL:
            return p[:-1], ROOT_LABEL
        return list(p), ROOT_LABEL

    l_body, _ = split(left)
    r_body, _ = split(right)
    n = max(len(l_body), len(r_body), 1)

    def pad(body: list[str], n: int) -> list[str]:
        if len(body) >= n:
            return body[:n]
        if not body:
            return ["↓"] * n
        return body + ["↓"] * (n - len(body))

    return pad(l_body, n) + [ROOT_LABEL], pad(r_body, n) + [ROOT_LABEL]


# ---------------------------------------------------------------------------
# Visualization
# ---------------------------------------------------------------------------

from matplotlib.font_manager import FontProperties

# Liberation Serif ≈ Times New Roman (metric-compatible on Linux).
_HEADER_FONT = FontProperties(family="Liberation Serif", weight="bold", size=10)
_TITLE_FONT = FontProperties(family="Liberation Serif", weight="normal", size=12.5)


def _cell_color(row: int, n_rows: int) -> tuple[float, float, float]:
    """Single-hue soft slate → paper white (smooth, low-chroma)."""
    if n_rows <= 1:
        t = 1.0
    else:
        t = row / (n_rows - 1)
    r = 0.788 + 0.164 * t
    g = 0.839 + 0.124 * t
    b = 0.886 + 0.086 * t
    return (r, g, b)


def _wrap_label(lab: str, width: int = 24) -> str:
    """Prefer fewer, wider lines so text fills horizontal space and stays in-cell."""
    if lab in ("↓", "▼", ROOT_LABEL) or len(lab) <= width:
        return lab
    lines = textwrap.wrap(
        lab, width=width, break_long_words=False, break_on_hyphens=True
    )
    out = []
    for line in lines:
        if len(line) <= width + 2:
            out.append(line)
        else:
            out.extend(
                textwrap.wrap(
                    line, width=width, break_long_words=True, break_on_hyphens=True
                )
            )
    return "\n".join(out)


def _label_font(lab: str, wrapped: str) -> FontProperties:
    if lab == "↓":
        # Compact arrow, not dominating the cell
        return FontProperties(family="DejaVu Sans", weight="normal", size=6)
    if lab == ROOT_LABEL:
        return FontProperties(
            family="Liberation Serif", style="normal", weight="normal", size=6.0
        )
    n_lines = wrapped.count("\n") + 1
    longest = max(len(line) for line in wrapped.split("\n"))
    if n_lines >= 3:
        size = 5.6
    elif n_lines == 2 or longest >= 22:
        size = 6.2
    else:
        size = 6.8
    return FontProperties(
        family="Liberation Serif", style="italic", weight="normal", size=size
    )


def plot_traversal_panel(
    image_paths: list[str],
    hyp_paths: list[list[str]],
    euc_paths: list[list[str]],
    out_path: str,
    left_name: str = "HMClip",
    right_name: str = "RadIR",
    keep: int = 6,
    keep_right: int | None = None,
    dpi: int = 180,
) -> None:
    n = len(image_paths)
    fig_w = max(4.0 * n, 9.0)
    keep_euc = keep if keep_right is None else keep_right
    aligned = []
    max_rows = 1
    for h, e in zip(hyp_paths, euc_paths):
        a, b = align_two_paths(
            compress_traversal(h, keep=keep),
            compress_traversal(e, keep=keep_euc),
        )
        aligned.append((a, b))
        max_rows = max(max_rows, len(a))

    row_h_in = 0.46  # taller cells so multi-line text has vertical room
    header_h_in = 0.32
    img_h_in = 1.85
    fig_h = img_h_in + header_h_in + row_h_in * max_rows + 0.5
    fig, axes = plt.subplots(1, n, figsize=(fig_w, fig_h))
    if n == 1:
        axes = [axes]

    from mpl_toolkits.axes_grid1.inset_locator import inset_axes

    for ax, ip, (left_labs, right_labs) in zip(axes, image_paths, aligned):
        ax.set_axis_off()
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)

        img = Image.open(ip).convert("L")
        img_frac = img_h_in / fig_h
        img_ax = inset_axes(
            ax,
            width="92%",
            height=f"{img_frac * 100 * 0.92:.1f}%",
            loc="upper center",
            borderpad=0.25,
        )
        img_ax.imshow(img, cmap="gray")
        img_ax.set_xticks([])
        img_ax.set_yticks([])
        for spine in img_ax.spines.values():
            spine.set_visible(True)
            spine.set_linewidth(0.6)
            spine.set_color("0.4")

        n_rows = len(left_labs)
        header_h = header_h_in / fig_h
        row_h = row_h_in / fig_h
        table_top = 1.0 - img_frac - 0.02
        col_w = 0.47
        left_x0, right_x0 = 0.025, 0.505

        for x0, name in ((left_x0, left_name), (right_x0, right_name)):
            ax.add_patch(
                FancyBboxPatch(
                    (x0, table_top - header_h),
                    col_w,
                    header_h,
                    boxstyle="square,pad=0",
                    facecolor="#A8B4C0",
                    edgecolor="#3D4650",
                    linewidth=0.8,
                    transform=ax.transAxes,
                    clip_on=False,
                )
            )
            ax.text(
                x0 + col_w / 2,
                table_top - header_h / 2,
                name,
                ha="center",
                va="center",
                fontproperties=_HEADER_FONT,
                color="#1A1A1A",
                transform=ax.transAxes,
            )

        for r, (ll, rl) in enumerate(zip(left_labs, right_labs)):
            y = table_top - header_h - (r + 1) * row_h
            color = _cell_color(r, n_rows)
            for x0, lab in ((left_x0, ll), (right_x0, rl)):
                ax.add_patch(
                    FancyBboxPatch(
                        (x0, y),
                        col_w,
                        row_h,
                        boxstyle="square,pad=0",
                        facecolor=color,
                        edgecolor="#6B7682",
                        linewidth=0.7,
                        transform=ax.transAxes,
                        clip_on=False,
                    )
                )
                # Use thin ↓ (not heavy ▼); wrap wide to use horizontal space
                if lab == "↓":
                    wrapped = "↓"
                else:
                    wrapped = _wrap_label(lab, width=24)
                fp = _label_font(lab, wrapped)
                # More vertical inset, less horizontal (use cell width)
                inset_x = 0.012
                inset_y = 0.035 * (row_h / (row_h_in / fig_h) if row_h > 0 else 1)
                inset_y = max(0.018, min(0.028, row_h * 0.12))
                txt = ax.text(
                    x0 + col_w / 2,
                    y + row_h / 2,
                    wrapped,
                    ha="center",
                    va="center",
                    fontproperties=fp,
                    color="#1F2933",
                    transform=ax.transAxes,
                    linespacing=0.98,
                    zorder=3,
                )
                clip = FancyBboxPatch(
                    (x0 + inset_x, y + inset_y),
                    col_w - 2 * inset_x,
                    row_h - 2 * inset_y,
                    boxstyle="square,pad=0",
                    transform=ax.transAxes,
                    facecolor="none",
                    edgecolor="none",
                )
                clip.set_visible(False)
                ax.add_patch(clip)
                txt.set_clip_path(clip)

    fig.suptitle(
        "Image traversals with HMClip (hyperbolic) and RadIR (Euclidean)",
        fontproperties=_TITLE_FONT,
        y=0.995,
    )
    os.makedirs(os.path.dirname(os.path.abspath(out_path)) or ".", exist_ok=True)
    fig.savefig(out_path, dpi=dpi, bbox_inches="tight", facecolor="white")
    plt.close(fig)
    print(f"Saved traversal panel to {out_path}")


def save_json_summary(
    out_json: str,
    image_paths: list[str],
    hyp_raw: list[list[str]],
    euc_raw: list[list[str]],
    keep: int = 6,
    keep_right: int | None = None,
) -> None:
    keep_euc = keep if keep_right is None else keep_right
    payload = []
    for ip, h, e in zip(image_paths, hyp_raw, euc_raw):
        hc = compress_traversal(h, keep=keep)
        ec = compress_traversal(e, keep=keep_euc)
        la, ra = align_two_paths(hc, ec)
        payload.append(
            {
                "image": ip,
                "radir_hyp": la,
                "radir_compare": ra,
                "hyp_unique_all": compress_traversal(h, keep=30),
                "compare_unique_all": compress_traversal(e, keep=30),
            }
        )
    os.makedirs(os.path.dirname(os.path.abspath(out_json)) or ".", exist_ok=True)
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)
    print(f"Saved JSON summary to {out_json}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser(description="MERU-style HMClip vs RadIR_compare traversals")
    p.add_argument("--images", nargs="*", default=None, help="Query image paths")
    p.add_argument("--auto_pick", type=int, default=0, help="Auto-pick N diverse CXRs")
    p.add_argument("--device", default="cuda:0")
    p.add_argument("--hyp_ckpt", default=str(ROOT / "outputs" / "radir_best.pt"))
    p.add_argument(
        "--euc_ckpt",
        default=str(
            COMPARE_ROOT / "outputs" / "logs" / "CXRCT_UNI_stage1" / "Rad_IR.12000.pt"
        ),
    )
    p.add_argument(
        "--hyp_bert",
        default=str(
            ROOT / "hf_models" / "microsoft" / "BiomedVLP-CXR-BERT-specialized"
        ),
    )
    p.add_argument(
        "--euc_bert",
        default=str(
            COMPARE_ROOT / "hf_models" / "microsoft" / "BiomedVLP-CXR-BERT-specialized"
        ),
    )
    p.add_argument("--dataset_root", default=str(ROOT / "dataset"))
    p.add_argument("--jsonl", default=str(ROOT / "dataset" / "valid.jsonl"))
    p.add_argument(
        "--labels_csv",
        default=str(ROOT / "outputs" / "tsne_disease_labels.csv"),
    )
    p.add_argument(
        "--concepts_file",
        default=str(ROOT / "scripts" / "traversal_concepts.txt"),
    )
    p.add_argument(
        "--extra_captions",
        type=int,
        default=0,
        help="Optionally add N short first-sentences from valid.jsonl (0 = concepts file only)",
    )
    p.add_argument("--steps", type=int, default=50)
    p.add_argument(
        "--keep",
        type=int,
        default=6,
        help="Max concepts shown per HMClip column including [ROOT] (subsample if longer)",
    )
    p.add_argument("--hyp_imsize", type=int, default=224)
    p.add_argument("--euc_imsize", type=int, default=480)
    p.add_argument(
        "--entailment",
        dest="entailment",
        action="store_true",
        default=True,
        help="Hyperbolic: filter pool by entailment cone before NN (MERU-style, default on)",
    )
    p.add_argument(
        "--no_entailment",
        dest="entailment",
        action="store_false",
        help="Hyperbolic: pure distance NN (no cone filter)",
    )
    p.add_argument(
        "--aperture_thresh",
        type=float,
        default=1.2,
        help="Entailment aperture multiplier (training uses ~0.7; higher = weaker filter)",
    )
    p.add_argument("--out", default=str(ROOT / "outputs" / "traversals_compare.png"))
    p.add_argument(
        "--out_json",
        default=str(ROOT / "outputs" / "traversals_compare.json"),
    )
    p.add_argument("--dpi", type=int, default=160)
    p.add_argument("--seed", type=int, default=42)
    return p.parse_args()


def main():
    args = parse_args()
    device = args.device
    if device.startswith("cuda") and not torch.cuda.is_available():
        print("CUDA not available, falling back to CPU")
        device = "cpu"

    print(f"COMPARE_ROOT = {COMPARE_ROOT}")
    if args.images:
        image_paths = [resolve_image_path(p, args.dataset_root) for p in args.images]
    else:
        n = args.auto_pick if args.auto_pick > 0 else 4
        print(f"Auto-picking {n} diverse CXRs ...")
        image_paths = auto_pick_images(
            args.dataset_root,
            args.labels_csv,
            args.jsonl,
            n=n,
            seed=args.seed,
        )
    print("Query images:")
    for p in image_paths:
        print(f"  - {p}")

    concepts = load_concept_pool(
        args.concepts_file,
        extra_captions=args.extra_captions,
        caption_jsonl=args.jsonl,
        seed=args.seed,
    )
    print(f"Concept pool size: {len(concepts)} (includes {ROOT_LABEL})")

    # Same keep for both sides (natural unique paths; ↓ only if one is shorter).
    keep = args.keep

    print("\n=== Hyperbolic RadIR/RadIR ===")
    hyp_paths, curv = run_hyperbolic(
        image_paths,
        concepts,
        ckpt=args.hyp_ckpt,
        model_path=args.hyp_bert,
        device=device,
        imsize=args.hyp_imsize,
        steps=args.steps,
        use_entailment=args.entailment,
        aperture_thresh=args.aperture_thresh,
    )
    print(f"curvature={curv:.6f}")
    for ip, labs in zip(image_paths, hyp_paths):
        print(
            f"  {os.path.basename(ip)}: "
            f"{' → '.join(compress_traversal(labs, keep=keep))}"
        )

    print("\n=== Euclidean RadIR_compare ===")
    euc_paths = run_euclidean(
        image_paths,
        concepts,
        ckpt=args.euc_ckpt,
        model_path=args.euc_bert,
        device=device,
        imsize=args.euc_imsize,
        steps=args.steps,
    )
    for ip, labs in zip(image_paths, euc_paths):
        print(
            f"  {os.path.basename(ip)}: "
            f"{' → '.join(compress_traversal(labs, keep=keep))}"
        )

    plot_traversal_panel(
        image_paths,
        hyp_paths,
        euc_paths,
        out_path=args.out,
        keep=keep,
        dpi=args.dpi,
    )
    save_json_summary(
        args.out_json,
        image_paths,
        hyp_paths,
        euc_paths,
        keep=keep,
    )


if __name__ == "__main__":
    main()