app.py

"""
Apiarist - Offline AI inspector for honeybee hive frames.

Stack:
- Qwen2.5-VL-3B on ZeroGPU for the narrative pass.
- (Coming) custom-trained YOLOv8s for queen / drone / bee detection.
- SQLite for hive registry + inspection history.
- Gradio with custom field-tool theme.
"""

import json
import re
import sqlite3
import time
from pathlib import Path

import gradio as gr
import torch
from PIL import Image
from transformers import AutoProcessor, AutoModelForImageTextToText

import db
import detector
import report
import cascade
# build trigger v3

# ZeroGPU integration, no-op outside HF Spaces
try:
    import spaces

    def gpu(fn):
        return spaces.GPU(duration=90)(fn)
except ImportError:
    def gpu(fn):
        return fn


# ---------------------------------------------------------------- model setup

MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct"

_model = None
_processor = None


def get_model():
    global _model, _processor
    if _model is None:
        print(f"Loading {MODEL_ID} ...")
        _processor = AutoProcessor.from_pretrained(MODEL_ID)
        _model = AutoModelForImageTextToText.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.float16,
        )
        _model.eval()
        print("Model loaded.")
    return _model, _processor


INSPECTION_PROMPT = """You are an expert beekeeper inspecting a honeycomb frame. Look at the image carefully and respond ONLY in this exact format, one field per line:

QUEEN: yes or no
MITES: a number from 0 to 20
SWARM_CELLS: yes or no
BROOD_PATTERN: solid, spotty, none, or uncertain
HEALTH: good, watch, or alarm
NOTES: one short sentence describing what you see

Definitions:
- Queens are noticeably larger bees with elongated abdomens.
- Varroa mites are small reddish-brown parasites on bees or comb cells.
- Swarm cells are peanut-shaped cells hanging from the bottom or edges of the comb.
- Brood pattern is solid when capped cells are tightly packed, spotty when scattered.
- Only say "yes" when you can clearly see the feature."""


# ---------------------------------------------------------------- parsing


def parse_response(text: str, hive_name: str) -> dict:
    def grab(field: str, default: str = "") -> str:
        m = re.search(rf"{field}\s*[:\-]\s*([^\n]+)", text, re.IGNORECASE)
        return m.group(1).strip() if m else default

    mites_raw = grab("MITES", "0")
    mite_match = re.search(r"\d+", mites_raw)
    mite_count = int(mite_match.group()) if mite_match else 0

    return {
        "hive": hive_name or "Unnamed Hive",
        "queen_detected": grab("QUEEN", "no").lower().startswith("y"),
        "varroa_mites_visible": mite_count,
        "swarm_cells_detected": grab("SWARM_CELLS", "no").lower().startswith("y"),
        "brood_pattern": grab("BROOD_PATTERN", "uncertain").lower(),
        "frame_health": grab("HEALTH", "uncertain").lower(),
        "notes": grab("NOTES", ""),
        "model": MODEL_ID,
    }


def build_dashboard(r: dict) -> str:
    """Render a polished inspection 'report card' as HTML: a health-verdict
    banner plus color-coded stat tiles. This is the visual hero of a result."""
    counts = r.get("yolo_counts") or {}
    top_probs = r.get("cascade_top_probs", [])
    queen_conf = int(top_probs[0] * 100) if top_probs else 0

    n_bees = counts.get("bee", 0)
    n_drones = counts.get("drone", 0)
    n_mites = counts.get("varroa", 0)
    queen = r.get("queen_detected", False)

    # Verdict logic: mites drive alarm, missing queen drives watch.
    if n_mites >= 3:
        verdict, vcolor, vsub = "ALERT", "#e3493b", f"{n_mites} varroa mites detected - treat soon"
    elif n_mites >= 1:
        verdict, vcolor, vsub = "WATCH", "#e8a317", f"{n_mites} varroa mite(s) - monitor"
    elif not queen and r.get("detector_used"):
        verdict, vcolor, vsub = "WATCH", "#e8a317", "No queen confirmed on this frame"
    elif r.get("detector_used"):
        verdict, vcolor, vsub = "HEALTHY", "#2faa55", "Queen present, no mites detected"
    else:
        verdict, vcolor, vsub = "ANALYZED", "#8a8a8a", "Narrative-only (detector offline)"

    def tile(label, value, sub, color):
        return f"""
        <div class="stat-tile" style="border-top:3px solid {color};">
          <div class="stat-val" style="color:{color};">{value}</div>
          <div class="stat-label">{label}</div>
          <div class="stat-sub">{sub}</div>
        </div>"""

    queen_val = "1" if queen else "0"
    queen_sub = f"{queen_conf}% confidence" if queen else "not found"

    return f"""
    <div id="dash">
      <div class="verdict" style="background:linear-gradient(90deg,{vcolor}33,transparent);border-left:5px solid {vcolor};">
        <div class="verdict-hive">{r.get('hive','Unnamed Hive')}</div>
        <div class="verdict-main" style="color:{vcolor};">{verdict}</div>
        <div class="verdict-sub">{vsub}</div>
      </div>
      <div class="stat-row">
        {tile("QUEEN", queen_val, queen_sub, "#32c864" if queen else "#8a8a8a")}
        {tile("WORKERS", n_bees, "detected", "#f4a300")}
        {tile("DRONES", n_drones, "detected", "#ff7d4d")}
        {tile("MITES", n_mites, "specialist", "#dc46dc" if n_mites else "#5f8a5f")}
      </div>
    </div>"""


