"""
Gradio app for CLIP Progressive Steering Pipeline.
Entry point for Hugging Face Spaces (HF looks for app.py).
"""

import time
import gradio as gr
import numpy as np

from backend2 import (
    load_clip_model,
    load_or_compute_embeddings,
    get_raw_embeddings,
    get_text_embedding,
    generate_feedback_with_weights,
    linear_steering,
    subspace_steering,
    energy_based_steering,
    weighted_energy_steering,
    switch_dataset,
    get_available_datasets,
    get_image_path,
    DATASETS,
)
from sae_backend import sae_prf_steering
from study_utils import (
    STUDY_QUERIES,
    NUM_QUERIES,
    MAX_ROUNDS,
    register_participant,
    participant_exists,
    validate_email,
    log_interaction,
    log_image_annotations,
    log_method_comparison,
    log_survey_responses,
    log_final_selections,
    log_final_survey,
    firestore_batch_add,
    _iso_ts,
)

EMPTY_ATTR = {"attribute": "", "weight": 0.7}


def _study_error(msg: str) -> str:
    """Format a user-facing error for study phases (Markdown)."""
    return f"**⚠️ {msg}**" if msg else ""


def _study_warning(msg: str) -> str:
    """Format a non-blocking warning (e.g. save issue)."""
    return f"*Note: {msg}*" if msg else ""


# ── helpers ──────────────────────────────────────────────────────────────────

def _dataset_choices() -> list[str]:
    info = get_available_datasets()
    choices = [k for k, m in info.items() if m["has_embeddings"] or m["has_images"]]
    return choices if choices else list(DATASETS.keys())


def _build_gallery(indices, sims, img_names, baseline_names=None):
    gallery = []
    for rank, idx in enumerate(indices, 1):
        img_path = get_image_path(img_names[idx])
        if img_path.exists():
            caption = f"#{rank}  |  {sims[idx] * 100:.1f}%"
            if baseline_names and img_names[idx] not in baseline_names:
                caption += "  | NEW"
            gallery.append((str(img_path), caption))
    return gallery


def _format_feedback_md(feedback: dict) -> str:
    lines = [
        f"**Alpha** = {feedback.get('alpha', 0.4)}  |  "
        f"**Beta** = {feedback.get('beta', 0.4)}\n",
    ]
    pos = feedback.get("positive", [])
    neg = feedback.get("negative", [])
    if pos:
        lines.append(
            "**Positive:** "
            + ", ".join(f"{a['attribute']} ({a.get('weight', 0.7)})" for a in pos)
        )
    if neg:
        lines.append(
            "**Negative:** "
            + ", ".join(f"{a['attribute']} ({a.get('weight', 0.7)})" for a in neg)
        )
    return "\n\n".join(lines)


def _parse_attr_comma(s: str) -> list:
    """Parse comma-separated attributes into list of {attribute, weight} dicts."""
    if not s or not str(s).strip():
        return []
    return [{"attribute": x.strip(), "weight": 0.7} for x in str(s).split(",") if x.strip()]


def study_run_query(query_text: str, dataset_key: str, pos_list: list, neg_list: list, alpha: float, beta: float, use_llm: bool):
    """
    Run one study query: baseline + linear (with LLM or user attributes).
    pos_list/neg_list: list of attribute strings (no weights).
    Returns: (baseline_gallery, linear_gallery, feedback_md, baseline_ids, linear_ids, attributes_str, pos_attrs, neg_attrs).
    """
    top_k = 5
    alpha = float(alpha) if alpha else 0.4
    beta = float(beta) if beta else 0.4
    try:
        switch_dataset(dataset_key)
    except Exception as e:
        return [], [], f"Dataset error: {e}", [], [], "", [], []
    embeddings, img_names = load_or_compute_embeddings()
    query_emb = get_text_embedding(query_text)
    base_sims = embeddings @ query_emb
    base_idx = np.argsort(base_sims)[::-1][:top_k]
    baseline_ids = [img_names[i] for i in base_idx]
    gal_baseline = _build_gallery(base_idx, base_sims, img_names)

    if use_llm:
        feedback = generate_feedback_with_weights(query_text)
        feedback["alpha"] = alpha
        feedback["beta"] = beta
    else:
        pos_dicts = [{"attribute": a, "weight": 1.0} for a in pos_list if a.strip()] if pos_list else []
        neg_dicts = [{"attribute": a, "weight": 1.0} for a in neg_list if a.strip()] if neg_list else []
        if not pos_dicts:
            pos_dicts = [{"attribute": query_text, "weight": 1.0}]
        feedback = {
            "positive": pos_dicts,
            "negative": neg_dicts,
            "alpha": alpha,
            "beta": beta,
        }
    llm_emb = linear_steering(
        query_emb, feedback["positive"], feedback["negative"], alpha, beta
    )
    llm_sims = embeddings @ llm_emb
    llm_idx = np.argsort(llm_sims)[::-1][:top_k]
    linear_ids = [img_names[i] for i in llm_idx]
    gal_linear = _build_gallery(llm_idx, llm_sims, img_names)
    pos_attrs = [a["attribute"] for a in feedback.get("positive", []) if a.get("attribute", "").strip()]
    neg_attrs = [a["attribute"] for a in feedback.get("negative", []) if a.get("attribute", "").strip()]
    pos_str = ",".join(pos_attrs)
    neg_str = ",".join(neg_attrs)
    attributes_used = f"pos:{pos_str}|neg:{neg_str}|a:{alpha}|b:{beta}"
    feedback_md = f"**Alpha** = {alpha}  |  **Beta** = {beta}\n\n"
    if pos_attrs:
        feedback_md += "**Positive:** " + ", ".join(pos_attrs) + "\n\n"
    if neg_attrs:
        feedback_md += "**Negative:** " + ", ".join(neg_attrs)
    return gal_baseline, gal_linear, feedback_md, baseline_ids, linear_ids, attributes_used, pos_attrs, neg_attrs


# ── core search function ────────────────────────────────────────────────────

