"""
HPO Mapper 2 — FULL SCRIPT (HF InferenceClient chat.completions + user-selectable LLM)

✅ Embedding-based HPO mapping (SQLite embeddings)
✅ LLM QC using huggingface_hub.InferenceClient.chat.completions (your example style)
✅ User-selectable QC LLM (dropdown + custom entry)
✅ Robust QC parsing (handles empty content + retries + diagnostics)
✅ HPO ontology tree (auto-download hp.json if missing)
✅ Gene enrichment via Enrichr public API (no enrichrpy dependency)
✅ Single + Bulk modes + Embedding iframe views
✅ Runtime status panel shows DB/HPO readiness and a QC test for the currently selected model

HF Space requirements:
  pip install gradio sentence-transformers huggingface_hub requests numpy pandas

Secrets:
  HF_TOKEN = your Hugging Face token

Notes:
- Some models/providers may not support chat.completions via HF Inference API for your account.
  If QC fails, you will see ERR: ... in the QC box and qc_raw column.
"""

import os
import io
import json
import html
import tempfile
import zipfile
import sqlite3
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Dict, List, Any, Optional, Tuple

import gradio as gr
import numpy as np
from numpy.linalg import norm
import pandas as pd
import requests

from huggingface_hub import hf_hub_download, InferenceClient
from sentence_transformers import SentenceTransformer


# =========================
# ENV & CONSTANTS
# =========================

HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("Missing HF_TOKEN. Add it as a Secret in your Space.")

# Matches your example usage: no provider param, use chat.completions
client = InferenceClient(api_key=HF_TOKEN)

# Default QC model
DEFAULT_QC_MODEL = os.getenv("QC_MODEL_NAME", "openai/gpt-oss-120b")

# Your requested models (plus default). User can also type any model id.
QC_MODEL_CHOICES = [
    DEFAULT_QC_MODEL,
    "google/gemma-3-27b-it",
    "meta-llama/Llama-3.1-8B-Instruct",
    "deepseek-ai/DeepSeek-V3.2",
    "Qwen/Qwen3-VL-8B-Instruct",
]

# Embedding model
EMBEDDING_MODEL = "nomic-ai/nomic-embed-text-v1.5"
embedder = SentenceTransformer(EMBEDDING_MODEL, trust_remote_code=True)

# HPO DB (with synonyms + embeddings)
db_filename = "hpo_genes_with_synonyms.db"
db_repo = "UoS-HGIG/HPOmapper2"

db_path = os.path.join(os.getcwd(), db_filename)
if not os.path.exists(db_path):
    db_path = hf_hub_download(
        repo_id=db_repo,
        filename=db_filename,
        repo_type="space",
        token=HF_TOKEN,
    )

# HPO ontology JSON file (auto-download if missing)
HPO_JSON_PATH = "hp.json"
HPO_JSON_URL = "https://raw.githubusercontent.com/obophenotype/human-phenotype-ontology/master/hp.json"

# Embedding view HTML files (these live in the Space repo itself, so they are
# on local disk in the container; no need to go through huggingface.co URLs,
# which are served with headers that stop browsers rendering them in iframes)
EMBED_FILES = {
    "2D": "hp00118_branch_2d.html",
    "3D": "hp00118_branch_3d.html",
}

# Optional image
VISUAL_ABSTRACT_PATH = "HPO Mapper Visual Abstract(2).jpeg"


# =========================
# UTILS
# =========================

def cosine_sim(a: np.ndarray, b: np.ndarray) -> float:
    denom = (norm(a) * norm(b))
    if denom == 0:
        return -1.0
    return float(np.dot(a, b) / denom)

def safe_read_gradio_file(uploaded_file) -> Tuple[Optional[str], Optional[bytes]]:
    if uploaded_file is None:
        return None, None

    # Gradio 4/5: gr.File passes a filepath string (sometimes a str subclass
    # carrying a .name attribute). Handle that first; the old file-like and
    # dict payloads are kept as fallbacks for older Gradio versions.
    path = None
    if isinstance(uploaded_file, (str, os.PathLike)):
        path = str(uploaded_file)
    else:
        candidate = getattr(uploaded_file, "name", None) or getattr(uploaded_file, "path", None)
        if isinstance(candidate, str):
            path = candidate
    if path and os.path.exists(path):
        try:
            with open(path, "rb") as f:
                return os.path.basename(path), f.read()
        except Exception:
            pass

    try:
        filename = getattr(uploaded_file, "name", None) or getattr(uploaded_file, "orig_name", None)
        file_bytes = uploaded_file.read()
        if filename and file_bytes:
            return filename, file_bytes
    except Exception:
        pass
    try:
        if isinstance(uploaded_file, dict):
            filename = uploaded_file.get("name")
            data = uploaded_file.get("data")
            if isinstance(data, bytes):
                return filename, data
            if isinstance(data, str) and os.path.exists(data):
                with open(data, "rb") as f:
                    return filename, f.read()
    except Exception:
        pass
    return None, None


# =========================
# HPO MAPPER CORE
# =========================

_emb_cache = None  # (ids, names, matrix) with matrix rows L2-normalised

