"""FleXray: Universal Clinical X-ray Segmentation — interactive Gradio demo.

Loads the flagship FleXray UNet (plus its 4 ensemble siblings for the higher
quality modes) from `VictorButoi/flexray` through the official `flexray` PyPI
package and segments 60 anatomical structures from any clinical radiograph.
"""

import os

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

import spaces  # noqa: E402  (must precede torch)
import colorsys  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from PIL import Image  # noqa: E402

from fxr.inference import FleXraySegmenter  # noqa: E402

MODEL_ID = "VictorButoi/flexray"

# ---------------------------------------------------------------------------
# Model loading (module scope, eager .to("cuda") — ZeroGPU intercepts this)
# ---------------------------------------------------------------------------

_t0 = time.perf_counter()
_flagship = FleXraySegmenter.from_pretrained(MODEL_ID, device="cuda")
_ensemble = FleXraySegmenter.from_pretrained(MODEL_ID, ensemble=True, device="cuda")
print(f"[flexray] loaded flagship + {len(_ensemble.models)} ensemble members "
      f"in {time.perf_counter() - _t0:.1f}s", flush=True)

LABEL_NAMES = [n for n in _flagship.label_names if n != "background"]
MODES = ["Low", "Normal", "High", "X-High"]


# ---------------------------------------------------------------------------
# Label colors — the official FleXray protocol colormap, gaps filled deterministically
# ---------------------------------------------------------------------------

_PROTOCOL_COLORS = {
    "skull": "#8B4EC6", "scapulae": "#00FF00", "clavicles": "#0000FF",
    "humeri": "#FFFF00", "phalanges": "#00FFFF", "metacarpals": "#FF00FF",
    "carpals": "#FF6B6B", "radii": "#0000CD", "ulnae": "#CD853F",
    "tibiae": "#D2B48C", "femurs": "#CC66A0", "patellae": "#FF0000",
    "fibulae": "#008B8B", "metatarsals": "#2E8B57", "tarsals": "#4D96FF",
    "toes": "#6A5ACD",
    "rib_1": "#FFE0B2", "rib_2": "#FFCC80", "rib_3": "#FFB74D",
    "rib_4": "#FFA726", "rib_5": "#FF9800", "rib_6": "#FB8C00",
    "rib_7": "#F57C00", "rib_8": "#EF6C00", "rib_9": "#E65100",
    "rib_10": "#BF360C", "rib_11": "#991D0A", "rib_12": "#7D0C08",
    "sternum": "#B5179E",
    "vertebra_c1": "#BBDEFB", "vertebra_c2": "#A9D2F6", "vertebra_c3": "#97C6F1",
    "vertebra_c4": "#85BAEC", "vertebra_c5": "#73AEE7", "vertebra_c6": "#61A2E2",
    "vertebra_c7": "#4F96DD", "vertebra_t1": "#3D8AD8", "vertebra_t2": "#2B7ED3",
    "vertebra_t3": "#1976D2", "vertebra_t4": "#176DC4", "vertebra_t5": "#1565C0",
    "vertebra_t6": "#145DB2", "vertebra_t7": "#1255A4", "vertebra_t8": "#104D96",
    "vertebra_t9": "#0E4588", "vertebra_t10": "#0C3D7A", "vertebra_t11": "#0A356C",
    "vertebra_t12": "#082D5E",
    "vertebra_l1": "#6D28D9", "vertebra_l2": "#8B5CF6", "vertebra_l3": "#A855F7",
    "vertebra_l4": "#D946EF", "vertebra_l5": "#EC4899",
    "hips": "#90EE90", "sacrum": "#ADFF2F",
    "kidneys": "#2A9D8F", "liver": "#800000", "spleen": "#556B2F",
    "lungs": "#32CD32", "heart": "#F4A460",
}


def _fallback_color(index: int) -> str:
    """Deterministic, well-spread color for labels outside the official map."""
    hue = (index * 0.61803398875) % 1.0  # golden-ratio spacing
    r, g, b = colorsys.hls_to_rgb(hue, 0.55, 0.75)
    return "#%02X%02X%02X" % (int(r * 255), int(g * 255), int(b * 255))


