app_ui.py

from __future__ import annotations

import gradio as gr

from config import (
    DEFAULT_FLIP_X,
    DEFAULT_FLIP_Y,
    DEFAULT_RERANK_ALPHA,
    DEFAULT_RERANK_TOPK,
    DEFAULT_ROTATION_PRESET,
    DINO_MODEL_LABEL,
    DINO_MODEL_URL,
    FIXED_STEP_VOX,
    INDEX_NUM_SLICES,
    INDEX_SEARCH_BACKEND,
    MOUSE_BRAIN_RESOLUTION,
    K_NORMALS,
    K_PER_ANGLE_DEFAULT,
    PATCH_OVERLAP,
    PATCH_SCALES,
    PROJECT_GITHUB_URL,
    RERANKER_MODEL_LABEL,
    RERANKER_MODEL_URL,
    ROTATION_PRESET_CHOICES,
    TOPK_DEFAULT,
)
from env import gpu_decorator
from runtime import RUNTIME
from search_service import run_search


SEARCH_STRATEGY_CHOICES = [
    ("Fast", "fast"),
    ("Smart (recommended)", "smart"),
    ("Enhanced", "enhanced"),
]
DEFAULT_SEARCH_STRATEGY = "smart"

SCALE_CATEGORY_CHOICES = ["Big", "Medium", "Small"]
DEFAULT_SCALE_CATEGORIES = ["Big", "Medium", "Small"]


def _subsection_title(text: str) -> str:
    return (
        "<div style='padding-left: 6px;'>"
        f"<h3 style='margin-top: 4px; margin-bottom: 8px;'>{text}</h3>"
        "</div>"
    )


def _small_note(text: str) -> str:
    return f"<div style='padding-left: 6px; font-size: 0.95em; opacity: 0.9;'>{text}</div>"


def _reranker_visibility(use_reranker: bool):
    return gr.update(visible=bool(use_reranker))


def _search_strategy_help(strategy: str) -> str:
    strategy = str(strategy or DEFAULT_SEARCH_STRATEGY)
    if strategy == "fast":
        return (
            "<div style='padding-left: 6px;'>"
            "<strong>Fast</strong> — Uses only the global query image with rotations and flips. "
            "Best for centered, clean, standard slices. Fastest mode."
            "</div>"
        )
    if strategy == "enhanced":
        return (
            "<div style='padding-left: 6px;'>"
            "<strong>Enhanced</strong> — Forces extra local crops in addition to the global query. "
            "Best for difficult queries, partial views, or when the first results are not satisfying. Slowest but most exhaustive mode."
            "</div>"
        )
    return (
        "<div style='padding-left: 6px;'>"
        "<strong>Smart</strong> — Starts from the global query and automatically adds local crops when the image shape looks difficult "
        "(for example long rectangular crops or partial-brain views). Best default balance."
        "</div>"
    )


def _rotation_help(rotation_preset, flip_x, flip_y, strategy) -> str:
    try:
        n = int(rotation_preset) if rotation_preset is not None else 1
    except (TypeError, ValueError):
        n = 4

    if n not in ROTATION_PRESET_CHOICES:
        n = DEFAULT_ROTATION_PRESET

    flip_multiplier = (2 if bool(flip_x) else 1) * (2 if bool(flip_y) else 1)
    global_variants = int(n) * int(flip_multiplier)

    strategy = str(strategy or DEFAULT_SEARCH_STRATEGY)
    if strategy == "fast":
        tail = "No extra local crops are added."
    elif strategy == "enhanced":
        tail = "The search may also add extra local crops internally, so the total number of evaluated views can be much higher."
    else:
        tail = "The search may automatically add local crops when needed, so the total number of evaluated views can be higher on difficult queries."

    return (
        "<div style='padding-left: 6px;'>"
        f"Base query views from rotations/flips: <strong>{global_variants}</strong>. "
        f"{tail}"
        "</div>"
    )


def _scale_help(selected_categories) -> str:
    selected = set(selected_categories or [])
    if selected == {"Big", "Medium", "Small"} or not selected:
        desc = "All indexed patch sizes are allowed."
    else:
        parts = []
        if "Big" in selected:
            parts.append("Big = whole slices")
        if "Medium" in selected:
            parts.append("Medium = scales 2 and 4")
        if "Small" in selected:
            parts.append("Small = scale 8")
        desc = "; ".join(parts) if parts else "All indexed patch sizes are allowed."

    return (
        "<div style='padding-left: 6px;'>"
        "<strong>Scale filter</strong> — Restrict results by indexed patch size. "
        f"{desc}"
        "</div>"
    )