def _load_embedding_matrix():
    """Load all HPO embeddings from SQLite ONCE and keep them in memory as a
    normalised float32 matrix. The old code re-read and re-parsed the entire
    table for every input row, which made bulk runs scale as rows x DB size."""
    global _emb_cache
    if _emb_cache is not None:
        return _emb_cache
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT hpo_id, hpo_name, embedding FROM hpo_embeddings")
    ids, names, vecs = [], [], []
    for hpo_id, hpo_name, embedding_str in cursor.fetchall():
        ids.append(hpo_id)
        names.append(hpo_name)
        vecs.append(np.asarray(json.loads(embedding_str), dtype=np.float32))
    conn.close()
    matrix = np.vstack(vecs)
    norms = np.linalg.norm(matrix, axis=1, keepdims=True)
    norms[norms == 0] = 1.0
    _emb_cache = (ids, names, matrix / norms)
    return _emb_cache

def find_best_hpo_matches_batch(
    pairs: List[Tuple[str, str]], threshold: float
) -> List[Optional[dict]]:
    """Match many (finding, region) pairs at once: batched encoding, then a
    single matrix multiply per chunk. Returns a list aligned with `pairs`
    (dict on a hit, None when the best similarity is below threshold)."""
    if not pairs:
        return []
    ids, names, matrix = _load_embedding_matrix()
    texts = [f"{f} in {r}" if r else f for f, r in pairs]

    results: List[Optional[dict]] = []
    chunk = 256  # bounds the size of the similarity matrix in memory
    for start in range(0, len(texts), chunk):
        batch = texts[start:start + chunk]
        q = np.asarray(
            embedder.encode(batch, batch_size=64, show_progress_bar=False),
            dtype=np.float32,
        )
        if q.ndim == 1:
            q = q[None, :]
        qnorms = np.linalg.norm(q, axis=1, keepdims=True)
        qnorms[qnorms == 0] = 1.0
        sims = (q / qnorms) @ matrix.T
        best_idx = sims.argmax(axis=1)
        for row_i, col_i in enumerate(best_idx):
            score = float(sims[row_i, col_i])
            if score >= threshold:
                results.append(
                    {"hpo_id": ids[col_i], "hpo_term": names[col_i], "similarity": score}
                )
            else:
                results.append(None)
    return results

def find_best_hpo_match(finding: str, region: str, threshold: float):
    return find_best_hpo_matches_batch([(finding, region or "")], threshold)[0]

def get_genes_for_hpo(hpo_id: str):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    cursor.execute("SELECT genes FROM hpo_gene WHERE hpo_id = ?", (hpo_id,))
    result = cursor.fetchone()
    conn.close()
    return result[0].split(", ") if result and result[0] else []


# =========================
# QC (InferenceClient.chat.completions) — Robust
# =========================

def build_qc_prompt(row: dict) -> str:
    return (
        "You validate a mapping from a clinical/pathology finding to an HPO term.\n"
        "Return EXACTLY ONE character: 1 or 0.\n\n"
        f"Finding: {row.get('finding','')}\n"
        f"Region: {row.get('region','')}\n"
        f"HPO ID: {row.get('hpo_id','')}\n"
        f"HPO Term: {row.get('hpo_term','')}\n"
        f"Similarity: {row.get('similarity','')}\n\n"
        "Output:\n"
        "1 = incorrect mapping\n"
        "0 = correct mapping\n"
    )

def _extract_text(msg) -> str:
    # Different providers may populate different fields; check a few.
    for attr in ("content", "text", "output_text"):
        v = getattr(msg, attr, None)
        if isinstance(v, str) and v.strip():
            return v.strip()
    return ""

def _parse_first_bit(text: str) -> Optional[str]:
    if not text:
        return None
    for ch in text:
        if ch in ("0", "1"):
            return ch
    return None

def call_qc_model(prompt: str, model_name: str) -> str:
    """
    Returns "0"/"1" or "ERR: ...".
    Retries with larger max_tokens if content is empty / length-truncated.
    """
    try:
        # Attempt 1: modest tokens
        comp = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": "Reply with exactly ONE character: 0 or 1. No other text."},
                {"role": "user", "content": prompt},
            ],
            temperature=0.0,
            max_tokens=128,
        )
        choice = comp.choices[0]
        msg = choice.message
        text = _extract_text(msg)
        bit = _parse_first_bit(text)
        if bit:
            return bit

        fr1 = getattr(choice, "finish_reason", None)

        # Attempt 2: bigger budget + no system prompt (some backends behave better)
        comp2 = client.chat.completions.create(
            model=model_name,
            messages=[{"role": "user", "content": "Return exactly one character: 0 or 1.\n\n" + prompt}],
            temperature=0.0,
            max_tokens=512,
        )
        choice2 = comp2.choices[0]
        msg2 = choice2.message
        text2 = _extract_text(msg2)
        bit2 = _parse_first_bit(text2)
        if bit2:
            return bit2

        fr2 = getattr(choice2, "finish_reason", None)
        return f"ERR: No 0/1 found. fr1={fr1!r} raw1={text!r} fr2={fr2!r} raw2={text2!r}"

    except Exception as e:
        return f"ERR: {type(e).__name__}: {e}"