def build_narrative(r: dict, raw: str) -> str:
    top_probs = r.get("cascade_top_probs", [])
    if r["queen_detected"]:
        queen_line = (
            f"**Queen detected** (specialist classifier, {int(top_probs[0]*100)}% confidence)"
            if top_probs else "**Queen detected** (specialist classifier)"
        )
    else:
        queen_line = (
            "No queen detected - she may be hidden or on another frame"
        )
    swarm_line = (
        " Swarm cells (VLM estimate)"
        if r["swarm_cells_detected"]
        else " No swarm cells (VLM estimate)"
    )

    yolo_block = ""
    if r.get("detector_used"):
        counts = r.get("yolo_counts") or {}
        max_conf = r.get("yolo_max_conf") or {}
        ordered = ["queen", "drone", "varroa", "bee", "pollenbee"]
        lines = []
        for c in ordered:
            if counts.get(c):
                conf_str = f" (top conf {max_conf[c]:.0%})" if c in max_conf else ""
                lines.append(f" - **{counts[c]}** {c}{conf_str}")
        if lines:
            yolo_block = "\n**Specialist detector counts:**\n" + "\n".join(lines)
        else:
            yolo_block = "\n*Specialist detector ran but found no high-confidence bees in this frame.*"
        footer = "*Powered by custom YOLOv8s + Qwen2.5-VL-3B on ZeroGPU.*"
    else:
        footer = "*Powered by Qwen2.5-VL-3B on ZeroGPU. YOLO weights not present.*"

    # Mite count is Qwen-only (no specialist), flag that clearly
    mite_line = (
        f" ~{r['varroa_mites_visible']} mite(s) (VLM estimate, not specialist)"
    )

    return f"""**Hive: {r['hive']}**

{queen_line}
{swarm_line}
{mite_line}
{yolo_block}

**Brood pattern:** {r['brood_pattern']} | **Overall:** {r['frame_health']}
**Notes:** {r['notes']}

---
{footer}

<details><summary>Raw model output</summary>

```
{raw}
```
</details>
"""


# ---------------------------------------------------------------- inference


@gpu
def _qwen_only(image: Image.Image, prompt_text: str) -> str:
    """GPU-only: runs Qwen-3B and returns the raw text response."""
    model, processor = get_model()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32
    model = model.to(device=device, dtype=dtype)

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": image},
                {"type": "text", "text": prompt_text},
            ],
        }
    ]
    inputs = processor.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    ).to(device)

    with torch.no_grad():
        generated = model.generate(
            **inputs, max_new_tokens=300, do_sample=False
        )
    prompt_len = inputs["input_ids"].shape[1]
    return processor.batch_decode(
        generated[:, prompt_len:], skip_special_tokens=True
    )[0].strip()


def analyze_frame(image: Image.Image, hive_name: str):
    if image is None:
        return None, "", "Upload a frame photo first.", "", _hives_table_state(), gr.update(), gr.update()

    # Step 1: YOLO on the MAIN container (file access + log visibility).
    detections, _ = detector.detect(image)
    yolo_active = detector.is_available()

    # Step 1b: VLM cascade - send candidate bee crops as a numbered grid
    # to Qwen-VL and let it identify the queen by side-by-side comparison.
    # Far more reliable than asking the VLM to find a queen in the full
    # wide-shot photo.
    cascade_info = {"grid_size": 0, "queen_indices": set()}
    if yolo_active and detections:
        try:
            detections, cascade_info = cascade.verify_queens(
                image, detections, _qwen_only
            )
        except Exception as e:
            print(f"[cascade] verify_queens failed: {e}")

    # Re-annotate with the (possibly updated) detection classes
    annotated = detector.draw_annotations(image, detections) if yolo_active else None
    counts = detector.summarize_counts(detections)
    display_image = annotated if annotated is not None else image

    # Step 2: build prompt with YOLO context if available
    if yolo_active and counts:
        detection_summary = (
            "A specialist detector has already analyzed the image and found:\n"
            + "\n".join(f"- {n} {cls}" for cls, n in counts.items())
            + "\n\nIncorporate these counts into your assessment.\n\n"
        )
    elif yolo_active:
        detection_summary = (
            "A specialist detector ran and found no clear bees.\n\n"
        )
    else:
        detection_summary = ""

    # Step 3: Qwen narrative pass on the GPU worker
    try:
        response = _qwen_only(image, detection_summary + INSPECTION_PROMPT)
    except Exception as e:
        return (
            display_image,
            "",
            f"Model inference failed: {type(e).__name__}: {e}",
            "",
            _hives_table_state(),
            gr.update(),
            gr.update(),
        )

    results = parse_response(response, hive_name)

    # Override hallucinated fields with hard YOLO evidence.
    # Two distinct, honest signals:
    #   queen_detected  - a CONFIRMED queen (only the trained classifier sets
    #                     class="queen"). Hard yes.
    #   queen_candidate - a best-effort geometric outlier flagged "confirm by
    #                     eye". Never a hard yes - the queen may not even be
    #                     in frame.
    if yolo_active:
        cand = next((d for d in detections if d.get("queen_candidate")), None)
        results["queen_detected"] = counts.get("queen", 0) > 0
        results["queen_candidate"] = cand is not None
        results["queen_standout"] = (
            cand.get("queen_standout", 0.0) if cand else 0.0
        )
        results["yolo_counts"] = counts
        # Track the top confidence per class so the UI can show how
        # confident the model was about each detection.
        max_conf: dict[str, float] = {}
        for d in detections:
            cls = d["class"]
            if d["confidence"] > max_conf.get(cls, 0):
                max_conf[cls] = d["confidence"]
        results["yolo_max_conf"] = max_conf
        results["detector_used"] = True
        results["cascade_grid_size"] = cascade_info.get("grid_size", 0)
        results["cascade_queen_idx"] = sorted(cascade_info.get("queen_indices", set()))
        results["cascade_top_probs"] = cascade_info.get("top_3_probs", [])
    else:
        results["detector_used"] = False

    narrative = build_narrative(results, response)
    dashboard = build_dashboard(results)

    # Persist the inspection
    hive_id = db.get_or_create_hive(results["hive"])
    db.add_inspection(hive_id, results, raw_response=response)

    hive_names = [h["name"] for h in db.list_hives()]
    return (
        display_image,
        dashboard,
        narrative,
        json.dumps(results, indent=2),
        _hives_table_state(),
        gr.update(choices=hive_names, value=results["hive"]),
        gr.update(choices=hive_names),
    )