LABEL_COLORS = {
    name: _PROTOCOL_COLORS.get(name, _fallback_color(i))
    for i, name in enumerate(LABEL_NAMES)
}
LABEL_HEX = [LABEL_COLORS[n][1:] for n in LABEL_NAMES]
_CHANNELS = {name: i + 1 for i, name in enumerate(LABEL_NAMES)}  # model channel idx (background=0)


def _hex_to_rgb(h: str) -> tuple[int, int, int]:
    return (int(h[1:3], 16), int(h[3:5], 16), int(h[5:7], 16))


LABEL_RGB = np.array([_hex_to_rgb(LABEL_COLORS[n]) for n in LABEL_NAMES], dtype=np.float32)
CHANNEL_TO_SLOT = np.array([_CHANNELS[name] for name in LABEL_NAMES], dtype=np.int64)


# ---------------------------------------------------------------------------
# Rendering
# ---------------------------------------------------------------------------

def _render_overlay(pil_image: Image.Image, masks: np.ndarray, selected: list[str]) -> Image.Image:
    """Blend a color-coded mask over the original image at original resolution.

    masks: uint8 array shaped (61, H, W) — channel 0 is background.
    selected: list of label names to draw; empty list draws all.
    """

    width, height = pil_image.size
    base_rgb = np.repeat(
        np.asarray(pil_image.convert("L").resize((width, height)), dtype=np.float32)[:, :, None],
        3,
        axis=2,
    )

    present = masks[CHANNEL_TO_SLOT].astype(bool)  # (60, h, w) ordered like LABEL_NAMES
    keep = np.ones(len(LABEL_NAMES), dtype=bool)
    if selected:
        keep = np.array([name in selected for name in LABEL_NAMES], dtype=bool)
    present = present & keep[:, None, None]

    # Class-index map: later labels overwrite earlier ones on overlap.
    h, w = masks.shape[-2:]
    idx = np.zeros((h, w), dtype=np.int64)
    for i in range(present.shape[0]):
        if keep[i]:
            idx[present[i]] = i + 1

    # One RGB color per class index (0 = background, never drawn).
    lookup = np.zeros((len(LABEL_NAMES) + 1, 3), dtype=np.float32)
    lookup[1:] = LABEL_RGB
    rgb = np.zeros((h, w, 3), dtype=np.float32)
    nonzero = idx > 0
    rgb[nonzero] = lookup[idx[nonzero]]

    # Nearest-neighbor upscale keeps color regions crisp; alpha comes from class presence.
    overlay_np = np.asarray(
        Image.fromarray(rgb.astype(np.uint8), mode="RGB").resize(
            (width, height), resample=Image.Resampling.NEAREST
        ),
        dtype=np.float32,
    )
    alpha_img = np.asarray(
        Image.fromarray((nonzero * 255).astype(np.uint8)).resize(
            (width, height), resample=Image.Resampling.NEAREST
        ),
        dtype=np.float32,
    ) / 255.0
    alpha3 = alpha_img[:, :, None]

    out = overlay_np * (0.55 * alpha3) + base_rgb * (1.0 - 0.55 * alpha3)
    return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8), mode="RGB")


def _build_selected_list(show_all: bool, label_selection: list[str]) -> list[str]:
    """Resolve the checkbox-group selection into a concrete label list."""
    if show_all or not label_selection:
        return list(LABEL_NAMES)
    return [n for n in label_selection if n in LABEL_NAMES]


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------