def qc_connectivity_check(model_name: str) -> str:
    resp = call_qc_model("Return exactly one character: 0", model_name)
    if resp == "0":
        return f"QC OK (model={model_name})"
    return f"QC ERROR (model={model_name}): {resp}"


# =========================
# HPO TREE VISUALISATION (hp.json)
# =========================

_hpo_graph_cache = None

def ensure_hp_json() -> Tuple[bool, str]:
    if os.path.exists(HPO_JSON_PATH):
        return True, "hp.json present."
    try:
        r = requests.get(HPO_JSON_URL, timeout=90)
        r.raise_for_status()
        with open(HPO_JSON_PATH, "wb") as f:
            f.write(r.content)
        return True, "hp.json downloaded."
    except Exception as e:
        return False, f"hp.json download failed: {type(e).__name__}: {e}"

def load_hpo_graph():
    global _hpo_graph_cache
    if _hpo_graph_cache is not None:
        return _hpo_graph_cache

    ok, _ = ensure_hp_json()
    if not ok:
        _hpo_graph_cache = None
        return None

    try:
        with open(HPO_JSON_PATH, "r", encoding="utf-8") as f:
            obj = json.load(f)
    except Exception:
        _hpo_graph_cache = None
        return None

    graphs = obj.get("graphs", [])
    if not graphs:
        _hpo_graph_cache = None
        return None

    g = graphs[0]
    nodes = g.get("nodes", [])
    edges = g.get("edges", [])

    def _norm_id(x: Optional[str]) -> Optional[str]:
        # obographs hp.json uses IRIs like
        # http://purl.obolibrary.org/obo/HP_0000118 while the DB stores
        # HP:0000118. Normalise so tree lookups actually connect.
        if not x:
            return x
        x = x.rsplit("/", 1)[-1]
        return x.replace("_", ":", 1) if x.startswith("HP_") else x

    id2label = {}
    for n in nodes:
        nid = _norm_id(n.get("id"))
        meta = n.get("meta") or {}
        lbl = meta.get("lbl") or n.get("lbl") or nid
        if nid:
            id2label[nid] = lbl

    parents = {}
    for e in edges:
        sub = _norm_id(e.get("sub"))
        obj_ = _norm_id(e.get("obj"))
        pred = e.get("pred", "")
        if "is_a" not in pred and "subClassOf" not in pred:
            continue
        if not sub or not obj_:
            continue
        parents.setdefault(sub, []).append(obj_)

    _hpo_graph_cache = {"id2label": id2label, "parents": parents, "root": "HP:0000118"}
    return _hpo_graph_cache

def get_paths_to_root(hpo_id: str, parents_map: Dict[str, List[str]], root: str) -> List[List[str]]:
    paths = []
    def dfs(curr, path):
        if curr == root:
            paths.append(path + [curr])
            return
        if curr not in parents_map:
            return
        for p in parents_map[curr]:
            if p in path:
                continue
            dfs(p, path + [curr])
    dfs(hpo_id, [])
    fixed = []
    for p in paths:
        if p and p[0] != hpo_id:
            fixed.append([hpo_id] + p)
        else:
            fixed.append(p)
    return fixed

def build_minimal_tree_for_matches(match_ids: List[str]):
    g = load_hpo_graph()
    if not g:
        return None
    root = g["root"]
    parents = g["parents"]

    keep_edges = set()
    for mid in match_ids:
        if not mid:
            continue
        for path in get_paths_to_root(mid, parents, root):
            for i in range(len(path) - 1):
                keep_edges.add((path[i + 1], path[i]))  # (parent, child)

    adj = {}
    for parent, child in keep_edges:
        adj.setdefault(parent, []).append(child)

    id2label = g["id2label"]
    for p in adj:
        adj[p].sort(key=lambda x: id2label.get(x, x))

    def make_node(nid):
        return {"id": nid, "children": [make_node(ch) for ch in adj.get(nid, [])]}

    return make_node(root)

def tree_to_html(tree: dict, id2label: dict, highlight: set) -> str:
    def any_descendant_hit(n):
        for ch in n.get("children", []):
            if ch["id"] in highlight or any_descendant_hit(ch):
                return True
        return False

    def node_html(n):
        nid = n["id"]
        label = id2label.get(nid, nid)
        cls = "hpo-hit" if nid in highlight else ""
        title = f"{label} ({nid})"

        if not n["children"]:
            return f"<div class='hpo-leaf {cls}'>{title}</div>"

        open_attr = " open" if (nid in highlight or any_descendant_hit(n)) else ""
        inner = "".join([node_html(ch) for ch in n["children"]])
        return (
            f"<details{open_attr} class='hpo-node {cls}'>"
            f"<summary>{title}</summary>"
            f"<div class='hpo-children'>{inner}</div>"
            f"</details>"
        )
    return node_html(tree)