# ---------------------------------------------------------------- comparison helpers


def _run_qwen(image: Image.Image, prefix: str = "") -> str:
    """Convenience wrapper that calls the GPU-only Qwen function."""
    return _qwen_only(image, prefix + INSPECTION_PROMPT)


def run_comparison(image: Image.Image):
    """A/B: Apiarist (YOLO + Qwen) vs raw Qwen alone, side-by-side."""
    if image is None:
        msg = "Upload a frame photo to compare."
        return None, msg, None, msg

    # YOLO on main container (CPU, fast)
    detections, annotated = detector.detect(image)
    counts = detector.summarize_counts(detections)
    yolo_active = detector.is_available()

    # Run the queen cascade so the left side shows the confirmed queen.
    if yolo_active and detections:
        try:
            detections, _ = cascade.verify_queens(image, detections, _qwen_only)
        except Exception:
            pass
        annotated = detector.draw_annotations(image, detections)
        counts = detector.summarize_counts(detections)

    if yolo_active:
        prefix = (
            "A specialist detector found: "
            + ", ".join(f"{n} {cls}" for cls, n in counts.items())
            + ". Incorporate these counts.\n\n"
        ) if counts else ""
        left_response = _run_qwen(image, prefix=prefix)
        left_results = parse_response(left_response, "")
        left_results["queen_detected"] = counts.get("queen", 0) > 0
        left_results["yolo_counts"] = counts
        left_results["detector_used"] = True
        left_html = _cmp_apiarist_card(left_results, counts)
        left_image = annotated if annotated is not None else image
    else:
        left_html = "<div class='cmp-card cmp-apiarist'><h3>Apiarist</h3><p>Detector offline.</p></div>"
        left_image = image

    # ----- RIGHT: Qwen alone (generalist, no specialist context) -----
    right_response = _run_qwen(image, prefix="")
    right_html = _cmp_raw_card(right_response)

    return left_image, left_html, image, right_html


def _grounded_report(counts: dict) -> str:
    """Build the report sentence from the specialist counts ONLY, so it can
    never contradict the count table (the VLM prose sometimes invents mites)."""
    q = counts.get("queen", 0)
    b = counts.get("bee", 0)
    d = counts.get("drone", 0)
    v = counts.get("varroa", 0)
    parts = []
    if q:
        parts.append(f"{q} queen confirmed")
    parts.append(f"{b} worker bee{'s' if b != 1 else ''}")
    if d:
        parts.append(f"{d} drone{'s' if d != 1 else ''}")
    sentence = ", ".join(parts) + " detected on this frame. "
    if v == 0:
        sentence += "No varroa mites found."
    else:
        sentence += f"{v} varroa mite{'s' if v != 1 else ''} flagged - monitor closely."
    return sentence