def progressive_search(query, dataset, top_k, alpha, beta, pos_attrs, neg_attrs, last_query):
    """
    pos_attrs / neg_attrs come from gr.State — list of dicts.
    last_query tracks what query generated the current attributes.
    Returns: 6 galleries + feedback_md + pos_state + neg_state + last_query_state
    """
    empty = []
    default = [dict(EMPTY_ATTR)]

    if not query.strip():
        return tuple([empty] * 6 + ["Please enter a search query.", default, default, ""])

    try:
        switch_dataset(dataset)
    except Exception as exc:
        return tuple([empty] * 6 + [f"Dataset error: {exc}", default, default, ""])

    embeddings, img_names = load_or_compute_embeddings()

    query_changed = query.strip() != (last_query or "").strip()

    # Collect non-empty attributes from state
    user_pos = [a for a in pos_attrs if a.get("attribute", "").strip()]
    user_neg = [a for a in neg_attrs if a.get("attribute", "").strip()]

    if (user_pos or user_neg) and not query_changed:
        # Same query, user has custom attributes → keep them
        feedback = {
            "positive": user_pos,
            "negative": user_neg,
            "alpha": alpha,
            "beta": beta,
        }
    else:
        # New query or no attributes → LLM generates fresh ones
        feedback = generate_feedback_with_weights(query)
        feedback["alpha"] = alpha
        feedback["beta"] = beta

    top_k = max(1, min(10, int(top_k)))
    query_emb = get_text_embedding(query)

    # ── Stage 1: Baseline CLIP ──
    base_sims = embeddings @ query_emb
    base_idx = np.argsort(base_sims)[::-1][:top_k]
    baseline_names = {img_names[i] for i in base_idx}
    gal_baseline = _build_gallery(base_idx, base_sims, img_names)

    # ── Stage 2: LLM Linear Steering ──
    llm_emb = linear_steering(
        query_emb, feedback["positive"], feedback["negative"], alpha, beta
    )
    llm_sims = embeddings @ llm_emb
    llm_idx = np.argsort(llm_sims)[::-1][:top_k]
    gal_llm = _build_gallery(llm_idx, llm_sims, img_names, baseline_names)

    # ── Stage 3: Contrastive Subspace ──
    sub_emb = subspace_steering(
        query_emb, feedback["positive"], feedback["negative"]
    )
    sub_sims = embeddings @ sub_emb
    sub_idx = np.argsort(sub_sims)[::-1][:top_k]
    gal_sub = _build_gallery(sub_idx, sub_sims, img_names, baseline_names)

    # ── Stage 4: Energy-Based ──
    eng_emb = energy_based_steering(
        query_emb, feedback["positive"], feedback["negative"]
    )
    eng_sims = embeddings @ eng_emb
    eng_idx = np.argsort(eng_sims)[::-1][:top_k]
    gal_energy = _build_gallery(eng_idx, eng_sims, img_names, baseline_names)

    # ── Stage 5: Per-Concept Weighted ──
    w_emb = weighted_energy_steering(
        query_emb, feedback["positive"], feedback["negative"]
    )
    w_sims = embeddings @ w_emb
    w_idx = np.argsort(w_sims)[::-1][:top_k]
    gal_weighted = _build_gallery(w_idx, w_sims, img_names, baseline_names)

    # ── Stage 6: SAE PRF Steering ──
    try:
        raw_embs = get_raw_embeddings()
        # SAE expects un-normalised CLIP embeddings (matches training distribution)
        query_emb_raw = get_text_embedding(query, normalize=False)
        # Get more baseline indices for PRF (up to 20)
        prf_base_idx = np.argsort(base_sims)[::-1][:max(top_k, 20)]
        sae_emb = sae_prf_steering(
            query_emb_raw, raw_embs, prf_base_idx,
            prf_k=10, top_feats=32, scale=1.0,
        )
        sae_sims = embeddings @ sae_emb
        sae_idx = np.argsort(sae_sims)[::-1][:top_k]
        gal_sae = _build_gallery(sae_idx, sae_sims, img_names, baseline_names)
    except Exception as exc:
        print(f"⚠️ SAE steering error: {exc}")
        gal_sae = []

    feedback_md = _format_feedback_md(feedback)

    # Update states with current feedback attributes
    pos_out = feedback.get("positive", [])
    neg_out = feedback.get("negative", [])
    if not pos_out:
        pos_out = [dict(EMPTY_ATTR)]
    if not neg_out:
        neg_out = [dict(EMPTY_ATTR)]

    return (
        gal_baseline, gal_llm, gal_sub, gal_energy, gal_weighted, gal_sae,
        feedback_md,
        pos_out, neg_out, query.strip(),
    )


# ── Gradio UI ────────────────────────────────────────────────────────────────

STAGE_INFO = {
    "baseline": ("1. Baseline CLIP", "Pure cosine similarity \u2014 no steering"),
    "llm": ("2. LLM Linear Steering", "q\u2032 = q + \u03b1\u00b7pos \u2212 \u03b2\u00b7neg"),
    "subspace": ("3. Contrastive Subspace", "Centroid-based steering direction"),
    "energy": ("4. Energy-Based", "Gradient-descent optimisation in embedding space"),
    "weighted": ("5. Per-Concept Weighted", "Normalised per-attribute weight steering"),
    "sae_prf": ("6. SAE PRF Steering", "Pseudo-relevance feedback in sparse autoencoder latent space"),
}

available = _dataset_choices()
default_ds = available[0] if available else "flickr"

