"""Jolia — zero-shot CT analysis demo.

Upload a chest / abdominal CT volume (NIfTI) and score it against free-text
findings, either against the whole volume (global CLIP head) or routed to a
specific organ query (ParallelOrganCLIP head).

Mirrors `example_zero_shot.py` from the raidium/Jolia repo 1:1: same
preprocessing (`JoliaPreprocessor`), same paired text encoder
(Qwen3-Embedding-8B with last-token pooling), same calibrated logits.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: F401  # must precede torch / any CUDA-touching import

import sys
import time

import gradio as gr
import nibabel as nib
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from transformers import AutoModel

JOLIA_ID = "raidium/Jolia"
TEXT_ID = "Qwen/Qwen3-Embedding-8B"
MAX_PROMPTS = 10
MAX_ORGANS = 16
CACHE_VERSION = "v1"

# The Jolia repo ships its own preprocessing / text-encoder helpers.
_repo = snapshot_download(JOLIA_ID)
if _repo not in sys.path:
    sys.path.insert(0, _repo)

from jolia_windowing import get_available_windows  # noqa: E402
from preprocessing_jolia import JoliaPreprocessor  # noqa: E402
from text_encoder_jolia import JoliaTextEncoder  # noqa: E402

PRE = JoliaPreprocessor()
CT_WINDOWS = get_available_windows("CT")  # channel order of the 11 windowing channels
PREVIEW_WINDOWS = ["auto", "lung", "mediastinum", "abdomen", "liver", "bone", "soft_tissue"]

print("[1/2] Loading Jolia vision backbone ...", flush=True)
JOLIA = AutoModel.from_pretrained(JOLIA_ID, trust_remote_code=True).eval().to("cuda")
print("[2/2] Loading paired text encoder Qwen3-Embedding-8B (~15 GB) ...", flush=True)
TEXT = (
    JoliaTextEncoder.from_pretrained(
        TEXT_ID,
        dtype=torch.bfloat16,
        context_length=JOLIA.config.text_context_length,
    )
    .eval()
    .to("cuda")
)
ORGAN_NAMES = list(JOLIA.organ_slot_names)
print(f"Ready — {len(ORGAN_NAMES)} organ slots available.", flush=True)

DEFAULT_ORGANS = [
    "lungs",
    "pleura",
    "heart",
    "mediastinum",
    "liver",
    "kidneys",
    "spleen",
    "pancreas",
    "spine",
]
DEFAULT_VOLUME_PROMPTS = "\n".join(
    [
        "a normal chest CT",
        "a chest CT showing a pulmonary nodule",
        "a chest CT showing pneumonia",
        "a chest CT showing pleural effusion",
        "a CT showing a liver lesion",
    ]
)
DEFAULT_ORGAN_PROMPTS = "\n".join(["looks normal", "a lesion", "a mass", "an enlarged organ"])


# ----------------------------------------------------------------------------
# CT loading / preprocessing
# ----------------------------------------------------------------------------
def _load_ct(path: str):
    """NIfTI file -> (volume (H, W, D) in HU, resolution (row, col, slice) mm, info)."""
    try:
        img = nib.load(path)
    except Exception as exc:  # noqa: BLE001
        raise gr.Error(f"Could not read this file as NIfTI ({exc}). Convert DICOM with dcm2niix first.")
    try:
        img = nib.as_closest_canonical(img)  # reorient to RAS+
    except Exception:  # noqa: BLE001
        pass

    arr = np.asanyarray(img.dataobj)
    while arr.ndim > 3:
        arr = arr[..., 0]
    if arr.ndim != 3:
        raise gr.Error(f"Expected a 3D volume, got shape {tuple(arr.shape)}.")
    arr = np.nan_to_num(arr.astype(np.float32), nan=-1024.0)
    zooms = [float(z) for z in img.header.get_zooms()[:3]]
    zooms = [z if z > 0 else 1.0 for z in zooms]

    # RAS+ (x->Right, y->Anterior, z->Superior) to the radiological axial layout
    # the checkpoint was trained on: rows anterior->posterior, columns
    # right->left, slices inferior->superior (PrepareVolume then flips depth).
    vol = np.ascontiguousarray(arr.transpose(1, 0, 2)[::-1, ::-1, :])
    resolution = (zooms[1], zooms[0], zooms[2])  # (row, col, slice) mm
    info = {
        "shape": tuple(int(s) for s in arr.shape),
        "spacing": tuple(round(z, 3) for z in zooms),
        "hu_range": (float(np.percentile(vol, 0.5)), float(np.percentile(vol, 99.5))),
        "z_coverage_mm": round(arr.shape[2] * zooms[2], 1),
    }
    return vol, resolution, info


def _auto_window(vol: np.ndarray) -> str:
    """Pick a sensible display window: lung if there is lung parenchyma, else abdomen."""
    lung_frac = float(np.mean((vol > -900.0) & (vol < -500.0)))
    return "lung" if lung_frac > 0.06 else "abdomen"


def _u8(plane: np.ndarray) -> np.ndarray:
    img = (np.clip(plane, 0.0, 1.0) * 255.0).astype(np.uint8)
    return np.repeat(np.repeat(img, 2, axis=0), 2, axis=1)  # 192 -> 384 px


def _preview_tiles(image: torch.Tensor, window: str) -> list:
    """Orthogonal previews of the exact 192**3 cube the model sees."""
    vol = image[CT_WINDOWS.index(window)].float().numpy()
    depth, height, width = vol.shape
    tiles = []
    for frac in (0.3, 0.5, 0.7):
        idx = int(round(frac * (depth - 1)))
        tiles.append((_u8(vol[idx]), f"axial · slice {idx}/{depth - 1}"))
    tiles.append((_u8(vol[:, height // 2, :]), "coronal · mid"))
    tiles.append((_u8(vol[:, :, width // 2]), "sagittal · mid"))
    return tiles


def _prep(path: str, preview_window: str):
    """Load + preprocess a CT and render previews. CPU only."""
    vol, resolution, info = _load_ct(path)
    window = _auto_window(vol) if preview_window == "auto" else preview_window
    image = PRE(vol, resolution=resolution)  # (11, 192, 192, 192) float32
    return image, _preview_tiles(image, window), info, window


def _volume_summary(info: dict, window: str, extra: str = "") -> str:
    sx, sy, sz = info["spacing"]
    lo, hi = info["hu_range"]
    return (
        f"**Volume** {info['shape'][0]}×{info['shape'][1]}×{info['shape'][2]} @ "
        f"{sx}×{sy}×{sz} mm · {info['z_coverage_mm']} mm cranio-caudal coverage · "
        f"HU p0.5–p99.5 {lo:.0f} → {hi:.0f}  \n"
        f"**Model input** 11×192×192×192 (1.5 mm isotropic, centre crop) · preview window `{window}`"
        + (f"  \n{extra}" if extra else "")
    )


def _parse_lines(text: str, limit: int) -> list:
    lines = [ln.strip() for ln in (text or "").splitlines()]
    return [ln for ln in lines if ln][:limit]


# ----------------------------------------------------------------------------
# Inference
# ----------------------------------------------------------------------------
@spaces.GPU(duration=45)
def analyze(
    ct_file: str,
    volume_prompts: str = DEFAULT_VOLUME_PROMPTS,
    organ_prompts: str = DEFAULT_ORGAN_PROMPTS,
    organs: list = DEFAULT_ORGANS,
    preview_window: str = "auto",
):
    """Zero-shot classify a CT volume against free-text findings with Jolia.

    Args:
        ct_file: Path to a chest / abdominal CT volume in NIfTI format (.nii or .nii.gz).
        volume_prompts: Whole-volume prompts, one per line (global CLIP head).
        organ_prompts: Short findings phrases, one per line (per-organ CLIP head).
        organs: Organ query slots to route the findings phrases to.
        preview_window: CT display window for the preview images ("auto" picks lung or abdomen).

    Returns:
        Orthogonal previews of the model input, a volume summary, the global
        zero-shot table and the per-organ probability matrix.
    """
    if not ct_file:
        raise gr.Error("Upload a CT volume (NIfTI .nii / .nii.gz) first.")
    vol_prompts = _parse_lines(volume_prompts, MAX_PROMPTS)
    org_prompts = _parse_lines(organ_prompts, MAX_PROMPTS)
    organs = [o for o in (organs or []) if o in ORGAN_NAMES][:MAX_ORGANS]
    if not vol_prompts and not org_prompts:
        raise gr.Error("Enter at least one prompt.")

    t0 = time.perf_counter()
    image, tiles, info, window = _prep(ct_file, preview_window)
    t_prep = time.perf_counter() - t0

    t0 = time.perf_counter()
    with torch.no_grad():
        x = image.unsqueeze(0).to("cuda")
        cls, organ_queries = JOLIA.forward_with_queries(x)  # (1, 576), (1, slots, 576)
        image_emb = F.normalize(cls.float(), dim=-1, eps=1e-6)

        global_rows = []
        if vol_prompts:
            text_features = TEXT(vol_prompts).to(image_emb.device)  # (N, 4096)
            text_emb = JOLIA.encode_text(text_features)  # (N, 576)
            cosine = (image_emb @ text_emb.t())[0]
            logits = JOLIA.zero_shot_logits(image_emb, text_emb)[0]
            probs = torch.sigmoid(logits)
            global_rows = [
                [p, round(float(lg), 4), round(float(pr), 4), round(float(cs), 4)]
                for p, lg, pr, cs in zip(vol_prompts, logits, probs, cosine)
            ]
            global_rows.sort(key=lambda r: -r[1])

        organ_rows = []
        if org_prompts and organs:
            organ_text = TEXT(org_prompts).to(image_emb.device)
            organ_text_emb = JOLIA.encode_organ_text(organ_text)  # (N, 576)
            for name in organs:
                idx = ORGAN_NAMES.index(name)
                emb = F.normalize(organ_queries[:, idx, :].float(), dim=-1, eps=1e-6)
                scale = JOLIA.organ_logit_scale[idx].float().exp()
                bias = JOLIA.organ_text_bias[idx].float()
                logits = (emb @ organ_text_emb.t())[0] * scale + bias
                organ_rows.append([name] + [round(float(v), 4) for v in torch.sigmoid(logits)])
    t_gpu = time.perf_counter() - t0

    global_df = pd.DataFrame(
        global_rows or [["—", 0.0, 0.0, 0.0]],
        columns=["prompt", "calibrated logit", "match probability", "cosine"],
    )
    organ_df = pd.DataFrame(
        organ_rows or [["—"] + [0.0] * max(1, len(org_prompts))],
        columns=["organ"] + (org_prompts or ["—"]),
    )
    best = f"**Top whole-volume match** · `{global_rows[0][0]}` (p={global_rows[0][2]:.3f})  \n" if global_rows else ""
    summary = _volume_summary(
        info,
        window,
        f"{best}*preprocess {t_prep:.1f}s · encode + score {t_gpu:.1f}s*",
    )
    return tiles, summary, global_df, organ_df


def preview(ct_file: str, preview_window: str = "auto"):
    """Render orthogonal previews of the preprocessed CT volume (no GPU).

    Args:
        ct_file: Path to a CT volume in NIfTI format (.nii or .nii.gz).
        preview_window: CT display window ("auto" picks lung or abdomen).

    Returns:
        Preview images of the 192**3 model input and a short volume summary.
    """
    if not ct_file:
        return [], "Upload a CT volume (NIfTI `.nii` / `.nii.gz`) to get started."
    _, tiles, info, window = _prep(ct_file, preview_window)
    return tiles, _volume_summary(info, window, "*Preview only — press **Analyze** to score prompts.*")


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1250px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Jolia — zero-shot CT") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# Jolia — zero-shot CT analysis\n"
            "[`raidium/Jolia`](https://huggingface.co/raidium/Jolia) is a 3D CT foundation model: it "
            "encodes a whole chest / abdominal CT into one global embedding **and** 102 named "
            "organ-query embeddings, both aligned with report text. Score any free-text finding "
            "against the whole volume, or route it to a single organ.\n\n"
            "⚠️ Research preview — **not a medical device, not for clinical use.**"
        )
        with gr.Row():
            with gr.Column(scale=4):
                ct_file = gr.File(
                    label="CT volume (NIfTI .nii / .nii.gz)",
                    file_types=[".nii", ".gz"],
                    type="filepath",
                )
                volume_prompts = gr.Textbox(
                    label="Whole-volume prompts (one per line)",
                    info="Scored against the global CLIP head — full sentences work best.",
                    value=DEFAULT_VOLUME_PROMPTS,
                    lines=5,
                )
                organ_prompts = gr.Textbox(
                    label="Per-organ findings phrases (one per line)",
                    info="Scored against the per-organ head — short phrases work best.",
                    value=DEFAULT_ORGAN_PROMPTS,
                    lines=4,
                )
                organs = gr.Dropdown(
                    label="Organ query slots",
                    choices=ORGAN_NAMES,
                    value=DEFAULT_ORGANS,
                    multiselect=True,
                    max_choices=MAX_ORGANS,
                )
                run = gr.Button("Analyze", variant="primary")
                with gr.Accordion("Advanced", open=False):
                    preview_window = gr.Dropdown(
                        label="Preview window",
                        choices=PREVIEW_WINDOWS,
                        value="auto",
                        info="Display only — the model always sees all 11 windowing channels.",
                    )
            with gr.Column(scale=6):
                gallery = gr.Gallery(
                    label="Model input (1.5 mm isotropic, 192³ centre crop)",
                    columns=3,
                    height=340,
                    object_fit="contain",
                )
                summary = gr.Markdown("Upload a CT volume (NIfTI `.nii` / `.nii.gz`) to get started.")
                global_df = gr.Dataframe(
                    label="Whole-volume zero-shot (global CLIP head)",
                    headers=["prompt", "calibrated logit", "match probability", "cosine"],
                    wrap=True,
                )
                organ_df = gr.Dataframe(
                    label="Per-organ zero-shot — match probability per (organ, phrase)",
                    wrap=True,
                )

        gr.Markdown(
            "### Examples\n"
            "Public CT volumes from the [TotalSegmentator dataset](https://zenodo.org/records/10047292) "
            "(Wasserthal et al., CC-BY-4.0), via "
            "[`YongchengYAO/TotalSegmentator-CT-Lite`](https://huggingface.co/datasets/YongchengYAO/TotalSegmentator-CT-Lite). "
            "Radiology labels in the file names come from that dataset's metadata."
        )
        gr.Examples(
            examples=[
                [
                    "examples/chest_ct_lung_tumor_s1173.nii.gz",
                    "\n".join(
                        [
                            "a normal chest CT",
                            "a chest CT showing a pulmonary nodule",
                            "a chest CT showing pneumonia",
                            "a chest CT showing pleural effusion",
                            "a chest CT showing emphysema",
                        ]
                    ),
                    "\n".join(["looks normal", "a nodule", "a mass", "an effusion"]),
                ],
                [
                    "examples/chest_ct_inflammation_s1353.nii.gz",
                    "\n".join(
                        [
                            "a normal chest CT",
                            "a chest CT showing pneumonia",
                            "a chest CT showing consolidation",
                            "a chest CT showing a pulmonary nodule",
                        ]
                    ),
                    "\n".join(["looks normal", "consolidation", "an infection", "a nodule"]),
                ],
                [
                    "examples/abdomen_pelvis_ct_normal_s0143.nii.gz",
                    "\n".join(
                        [
                            "a normal abdominal CT",
                            "an abdominal CT showing a liver lesion",
                            "an abdominal CT showing hepatic steatosis",
                            "an abdominal CT showing bowel obstruction",
                        ]
                    ),
                    "\n".join(["looks normal", "a lesion", "an enlarged organ"]),
                ],
                [
                    "examples/abdomen_ct_tumor_s0168.nii.gz",
                    "\n".join(
                        [
                            "a normal abdominal CT",
                            "an abdominal CT showing a tumour",
                            "an abdominal CT showing a liver lesion",
                            "an abdominal CT showing enlarged lymph nodes",
                        ]
                    ),
                    "\n".join(["looks normal", "a lesion", "a mass", "an enlarged organ"]),
                ],
            ],
            inputs=[ct_file, volume_prompts, organ_prompts],
            outputs=[gallery, summary, global_df, organ_df],
            fn=analyze,
            cache_examples=True,
            cache_mode="lazy",
            label=f"Example CT volumes ({CACHE_VERSION})",
        )

        gr.Markdown(
            "Whole-volume scores use Jolia's global CLIP head; per-organ scores route the phrase to "
            "one organ query through the ParallelOrganCLIP head (each organ has its own trained "
            "temperature and bias). Probabilities are `sigmoid(calibrated logit)` — a per-pair "
            '"is this a match?" score, not a softmax over prompts, so they do not sum to 1. '
            "Text is encoded with the paired [`Qwen/Qwen3-Embedding-8B`](https://huggingface.co/Qwen/Qwen3-Embedding-8B) "
            "(last-token pooling, context length 512). DICOM series can be converted with `dcm2niix`."
        )

    ct_file.change(preview, inputs=[ct_file, preview_window], outputs=[gallery, summary])
    preview_window.change(preview, inputs=[ct_file, preview_window], outputs=[gallery, summary])
    run.click(
        analyze,
        inputs=[ct_file, volume_prompts, organ_prompts, organs, preview_window],
        outputs=[gallery, summary, global_df, organ_df],
        api_name="analyze",
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)