def _cmp_apiarist_card(r: dict, counts: dict) -> str:
    q = counts.get("queen", 0)
    rows = "".join(
        f"<tr><td>{cls.capitalize()}</td><td style='text-align:right;font-weight:700;color:#ffd066'>{n}</td></tr>"
        for cls, n in (("queen", counts.get("queen", 0)),
                       ("bee", counts.get("bee", 0)),
                       ("drone", counts.get("drone", 0)),
                       ("varroa", counts.get("varroa", 0)))
    )
    queen_badge = (
        "<span style='color:#32ff64;font-weight:700'>QUEEN CONFIRMED</span>"
        if q else "<span style='color:#c9b285'>no queen this frame</span>"
    )
    return f"""
    <div class="cmp-card cmp-apiarist">
      <h3 style="color:#32ff64;">Apiarist (specialist + generalist)</h3>
      <p style="margin:2px 0 8px 0;">{queen_badge}</p>
      <table style="width:100%;border-collapse:collapse;font-size:0.9rem;">
        <thead><tr><th style="text-align:left;color:#c9b285">Class</th>
        <th style="text-align:right;color:#c9b285">Exact count</th></tr></thead>
        <tbody>{rows}</tbody>
      </table>
      <p style="margin-top:10px;font-size:0.84rem;color:#d8c298;">
        <b>Report:</b> {_grounded_report(counts)}</p>
      <p style="font-size:0.74rem;color:#8fae8f;margin-top:6px;">
        Grounded in pixel-level detections from a custom-trained model.</p>
    </div>"""


def _cmp_raw_card(response: str) -> str:
    text = (response or "").strip()
    # Show the generalist's raw prose - typically vague, no hard counts.
    if len(text) > 420:
        text = text[:420] + "..."
    return f"""
    <div class="cmp-card cmp-raw">
      <h3 style="color:#bbb;">Raw generalist VLM (no specialist)</h3>
      <p style="margin:2px 0 8px 0;color:#999;">prose only - no reliable counts</p>
      <p style="font-size:0.86rem;color:#cfcfcf;white-space:pre-wrap;">{text}</p>
      <p style="font-size:0.74rem;color:#9a7d7d;margin-top:8px;">
        A generalist describes the scene but hallucinates or omits exact
        queen / mite / drone counts.</p>
    </div>"""


# ---------------------------------------------------------------- Hives tab helpers


def _hives_table_state() -> list[list]:
    rows = db.list_hives()
    out = []
    for h in rows:
        last = (
            time.strftime("%Y-%m-%d %H:%M", time.localtime(h["last_inspected"]))
            if h["last_inspected"]
            else "-"
        )
        out.append(
            [
                h["name"],
                h.get("location") or "",
                h.get("queen_marker") or "",
                h["inspection_count"],
                last,
            ]
        )
    return out


def add_hive_action(name, location, marker, notes):
    name = (name or "").strip()
    if not name:
        return (
            _hives_table_state(),
            gr.update(),
            gr.update(),
            " Name required.",
        )
    try:
        db.add_hive(name, location or "", marker or "", notes or "")
        msg = f" Added hive '{name}'."
    except sqlite3.IntegrityError:
        msg = f" Hive '{name}' already exists."
    hive_names = [h["name"] for h in db.list_hives()]
    return (
        _hives_table_state(),
        gr.update(choices=hive_names),
        gr.update(choices=hive_names),
        msg,
    )


def _hidden_trend():
    return gr.update(visible=False)


def view_hive_history(hive_name):
    import pandas as pd
    if not hive_name:
        return [], "_Pick a hive above to see its inspection history._", _hidden_trend()
    hive = next((h for h in db.list_hives() if h["name"] == hive_name), None)
    if not hive:
        return [], "_Hive not found._", _hidden_trend()
    inspections = db.get_inspections_for_hive(hive["id"])
    if not inspections:
        return [], f"_No inspections recorded for **{hive_name}** yet._", _hidden_trend()
    rows = []
    for i in inspections:
        rows.append(
            [
                time.strftime("%Y-%m-%d %H:%M", time.localtime(i["created_at"])),
                "Y" if i["queen_detected"] else "N",
                i["varroa_mites_visible"],
                "Y" if i["swarm_cells_detected"] else "N",
                i["frame_health"],
                (i["notes"] or "")[:60],
            ]
        )
    summary = (
        f"### {hive_name}\n"
        f"**Total inspections:** {len(inspections)}  |  "
        f"**Last inspected:** "
        f"{time.strftime('%Y-%m-%d %H:%M', time.localtime(inspections[0]['created_at']))}"
    )

    # Build a long-format trend frame (oldest -> newest) for the LinePlot.
    chrono = list(reversed(inspections))
    trend_rows = []
    for n, i in enumerate(chrono, start=1):
        mites = int(i["varroa_mites_visible"] or 0)
        # bee count isn't stored as a column; derive from structured_json if present
        bees = 0
        try:
            sj = json.loads(i.get("structured_json") or "{}")
            bees = int((sj.get("yolo_counts") or {}).get("bee", 0))
        except Exception:
            bees = 0
        trend_rows.append({"inspection": n, "count": bees, "metric": "bees"})
        trend_rows.append({"inspection": n, "count": mites, "metric": "varroa mites"})

    # Only show the trend chart when there are at least 2 inspections to plot.
    # With fewer, a line chart is just an awkward empty/single-dot box.
    if len(chrono) >= 2:
        trend_update = gr.update(value=pd.DataFrame(trend_rows), visible=True)
    else:
        trend_update = gr.update(visible=False)
    return rows, summary, trend_update


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


custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Fraunces:opsz,wght@9..144,600;9..144,800&family=JetBrains+Mono:wght@400;600&display=swap');