def _build_about_markdown() -> str:
    return f"""
<div style="padding: 14px 16px;">
  <h3 style="margin-top: 0; margin-bottom: 10px;">About this demo</h3>
  <p style="margin-bottom: 12px;">
    This demo retrieves 2D mouse-brain query slices against a precomputed patch index.
  </p>

  <p style="margin: 0 0 6px 0;"><strong>Current index</strong></p>
  <ul style="margin-top: 0; margin-bottom: 12px;">
    <li>Data: Allen average mouse brain <code>{MOUSE_BRAIN_RESOLUTION}um</code></li>
    <li>Normals: <code>{K_NORMALS}</code></li>
    <li>Patch scales: <code>{PATCH_SCALES}</code></li>
    <li>Fixed step (voxels): <code>{FIXED_STEP_VOX}</code></li>
    <li>Sliding window overlap: <code>{PATCH_OVERLAP}</code></li>
    <li>Indexed slices / candidates: <code>{INDEX_NUM_SLICES}</code></li>
    <li>Index search backend: <code>{INDEX_SEARCH_BACKEND}</code></li>
  </ul>

  <p style="margin: 0 0 6px 0;"><strong>Links</strong></p>
  <ul style="margin-top: 0; margin-bottom: 0;">
    <li>Paper: <em>Coming soon</em></li>
    <li>Code: <a href="{PROJECT_GITHUB_URL}" target="_blank">{PROJECT_GITHUB_URL}</a></li>
  </ul>
</div>
"""