def build_hpo_paths_html(match_ids: List[str]) -> str:
    ok, msg = ensure_hp_json()
    if not ok:
        return f"<div><b>Ontology view unavailable</b>: {msg}</div>"

    g = load_hpo_graph()
    if not g:
        return "<div><b>Ontology view unavailable</b>: could not parse hp.json.</div>"

    match_ids = [m for m in match_ids if m]
    if not match_ids:
        return "<div>No matched HPO IDs to display.</div>"

    tree = build_minimal_tree_for_matches(match_ids)
    if not tree:
        return "<div>Could not build ontology tree from hp.json.</div>"

    style = """
    <style>
      .hpo-wrap { font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Arial; font-size: 14px; line-height: 1.25; }
      .hpo-node { margin-left: 10px; padding: 2px 0; }
      .hpo-children { margin-left: 14px; border-left: 1px solid #ddd; padding-left: 10px; }
      .hpo-leaf { margin-left: 24px; padding: 2px 0; }
      .hpo-hit > summary, .hpo-leaf.hpo-hit { background: #fff3bf; border-radius: 6px; padding: 2px 6px; display: inline-block; }
      summary { cursor: pointer; }
      .hpo-note { color: #666; margin-bottom: 8px; }
    </style>
    """

    return (
        style
        + "<div class='hpo-wrap'>"
        + "<div class='hpo-note'>Ontology paths up to HP:0000118 (Phenotypic abnormality).</div>"
        + tree_to_html(tree, g["id2label"], set(match_ids))
        + "</div>"
    )


# =========================
# ENRICHR (API)
# =========================

def enrichr_query(genes: List[str], library: str = "GO_Biological_Process_2021") -> List[dict]:
    genes = [g for g in genes if g]
    if len(genes) < 3:
        return []

    add_url = "https://maayanlab.cloud/Enrichr/addList"
    enrich_url = "https://maayanlab.cloud/Enrichr/enrich"

    payload = {"list": "\n".join(genes), "description": "HPO Mapper gene list"}
    r1 = requests.post(add_url, files=payload, timeout=30)
    r1.raise_for_status()
    user_list_id = r1.json()["userListId"]

    r2 = requests.get(enrich_url, params={"userListId": user_list_id, "backgroundType": library}, timeout=30)
    r2.raise_for_status()
    data = r2.json()

    out = []
    rows = data.get(library, [])
    for row in rows[:20]:
        out.append({
            "rank": row[0],
            "term": row[1],
            "p_value": row[2],
            "z_score": row[3],
            "combined_score": row[4],
            "genes": row[5],
            "adj_p_value": row[6] if len(row) > 6 else None,
        })
    return out

def build_enrichr_html_from_rows(rows: List[Dict[str, Any]]) -> str:
    genes = []
    for r in rows or []:
        g = (r.get("genes") or "").strip()
        if not g:
            continue
        genes.extend([x.strip() for x in g.split(";") if x.strip()])

    genes = sorted(set(genes))
    if len(genes) < 3:
        return "<div>Not enough genes for enrichment (need ~3).</div>"

    try:
        top = enrichr_query(genes)
        if not top:
            return "<div>No enrichment results returned.</div>"

        rows_html = []
        for t in top:
            adj = t["adj_p_value"]
            adj_str = f"{adj:.2e}" if isinstance(adj, (int, float)) else ""
            rows_html.append(
                "<tr>"
                f"<td>{t['rank']}</td>"
                f"<td>{t['term']}</td>"
                f"<td>{t['p_value']:.2e}</td>"
                f"<td>{adj_str}</td>"
                f"<td>{t['combined_score']:.2f}</td>"
                f"<td style='max-width:420px; white-space:nowrap; overflow:hidden; text-overflow:ellipsis;' title='{t['genes']}'>{t['genes']}</td>"
                "</tr>"
            )

        style = """
        <style>
          .enrichr-wrap { font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Arial; font-size: 13px; }
          .enrichr-wrap table { border-collapse: collapse; width: 100%; }
          .enrichr-wrap th, .enrichr-wrap td { border: 1px solid #e5e5e5; padding: 6px 8px; vertical-align: top; }
          .enrichr-wrap th { background: #f7f7f7; text-align: left; }
          .enrichr-note { color: #666; margin: 6px 0 10px; }
        </style>
        """
        return (
            style
            + "<div class='enrichr-wrap'>"
            + "<div class='enrichr-note'>Top GO Biological Process enrichments (Enrichr API).</div>"
            + "<table><thead><tr><th>Rank</th><th>Term</th><th>P</th><th>Adj P</th><th>Combined</th><th>Genes</th></tr></thead>"
            + "<tbody>" + "".join(rows_html) + "</tbody></table></div>"
        )
    except Exception as e:
        return f"<div>Enrichment failed: {type(e).__name__}: {e}</div>"


# =========================
# EMBEDDING IFRAMES
# =========================

def _load_embed_html(view: str) -> Optional[str]:
    fname = EMBED_FILES.get(view, EMBED_FILES["2D"])
    path = os.path.join(os.getcwd(), fname)
    if not os.path.exists(path):
        # Fallback: pull it from the Space repo (e.g. if cwd differs)
        try:
            path = hf_hub_download(
                repo_id=db_repo, filename=fname, repo_type="space", token=HF_TOKEN
            )
        except Exception:
            return None
    try:
        with open(path, "r", encoding="utf-8") as f:
            return f.read()
    except Exception:
        return None