.gradio-container {
    background:
        radial-gradient(1200px 500px at 80% -10%, rgba(244,163,0,0.10), transparent 60%),
        linear-gradient(180deg, #17110c 0%, #221913 55%, #2a1f15 100%) !important;
    color: #f4e4bc !important;
    font-family: 'JetBrains Mono', ui-monospace, 'Courier New', monospace !important;
    max-width: 1600px !important;
    margin: 0 auto !important;
    padding-left: 24px !important;
    padding-right: 24px !important;
}

/* ----- Hero header ----- */
#apiarist-hero {
    border: 1px solid rgba(244,163,0,0.30);
    border-radius: 18px;
    padding: 22px 26px;
    margin: 6px 0 14px 0;
    background:
        linear-gradient(135deg, rgba(244,163,0,0.14) 0%, rgba(42,31,21,0.20) 55%);
    box-shadow: 0 8px 30px rgba(0,0,0,0.35), inset 0 1px 0 rgba(255,214,102,0.15);
}
#apiarist-hero h1 {
    font-family: 'Fraunces', Georgia, serif !important;
    font-weight: 800 !important;
    font-size: 2.5rem !important;
    letter-spacing: 0.04em;
    margin: 0 0 4px 0 !important;
    color: #ffc23d !important;
    text-shadow: 0 2px 12px rgba(244,163,0,0.25);
}
#apiarist-hero .tagline {
    color: #e8d2a0 !important;
    font-size: 0.98rem;
    margin: 0;
}
#apiarist-hero .pills { margin-top: 12px; }
#apiarist-hero .pill {
    display: inline-block;
    border: 1px solid rgba(244,163,0,0.40);
    border-radius: 999px;
    padding: 3px 11px;
    margin: 3px 6px 0 0;
    font-size: 0.74rem;
    color: #ffd066;
    background: rgba(244,163,0,0.08);
    white-space: nowrap;
}

h1, h2, h3 { color: #f4a300 !important; font-family: 'Fraunces', Georgia, serif !important; }
h4 { color: #ffd066 !important; }

button.primary, button[variant="primary"] {
    background: linear-gradient(180deg, #ffc23d 0%, #f4a300 100%) !important;
    color: #1a1410 !important;
    font-weight: 700 !important;
    border: none !important;
    border-radius: 12px !important;
    box-shadow: 0 3px 12px rgba(244,163,0,0.30) !important;
    transition: transform 80ms ease, box-shadow 200ms ease !important;
}
button.primary:hover {
    transform: translateY(-1px);
    box-shadow: 0 6px 20px rgba(244,163,0,0.45) !important;
}

.gr-box, .block, .gradio-container .block {
    border-color: rgba(244,163,0,0.28) !important;
    background: rgba(36,26,18,0.55) !important;
    border-radius: 14px !important;
}
.tabs > .tab-nav > button.selected {
    color: #ffc23d !important;
    border-bottom: 2px solid #f4a300 !important;
}
.gradio-container a { color: #ffd066 !important; }
.gradio-container a:hover { color: #ffe599 !important; text-decoration: underline; }
table { border-color: rgba(244,163,0,0.25) !important; }

/* ----- Detection legend ----- */
#legend {
    display: flex; flex-wrap: wrap; gap: 14px;
    font-size: 0.82rem; color: #e8d2a0;
    padding: 4px 2px 2px 2px;
}
#legend .chip { display: inline-flex; align-items: center; gap: 6px; }
#legend .sw {
    width: 13px; height: 13px; border-radius: 3px; display: inline-block;
    border: 1px solid rgba(255,255,255,0.25);
}
.sw-queen  { background: #32ff64; }
.sw-bee    { background: #f4a300; }
.sw-drone  { background: #ff5050; }
.sw-mite   { background: #dc32dc; }

/* ----- Inspection dashboard ----- */
#dash { margin: 6px 0 4px 0; }
#dash .verdict {
    border-radius: 12px; padding: 12px 16px; margin-bottom: 12px;
}
#dash .verdict-hive {
    font-size: 0.8rem; color: #c9b285; letter-spacing: 0.04em;
    text-transform: uppercase;
}
#dash .verdict-main {
    font-family: 'Fraunces', Georgia, serif; font-weight: 800;
    font-size: 1.7rem; line-height: 1.1;
}
#dash .verdict-sub { font-size: 0.86rem; color: #d8c298; margin-top: 2px; }
#dash .stat-row {
    display: grid; grid-template-columns: repeat(4, 1fr); gap: 10px;
}
#dash .stat-tile {
    background: rgba(20,15,10,0.55); border-radius: 10px;
    padding: 12px 8px 10px 8px; text-align: center;
    border: 1px solid rgba(244,163,0,0.15);
}
#dash .stat-val {
    font-family: 'Fraunces', Georgia, serif; font-weight: 800;
    font-size: 1.7rem; line-height: 1;
}
#dash .stat-label {
    font-size: 0.72rem; letter-spacing: 0.08em;
    color: #c9b285; margin-top: 4px; text-transform: uppercase;
}
#dash .stat-sub { font-size: 0.68rem; color: #9c8662; margin-top: 1px; }