@spaces.GPU(duration=25)
def segment(
    image: Image.Image,
    mode: str = "Low",
    show_all_labels: bool = True,
    label_selection: list[str] | None = None,
    threshold: float = 0.5,
) -> tuple[Image.Image, str, str]:
    """Segment anatomical structures in a clinical X-ray.

    Args:
        image: Grayscale or color X-ray image (PNG/JPG; DICOM works via the CLI).
        mode: Quality mode — Low (flagship, 1 pass), Normal (flagship, 8-pass TTA),
            High (5-model ensemble), X-High (ensemble + 8-pass TTA).
        show_all_labels: When True, draw every predicted structure.
        label_selection: Specific structures to draw when show_all_labels is False.
        threshold: Probability threshold for a structure to count as present.

    Returns:
        (color overlay, detected-structure summary, timing note).
    """
    if image is None:
        raise gr.Error("Please upload an X-ray image first.")
    start = time.perf_counter()
    selected = _build_selected_list(show_all_labels, label_selection or [])

    segmenter = _ensemble if mode in ("High", "X-High") else _flagship
    tta_samples = 8 if mode in ("Normal", "X-High") else 1

    prediction = segmenter.predict(image, threshold=float(threshold), tta_samples=tta_samples)
    probs = prediction.probabilities[0].detach().cpu().numpy()  # (61, 256, 256)
    masks = (probs >= float(threshold)).astype(np.uint8)

    elapsed = time.perf_counter() - start

    overlay = _render_overlay(image, masks, selected)

    found = []
    for i, name in enumerate(LABEL_NAMES):
        frac = float(masks[_CHANNELS[name]].mean())
        if frac > 0.001:
            found.append((frac, name, i))
    found.sort(reverse=True)
    if found:
        summary = "<br>".join(
            f'<span style="display:inline-block;width:12px;height:12px;'
            f'background:#{LABEL_HEX[i]};margin-right:6px;border-radius:3px;"></span>'
            f"{name} — {frac * 100:.1f}% of image"
            for frac, name, i in found[:12]
        )
        if len(found) > 12:
            summary += f"<br>… and {len(found) - 12} more"
    else:
        summary = "No structures detected above threshold — try another image or lower the threshold."

    note = f"mode: {mode} · {len(segmenter.models)} model(s) · TTA×{tta_samples} · {elapsed:.1f}s on GPU"
    return overlay, summary, note


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------

CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.legend { display: flex; flex-wrap: wrap; gap: 6px 14px; font-size: 12px; }
.legend span { display: inline-flex; align-items: center; gap: 5px; }
.legend i { display: inline-block; width: 10px; height: 10px; border-radius: 2px; }
"""

legend_html = '<div class="legend">' + "".join(
    f'<span><i style="background:#{LABEL_HEX[i]}"></i>{name}</span>'
    for i, name in enumerate(LABEL_NAMES)
) + "</div>"


def _change_mode(mode: str):
    """Placeholder kept for API stability; currently unused."""
    return gr.update()


with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            # 🦴 FleXray: Universal Clinical X-ray Segmentation
            One 2D UNet that segments **60 anatomical structures** from any clinical
            radiograph — chest, spine, pelvis, hands, knees, and more.
            [Paper](https://arxiv.org/abs/2609.26756) ·
            [Code](https://github.com/VictorButoi/FleXray) ·
            [Weights](https://huggingface.co/VictorButoi/flexray) ·
            [Website](https://flexray.csail.mit.edu/)
            """
        )

        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(
                    label="X-ray image", type="pil", image_mode="L", height=420
                )
                mode = gr.Radio(
                    MODES, value="Low", label="Quality mode",
                    info="Low: flagship · Normal: +8-pass TTA · High: 5-model ensemble · X-High: both",
                )
                threshold = gr.Slider(
                    0.1, 0.9, value=0.5, step=0.05, label="Probability threshold"
                )
                run_btn = gr.Button("Segment", variant="primary")
                with gr.Accordion("Focus on specific structures", open=False):
                    show_all = gr.Checkbox(value=True, label="Show all structures")
                    label_sel = gr.CheckboxGroup(
                        LABEL_NAMES, label="Structures to display (when 'show all' is off)"
                    )
            with gr.Column(scale=1):
                overlay_out = gr.Image(
                    label="Segmentation overlay", type="pil", format="png", height=420
                )
                stats_out = gr.HTML(label="Detected structures")
                note_out = gr.Textbox(label="Run info", interactive=False)

        with gr.Accordion("Label legend (60 structures)", open=False):
            gr.HTML(legend_html)

        gr.Examples(
            examples=[["examples/chest_lateral.png"], ["examples/pelvis_ap.png"]],
            inputs=[image_in],
            outputs=[overlay_out, stats_out, note_out],
            fn=segment,
            cache_examples=True,
            cache_mode="lazy",
            label="Try a sample X-ray (radiographs courtesy of Radiopaedia.org, used with permission)",
        )

        gr.Markdown(
            """
            ⚠️ **Research demo** — not a medical device. Weights are CC-BY-NC-4.0.
            """
        )

    run_btn.click(
        segment,
        inputs=[image_in, mode, show_all, label_sel, threshold],
        outputs=[overlay_out, stats_out, note_out],
        api_name="segment",
    )

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