def get_embedding_iframe(view: str) -> str:
    content = _load_embed_html(view)
    if content is None:
        return (
            "<div>Embedding HTML not found. Make sure "
            f"<code>{EMBED_FILES.get(view)}</code> is committed to the Space repo.</div>"
        )
    # srcdoc inlines the document, so no cross-origin request is made and the
    # plot's JS runs inside the sandboxed frame.
    return (
        "<div style='width: 100%; height: 720px;'>"
        f"<iframe srcdoc=\"{html.escape(content, quote=True)}\" "
        "style='width: 100%; height: 100%; border: none;' "
        "sandbox='allow-scripts allow-downloads'></iframe>"
        "</div>"
    )


# =========================
# BULK PIPELINE
# =========================

def run_qc_on_df(df: pd.DataFrame, qc_model_name: str, progress_cb=None) -> pd.DataFrame:
    df = df.copy()

    # Identical finding/region/HPO mappings need only ONE QC verdict, and the
    # network-bound calls can run concurrently.
    prompts: Dict[Tuple[str, str, str], str] = {}
    for _, row in df.iterrows():
        if row.get("hpo_id") and row.get("hpo_term"):
            key = (row.get("finding", ""), row.get("region", ""), row.get("hpo_id", ""))
            if key not in prompts:
                prompts[key] = build_qc_prompt(row.to_dict())

    results: Dict[Tuple[str, str, str], str] = {}
    if prompts:
        total = len(prompts)
        done = 0
        with ThreadPoolExecutor(max_workers=6) as ex:
            futures = {
                ex.submit(call_qc_model, p, qc_model_name): k for k, p in prompts.items()
            }
            for fut in as_completed(futures):
                key = futures[fut]
                try:
                    results[key] = fut.result()
                except Exception as e:
                    results[key] = f"ERR: {type(e).__name__}: {e}"
                done += 1
                if progress_cb:
                    progress_cb(done, total)

    flags, raws = [], []
    for _, row in df.iterrows():
        if row.get("hpo_id") and row.get("hpo_term"):
            key = (row.get("finding", ""), row.get("region", ""), row.get("hpo_id", ""))
            qc = results.get(key, "")
            raws.append(qc)
            flags.append("1" if qc == "1" else "")
        else:
            raws.append("")
            flags.append("")
    df["qc_raw"] = raws
    df["flag"] = flags
    return df

def run_mapping_on_rows(rows: List[Dict[str, str]], threshold: float) -> pd.DataFrame:
    # Pathology tables repeat findings heavily, so map each unique
    # finding/region pair once and reuse the result for every occurrence.
    unique_pairs = sorted({(r["finding"], r.get("region", "") or "") for r in rows})
    matches = find_best_hpo_matches_batch(unique_pairs, threshold)
    match_by_pair = dict(zip(unique_pairs, matches))

    gene_cache: Dict[str, List[str]] = {}
    out = []
    for r in rows:
        finding = r["finding"]
        region = r.get("region", "") or ""
        sid = r.get("subject_id", "") or ""

        match = match_by_pair.get((finding, region))
        if match:
            hpo_id = match["hpo_id"]
            if hpo_id not in gene_cache:
                gene_cache[hpo_id] = get_genes_for_hpo(hpo_id)
            genes = gene_cache[hpo_id]
            out.append({
                "subject_id": sid,
                "finding": finding,
                "region": region,
                "hpo_id": hpo_id,
                "hpo_term": match["hpo_term"],
                "similarity": match.get("similarity", ""),
                "genes": ";".join(genes) if genes else "",
                "flag": "",
                "qc_raw": "",
            })
        else:
            out.append({
                "subject_id": sid,
                "finding": finding,
                "region": region,
                "hpo_id": "",
                "hpo_term": "",
                "similarity": "",
                "genes": "",
                "flag": "",
                "qc_raw": "",
            })
    return pd.DataFrame(out)

def _read_table_file_to_rows(file_bytes: bytes, filename: str) -> List[Dict[str, str]]:
    suffix = (filename or "").lower()
    sep = "\t" if suffix.endswith(".tsv") else None
    df = pd.read_csv(io.BytesIO(file_bytes), sep=sep, engine="python")

    col_map = {c.lower().strip(): c for c in df.columns}

    def pick_col(cands):
        for c in cands:
            if c in col_map:
                return col_map[c]
        return None

    finding_col = pick_col(["finding", "pathological_finding", "findings"])
    if finding_col is None:
        raise ValueError(f"No finding column found in {filename}")

    region_col = pick_col(["region", "anatomical_region", "location", "site"])
    subject_col = pick_col(["subject_id", "patient", "patient_id", "id"])

    rows = []
    for _, r in df.iterrows():
        finding = "" if pd.isna(r[finding_col]) else str(r[finding_col]).strip()
        if not finding:
            continue
        region = "" if (not region_col or pd.isna(r[region_col])) else str(r[region_col]).strip()
        subject_id = "" if (not subject_col or pd.isna(r[subject_col])) else str(r[subject_col]).strip()
        rows.append({"subject_id": subject_id, "finding": finding, "region": region})
    return rows