@gpu_decorator(duration=120)
def run_search_entry(
    image,
    search_strategy,
    rotation_preset,
    flip_x,
    flip_y,
    crop_foreground,
    scale_categories,
    top_k,
    k_per_angle,
    use_reranker,
    rerank_topk,
    rerank_alpha,
):
    return run_search(
        image=image,
        top_k=top_k,
        k_per_angle=k_per_angle,
        crop_foreground=crop_foreground,
        use_reranker=use_reranker,
        rerank_topk=rerank_topk,
        rerank_alpha=rerank_alpha,
        rotation_preset=rotation_preset,
        flip_x=flip_x,
        flip_y=flip_y,
        search_strategy=search_strategy,
        scale_categories=scale_categories,
    )


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="MouseBrain 2D→3D Demo") as demo:
        gr.Markdown(
            """
# MouseBrain 2D→3D Demo
Upload a query image, choose how aggressive the search should be, and inspect the retrieved patch matches.
"""
        )

        status = gr.Textbox(
            label="Runtime status",
            value=RUNTIME.status_text(),
            interactive=False,
        )

        with gr.Row():
            with gr.Column(scale=1):
                query_image = gr.Image(type="numpy", label="Query image")
                search_btn = gr.Button("Search", variant="primary")

                gr.Markdown("## Search strategy")
                search_strategy = gr.Radio(
                    choices=SEARCH_STRATEGY_CHOICES,
                    value=DEFAULT_SEARCH_STRATEGY,
                    label="Search aggressiveness",
                    info="How much query expansion should be used.",
                )
                search_strategy_help = gr.Markdown(value=_search_strategy_help(DEFAULT_SEARCH_STRATEGY))

                with gr.Group():
                    gr.Markdown(_subsection_title("Base query transforms"))
                    rotation_help = gr.Markdown(
                        value=_rotation_help(
                            DEFAULT_ROTATION_PRESET,
                            DEFAULT_FLIP_X,
                            DEFAULT_FLIP_Y,
                            DEFAULT_SEARCH_STRATEGY,
                        )
                    )

                    rotation_preset = gr.Radio(
                        choices=ROTATION_PRESET_CHOICES,
                        value=DEFAULT_ROTATION_PRESET,
                        label="Rotation preset",
                        info="Number of evenly spaced base query rotations.",
                    )

                    flip_x = gr.Checkbox(
                        value=DEFAULT_FLIP_X,
                        label="Apply horizontal flip (flip_x)",
                        info="Mirror the query left-right.",
                    )

                    flip_y = gr.Checkbox(
                        value=DEFAULT_FLIP_Y,
                        label="Apply vertical flip (flip_y)",
                        info="Mirror the query top-bottom.",
                    )

                    crop_foreground = gr.Checkbox(
                        value=True,
                        label="Auto-crop foreground",
                        info="Crop empty background before the search pipeline analyzes the query.",
                    )

                gr.Markdown("## Search settings")

                with gr.Group():
                    gr.Markdown(
                        _subsection_title(
                            f"Base search — <a href='{DINO_MODEL_URL}' target='_blank'>{DINO_MODEL_LABEL}</a>"
                        )
                    )
                    gr.Markdown(
                        _small_note(
                            "Fast uses only the global query. Smart may add local crops automatically. Enhanced forces stronger local query expansion."
                        )
                    )

                    scale_help = gr.Markdown(value=_scale_help(DEFAULT_SCALE_CATEGORIES))
                    scale_categories = gr.CheckboxGroup(
                        choices=SCALE_CATEGORY_CHOICES,
                        value=DEFAULT_SCALE_CATEGORIES,
                        label="Slice size filter",
                        info="Choose which indexed patch sizes are allowed in the results.",
                    )

                    top_k = gr.Slider(
                        1,
                        50,
                        value=TOPK_DEFAULT,
                        step=1,
                        label="Top-k",
                        info="Final number of results returned after merge and reranking.",
                    )

                    k_per_angle = gr.Slider(
                        10,
                        500,
                        value=K_PER_ANGLE_DEFAULT,
                        step=1,
                        label="k per base view",
                        info="Number of coarse neighbors fetched for each base query view before merging.",
                    )

                    gr.Markdown(
                        _subsection_title(
                            f"Reranker — <a href='{RERANKER_MODEL_URL}' target='_blank'>{RERANKER_MODEL_LABEL}</a>"
                        )
                    )

                    use_reranker = gr.Checkbox(
                        value=True,
                        label="Use reranker",
                        info="Apply a second-stage neural reranker to improve ordering.",
                    )

                    with gr.Column(visible=True) as reranker_section:
                        rerank_topk = gr.Slider(
                            1,
                            1000,
                            value=DEFAULT_RERANK_TOPK,
                            step=1,
                            label="Rerank top-k",
                            info="How many coarse candidates are sent to the reranker.",
                        )

                        rerank_alpha = gr.Slider(
                            0.0,
                            1.0,
                            value=DEFAULT_RERANK_ALPHA,
                            step=0.05,
                            label="Rerank alpha",
                            info="Blend between coarse score and reranker score. 0 = coarse only, 1 = reranker only.",
                        )

            with gr.Column(scale=2):
                gallery = gr.Gallery(
                    label="Top matches",
                    columns=4,
                    height=600,
                    preview=True,
                )
                results_df = gr.Dataframe(label="Search results")

        search_strategy.change(
            fn=_search_strategy_help,
            inputs=[search_strategy],
            outputs=search_strategy_help,
        )
        search_strategy.change(
            fn=_rotation_help,
            inputs=[rotation_preset, flip_x, flip_y, search_strategy],
            outputs=rotation_help,
        )
        rotation_preset.change(
            fn=_rotation_help,
            inputs=[rotation_preset, flip_x, flip_y, search_strategy],
            outputs=rotation_help,
        )
        flip_x.change(
            fn=_rotation_help,
            inputs=[rotation_preset, flip_x, flip_y, search_strategy],
            outputs=rotation_help,
        )
        flip_y.change(
            fn=_rotation_help,
            inputs=[rotation_preset, flip_x, flip_y, search_strategy],
            outputs=rotation_help,
        )

        scale_categories.change(
            fn=_scale_help,
            inputs=[scale_categories],
            outputs=scale_help,
        )

        use_reranker.change(
            fn=_reranker_visibility,
            inputs=[use_reranker],
            outputs=[reranker_section],
        )

        search_btn.click(
            fn=run_search_entry,
            inputs=[
                query_image,
                search_strategy,
                rotation_preset,
                flip_x,
                flip_y,
                crop_foreground,
                scale_categories,
                top_k,
                k_per_angle,
                use_reranker,
                rerank_topk,
                rerank_alpha,
            ],
            outputs=[gallery, results_df, status],
        )

        with gr.Group():
            gr.Markdown(_build_about_markdown())

    return demo