/* ----- A/B compare cards ----- */
.cmp-card { border-radius: 12px; padding: 14px 16px; height: 100%; }
.cmp-apiarist {
    border: 1px solid rgba(50,200,100,0.45);
    background: linear-gradient(160deg, rgba(50,200,100,0.10), transparent);
}
.cmp-raw {
    border: 1px dashed rgba(150,150,150,0.45);
    background: rgba(40,40,40,0.25);
}
.cmp-card h3 { margin-top: 0 !important; }

/* ----- Footer ----- */
#apiarist-footer {
    margin-top: 22px; padding-top: 14px;
    border-top: 1px solid rgba(244,163,0,0.20);
    color: #b8a079; font-size: 0.8rem; text-align: center;
}

/* ============================================================
   OFF-BRAND PASS: strip the remaining 'default Gradio' tells
   ============================================================ */

/* Custom tab bar -> bespoke segmented nav */
.tabs > .tab-nav, .tab-nav {
    border-bottom: 1px solid rgba(244,163,0,0.20) !important;
    gap: 6px !important;
    padding-bottom: 0 !important;
}
.tab-nav button {
    background: transparent !important;
    border: none !important;
    color: #b9a079 !important;
    font-family: 'JetBrains Mono', monospace !important;
    font-weight: 600 !important;
    letter-spacing: 0.04em;
    text-transform: uppercase;
    font-size: 0.82rem !important;
    padding: 10px 16px !important;
    border-radius: 10px 10px 0 0 !important;
    transition: color 120ms ease, background 120ms ease;
}
.tab-nav button:hover { color: #ffd066 !important; background: rgba(244,163,0,0.06) !important; }
.tab-nav button.selected {
    color: #1a1410 !important;
    background: linear-gradient(180deg,#ffc23d,#f4a300) !important;
    border: none !important;
}

/* Component labels -> small amber field-tags instead of grey Gradio labels */
.block label > span,
span[data-testid="block-info"],
.gr-box label span {
    color: #d8b66a !important;
    font-size: 0.72rem !important;
    letter-spacing: 0.05em;
    text-transform: uppercase;
    font-weight: 600 !important;
}

/* Soften component chrome so it reads as one bespoke surface, not stacked Gradio boxes */
.block, .gr-box, .form, .gr-form {
    box-shadow: none !important;
}
.gradio-container .block {
    backdrop-filter: blur(2px);
}

/* Inputs / dropdowns -> dark themed, not white Gradio defaults */
input, textarea, select,
.gr-input, .gr-text-input, .gr-dropdown,
[data-testid="textbox"] textarea,
[data-testid="textbox"] input {
    background: rgba(20,15,10,0.6) !important;
    color: #f4e4bc !important;
    border: 1px solid rgba(244,163,0,0.22) !important;
    border-radius: 10px !important;
}

/* Accordions -> themed headers */
.label-wrap, .gr-accordion > .label-wrap {
    color: #ffd066 !important;
}

/* Dataframe / table -> dark themed */
.gr-dataframe table, table.dataframe {
    background: rgba(20,15,10,0.4) !important;
    color: #f4e4bc !important;
}
.gr-dataframe thead th, table.dataframe thead th {
    background: rgba(244,163,0,0.12) !important;
    color: #ffd066 !important;
}

/* Custom scrollbar (a small but distinct off-brand touch) */
*::-webkit-scrollbar { width: 10px; height: 10px; }
*::-webkit-scrollbar-track { background: rgba(20,15,10,0.4); }
*::-webkit-scrollbar-thumb {
    background: rgba(244,163,0,0.35); border-radius: 8px;
    border: 2px solid rgba(20,15,10,0.4);
}
*::-webkit-scrollbar-thumb:hover { background: rgba(244,163,0,0.55); }

/* Tone down the stock 'Built with Gradio' footer chrome so it doesn't clash */
footer { opacity: 0.45 !important; }
footer:hover { opacity: 0.8 !important; }
"""


def _seed_sample_data():
    """Drop a handful of plausible hives into the DB on first boot
    so the Hives tab isn't empty for first-visit judges."""
    if db.list_hives():
        return # someone already has data
    samples = [
        ("Hive #1, Apricot tree", "South corner of yard", "yellow",
         "Italian queen, marked yellow"),
        ("Hive #2, Cedar fence", "East side, near cedar", "red",
         "2024 queen, watch for supersedure cells"),
        ("Hive #3, Vegetable garden", "By the tomato beds", "white",
         "Newly split from Hive #1 in May"),
    ]
    for name, loc, marker, notes in samples:
        try:
            db.add_hive(name, loc, marker, notes)
        except Exception:
            pass


def build_ui() -> gr.Blocks:
    db.init_db()
    _seed_sample_data()

    with gr.Blocks(title="Apiarist - Hive Frame Inspector", fill_width=True) as app:
        gr.HTML(
            """
            <div id="apiarist-hero">
              <h1>&#128029; APIARIST</h1>
              <p class="tagline">Offline AI hive-frame inspector for backyard
              beekeepers &mdash; finds the queen, counts bees &amp; drones,
              flags varroa mites, all on a laptop with no cloud.</p>
              <div class="pills">
                <span class="pill">Custom YOLOv8s detector</span>
                <span class="pill">EfficientNet-B0 queen classifier</span>
                <span class="pill">Qwen2.5-VL-3B narrator</span>
                <span class="pill">Runs on ZeroGPU</span>
                <span class="pill">No cloud APIs</span>
              </div>
            </div>
            """
        )

        with gr.Tabs():
            # ------- INSPECT TAB -------
            with gr.Tab("Inspect"):
                with gr.Row():
                    with gr.Column():
                        hive_input = gr.Dropdown(
                            label="Hive",
                            choices=[h["name"] for h in db.list_hives()],
                            allow_custom_value=True,
                            info="Pick an existing hive or type a new name.",
                        )
                        image_input = gr.Image(
                            label="Frame Photo",
                            type="pil",
                            sources=["upload", "webcam"],
                        )
                        gr.Markdown(
                            "*Best with close-up macro shots of frame inspections "
                            "(a single bee or small cluster filling the frame). "
                            "Wide-angle photos with hands and background may not "
                            "reliably detect the queen.*"
                        )
                        analyze_btn = gr.Button(
                            " Analyze Frame", variant="primary"
                        )

                        sample_dir = Path(__file__).parent / "sample_photos"
                        sample_files = (
                            sorted(sample_dir.glob("*.jpg"))
                            if sample_dir.exists() else []
                        )
                        if sample_files:
                            gr.Markdown("### Or try a sample photo")
                            sample_gallery = gr.Gallery(
                                value=[str(p) for p in sample_files],
                                label="Click any sample to load it",
                                columns=3, height="auto",
                                allow_preview=False, show_label=False,
                                interactive=False,
                            )

                            def _load_sample(evt: gr.SelectData):
                                if evt is None or evt.index is None:
                                    return gr.update()
                                idx = (
                                    evt.index
                                    if isinstance(evt.index, int)
                                    else evt.index[0]
                                )
                                if 0 <= idx < len(sample_files):
                                    return Image.open(str(sample_files[idx])).convert("RGB")
                                return gr.update()

                            sample_gallery.select(
                                fn=_load_sample,
                                outputs=[image_input],
                            )
                    with gr.Column():
                        annotated_output = gr.Image(label="Annotated Frame")
                        gr.HTML(
                            """
                            <div id="legend">
                              <span class="chip"><span class="sw sw-queen"></span>Queen</span>
                              <span class="chip"><span class="sw sw-bee"></span>Worker bee</span>
                              <span class="chip"><span class="sw sw-drone"></span>Drone</span>
                              <span class="chip"><span class="sw sw-mite"></span>Varroa mite</span>
                            </div>
                            """
                        )
                        dashboard_output = gr.HTML()
                        with gr.Accordion("Full report + raw model output", open=False):
                            narrative_output = gr.Markdown()
                        with gr.Accordion("Raw JSON", open=False):
                            json_output = gr.Code(language="json")

            # ------- HIVES TAB -------
            with gr.Tab("Hives") as hives_tab:
                with gr.Row():
                    with gr.Column(scale=1):
                        gr.Markdown("### Add a Hive")
                        new_name = gr.Textbox(
                            label="Name", placeholder="Hive #7"
                        )
                        new_location = gr.Textbox(
                            label="Location (optional)",
                            placeholder="South corner of yard",
                        )
                        new_marker = gr.Dropdown(
                            label="Queen marker color (optional)",
                            choices=["", "white", "yellow", "red", "green", "blue"],
                            value="",
                        )
                        new_notes = gr.Textbox(
                            label="Notes (optional)", lines=2
                        )
                        add_btn = gr.Button(" Add Hive", variant="primary")
                        add_msg = gr.Markdown()

                    with gr.Column(scale=2):
                        gr.Markdown("### Your Apiary")
                        hives_table = gr.Dataframe(
                            headers=[
                                "Name", "Location", "Queen marker",
                                "Inspections", "Last inspected",
                            ],
                            datatype=["str", "str", "str", "number", "str"],
                            interactive=False,
                            value=_hives_table_state(),
                            wrap=True,
                        )
                        refresh_btn = gr.Button(" Refresh")

                gr.Markdown("---\n### Inspection history")
                history_select = gr.Dropdown(
                    label="Select a hive",
                    choices=[h["name"] for h in db.list_hives()],
                )
                history_summary = gr.Markdown()
                history_trend = gr.LinePlot(
                    x="inspection",
                    y="count",
                    color="metric",
                    title="Colony trend (bees & varroa mites over inspections)",
                    height=240,
                    visible=False,
                )
                history_table = gr.Dataframe(
                    headers=[
                        "When", "Queen?", "Mites", "Swarm?",
                        "Health", "Notes",
                    ],
                    datatype=["str", "str", "number", "str", "str", "str"],
                    interactive=False,
                    wrap=True,
                )

                gr.Markdown("---\n### Weekly report")
                with gr.Row():
                    report_btn = gr.Button(
                        " Generate PDF report", variant="primary"
                    )
                    report_file = gr.File(
                        label="Latest report", interactive=False
                    )

            # ------- COMPARE TAB -------
            with gr.Tab("Compare"):
                gr.Markdown(
                    "### Apiarist vs raw generalist VLM\n"
                    "Same image, two pipelines. Apiarist combines a "
                    "**custom-trained YOLO specialist** with a generalist VLM. "
                    "The right column shows what happens when you ask the "
                    "generalist alone, same model, no specialist."
                )
                with gr.Row():
                    cmp_image = gr.Image(
                        label="Frame Photo",
                        type="pil",
                        sources=["upload", "webcam"],
                    )
                    cmp_btn = gr.Button(
                        " Run Comparison", variant="primary", scale=0
                    )

                with gr.Row():
                    with gr.Column():
                        cmp_left_image = gr.Image(
                            label="Annotated by YOLO + classifier"
                        )
                        cmp_left_text = gr.HTML()
                    with gr.Column():
                        cmp_right_image = gr.Image(label="No annotations")
                        cmp_right_text = gr.HTML()

            # ------- ABOUT TAB -------
            with gr.Tab("About"):
                gr.Markdown(
                    """
## Apiarist

An offline AI hive frame inspector for backyard beekeepers. Built in
10 days for the [Build Small Hackathon](https://huggingface.co/build-small-hackathon).

### The problem

Frank, a beekeeper down the road, keeps 14 hives and inspects each
weekend by hand. Last summer he lost four colonies to queen failures
he didn't spot in time. Apiarist is the assistant he should have had:
fast, focused, runs on a laptop in a field, never sends his hive data
anywhere.

### How it works

Two models, two jobs:

- **YOLOv8s, custom-trained** (22 MB), finds bees, drones, queens, and
  varroa mites with bounding boxes. CPU-fast.
- **Qwen2.5-VL-3B-Instruct**, takes the image + YOLO detection counts
  and writes a narrative inspection report. Runs on ZeroGPU.

Together: structured detection + narrative report in ~5 seconds, no
internet required.

### What's in the box

- **Inspect**, upload a frame, get an annotated image + report
- **Hives**, registry, per-hive inspection history, weekly PDF reports
- **Compare**, Apiarist vs raw generalist VLM, side-by-side
- SQLite persistence for the whole apiary

### Stack

- Vision-language model: [Qwen2.5-VL-3B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct)
- Specialist detector: [Apiarist Honey-Bee Detector (custom YOLOv8s)](https://huggingface.co/maryammeda/apiarist-honey-bee-detector)
- Training data: [Hendricks Ricky bee-project](https://universe.roboflow.com/hendricks_ricky-hotmail-de/bee-project) on Roboflow Universe (3,308 labeled images, 892 queens)
- Bee imagery: [Apiarist iNaturalist bees dataset](https://huggingface.co/datasets/maryammeda/apiarist-inaturalist-bees)
- Compute: Modal (training), Hugging Face ZeroGPU (inference)
- Frontend: Gradio with custom field-tool theme

### Badges chased

Off the Grid - Well-Tuned - Off-Brand - Sharing is Caring
"""
                )

        gr.HTML(
            """
            <div id="apiarist-footer">
              Apiarist - built for the Build Small Hackathon -
              custom YOLOv8s + EfficientNet-B0 + Qwen2.5-VL-3B -
              trained on <a href="https://modal.com">Modal</a>,
              served on <a href="https://huggingface.co/docs/hub/spaces-zerogpu">ZeroGPU</a>.
              Fully offline at inference.
            </div>
            """
        )

        # ---- wiring ----

        analyze_btn.click(
            fn=analyze_frame,
            inputs=[image_input, hive_input],
            outputs=[
                annotated_output,
                dashboard_output,
                narrative_output,
                json_output,
                hives_table,
                hive_input,
                history_select,
            ],
        )

        add_btn.click(
            fn=add_hive_action,
            inputs=[new_name, new_location, new_marker, new_notes],
            outputs=[hives_table, hive_input, history_select, add_msg],
        )

        def _refresh_all():
            names = [h["name"] for h in db.list_hives()]
            return (
                _hives_table_state(),
                gr.update(choices=names),
                gr.update(choices=names),
            )

        refresh_btn.click(
            fn=_refresh_all,
            outputs=[hives_table, hive_input, history_select],
        )

        cmp_btn.click(
            fn=run_comparison,
            inputs=[cmp_image],
            outputs=[cmp_left_image, cmp_left_text, cmp_right_image, cmp_right_text],
        )

        history_select.change(
            fn=view_hive_history,
            inputs=[history_select],
            outputs=[history_table, history_summary, history_trend],
        )

        def _build_report():
            out = Path(__file__).parent / "reports"
            out.mkdir(exist_ok=True)
            ts = time.strftime("%Y%m%d-%H%M%S")
            target = out / f"apiarist_report_{ts}.pdf"
            report.save_report(target)
            return str(target)

        report_btn.click(fn=_build_report, outputs=[report_file])

    return app


app = build_ui()


if __name__ == "__main__":
    app.launch(css=custom_css)