def _read_subject_json_to_rows(obj: Any) -> List[Dict[str, str]]:
    def one_subject(o):
        sid = str(o.get("subject_id", "") or "")
        out = []
        for item in (o.get("data") or []):
            finding = str(item.get("finding", "") or "").strip()
            if not finding:
                continue
            region = str(item.get("anatomical_region", "") or "").strip()
            out.append({"subject_id": sid, "finding": finding, "region": region})
        return out

    if isinstance(obj, list):
        rows = []
        for o in obj:
            if isinstance(o, dict):
                rows.extend(one_subject(o))
        return rows

    if isinstance(obj, dict):
        if "subject_id" in obj and "data" in obj:
            return one_subject(obj)
        if "finding" in obj:
            return [{
                "subject_id": str(obj.get("subject_id", "") or ""),
                "finding": str(obj.get("finding", "") or "").strip(),
                "region": str(obj.get("anatomical_region", obj.get("region", "")) or "").strip(),
            }]
    return []

def parse_bulk_inputs(uploaded_file, pasted_table: Optional[pd.DataFrame]) -> List[Dict[str, str]]:
    rows: List[Dict[str, str]] = []

    if pasted_table is not None and isinstance(pasted_table, pd.DataFrame) and len(pasted_table) > 0:
        df = pasted_table.copy()
        df.columns = [str(c).strip() for c in df.columns]
        lower_cols = {c.lower(): c for c in df.columns}

        def getcol(name, fallback_idx=None):
            if name in lower_cols:
                return lower_cols[name]
            if fallback_idx is not None and fallback_idx < len(df.columns):
                return df.columns[fallback_idx]
            return None

        sub_c = getcol("subject_id", 0)
        fin_c = getcol("finding", 1)
        reg_c = getcol("region", 2)

        for _, r in df.iterrows():
            finding = "" if pd.isna(r[fin_c]) else str(r[fin_c]).strip()
            if not finding:
                continue
            subject_id = "" if (not sub_c or pd.isna(r[sub_c])) else str(r[sub_c]).strip()
            region = "" if (not reg_c or pd.isna(r[reg_c])) else str(r[reg_c]).strip()
            rows.append({"subject_id": subject_id, "finding": finding, "region": region})

    filename, file_bytes = safe_read_gradio_file(uploaded_file)
    if filename and file_bytes:
        if filename.lower().endswith(".zip"):
            with zipfile.ZipFile(io.BytesIO(file_bytes), "r") as zf:
                for info in zf.infolist():
                    if info.is_dir():
                        continue
                    inner_name = info.filename
                    inner_bytes = zf.read(inner_name)
                    lower = inner_name.lower()
                    try:
                        if lower.endswith(".json"):
                            obj = json.loads(inner_bytes.decode("utf-8"))
                            rows.extend(_read_subject_json_to_rows(obj))
                        elif lower.endswith((".csv", ".tsv", ".txt")):
                            rows.extend(_read_table_file_to_rows(inner_bytes, inner_name))
                    except Exception:
                        continue
        elif filename.lower().endswith(".json"):
            obj = json.loads(file_bytes.decode("utf-8"))
            rows.extend(_read_subject_json_to_rows(obj))
        elif filename.lower().endswith((".csv", ".tsv", ".txt")):
            rows.extend(_read_table_file_to_rows(file_bytes, filename))

    cleaned = []
    for r in rows:
        finding = str(r.get("finding", "") or "").strip()
        if not finding:
            continue
        cleaned.append({
            "subject_id": str(r.get("subject_id", "") or "").strip(),
            "finding": finding,
            "region": str(r.get("region", "") or "").strip(),
        })
    return cleaned

def process_bulk(uploaded_file, threshold, pasted_table, qc_model_name, run_qc=True, progress=gr.Progress()):
    try:
        progress(0.0, desc="Parsing inputs")
        rows = parse_bulk_inputs(uploaded_file, pasted_table)
        if not rows:
            return (
                "No valid rows found. The file needs a 'finding' column "
                "(optional: 'region', 'subject_id').",
                None, None, "<div>No ontology view.</div>", "<div>No enrichment.</div>",
            )

        n_unique = len({(r["finding"], r.get("region", "") or "") for r in rows})
        progress(0.1, desc=f"Mapping {len(rows)} rows ({n_unique} unique finding/region pairs)")
        df_map = run_mapping_on_rows(rows, threshold)

        if run_qc:
            def qc_progress(done, total):
                progress(0.3 + 0.6 * (done / max(total, 1)), desc=f"LLM QC {done}/{total}")
            df_out = run_qc_on_df(df_map, qc_model_name, progress_cb=qc_progress)
        else:
            df_out = df_map

        progress(0.92, desc="Building outputs")
        tmp_dir = tempfile.mkdtemp(prefix="hpo_mapper_qc_")
        out_path = os.path.join(tmp_dir, "hpo_mapper_qc_results.csv")
        df_out.to_csv(out_path, index=False)

        hpo_ids = [x for x in df_out["hpo_id"].dropna().astype(str).tolist() if x.strip()]
        ontology_html = build_hpo_paths_html(hpo_ids)
        enrichment_html = build_enrichr_html_from_rows(df_out.to_dict(orient="records"))

        n_mapped = int((df_out["hpo_id"].astype(str).str.strip() != "").sum())
        n_flagged = int((df_out["flag"] == "1").sum())
        n_err = int(df_out["qc_raw"].astype(str).str.startswith("ERR").sum())
        qc_part = (
            f"QC flagged: {n_flagged}. QC errors: {n_err}." if run_qc else "QC skipped."
        )
        status = (
            f"Processed {len(df_out)} rows ({n_unique} unique finding/region pairs). "
            f"Mapped: {n_mapped}. Below threshold: {len(df_out) - n_mapped}. " + qc_part
        )
        return status, out_path, df_out.head(500), ontology_html, enrichment_html

    except Exception as e:
        return f"Bulk processing failed: {type(e).__name__}: {e}", None, None, "<div>No ontology view.</div>", "<div>No enrichment.</div>"