with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="purple"),
    title="CLIP Progressive Steering Pipeline",
    css="""
        .stage-title { margin-bottom: 0 !important; }
        footer { display: none !important; }
        .gallery-container { overflow: visible !important; }
        .grid-wrap { overflow-y: visible !important; max-height: none !important; }
        .grid-container { max-height: none !important; }
        .thumbnail-item .caption-label { font-size: 0.85em !important; padding: 2px 4px !important; }
        /* Floating Back-to-Top button */
        #back-to-top-btn {
            position: fixed; bottom: 24px; right: 24px; z-index: 9999;
            width: 48px; height: 48px; border-radius: 50%;
            background: #6366f1; color: #fff; border: none;
            font-size: 22px; line-height: 48px; text-align: center;
            cursor: pointer; box-shadow: 0 3px 10px rgba(0,0,0,0.25);
            opacity: 0; pointer-events: none;
            transition: opacity 0.25s ease;
        }
        #back-to-top-btn.show { opacity: 1; pointer-events: auto; }
        #back-to-top-btn:hover { background: #4f46e5; }
    """,
) as demo:

    # Inject layout-fix + back-to-top JS via gr.HTML (visible=True so the script runs)
    gr.HTML(value="""<div id="study-js-anchor" style="display:none"></div>
    <script>
    (function(){
        if(window.__studyFixLoaded) return;
        window.__studyFixLoaded = true;

        /* ── Back-to-top floating button ── */
        var b = document.createElement('button');
        b.id = 'back-to-top-btn';
        b.title = 'Back to top';
        b.innerHTML = '&#8593;';
        b.addEventListener('click', function(){
            window.scrollTo({top:0, behavior:'smooth'});
            document.documentElement.scrollTop = 0;
        });
        document.body.appendChild(b);

        function _checkScroll(){
            var y = window.scrollY || window.pageYOffset || document.documentElement.scrollTop || 0;
            b.classList.toggle('show', y > 300);
        }
        window.addEventListener('scroll', _checkScroll, {passive:true});
        document.addEventListener('scroll', _checkScroll, {passive:true});

        /* ── Layout fixer ── */
        var PHASE_IDS = ['study_phase0','study_phase1','study_phase1_success',
                         'study_phase2_anchor','study_phase3','study_thanks'];

        window.__fixLayout = function(){
            PHASE_IDS.forEach(function(id){
                var el = document.getElementById(id);
                if(!el) return;
                var t = el;
                /* Walk up through wrappers that Gradio may add around the elem_id div */
                for(var i=0;i<3;i++){
                    if(t.parentElement && t.parentElement !== document.body){
                        var p = t.parentElement;
                        if(p.children.length === 1) t = p;
                        else break;
                    }
                }
                var isHidden = (el.style.display === 'none' || t.style.display === 'none' ||
                                el.offsetParent === null);
                if(isHidden){
                    [el, t].forEach(function(n){
                        n.style.setProperty('height','0','important');
                        n.style.setProperty('max-height','0','important');
                        n.style.setProperty('overflow','hidden','important');
                        n.style.setProperty('padding','0','important');
                        n.style.setProperty('margin','0','important');
                        n.style.setProperty('border','none','important');
                    });
                } else {
                    [el, t].forEach(function(n){
                        ['height','max-height','overflow','padding','margin','border'].forEach(function(p){
                            n.style.removeProperty(p);
                        });
                    });
                }
            });
            /* Force every possible ancestor to shrink-wrap */
            document.querySelectorAll(
                '.gradio-container, .gradio-container > div, ' +
                '[role=\"tabpanel\"], .contain, .main, .wrap, .tabitem, ' +
                '.tab-content, .tabs > div'
            ).forEach(function(c){
                c.style.setProperty('min-height','0','important');
                c.style.setProperty('height','auto','important');
            });
        };

        /* ── Scroll helper ── */
        window.__scrollTop = function(targetId){
            /* Scroll every possible container to top */
            window.scrollTo(0,0);
            document.documentElement.scrollTop = 0;
            document.body.scrollTop = 0;
            var gc = document.querySelector('.gradio-container');
            if(gc){ gc.scrollTop = 0; }
            /* Walk all ancestors of gradio-container */
            var p = gc;
            while(p){ if(p.scrollTop > 0) p.scrollTop = 0; p = p.parentElement; }
            /* scrollIntoView on the visible target — this works even inside iframes */
            if(targetId){
                var te = document.getElementById(targetId);
                if(te && te.offsetParent !== null){
                    te.scrollIntoView({block:'start', behavior:'instant'});
                }
            }
        };

        /* ── Combined fix+scroll ── */
        window.__fixAndScroll = function(targetId){
            window.__fixLayout();
            window.__scrollTop(targetId);
        };

        /* MutationObserver with debounce */
        var _t = null;
        new MutationObserver(function(){
            if(_t) clearTimeout(_t);
            _t = setTimeout(window.__fixLayout, 80);
        }).observe(document.body, {
            childList:true, subtree:true,
            attributes:true, attributeFilter:['style','class']
        });

        /* Periodic fix during initial load */
        [200,500,1000,2000,4000].forEach(function(d){
            setTimeout(window.__fixLayout, d);
        });
    })();
    </script>""", visible=True)

    # ── Header with User Study entry ──
    with gr.Row():
        gr.Markdown(
            "# CLIP Progressive Steering Pipeline\n"
            "Compare 6 embedding-steering methods side-by-side on CLIP ViT-B/16 image retrieval."
        )
        gr.Markdown("")  # spacer
        study_tab_btn = gr.Button("👉 User Study", size="sm", variant="secondary")

    main_tabs = gr.Tabs(selected="retrieval_demo_tab")
    with main_tabs:
        # ─── Tab 1: Retrieval Demo ───
        with gr.TabItem("Retrieval Demo", id="retrieval_demo_tab"):
            # ── Controls ──
            with gr.Row():
                query_input = gr.Textbox(
                    label="Search Query",
                    placeholder="e.g. a guilty dog, a happy person wearing sunglasses",
                    value="a guilty dog",
                    scale=3,
                )
                dataset_dd = gr.Radio(
                    choices=available,
                    value=default_ds,
                    label="Dataset",
                    scale=1,
                )
                top_k_slider = gr.Slider(
                    minimum=1, maximum=10, value=5, step=1, label="Top K", scale=1,
                )

            with gr.Row():
                alpha_slider = gr.Slider(
                    minimum=0.1, maximum=0.8, value=0.4, step=0.05,
                    label="Alpha (\u03b1) \u2014 positive steering strength",
                )
                beta_slider = gr.Slider(
                    minimum=0.1, maximum=0.8, value=0.4, step=0.05,
                    label="Beta (\u03b2) \u2014 negative steering strength",
                )

            # ── Dynamic attribute section ──
            gr.Markdown(
                "### Steering Attributes\n"
                "Use **+** to add attributes, **\u2715** to remove, and drag sliders to set weights. "
                "Leave all empty to let the LLM auto-generate on the next search."
            )

            pos_state = gr.State([dict(EMPTY_ATTR)])
            neg_state = gr.State([dict(EMPTY_ATTR)])
            last_query_state = gr.State("")

            with gr.Row(equal_height=True):
                # ────── Positive column ──────
                with gr.Column():
                    gr.Markdown("**Positive** (steer toward)")

                    @gr.render(inputs=[pos_state])
                    def render_pos(attrs):
                        for i, attr in enumerate(attrs):
                            with gr.Row():
                                t = gr.Textbox(
                                    value=attr["attribute"],
                                    show_label=False,
                                    placeholder=f"positive attribute {i + 1}",
                                    scale=3,
                                    max_lines=1,
                                )
                                s = gr.Slider(
                                    minimum=0.0, maximum=1.0,
                                    value=attr["weight"], step=0.05,
                                    show_label=False, scale=2,
                                )
                                rm = gr.Button(
                                    "\u2715", variant="stop", size="sm",
                                    scale=0, min_width=40,
                                )

                                # Save text when user clicks away / tabs out
                                def _pos_blur(val, state, idx=i):
                                    ns = [dict(a) for a in state]
                                    if idx < len(ns):
                                        ns[idx]["attribute"] = val
                                    return ns
                                t.blur(_pos_blur, [t, pos_state], [pos_state])

                                # Save weight when slider is released
                                def _pos_release(val, state, idx=i):
                                    ns = [dict(a) for a in state]
                                    if idx < len(ns):
                                        ns[idx]["weight"] = val
                                    return ns
                                s.release(_pos_release, [s, pos_state], [pos_state])

                                # Remove this attribute
                                def _pos_remove(state, idx=i):
                                    ns = [a for j, a in enumerate(state) if j != idx]
                                    return ns if ns else [dict(EMPTY_ATTR)]
                                rm.click(_pos_remove, [pos_state], [pos_state])

                    add_pos_btn = gr.Button("+ Add Positive", size="sm")

                    def _add_pos(state):
                        return state + [dict(EMPTY_ATTR)]
                    add_pos_btn.click(_add_pos, [pos_state], [pos_state])

                # ────── Negative column ──────
                with gr.Column():
                    gr.Markdown("**Negative** (steer away from)")

                    @gr.render(inputs=[neg_state])
                    def render_neg(attrs):
                        for i, attr in enumerate(attrs):
                            with gr.Row():
                                t = gr.Textbox(
                                    value=attr["attribute"],
                                    show_label=False,
                                    placeholder=f"negative attribute {i + 1}",
                                    scale=3,
                                    max_lines=1,
                                )
                                s = gr.Slider(
                                    minimum=0.0, maximum=1.0,
                                    value=attr["weight"], step=0.05,
                                    show_label=False, scale=2,
                                )
                                rm = gr.Button(
                                    "\u2715", variant="stop", size="sm",
                                    scale=0, min_width=40,
                                )

                                def _neg_blur(val, state, idx=i):
                                    ns = [dict(a) for a in state]
                                    if idx < len(ns):
                                        ns[idx]["attribute"] = val
                                    return ns
                                t.blur(_neg_blur, [t, neg_state], [neg_state])

                                def _neg_release(val, state, idx=i):
                                    ns = [dict(a) for a in state]
                                    if idx < len(ns):
                                        ns[idx]["weight"] = val
                                    return ns
                                s.release(_neg_release, [s, neg_state], [neg_state])

                                def _neg_remove(state, idx=i):
                                    ns = [a for j, a in enumerate(state) if j != idx]
                                    return ns if ns else [dict(EMPTY_ATTR)]
                                rm.click(_neg_remove, [neg_state], [neg_state])

                    add_neg_btn = gr.Button("+ Add Negative", size="sm")

                    def _add_neg(state):
                        return state + [dict(EMPTY_ATTR)]
                    add_neg_btn.click(_add_neg, [neg_state], [neg_state])

            # Clear all
            def _clear_all():
                return [dict(EMPTY_ATTR)], [dict(EMPTY_ATTR)]

            clear_btn = gr.Button("Clear All Attributes", variant="secondary", size="sm")
            clear_btn.click(_clear_all, [], [pos_state, neg_state])

            # Search button
            search_btn = gr.Button(
                "Run Progressive Comparison", variant="primary", size="lg",
            )

            # ── Feedback summary ──
            feedback_md = gr.Markdown(value="*Run a search to see feedback*")

            # ── Stage galleries ──
            galleries = {}
            for key, (title, desc) in STAGE_INFO.items():
                gr.Markdown(f"### {title}\n_{desc}_", elem_classes=["stage-title"])
                galleries[key] = gr.Gallery(
                    label=title,
                    columns=5,
                    rows=2,
                    height=600,
                    object_fit="contain",
                    show_label=False,
                )

            # ── Wiring ──
            all_inputs = [
                query_input, dataset_dd, top_k_slider,
                alpha_slider, beta_slider,
                pos_state, neg_state, last_query_state,
            ]
            all_outputs = [
                galleries["baseline"],
                galleries["llm"],
                galleries["subspace"],
                galleries["energy"],
                galleries["weighted"],
                galleries["sae_prf"],
                feedback_md,
                pos_state,
                neg_state,
                last_query_state,
            ]

            search_btn.click(fn=progressive_search, inputs=all_inputs, outputs=all_outputs)

        # ─── Tab 2: User Study ───
        with gr.TabItem("👉 User Study", id="user_study_tab"):
            gr.Markdown("## User Study\nComplete the steps below in order. Progress is saved as you go.")

            # Phase 0: How to participate
            study_phase0 = gr.Column(visible=True, elem_id="study_phase0")
            with study_phase0:
                study_accordion = gr.Accordion("How to Participate", open=True)
                with study_accordion:
                    gr.Markdown(
                        "You will complete **22 image retrieval tasks**.\n\n"
                        "**Important:** Please complete this study **on your own** — we need your personal opinion, not anyone else's.\n\n"
                        "---\n"
                        "#### How each query works\n\n"
                        "For each query you will see two sets of 5 images:\n"
                        "- **Baseline (CLIP)** — top 5 results from plain CLIP similarity (top of the page). These **never change** when you edit attributes.\n"
                        "- **Linear Steering** — top 5 results after steering (below). **Only this section updates** when you refine attributes.\n\n"
                        "---\n"
                        "#### Attributes and refinement\n\n"
                        "When a query first loads, the system uses an **LLM to auto-generate** a set of positive and negative attributes. "
                        "You are free to **edit, add, or remove** them.\n\n"
                        "- **Positive attributes** steer the results *toward* that concept.\n"
                        "- **Negative attributes** steer the results *away* from that concept.\n\n"
                        "You can refine up to **3 rounds** per query. After 3 rounds, only \"Satisfied\" is available.\n\n"
                        "---\n"
                        "#### Alpha (α) and Beta (β)\n\n"
                        "- **Alpha (α)** controls how strongly positive attributes pull the results (default 0.4).\n"
                        "- **Beta (β)** controls how strongly negative attributes push the results away (default 0.4).\n"
                        "- You can adjust these sliders before clicking *Apply Refinement*.\n\n"
                        "---\n"
                        "#### Example walk-through\n\n"
                        "Suppose the query is **\"a cozy living room\"**:\n\n"
                        "1. The LLM might suggest positive = *warm lighting, soft furniture* and negative = *cluttered, dark*.\n"
                        "2. You look at the Linear Steering results. They look warm but too modern.\n"
                        "3. **Round 2:** You remove *soft furniture*, add *rustic* to positive, and add *modern* to negative. Click **Apply Refinement**.\n"
                        "4. The Linear Steering images update — now they look more rustic. Baseline stays the same.\n"
                        "5. **Round 3:** You tweak alpha to 0.6 for stronger pull. Click **Apply Refinement** one last time.\n"
                        "6. Happy with the results → click **Satisfied**.\n\n"
                        "---\n"
                        "#### After each query\n\n"
                        "- You will label whether each of the 10 retrieved images matches your intended meaning (Yes / No).\n"
                        "- You will answer a short comparison question and four rating questions.\n"
                        "- Then you move to the next query.\n\n"
                        "**There are no correct answers.** We are studying how people interpret subjective concepts.\n\n"
                        "**Progress:** Query 1 / 22 → Query 22 / 22  \n"
                        "**Estimated time:** 20–30 minutes."
                    )
                continue_phase0_btn = gr.Button("Continue to Participant Form", variant="primary")

            # Phase 1: Participant form
            study_phase1 = gr.Column(visible=False, elem_id="study_phase1")
            with study_phase1:
                gr.Markdown("### Participant information")
                study_email = gr.Textbox(label="Email (required)", placeholder="you@example.com")
                study_gender = gr.Radio(
                    choices=["Male", "Female", "Non-binary", "Prefer not to say"],
                    value=None,
                    label="Gender",
                )
                study_age = gr.Radio(
                    choices=["18–24", "25–34", "35–44", "45+"],
                    value=None,
                    label="Age range",
                )
                study_consent = gr.Checkbox(label="I consent to participate (required)", value=False)
                study_form_msg = gr.Markdown(value="")
                submit_participant_btn = gr.Button("Submit & Continue to Query 1", variant="primary")

            # After successful registration: show "Continue to Query 1"
            study_phase1_success = gr.Column(visible=False, elem_id="study_phase1_success")
            with study_phase1_success:
                gr.Markdown("**Registered.** Click below to start the first query.")
                continue_to_query1_btn = gr.Button("Continue to Query 1", variant="primary")

            # Phase 2: Query loop — state
            study_query_idx = gr.State(0)
            study_round = gr.State(0)
            study_query_start_time = gr.State(0.0)
            study_baseline_ids = gr.State([])
            study_linear_ids = gr.State([])
            study_attributes_used = gr.State("")
            study_round_satisfied = gr.State(-1)
            study_time_elapsed = gr.State(0.0)

            # Phase 2: Query loop UI
            study_phase2 = gr.Column(visible=False, elem_id="study_phase2_anchor")
            with study_phase2:
                study_progress_md = gr.Markdown(value="Query 1 / 22")
                study_round_md = gr.Markdown(value="Round 1 / 3")
                study_query_text_md = gr.Markdown(value="")
                gr.Markdown("### 🔵 Baseline (CLIP) Results")
                study_baseline_gallery = gr.Gallery(
                    label="Baseline", columns=5, rows=1, height=350, object_fit="contain", show_label=False,
                )
                gr.Markdown("### 🟢 Linear Steering Results *(only this section changes when you refine)*")
                study_linear_gallery = gr.Gallery(
                    label="Linear", columns=5, rows=1, height=350, object_fit="contain", show_label=False,
                )
                gr.Markdown("---")
                gr.Markdown("**Steering Attributes** — Add or remove attributes, then click *Apply Refinement*.")

                study_pos_state = gr.State([""])
                study_neg_state = gr.State([""])

                with gr.Row(equal_height=True):
                    with gr.Column():
                        gr.Markdown("**Positive** (steer toward)")

                        @gr.render(inputs=[study_pos_state])
                        def render_study_pos(attrs):
                            for i, attr in enumerate(attrs):
                                with gr.Row():
                                    t = gr.Textbox(
                                        value=attr,
                                        show_label=False,
                                        placeholder=f"positive attribute {i + 1}",
                                        scale=3,
                                        max_lines=1,
                                    )
                                    rm = gr.Button(
                                        "\u2715", variant="stop", size="sm",
                                        scale=0, min_width=40,
                                    )

                                    def _sp_blur(val, state, idx=i):
                                        ns = list(state)
                                        if idx < len(ns):
                                            ns[idx] = val
                                        return ns
                                    t.blur(_sp_blur, [t, study_pos_state], [study_pos_state])

                                    def _sp_remove(state, idx=i):
                                        ns = [a for j, a in enumerate(state) if j != idx]
                                        return ns if ns else [""]
                                    rm.click(_sp_remove, [study_pos_state], [study_pos_state])

                        study_add_pos_btn = gr.Button("+ Add Positive", size="sm")
                        def _sp_add(state):
                            return state + [""]
                        study_add_pos_btn.click(_sp_add, [study_pos_state], [study_pos_state])

                    with gr.Column():
                        gr.Markdown("**Negative** (steer away from)")

                        @gr.render(inputs=[study_neg_state])
                        def render_study_neg(attrs):
                            for i, attr in enumerate(attrs):
                                with gr.Row():
                                    t = gr.Textbox(
                                        value=attr,
                                        show_label=False,
                                        placeholder=f"negative attribute {i + 1}",
                                        scale=3,
                                        max_lines=1,
                                    )
                                    rm = gr.Button(
                                        "\u2715", variant="stop", size="sm",
                                        scale=0, min_width=40,
                                    )

                                    def _sn_blur(val, state, idx=i):
                                        ns = list(state)
                                        if idx < len(ns):
                                            ns[idx] = val
                                        return ns
                                    t.blur(_sn_blur, [t, study_neg_state], [study_neg_state])

                                    def _sn_remove(state, idx=i):
                                        ns = [a for j, a in enumerate(state) if j != idx]
                                        return ns if ns else [""]
                                    rm.click(_sn_remove, [study_neg_state], [study_neg_state])

                        study_add_neg_btn = gr.Button("+ Add Negative", size="sm")
                        def _sn_add(state):
                            return state + [""]
                        study_add_neg_btn.click(_sn_add, [study_neg_state], [study_neg_state])
                with gr.Row():
                    study_alpha = gr.Slider(0.1, 0.8, value=0.4, step=0.05, label="Alpha (α) — positive steering strength")
                    study_beta = gr.Slider(0.1, 0.8, value=0.4, step=0.05, label="Beta (β) — negative steering strength")
                study_feedback_md = gr.Markdown(value="")
                with gr.Row():
                    apply_refinement_btn = gr.Button("Apply Refinement", variant="secondary")
                    satisfied_btn = gr.Button("Satisfied", variant="primary")

            def _make_scroll_js(target_id=None):
                """Generate JS that fixes layout + scrolls to top (or to a specific phase element)."""
                tid = f"'{target_id}'" if target_id else "null"
                return f"""() => {{
                    function _go(){{
                        if(window.__fixAndScroll) window.__fixAndScroll({tid});
                        else {{
                            window.scrollTo(0,0);
                            document.documentElement.scrollTop=0;
                            document.body.scrollTop=0;
                            if({tid}){{
                                var e=document.getElementById({tid});
                                if(e) e.scrollIntoView({{block:'start'}});
                            }}
                        }}
                    }}
                    _go();
                    setTimeout(_go, 150);
                    setTimeout(_go, 400);
                    setTimeout(_go, 800);
                }}"""
            def _to_phase1():
                return gr.update(visible=False), gr.update(visible=True)
            continue_phase0_btn.click(
                fn=_to_phase1, inputs=[], outputs=[study_phase0, study_phase1],
            ).then(
                fn=lambda: None, inputs=[], outputs=[], js=_make_scroll_js("study_phase1"),
            )

            def _submit_participant(email, gender, age, consent):
                email = (email or "").strip()
                if not email:
                    return gr.update(visible=True), gr.update(visible=False), gr.update(value=_study_error("Please enter your email address.")), None
                if not consent:
                    return gr.update(visible=True), gr.update(visible=False), gr.update(value=_study_error("You must check the consent box to participate.")), None
                if not validate_email(email):
                    return gr.update(visible=True), gr.update(visible=False), gr.update(value=_study_error("Please enter a valid email address (e.g. name@example.com).")), None
                try:
                    if participant_exists(email):
                        return gr.update(visible=True), gr.update(visible=False), gr.update(value=_study_error("This email is already registered. Use a different one or contact the researchers.")), None
                except Exception as e:
                    return gr.update(visible=True), gr.update(visible=False), gr.update(value=_study_error(f"Could not verify participant: {str(e)}. Please try again.")), None
                try:
                    ok, err_msg = register_participant(email, (gender or "").strip(), (age or "").strip())
                    if not ok:
                        return gr.update(visible=True), gr.update(visible=False), gr.update(value=_study_error(err_msg or "Registration failed.")), None
                except Exception as e:
                    return gr.update(visible=True), gr.update(visible=False), gr.update(value=_study_error(f"Registration could not be saved: {str(e)}. Please try again.")), None
                return gr.update(visible=False), gr.update(visible=True), gr.update(value=""), email.lower()

            study_participant_id = gr.State(value=None)
            # NOTE: do NOT add js= here — on some browsers the scroll
            # fires before Gradio captures the textbox value, causing
            # the email to arrive as empty.  Scroll happens in .then().
            submit_participant_btn.click(
                fn=_submit_participant,
                inputs=[study_email, study_gender, study_age, study_consent],
                outputs=[study_phase1, study_phase1_success, study_form_msg, study_participant_id],
            ).then(
                fn=lambda: None, inputs=[], outputs=[], js=_make_scroll_js("study_phase1_success"),
            )

            def _run_study_query(qidx, rnd, pid, start_time, pos_list, neg_list, alpha, beta, use_llm):
                if qidx is None or qidx < 0 or qidx >= NUM_QUERIES:
                    return [], [], _study_error("Invalid query index."), [], [], "", [], [], 0, 0, 0.0
                qtext, dkey = STUDY_QUERIES[qidx]
                try:
                    gal_b, gal_l, feedback_md, base_ids, linear_ids, attrs_used, pos_attrs, neg_attrs = study_run_query(
                        qtext, dkey, pos_list or [], neg_list or [], alpha, beta, use_llm
                    )
                    base_ids = list(base_ids or [])
                    linear_ids = list(linear_ids or [])
                except Exception as e:
                    gal_b, gal_l, base_ids, linear_ids = [], [], [], []
                    err_msg = str(e).strip() or "Query failed."
                    feedback_md = _study_error(f"Retrieval error: {err_msg}. Please try again or click Satisfied to continue.")
                    attrs_used, pos_attrs, neg_attrs = "", [], []
                elapsed = (time.time() - start_time) if start_time else 0.0
                progress = f"Query {qidx + 1} / {NUM_QUERIES}"
                round_label = f"Round {rnd + 1} / {MAX_ROUNDS}"
                query_display = f"## {qtext}"
                return (
                    gal_b, gal_l, feedback_md, base_ids, linear_ids, attrs_used, pos_attrs, neg_attrs,
                    qidx, rnd, progress, round_label, query_display, elapsed,
                )

            def _start_first_query(pid):
                start = time.time()
                res = _run_study_query(0, 0, pid, start, [], [], 0.4, 0.4, True)
                gal_b, gal_l, feedback_md, base_ids, linear_ids, attrs_used, pos_attrs, neg_attrs, qidx, rnd, progress, round_label, query_display, elapsed = res
                pos_out = pos_attrs if pos_attrs else [""]
                neg_out = neg_attrs if neg_attrs else [""]
                return (
                    gr.update(visible=False), gr.update(visible=True),
                    gr.update(value=progress), gr.update(value=round_label), gr.update(value=query_display),
                    gal_b, gal_l, gr.update(value=feedback_md),
                    pos_out, neg_out,
                    gr.update(value=0.4), gr.update(value=0.4),
                    gr.update(interactive=True), gr.update(visible=True),
                    0, 0, start, base_ids, linear_ids, attrs_used, -1, elapsed,
                )

            def _apply_refinement(pid, qidx, rnd, start_time, base_ids, linear_ids, attrs_used, pos_state_val, neg_state_val, alpha, beta):
                if qidx is None or qidx < 0 or qidx >= NUM_QUERIES:
                    return (gr.update(), gr.update(), gr.update(), [], [], gr.update(), gr.update(), gr.update(),
                            gr.update(), gr.update(), gr.update(), gr.update(),
                            qidx, rnd, start_time, base_ids or [], linear_ids or [], attrs_used, -1, 0.0)
                if rnd >= MAX_ROUNDS:
                    return (gr.update(), gr.update(value=_study_error("Maximum 3 rounds reached. Click Satisfied to continue.")), gr.update(), [], [], gr.update(), gr.update(), gr.update(),
                            gr.update(), gr.update(), gr.update(), gr.update(visible=False),
                            qidx, rnd, start_time, base_ids or [], linear_ids or [], attrs_used, -1, 0.0)
                next_round = rnd + 1
                pos_list = [s.strip() for s in (pos_state_val or []) if s.strip()]
                neg_list = [s.strip() for s in (neg_state_val or []) if s.strip()]
                res = _run_study_query(qidx, next_round, pid, start_time, pos_list, neg_list, alpha, beta, False)
                gal_b, gal_l, feedback_md, new_base, new_linear, attrs_used_new, pos_attrs, neg_attrs, _qi, _rn, progress, round_label, query_display, elapsed = res
                # After 3 applies (next_round >= MAX_ROUNDS), hide apply button
                can_refine = next_round < MAX_ROUNDS
                pos_out = pos_attrs if pos_attrs else [""]
                neg_out = neg_attrs if neg_attrs else [""]
                return (
                    gr.update(value=progress), gr.update(value=round_label), gr.update(value=query_display),
                    gal_b, gal_l, gr.update(value=feedback_md),
                    pos_out, neg_out,
                    gr.update(value=alpha), gr.update(value=beta),
                    gr.update(interactive=True), gr.update(visible=can_refine),
                    qidx, next_round, start_time, new_base, new_linear, attrs_used_new, -1, elapsed,
                )

            def _on_satisfied(pid, qidx, rnd, start_time, base_ids, linear_ids, attrs_used):
                warning = ""
                if pid is None or qidx is None:
                    return (gr.update(), gr.update(), gr.update(), [], [], gr.update(), gr.update(), gr.update(),
                            gr.update(), gr.update(), gr.update(), gr.update(),
                            qidx, rnd, start_time, base_ids or [], linear_ids or [], attrs_used, rnd if rnd is not None else 0, 0.0), _study_error("Session state missing; interaction was not logged.")
                elapsed = time.time() - start_time
                base_str = ",".join(base_ids) if base_ids else ""
                linear_str = ",".join(linear_ids) if linear_ids else ""
                try:
                    log_interaction(pid, qidx, STUDY_QUERIES[qidx][0], "baseline", rnd, attrs_used, base_str, elapsed, 1)
                    log_interaction(pid, qidx, STUDY_QUERIES[qidx][0], "linear", rnd, attrs_used, linear_str, elapsed, 1)
                except Exception as e:
                    warning = _study_warning(f"Interaction log could not be saved: {str(e)}. Please contact the researchers.")
                return (
                    gr.update(), gr.update(), gr.update(), [], [], gr.update(), gr.update(), gr.update(),
                    gr.update(), gr.update(), gr.update(), gr.update(),
                    qidx, rnd, start_time, base_ids or [], linear_ids or [], attrs_used, rnd, elapsed,
                ), warning

            # ── Phase 3: ALL annotations on one scrollable page ────────
            #
            # After every query is completed the user sees ALL queries'
            # images on a single page and annotates them.  Each query
            # has its OWN unique radio/slider components — nothing is
            # reused or reset, so Gradio state-tracking bugs are avoided.
            study_all_images = gr.State({})   # {qidx: [10 paths]}
            study_all_meta   = gr.State({})   # {qidx: {base_ids, linear_ids, …}}

            study_phase3 = gr.Column(visible=False, elem_id="study_phase3")
            _p3_images  = []   # flat: NUM_QUERIES × 10 Image components
            _p3_radios  = []   # flat: NUM_QUERIES × 10 Radio components
            _p3_comps   = []   # NUM_QUERIES comparison Radios
            _p3_likerts = []   # 4 global Likert sliders (asked once, not per-query)
            _p3_done_btns = []   # NUM_QUERIES "Done with Query X" buttons
            _p3_update_btns = [] # NUM_QUERIES "Update Query X" buttons
            _p3_status_mds = []  # NUM_QUERIES status Markdown components
            study_saved_queries = gr.State({})  # {qidx: True} tracks saved queries
            with study_phase3:
                gr.Markdown(
                    "# Image Annotations\n\n"
                    "For each query, annotate all images then click "
                    "**Done with Query X** to save. You can update later if needed.\n\n"
                    "After all queries are saved, fill in the **Final Survey** "
                    "at the bottom and click **Submit & Finish**."
                )
                for _q in range(NUM_QUERIES):
                    _qtxt = STUDY_QUERIES[_q][0]
                    gr.Markdown(f"---\n### Query {_q + 1} / {NUM_QUERIES}: \"{_qtxt}\"")
                    for _img_i in range(10):
                        _method = "Baseline" if _img_i < 5 else "Linear"
                        with gr.Row():
                            _im = gr.Image(height=280, show_label=False, interactive=False)
                            _r  = gr.Radio(
                                ["Yes", "No"], value=None,
                                label=f"{_method} — Matches query?", scale=1,
                            )
                        _p3_images.append(_im)
                        _p3_radios.append(_r)
                    gr.Markdown(
                        "**Overall, did Linear Steering give you better top-5 "
                        "results than Baseline for this query?**\n\n"
                        "*Consider: even if both returned correct images, did "
                        "Linear provide more relevant or better-matching ones?*"
                    )
                    _comp = gr.Radio(
                        ["Yes", "No", "About the same"], value=None,
                        label="Linear gave better results than Baseline?",
                    )
                    _p3_comps.append(_comp)
                    # Per-query save controls
                    _st = gr.Markdown(value="", elem_id=f"q_status_{_q}")
                    with gr.Row():
                        _done = gr.Button(
                            f"✅ Done with Query {_q + 1}", variant="primary", size="lg",
                        )
                        _upd = gr.Button(
                            f"🔄 Update Query {_q + 1}", variant="secondary", size="lg",
                            visible=False,
                        )
                    _p3_done_btns.append(_done)
                    _p3_update_btns.append(_upd)
                    _p3_status_mds.append(_st)

                # ── Overall experience ratings (asked once, not per query) ──
                gr.Markdown(
                    "---\n# Overall Experience Ratings\n\n"
                    "Think about the **entire study** across all 22 queries.\n\n"
                    "**Rate each statement (1 = strongly disagree, 7 = strongly agree):**"
                )
                _p3_likerts.append(gr.Slider(1, 7, value=1, step=1,
                                             label="The system understood what I meant"))
                _p3_likerts.append(gr.Slider(1, 7, value=1, step=1,
                                             label="I felt in control of the results"))
                _p3_likerts.append(gr.Slider(1, 7, value=1, step=1,
                                             label="I am satisfied with the results"))
                _p3_likerts.append(gr.Slider(1, 7, value=1, step=1,
                                             label="The interaction was frustrating"))

                # ── Final Survey (at the bottom of the annotation page) ──
                gr.Markdown("---\n# Final Survey")
                study_final_preferred = gr.Dropdown(
                    choices=["Baseline", "Linear", "No preference"],
                    value=None,
                    label="Overall, which system did you prefer?",
                    interactive=True,
                )
                study_final_concept = gr.Dropdown(
                    choices=["Yes", "No"],
                    value=None,
                    label="Did the interaction change how you think about some concepts?",
                    interactive=True,
                )
                study_final_concept_text = gr.Textbox(
                    label="Optional explanation (if you said Yes)",
                    placeholder="Optional…",
                    lines=2,
                )
                study_final_feedback = gr.Textbox(
                    label="Open feedback",
                    placeholder="What worked well? What frustrated you?",
                    lines=4,
                )
                study_phase3_msg = gr.Markdown(value="", elem_id="study_phase3_msg")
                study_phase3_submit_btn = gr.Button(
                    "Submit & Finish", variant="primary", size="lg",
                    elem_id="study_submit_btn",
                )

            # State variables to capture final survey values via .change()
            # This avoids the Gradio bug where js= on .click() prevents
            # input values from being captured.
            _st_pref = gr.State(value=None)
            _st_conc = gr.State(value=None)
            _st_conc_text = gr.State(value="")
            _st_feedback = gr.State(value="")
            study_final_preferred.change(lambda v: v, [study_final_preferred], [_st_pref])
            study_final_concept.change(lambda v: v, [study_final_concept], [_st_conc])
            study_final_concept_text.change(lambda v: v, [study_final_concept_text], [_st_conc_text])
            study_final_feedback.change(lambda v: v, [study_final_feedback], [_st_feedback])

            # Capture 4 global Likert sliders via State
            _st_likert = [gr.State(value=1) for _ in range(4)]
            for _li in range(4):
                _p3_likerts[_li].change(lambda v: v, [_p3_likerts[_li]], [_st_likert[_li]])

            # ── Per-query save handler ──
            def _make_save_query_fn(qidx):
                """Create a handler that validates + saves query qidx's annotations."""
                def _save_query(pid, all_meta, saved_dict, *vals):
                    if pid is None:
                        return (
                            gr.update(value=_study_error("Session expired.")),
                            gr.update(), gr.update(), saved_dict,
                        )
                    # vals = 10 radios + 1 comparison = 11 values
                    radios = vals[:10]
                    comp = vals[10]

                    # Validate
                    errors = []
                    missing = [i + 1 for i, rv in enumerate(radios)
                               if rv is None or str(rv).strip() not in ("Yes", "No")]
                    if missing:
                        errors.append(f"Answer Yes/No for image(s) {', '.join(str(m) for m in missing)}")
                    if comp is None or str(comp).strip() not in ("Yes", "No", "About the same"):
                        errors.append("Missing comparison answer")
                    if errors:
                        return (
                            gr.update(value=_study_error(f"Query {qidx + 1}: " + "; ".join(errors))),
                            gr.update(), gr.update(), saved_dict,
                        )

                    # Build writes
                    qmeta = (all_meta or {}).get(qidx, {})
                    base_ids = qmeta.get("base_ids", [])
                    linear_ids = qmeta.get("linear_ids", [])
                    all_ids = base_ids[:5] + linear_ids[:5]
                    all_ids += [None] * (10 - len(all_ids))
                    ts = _iso_ts()
                    cv = str(comp).strip()

                    writes = []
                    for i, (img_id, rv) in enumerate(zip(all_ids, radios)):
                        if img_id:
                            method = "baseline" if i < 5 else "linear"
                            writes.append(("image_annotations", {
                                "participant_id": pid, "query_id": qidx,
                                "image_id": img_id, "method": method,
                                "meets_intent": 1 if str(rv).strip() == "Yes" else 0,
                                "timestamp": ts,
                            }))
                    writes.append(("method_comparison", {
                        "participant_id": pid, "query_id": qidx,
                        "linear_better": cv, "timestamp": ts,
                    }))
                    writes.append(("final_selections", {
                        "participant_id": pid, "query_id": qidx,
                        "baseline_final_image_ids": ",".join(base_ids[:5]),
                        "linear_final_image_ids": ",".join(linear_ids[:5]),
                        "round_satisfied": qmeta.get("round_sat", 0),
                        "time_elapsed": qmeta.get("time_elapsed", 0),
                        "timestamp": ts,
                    }))

                    try:
                        firestore_batch_add(writes)
                    except Exception as e:
                        return (
                            gr.update(value=_study_error(f"Save error: {str(e)}")),
                            gr.update(), gr.update(), saved_dict,
                        )

                    new_saved = dict(saved_dict or {})
                    new_saved[qidx] = True
                    n_saved = len(new_saved)
                    return (
                        gr.update(value=f"<p style='color:green;font-weight:bold'>"
                                        f"✅ Query {qidx + 1} saved! "
                                        f"({n_saved}/{NUM_QUERIES} done)</p>"),
                        gr.update(visible=False),   # hide Done button
                        gr.update(visible=True),     # show Update button
                        new_saved,
                    )
                return _save_query

            # Wire per-query Done / Update buttons
            for _q in range(NUM_QUERIES):
                _q_radios  = _p3_radios[_q * 10:(_q + 1) * 10]
                _q_inputs  = ([study_participant_id, study_all_meta, study_saved_queries]
                              + _q_radios + [_p3_comps[_q]])
                _q_outputs = [_p3_status_mds[_q], _p3_done_btns[_q],
                              _p3_update_btns[_q], study_saved_queries]

                _p3_done_btns[_q].click(
                    fn=_make_save_query_fn(_q),
                    inputs=_q_inputs,
                    outputs=_q_outputs,
                )
                _p3_update_btns[_q].click(
                    fn=_make_save_query_fn(_q),
                    inputs=_q_inputs,
                    outputs=_q_outputs,
                )

            def _annotation_image_paths(qidx, base_ids, linear_ids):
                """Return 10 image paths for Phase 3; ensure dataset is switched."""
                if qidx is None or qidx < 0 or qidx >= NUM_QUERIES:
                    return [None] * 10
                try:
                    switch_dataset(STUDY_QUERIES[qidx][1])
                except Exception:
                    return [None] * 10
                paths = []
                for img_id in (base_ids or [])[:5]:
                    try:
                        p = get_image_path(img_id)
                        paths.append(str(p) if p and p.exists() else None)
                    except Exception:
                        paths.append(None)
                while len(paths) < 5:
                    paths.append(None)
                for img_id in (linear_ids or [])[:5]:
                    try:
                        p = get_image_path(img_id)
                        paths.append(str(p) if p and p.exists() else None)
                    except Exception:
                        paths.append(None)
                while len(paths) < 10:
                    paths.append(None)
                return paths

            def _satisfied_and_advance(
                pid, qidx, rnd, start_time, base_ids, linear_ids,
                attrs_used, all_images, all_meta,
            ):
                """Log interaction, store images/meta, advance to next query or Phase 3."""
                _noop24 = tuple(
                    [gr.update()] * 20
                    + [gr.update(), gr.update(), all_images or {}, all_meta or {}]
                )
                if pid is None or qidx is None:
                    return _noop24

                elapsed = time.time() - start_time if start_time else 0.0
                base_str = ",".join(base_ids) if base_ids else ""
                linear_str = ",".join(linear_ids) if linear_ids else ""
                try:
                    log_interaction(pid, qidx, STUDY_QUERIES[qidx][0],
                                    "baseline", rnd, attrs_used, base_str, elapsed, 1)
                    log_interaction(pid, qidx, STUDY_QUERIES[qidx][0],
                                    "linear", rnd, attrs_used, linear_str, elapsed, 1)
                except Exception:
                    pass  # non-fatal

                # Store images + metadata for Phase 3
                paths = _annotation_image_paths(qidx, base_ids, linear_ids)
                new_images = dict(all_images or {})
                new_images[qidx] = paths
                new_meta = dict(all_meta or {})
                new_meta[qidx] = {
                    "base_ids": list(base_ids or []),
                    "linear_ids": list(linear_ids or []),
                    "round_sat": rnd if rnd is not None else 0,
                    "time_elapsed": elapsed,
                }

                next_qidx = qidx + 1
                if next_qidx >= NUM_QUERIES:
                    # Last query done → show Phase 3 (all-annotation page)
                    return tuple(
                        [gr.update()] * 20
                        + [gr.update(visible=False), gr.update(visible=True),
                           new_images, new_meta]
                    )

                # Load next query (Phase 2 continues)
                start = time.time()
                qtext = STUDY_QUERIES[next_qidx][0]
                try:
                    res = _run_study_query(next_qidx, 0, pid, start, [], [], 0.4, 0.4, True)
                    (gal_b, gal_l, feedback_md, new_base, new_linear,
                     attrs_used_new, pos_attrs, neg_attrs,
                     _qi, _rn, progress, round_label, query_display, _el) = res
                except Exception:
                    gal_b, gal_l, feedback_md = [], [], ""
                    new_base, new_linear, attrs_used_new = [], [], ""
                    pos_attrs, neg_attrs = [], []
                    progress = f"Query {next_qidx + 1} / {NUM_QUERIES}"
                    round_label = f"Round 1 / {MAX_ROUNDS}"
                    query_display = f"## {qtext}"
                pos_out = pos_attrs if pos_attrs else [""]
                neg_out = neg_attrs if neg_attrs else [""]
                return (
                    # phase2_outputs (20):
                    gr.update(value=progress), gr.update(value=round_label),
                    gr.update(value=query_display),
                    gal_b, gal_l, gr.update(value=feedback_md),
                    pos_out, neg_out,
                    gr.update(value=0.4), gr.update(value=0.4),
                    gr.update(interactive=True), gr.update(visible=True),
                    next_qidx, 0, start, new_base, new_linear, attrs_used_new, -1, 0.0,
                    # phase2 / phase3 visibility + state:
                    gr.update(visible=True), gr.update(visible=False),
                    new_images, new_meta,
                )

            def _populate_phase3_images(all_images):
                """Populate all Image components when Phase 3 becomes visible."""
                if len(all_images or {}) < NUM_QUERIES:
                    return [gr.update()] * (NUM_QUERIES * 10)
                result = []
                for q in range(NUM_QUERIES):
                    paths = all_images.get(q, [None] * 10)
                    for p in paths:
                        result.append(gr.update(value=p))
                return result

            phase2_outputs = [
                study_progress_md, study_round_md, study_query_text_md,
                study_baseline_gallery, study_linear_gallery, study_feedback_md,
                study_pos_state, study_neg_state,
                study_alpha, study_beta,
                satisfied_btn, apply_refinement_btn,
                study_query_idx, study_round, study_query_start_time,
                study_baseline_ids, study_linear_ids, study_attributes_used,
                study_round_satisfied, study_time_elapsed,
            ]
            continue_to_query1_btn.click(
                fn=_start_first_query,
                inputs=[study_participant_id],
                outputs=[study_phase1_success, study_phase2] + phase2_outputs,
            ).then(
                fn=lambda: None, inputs=[], outputs=[], js=_make_scroll_js("study_phase2_anchor"),
            )
            apply_refinement_btn.click(
                fn=_apply_refinement,
                inputs=[
                    study_participant_id, study_query_idx, study_round, study_query_start_time,
                    study_baseline_ids, study_linear_ids, study_attributes_used,
                    study_pos_state, study_neg_state, study_alpha, study_beta,
                ],
                outputs=phase2_outputs,
            )
            satisfied_btn.click(
                fn=_satisfied_and_advance,
                inputs=[
                    study_participant_id, study_query_idx, study_round,
                    study_query_start_time, study_baseline_ids, study_linear_ids,
                    study_attributes_used, study_all_images, study_all_meta,
                ],
                outputs=phase2_outputs + [study_phase2, study_phase3,
                                          study_all_images, study_all_meta],
            ).then(
                # Populate all 220 images once Phase 3 is shown (no-op otherwise)
                fn=_populate_phase3_images,
                inputs=[study_all_images],
                outputs=_p3_images,
            ).then(
                # Scroll to top AFTER content is fully loaded
                fn=lambda: None, inputs=[], outputs=[], js=_make_scroll_js("study_phase2_anchor"),
            )

            # Thank-you page (shown after successful submit)
            study_thanks = gr.Column(visible=False, elem_id="study_thanks")
            with study_thanks:
                gr.Markdown("## Thank you!\n\nYour responses have been saved. We appreciate your participation.")

            def _submit_all(pid, saved_dict, preferred, concept, concept_text, feedback,
                            lik_align, lik_agency, lik_satis, lik_frust):
                """Validate that all queries are saved, save Likert ratings + final survey."""
                _err = lambda msg: (gr.update(), gr.update(value=_study_error(msg)), gr.update())

                if pid is None:
                    return _err("Session expired. Please restart.")

                # Check all queries are saved
                saved = saved_dict or {}
                unsaved = [q + 1 for q in range(NUM_QUERIES) if q not in saved]
                if unsaved:
                    return _err(
                        f"Please click **Done with Query X** for: "
                        f"{', '.join(str(u) for u in unsaved[:10])}.\n\n"
                        f"({len(unsaved)} of {NUM_QUERIES} queries not yet saved.)"
                    )

                # Validate Likert ratings
                errors = []
                lik_vals = []
                for label, val in [("alignment", lik_align), ("agency", lik_agency),
                                   ("satisfaction", lik_satis), ("frustration", lik_frust)]:
                    try:
                        x = int(round(float(val)))
                        if not (1 <= x <= 7):
                            errors.append(f"Rating '{label}' out of range.")
                        lik_vals.append(x)
                    except (TypeError, ValueError):
                        errors.append(f"Invalid value for '{label}'.")
                        lik_vals.append(1)

                # Validate final survey
                pref_val = str(preferred or "").strip()
                if pref_val not in ("Baseline", "Linear", "No preference"):
                    errors.append(f"Please choose Baseline, Linear, or No preference (received: '{pref_val}').")
                conc_val = str(concept or "").strip()
                if conc_val not in ("Yes", "No"):
                    errors.append(f"Please answer Yes or No for the concept question (received: '{conc_val}').")
                if errors:
                    return _err("\n\n".join(errors))

                # Save survey_responses (1 global document) + final_survey
                try:
                    open_fb = (feedback or "").strip()
                    if conc_val == "Yes" and (concept_text or "").strip():
                        open_fb = (concept_text or "").strip() + "\n\n" + open_fb
                    ts = _iso_ts()
                    firestore_batch_add([
                        ("survey_responses", {
                            "participant_id": pid,
                            "alignment_score": lik_vals[0],
                            "agency_score": lik_vals[1],
                            "satisfaction_score": lik_vals[2],
                            "frustration_score": lik_vals[3],
                            "scope": "global",
                            "timestamp": ts,
                        }),
                        ("final_survey", {
                            "participant_id": pid,
                            "preferred_system": pref_val,
                            "concept_changed": conc_val,
                            "open_feedback": open_fb,
                            "timestamp": ts,
                        }),
                    ])
                except Exception as e:
                    return _err(f"Save error: {str(e)}. Please try again.")

                return gr.update(visible=False), gr.update(value=""), gr.update(visible=True)

            study_phase3_submit_btn.click(
                fn=_submit_all,
                inputs=[study_participant_id, study_saved_queries,
                        _st_pref, _st_conc, _st_conc_text, _st_feedback]
                       + _st_likert,
                outputs=[study_phase3, study_phase3_msg, study_thanks],
            ).then(
                fn=lambda: None, inputs=[], outputs=[], js=_make_scroll_js("study_thanks"),
            )

    # Switch to User Study tab when button is clicked
    def _go_to_study_tab():
        return gr.update(selected="user_study_tab")
    study_tab_btn.click(fn=_go_to_study_tab, inputs=[], outputs=[main_tabs])


# ── Launch ───────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    print("\U0001f504 Pre-loading model and embeddings\u2026")
    try:
        load_clip_model()
        switch_dataset(default_ds)
        print("\u2705 Ready")
    except Exception as exc:
        print(f"\u26a0\ufe0f Pre-load warning: {exc}")

    demo.launch(share=True)