"""Neuron Quantification using AI — MedCLIPSeg variant.

Applies the MedCLIPSeg (CVPR 2026) vision-language approach: a text prompt
describes the target ("nerve fibers ..."), a frozen CLIP backbone produces a
fiber-probability map, and that map is traced into a skeleton. The app reports a
white-on-black **tracked skeleton** and its **total trace length** in microns.

Zero-shot general CLIP is only a coarse prior on this domain — train the bundled
model (reference_medclipseg/) on a GPU with masks for real segmentation quality.
"""
import traceback
import gradio as gr

import processing as P
import medclipseg as MC


def analyze(file_obj, fg_prompt, threshold):
    if file_obj is None:
        return None, None, "Upload a CZI or TIFF z-stack."
    try:
        img = P.load_image(file_obj)
        nf, _ = P.guess_channels(img)
        nf_mip = P.channel_preview(img.data[nf])

        if fg_prompt and fg_prompt.strip():
            MC.FG_PROMPTS = [p.strip() for p in fg_prompt.split("|") if p.strip()]

        # --- MedCLIP text-prompted fiber-probability map -> tracked skeleton ---
        prob, _ = MC.segment_best(nf_mip)
        trace = P.trace_from_probability(prob, img.voxel,
                                         threshold=float(threshold), prune_um=3.0)
        m = P.compute_metrics(trace, "MedCLIP", min_fiber_um=5.0)
        skel_img = P.skeleton_image(trace.skeleton, dilate=1)   # white on black

        model = ("few-shot trained decoder" if MC.has_trained_model()
                 else "zero-shot CLIP prior")
        status = (f"**MedCLIP total trace length: "
                  f"{m.total_length_um:,.1f} µm.**\n\n"
                  f"Tracked from the MedCLIP fiber-probability map "
                  f"({model}, text-prompted) at threshold {float(threshold):.2f}. "
                  f"Prompt(s): *{', '.join(MC.FG_PROMPTS)}*.")
        return nf_mip, skel_img, status
    except Exception as e:  # noqa: BLE001
        return None, None, f"Error:\n{e}\n{traceback.format_exc()}"


HEADER_HTML = """
<div style="text-align:center; margin: 0.2rem 0 0.7rem;">
  <h1 style="font-size:2.1rem; font-weight:750; letter-spacing:-0.01em;
             margin:0 0 0.4rem;">Neuron Quantification using AI</h1>
  <div style="font-weight:700; font-size:1.15rem; color:#1a1a1a;">
    Iman Sabir Ezzat, Randa K Ismail, Ayden Chavez, Marisa Zallocchi, PhD, Steven Fernandes, PhD
  </div>
  <div style="font-weight:600; font-size:0.95rem; color:#4b5563; margin-top:0.25rem;">
    MedCLIPSeg variant — text-prompted vision-language fiber tracking
  </div>
</div>
"""

THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.slate, secondary_hue=gr.themes.colors.slate,
    neutral_hue=gr.themes.colors.gray,
    font=["system-ui", "-apple-system", "Segoe UI", "Roboto", "sans-serif"],
).set(
    body_background_fill="#ffffff", body_text_color="#1a1a1a",
    background_fill_primary="#ffffff", background_fill_secondary="#f7f7f8",
    block_background_fill="#ffffff", block_border_color="#e5e7eb",
    block_label_text_color="#1a1a1a", block_title_text_color="#1a1a1a",
    border_color_primary="#e5e7eb", button_primary_background_fill="#1f2937",
    button_primary_text_color="#ffffff", input_background_fill="#ffffff",
    input_border_color="#c0c5cc",
    body_background_fill_dark="#ffffff", body_text_color_dark="#1a1a1a",
    background_fill_primary_dark="#ffffff", background_fill_secondary_dark="#f7f7f8",
    block_background_fill_dark="#ffffff", block_border_color_dark="#e5e7eb",
    block_label_text_color_dark="#1a1a1a", block_title_text_color_dark="#1a1a1a",
    panel_background_fill_dark="#ffffff", border_color_primary_dark="#e5e7eb",
    button_primary_background_fill_dark="#1f2937",
    button_primary_text_color_dark="#ffffff", input_background_fill_dark="#ffffff",
)
CSS = """
.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }
:root, .dark {
    color-scheme: light; --body-background-fill:#ffffff;
    --background-fill-primary:#ffffff; --block-background-fill:#ffffff;
    --body-text-color:#1a1a1a; --block-label-text-color:#1a1a1a;
    --block-title-text-color:#1a1a1a; --border-color-primary:#e5e7eb;
    --input-background-fill:#ffffff; --neutral-950:#1a1a1a;
}
body, gradio-app, .gradio-container, .dark { background:#ffffff !important; color:#1a1a1a !important; }
"""

with gr.Blocks(title="Neuron Quantification using AI — MedCLIPSeg",
               theme=THEME, css=CSS) as demo:
    gr.HTML(HEADER_HTML)
    with gr.Row():
        with gr.Column(scale=1):
            file_in = gr.File(label="Neurofilament z-stack (.czi / .tif)",
                              type="filepath")
            prompt = gr.Textbox(
                label="Text prompt(s) for the target (separate with | )",
                value="a fluorescence microscopy image of nerve fibers | "
                      "neurofilament nerve fibers and axons")
            thresh = gr.Slider(0.1, 0.9, value=0.5, step=0.05,
                               label="Fiber-probability threshold")
            btn = gr.Button("Analyze", variant="primary")
        with gr.Column(scale=2):
            status = gr.Markdown()
            with gr.Row():
                out_orig = gr.Image(label="Neurofilament (MIP)", height=280)
                out_skel = gr.Image(label="Tracked skeleton (white on black)",
                                    height=280)
    btn.click(analyze, [file_in, prompt, thresh],
              [out_orig, out_skel, status])

if __name__ == "__main__":
    demo.launch()