# =========================
# SINGLE QUERY UI FUNCTIONS
# =========================

def map_only_ui(finding, region, threshold, qc_model_name):
    # mapping step; qc runs in next step but we keep qc model in state
    if not finding or not str(finding).strip():
        return "", "", "", "", "Enter a finding.", [], "<div>No ontology view.</div>", "<div>No enrichment.</div>"

    finding = str(finding).strip()
    region = str(region).strip() if region else ""

    match = find_best_hpo_match(finding, region, threshold)
    if not match:
        return "", "", "", "", "No match found.", [], "<div>No matched HPO IDs to display.</div>", "<div>Not enough genes for enrichment.</div>"

    genes = get_genes_for_hpo(match["hpo_id"])
    row_state = {
        "finding": finding,
        "region": region,
        "hpo_id": match["hpo_id"],
        "hpo_term": match["hpo_term"],
        "similarity": match.get("similarity", ""),
        "genes": ";".join(genes) if genes else "",
        "flag": "",
        "qc_raw": "",
        "subject_id": "",
        "qc_model": qc_model_name,
    }

    ontology_html = build_hpo_paths_html([match["hpo_id"]])
    enrichment_html = build_enrichr_html_from_rows([row_state])

    return (
        match["hpo_id"],
        match["hpo_term"],
        ", ".join(genes),
        f"{match.get('similarity', 0):.4f}",
        "QC running…",
        [row_state],
        ontology_html,
        enrichment_html,
    )

def qc_only_ui(row_state_records):
    if not row_state_records or not isinstance(row_state_records, list):
        return "—", row_state_records, "<div>No ontology view.</div>", "<div>No enrichment.</div>"

    row = row_state_records[0]
    if not row.get("hpo_id") or not row.get("hpo_term"):
        return "—", row_state_records, "<div>No ontology view.</div>", "<div>No enrichment.</div>"

    qc_model_name = row.get("qc_model") or DEFAULT_QC_MODEL
    qc = call_qc_model(build_qc_prompt(row), qc_model_name)
    row["qc_raw"] = qc

    if qc.startswith("ERR:"):
        row["flag"] = ""
        qc_display = qc
    else:
        row["flag"] = "1" if qc == "1" else ""
        qc_display = row["flag"]

    ontology_html = build_hpo_paths_html([row["hpo_id"]])
    enrichment_html = build_enrichr_html_from_rows([row])

    return qc_display, row_state_records, ontology_html, enrichment_html


# =========================
# GRADIO APP
# =========================

custom_css = "#visual-abstract img { max-width: 600px; width: 100%; height: auto; display: block; margin-left: auto; margin-right: auto; }"

description_md = (
    "Enter a pathological finding (e.g., *chronic inflammation*) and optional anatomical region "
    "(e.g., *terminal ileum*) to map it to the closest HPO term using Nomic embeddings.\n\n"
    "QC is performed by **InferenceClient.chat.completions**. Choose the QC model from the dropdown "
    "(you can also type any HF model id).\n\n"
    "If QC fails, you will see **ERR:** in the QC box and in the bulk output `qc_raw` column.\n"
)

example_json_md = (
    "Example JSON:\n\n"
    "```json\n"
    "{\n"
    '  \"subject_id\": \"U099\",\n'
    '  \"data\": [\n'
    "    {\n"
    '      \"finding\": \"Post COVID-19 condition, unspecified\",\n'
    '      \"anatomical_region\": \"\"\n'
    "    }\n"
    "  ]\n"
    "}\n"
    "```"
)

with gr.Blocks(title="HPO Mapper 2 — Selectable LLM QC", css=custom_css) as demo:
    gr.Markdown("# HPO Mapper 2 — Selectable LLM QC (Hugging Face InferenceClient.chat.completions)")

    gr.Markdown(
        "**Preprint:** Kadhim AZ, Green Z, Boags A, et al. "
        "[Human Phenotype Ontology (HPO) Mapper: Semantic Mapping of Clinical Findings "
        "to the Human Phenotype Ontology Using AI-Powered Embeddings and LLM-Based "
        "Quality Control](https://www.medrxiv.org/content/10.64898/2025.12.20.25342726v2). "
        "*medRxiv* (2025). "
        "doi: [10.64898/2025.12.20.25342726](https://doi.org/10.64898/2025.12.20.25342726)"
    )

    if os.path.exists(VISUAL_ABSTRACT_PATH):
        gr.Image(value=VISUAL_ABSTRACT_PATH, show_label=False, elem_id="visual-abstract")

    gr.Markdown(description_md)

    mapped_state = gr.State([])

    with gr.Tab("Single Query"):
        qc_model_in = gr.Dropdown(
            label="QC LLM (Hugging Face model id)",
            choices=QC_MODEL_CHOICES,
            value=DEFAULT_QC_MODEL,
            allow_custom_value=True,
        )

        finding_in = gr.Textbox(label="Pathological Finding")
        region_in = gr.Textbox(label="Anatomical Region (optional)")
        threshold_in = gr.Slider(0.0, 1.0, step=0.01, value=0.76, label="Similarity Threshold")

        hpo_id_out = gr.Textbox(label="HPO ID")
        hpo_term_out = gr.Textbox(label="HPO Term")
        genes_out = gr.Textbox(label="Genes annotated to this HPO term")
        sim_out = gr.Textbox(label="Similarity")
        qc_flag_out = gr.Textbox(label="LLM QC (blank=OK, 1=incorrect, ERR=API issue)")

        run_btn = gr.Button("Run Mapping → QC")

        gr.Markdown("### Visual outputs")
        ontology_html_out = gr.HTML(label="HPO Ontology Tree")
        enrichr_html_out = gr.HTML(label="Gene Enrichment (Enrichr)")

        event = run_btn.click(
            fn=map_only_ui,
            inputs=[finding_in, region_in, threshold_in, qc_model_in],
            outputs=[hpo_id_out, hpo_term_out, genes_out, sim_out, qc_flag_out, mapped_state, ontology_html_out, enrichr_html_out],
        )
        event.then(
            fn=qc_only_ui,
            inputs=[mapped_state],
            outputs=[qc_flag_out, mapped_state, ontology_html_out, enrichr_html_out],
        )

    with gr.Tab("Bulk Mapping + QC"):
        gr.Markdown("Upload CSV/TSV/JSON or ZIP, and/or paste rows below. Output includes qc_raw for debugging.")
        gr.Markdown(example_json_md)

        bulk_qc_model = gr.Dropdown(
            label="QC LLM (Hugging Face model id)",
            choices=QC_MODEL_CHOICES,
            value=DEFAULT_QC_MODEL,
            allow_custom_value=True,
        )

        bulk_file = gr.File(label="Upload CSV/TSV/JSON/ZIP", file_types=[".csv", ".tsv", ".txt", ".json", ".zip"])
        bulk_table = gr.Dataframe(
            headers=["subject_id", "finding", "region"],
            datatype=["str", "str", "str"],
            row_count=12,
            col_count=3,
            label="Paste or type rows here (optional)",
        )
        bulk_threshold = gr.Slider(0.0, 1.0, step=0.01, value=0.76, label="Similarity Threshold")
        bulk_run_qc = gr.Checkbox(value=True, label="Run LLM QC (uncheck for fast mapping-only runs)")
        bulk_btn = gr.Button("Run Bulk Mapping → QC")

        bulk_status = gr.Textbox(label="Status / Summary", interactive=False)
        bulk_file_out = gr.File(label="Download Results CSV")
        bulk_preview_out = gr.Dataframe(label="Preview Results (first 500 rows)", interactive=False, wrap=True)

        bulk_ontology_html = gr.HTML(label="HPO Ontology Tree")
        bulk_enrichr_html = gr.HTML(label="Gene Enrichment (Enrichr)")

        bulk_btn.click(
            fn=process_bulk,
            inputs=[bulk_file, bulk_threshold, bulk_table, bulk_qc_model, bulk_run_qc],
            outputs=[bulk_status, bulk_file_out, bulk_preview_out, bulk_ontology_html, bulk_enrichr_html],
        )

    with gr.Tab("Embedding Visualisation"):
        view_mode = gr.Dropdown(choices=["2D", "3D"], value="3D", label="View mode")
        html_view = gr.HTML(value=get_embedding_iframe("3D"), label="Embedding Viewer")
        view_mode.change(fn=get_embedding_iframe, inputs=view_mode, outputs=html_view)

    with gr.Accordion("Runtime status", open=True):
        runtime_status = gr.Textbox(label="Runtime status", interactive=False)
        runtime_model = gr.Dropdown(
            label="Model to test in Runtime status",
            choices=QC_MODEL_CHOICES,
            value=DEFAULT_QC_MODEL,
            allow_custom_value=True,
        )
        test_btn = gr.Button("Test QC model now")

        gr.Markdown(
            f"- HF token env: `HF_TOKEN`\n"
            f"- hp.json: `{HPO_JSON_PATH}` (auto-download)\n"
            f"- DB: `{db_path}`\n"
        )

        def run_runtime_test(model_name: str):
            qc_line = qc_connectivity_check(model_name)
            hpo_ok, hpo_msg = ensure_hp_json()
            hpo_line = "HPO OK (hp.json ready)" if hpo_ok else f"HPO ERROR: {hpo_msg}"
            db_line = "DB OK" if os.path.exists(db_path) else "DB ERROR: missing db"
            return f"{qc_line}\n{hpo_line}\n{db_line}"

        test_btn.click(fn=run_runtime_test, inputs=[runtime_model], outputs=[runtime_status])

        # initial load
        demo.load(fn=run_runtime_test, inputs=[runtime_model], outputs=runtime_status)

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