#!/usr/bin/env python3
from __future__ import annotations

# ================== STD / CORE ==================
import os, re, json, time, sys, csv, uuid, datetime
from typing import List, Dict, Any, Optional
from functools import lru_cache
from xml.etree import ElementTree as ET

# ================== TRANSFORMERS / TORCH ==================
from transformers import AutoTokenizer, AutoModelForCausalLM
try:
    from transformers import BitsAndBytesConfig  # may exist even if bnb isn't installed
except Exception:
    BitsAndBytesConfig = None

import numpy as np
import requests
import gradio as gr

# ================== CONFIG / PATHS ==================
ASSETS_DIR = os.environ.get("ASSETS_DIR", "assets")
FAISS_PATH = os.environ.get("FAISS_PATH", f"{ASSETS_DIR}/index.faiss")
META_PATH  = os.environ.get("META_PATH",  f"{ASSETS_DIR}/index_meta.filtered.jsonl")
REL_CONFIG_PATH = os.environ.get("REL_CONFIG_PATH", f"{ASSETS_DIR}/relevance_config.json")

# Normalize QUANTIZE env (default: no quantization)
QUANTIZE = os.environ.get("QUANTIZE", "none").strip().lower()  # "none" | "8bit" | "4bit"

# Turn logging
TRANSCRIPT_PATH   = os.environ.get("TRANSCRIPT_PATH", "transcripts.jsonl")
PUSH_TRANSCRIPTS  = os.environ.get("PUSH_TRANSCRIPTS", "1") == "1"  # set "0" to disable

# Models
BASE_MODEL   = os.environ.get("BASE_MODEL", "mistralai/Mistral-7B-Instruct-v0.2")
ADAPTER_PATH = os.environ.get("ADAPTER_PATH", f"{ASSETS_DIR}/lora_adapter")  # local folder if present
ADAPTER_REPO = os.environ.get("ADAPTER_REPO", "")  # optional HF repo id for adapter

# Embeddings used to build FAISS
EMBED_MODEL_NAME = os.environ.get("EMBED_MODEL_NAME", "pritamdeka/S-PubMedBERT-MS-MARCO")

# PubMed / NCBI (set in Space Secrets if you have them)
NCBI_EMAIL  = os.environ.get("NCBI_EMAIL", "")
NCBI_TOOL   = os.environ.get("NCBI_TOOL", "askstein")
NCBI_APIKEY = os.environ.get("NCBI_APIKEY", "")

# Feedback logging
FEEDBACK_PATH   = os.environ.get("FEEDBACK_PATH", "feedback.csv")
PUSH_FEEDBACK   = os.environ.get("PUSH_FEEDBACK", "0") == "1"  # set "1" to enable Hub upload
HF_READ_TOKEN   = os.environ.get("HF_READ_TOKEN", os.environ.get("HF_TOKEN", ""))
HF_WRITE_TOKEN  = os.environ.get("HF_WRITE_TOKEN", HF_READ_TOKEN)
SPACE_REPO_ID   = os.environ.get("SPACE_REPO_ID", "")

# Generation / toggles
ALLOW_WIKIPEDIA = False
DEBUG = True
MAX_NEW_TOKENS_GROUNDED = 512
MAX_NEW_TOKENS_FALLBACK = 256
MIN_USEFUL_CHARS = 260

# Auto-continue if we hit the cap without EOS
AUTO_CONTINUE = True
AUTO_CONT_MAX_STEPS = 2        # continue up to 2 extra chunks
AUTO_CONT_NEW_TOKENS = 256     # tokens per continuation step

def dlog(tag, msg):
    if DEBUG: print(f"[{tag}] {msg}")

    
# --- Logging/upload config (OFF by default unless env says "1")
PUSH_TRANSCRIPTS  = os.environ.get("PUSH_TRANSCRIPTS", "0") == "1"
PUSH_FEEDBACK     = os.environ.get("PUSH_FEEDBACK", "0") == "1"
HF_WRITE_TOKEN    = os.environ.get("HF_WRITE_TOKEN", os.environ.get("HF_TOKEN", ""))
LOGS_DATASET_ID   = os.environ.get("LOGS_DATASET_ID", "")

def _log_path(prefix: str, ext: str) -> str:
    # rotate daily, e.g., transcripts-20250929.jsonl
    today = datetime.datetime.utcnow().strftime("%Y%m%d")
    return f"{prefix}-{today}.{ext}"

# rotate by default if user didn't override via env
TRANSCRIPT_PATH = os.environ.get("TRANSCRIPT_PATH", _log_path("transcripts", "jsonl"))
FEEDBACK_PATH   = os.environ.get("FEEDBACK_PATH",   _log_path("feedback",   "csv"))

# ================== HEAVY IMPORTS ==================
import faiss
from sentence_transformers import SentenceTransformer
import torch
from peft import PeftModel
import wikipedia
from wikipedia.exceptions import DisambiguationError, PageError
from huggingface_hub import login, snapshot_download, HfApi

# ================== LOW-VRAM RUNTIME KNOBS ==================
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "max_split_size_mb:128,expandable_segments:True")

# ================== GPU CHECK ==================
if not torch.cuda.is_available():
    with gr.Blocks() as demo:
        gr.Markdown("### 🚦 GPU not detected.\nThis Space needs a GPU (A10G or better). Go to **Settings → Hardware** and select a GPU.")
    demo.launch()
    sys.exit(0)

device = "cuda"
dtype = torch.float16
torch.manual_seed(42)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

# ================== RELEVANCE CONFIG ==================
DEFAULT_REL_CONFIG = {
    "positive_terms": [
        "ctra","rigidity","ct-based","qct","micro-ct","hounsfield",
        "femur","femoral","hip","proximal femur",
        "bending","torsional","axial","failure load","modulus",
        "nazarian","freedman","alboro"
    ],
    "negative_terms": [
        "t cell","lymph","immunolog","synapse","receptor","egfr",
        "tumor","oncolog","immune","lymph node","cardio","myocard","neuro","skull","heart","brain"
    ],
    "weights": {"positive": 2, "negative": 1},
    "author_whitelist": ["Nazarian","Freedman","Alboro"],
    "mesh_positive": [
        "Femur", "Femoral Neck", "Hip", "Bone Density",
        "Tomography, X-Ray Computed", "Finite Element Analysis",
        "Bone and Bones", "Elastic Modulus", "Biomechanical Phenomena"
    ],
    "mesh_weight": 2,
    "author_weight": 3,
    "min_rel_to_use_faiss": 3,
    "ncbi_email": NCBI_EMAIL,
    "ncbi_tool":  NCBI_TOOL,
    "ncbi_apikey": NCBI_APIKEY,
}

def load_rel_config(path: str) -> Dict[str, Any]:
    cfg = DEFAULT_REL_CONFIG.copy()
    try:
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                user_cfg = json.load(f)
            cfg.update(user_cfg)
    except Exception as e:
        dlog("rel-config", f"using defaults ({e})")
    return cfg

REL_CFG = load_rel_config(REL_CONFIG_PATH)
dlog("rel-config", f"Loaded keys: {list(REL_CFG.keys())}")

# ================== HTTP / NCBI HELPERS ==================
class _Http:
    session = requests.Session()
    session.headers.update({"User-Agent": "Askstein/1.0 (https://huggingface.co)"})

def _ncbi_params(extra: Dict[str, Any] | None = None) -> Dict[str, Any]:
    p = {"retmode": "xml"}
    if REL_CFG.get("ncbi_email"): p["email"] = REL_CFG["ncbi_email"]
    if REL_CFG.get("ncbi_tool"):  p["tool"]  = REL_CFG["ncbi_tool"]
    if REL_CFG.get("ncbi_apikey"): p["api_key"] = REL_CFG["ncbi_apikey"]
    if extra: p.update(extra)
    return p

def _get_with_backoff(url: str, params: Dict[str, Any], tries: int = 2, base_sleep: float = 0.5, timeout: int = 6) -> str:
    for i in range(tries):
        try:
            if not REL_CFG.get("ncbi_apikey"):
                time.sleep(0.34)
            r = _Http.session.get(url, params=params, timeout=timeout)
            r.raise_for_status()
            return r.text
        except Exception:
            if i == tries - 1: raise
            time.sleep(base_sleep * (2 ** i))

# ================== WIKIPEDIA ==================
def wiki_summary_allow(query: str, sentences: int = 2) -> Optional[str]:
    prev = globals().get("ALLOW_WIKIPEDIA", False)
    globals()["ALLOW_WIKIPEDIA"] = True
    try:
        q = re.sub(r'^(what is|what are|define|where is|where are)\s+', '', query, flags=re.IGNORECASE)
        q = re.sub(r'\s+(located|location)\s*\?*$', '', q, flags=re.IGNORECASE).strip('?').strip()
        return wikipedia.summary(q, sentences=sentences)
    except (DisambiguationError, PageError, Exception):
        return None
    finally:
        globals()["ALLOW_WIKIPEDIA"] = prev

# ================== LOAD FAISS + META + EMBEDDER ==================
dlog("FAISS", "Loading index…")
index = faiss.read_index(FAISS_PATH)
dlog("FAISS", f"ntotal={index.ntotal}")

dlog("FAISS", "Loading metadata…")
all_chunks: List[Dict[str, Any]] = []
with open(META_PATH, "r", encoding="utf-8") as f:
    for line in f:
        try:
            all_chunks.append(json.loads(line))
        except Exception:
            pass
dlog("FAISS", f"Metadata records: {len(all_chunks)}")

if len(all_chunks) != index.ntotal:
    raise RuntimeError(f"[ALIGNMENT] Metadata rows ({len(all_chunks)}) != FAISS ntotal ({index.ntotal}).")

dlog("EMBED", f"Loading embedder: {EMBED_MODEL_NAME}")
embed_model = SentenceTransformer(EMBED_MODEL_NAME)

try:
    _probe = embed_model.encode(["__dimcheck__"], convert_to_numpy=True).astype("float32")
    _dim = _probe.shape[1] if _probe.ndim == 2 else len(_probe)
    assert index.d == _dim, f"FAISS dim {index.d} != embed dim {_dim} (model={EMBED_MODEL_NAME}). Rebuild index."
except Exception as e:
    raise RuntimeError(f"[FAISS] Dimension check failed: {e}")

_IS_IP = isinstance(index, faiss.IndexFlatIP) or "IndexFlatIP" in type(index).__name__

# ================== HUGGING FACE LOGIN & ADAPTER PATH ==================
if HF_READ_TOKEN:
    try:
        login(token=HF_READ_TOKEN)
        dlog("HF", "Login (read) OK")
    except Exception as e:
        dlog("HF", f"Login issue: {e}")

if ADAPTER_REPO:
    ADAPTER_PATH = snapshot_download(repo_id=ADAPTER_REPO, allow_patterns=["*"])

# ================== QUANTIZATION AVAILABILITY ==================
try:
    import bitsandbytes as _bnb  # noqa: F401
    _BNB_AVAILABLE = True
except Exception:
    _BNB_AVAILABLE = False

# ================== LLM (BASE + LORA) ==================
dlog("LLM", f"Loading base model: {BASE_MODEL}")
tokenizer_lm = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)

use_bnb = QUANTIZE in {"8bit", "4bit"} and BitsAndBytesConfig is not None and _BNB_AVAILABLE

if use_bnb:
    # Quantized path (only if explicitly requested and bnb is installed)
    bnb_config = BitsAndBytesConfig(
        load_in_8bit=(QUANTIZE == "8bit"),
        load_in_4bit=(QUANTIZE == "4bit"),
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=torch.float16,
    )
    base_model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        device_map="auto",
        quantization_config=bnb_config,
    )
else:
    # fp16 path with SDPA attention (lower VRAM). Fallback if not supported.
    try:
        base_model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL,
            torch_dtype=dtype,
            device_map="auto",
            attn_implementation="sdpa",
        )
    except TypeError:
        base_model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL,
            torch_dtype=dtype,
            device_map="auto",
        )

dlog("LLM", f"Loading LoRA adapter from: {ADAPTER_PATH}")
model_lm = PeftModel.from_pretrained(base_model, ADAPTER_PATH)
model_lm.eval()

GEN_ARGS_GROUNDED = dict(
    max_new_tokens=MAX_NEW_TOKENS_GROUNDED,
    min_new_tokens=220,
    do_sample=False,
    num_beams=1,
    no_repeat_ngram_size=3,
    repetition_penalty=1.08,
    eos_token_id=tokenizer_lm.eos_token_id,
)
GEN_ARGS_FALLBACK = dict(
    max_new_tokens=MAX_NEW_TOKENS_FALLBACK,
    min_new_tokens=120,
    do_sample=False,
    num_beams=1,
    no_repeat_ngram_size=3,
    repetition_penalty=1.05,
    eos_token_id=tokenizer_lm.eos_token_id,
)

# ================== UTILITIES ==================
_SANITIZE = re.compile(r"```.*?```|<\s*script[^>]*>.*?<\s*/\s*script\s*>", re.DOTALL|re.IGNORECASE)
def _to_text(rec: Any) -> str:
    if isinstance(rec, str): return rec.strip()
    for k in ("text","chunk_text","content","body","passage","raw_text","section_text","abstract"):
        v = rec.get(k)
        if isinstance(v, str) and v.strip():
            return _SANITIZE.sub("", v.strip())
    segs = rec.get("segments")
    if isinstance(segs, list):
        return _SANITIZE.sub("", " ".join(s.get("text","").strip() for s in segs if isinstance(s, dict)).strip())
    return ""

def _hydrate_text(item: Dict[str, Any]) -> str:
    t = _to_text(item)
    if t: return t
    pmid = item.get("pmid")
    if pmid:
        abs_chunks = fetch_pubmed_chunks(str(pmid), max_papers=1)
        if abs_chunks: return abs_chunks[0].get("text", "").strip()
    title = item.get("title")
    if isinstance(title, str) and title.strip():
        return title.strip()
    return ""

def _dedup_by_pmid_or_title(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    seen, out = set(), []
    for r in records:
        key = str(r.get("pmid") or "").strip()
        if not key: key = (r.get("title") or "").strip().lower()[:120]
        if not key: key = (r.get("text") or "").strip().lower()[:200]
        if key in seen: continue
        seen.add(key); out.append(r)
    return out

def _split_sentences(s: str) -> List[str]:
    s = s.replace("\r"," ").replace("\n"," ")
    parts = re.split(r"(?<=[\.\?\!])\s+", s)
    return [p.strip() for p in parts if p.strip()]

_BAD_BULLETS = re.compile(r"^\s*(?:\d+\s*\)|[•\-\*])\s*$", re.M)
_DANGLING = re.compile(r"[\[\(][^\]\)]$")
def _post_clean(text: str) -> str:
    t = re.sub(r"[ \t]+\n", "\n", text)
    t = _BAD_BULLETS.sub("", t)
    t = re.sub(r"\n{3,}", "\n\n", t).strip()
    sents = _split_sentences(t)
    seen, kept = set(), []
    for s in sents:
        key = s.lower()
        if key in seen: continue
        seen.add(key); kept.append(s)
    t = " ".join(kept)
    t = re.sub(_DANGLING, "", t).strip(" -,:;")
    return t

def _ensure_min_answer(ans: str) -> str:
    if len(ans) >= MIN_USEFUL_CHARS: return ans
    tail = " If you want, I can add a short checklist of assumptions, units, and typical parameter ranges."
    return (ans + tail) if not ans.endswith(".") else (ans + tail)

# ================== RELEVANCE / GATING ==================
def _rel_score(text: str, title: str = "", cfg: Dict[str, Any] | None = None) -> int:
    cfg = cfg or REL_CFG
    blob = (title + " " + text).lower()
    pos = sum(1 for k in cfg.get("positive_terms", []) if k.lower() in blob)
    neg = sum(1 for k in cfg.get("negative_terms", []) if k.lower() in blob)
    w_pos = int(cfg.get("weights", {}).get("positive", 2))
    w_neg = int(cfg.get("weights", {}).get("negative", 1))
    return pos * w_pos - neg * w_neg

@lru_cache(maxsize=4096)
def _mesh_by_pmid(pmid: str) -> List[str]:
    try:
        xml = _get_with_backoff(
            "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi",
            _ncbi_params({"db":"pubmed","id":str(pmid)})
        )
        root = ET.fromstring(xml)
        heads = []
        for mh in root.findall(".//MeshHeading"):
            dn = mh.find("DescriptorName")
            if dn is not None and dn.text:
                heads.append(dn.text.strip())
        return heads
    except Exception:
        return []

@lru_cache(maxsize=4096)
def _authors_by_pmid(pmid: str) -> List[str]:
    try:
        xml = _get_with_backoff(
            "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi",
            _ncbi_params({"db":"pubmed","id":str(pmid)})
        )
        root = ET.fromstring(xml)
        names = []
        for docsum in root.findall(".//DocSum"):
            for item in docsum.findall("Item"):
                if item.get("Name") == "AuthorList":
                    for au in item.findall("Item"):
                        if au.text:
                            last = au.text.split(",")[0].split(" ")[-1]
                            names.append(last)
        return names
    except Exception:
        return []

def _boost_by_author(pmid: str | int, rel_base: int, cfg: Dict[str, Any] | None = None) -> int:
    cfg = cfg or REL_CFG
    wl = set(cfg.get("author_whitelist", []))
    if not pmid or not wl: return rel_base
    names = _authors_by_pmid(str(pmid))
    return rel_base + int(cfg.get("author_weight", 3)) if any(n in wl for n in names) else rel_base

def _mesh_boost(pmid: str | int, rel_base: int, cfg: Dict[str, Any] | None = None) -> int:
    cfg = cfg or REL_CFG
    if not pmid: return rel_base
    targets = set(x.lower() for x in cfg.get("mesh_positive", []))
    weight  = int(cfg.get("mesh_weight", 2))
    heads   = [h.lower() for h in _mesh_by_pmid(str(pmid))]
    hit     = sum(1 for h in heads if h in targets)
    return rel_base + hit * weight

_MSK_MUST = re.compile(
    r"\b(femur|femoral|hip|proximal\s+femur|ctra|qct|ct-?based|rigidity|bending|torsional|axial|failure\s+load)\b",
    re.I
)
_CT_RIGIDITY_TOKENS = re.compile(r"\b(qct|ct[-\s]?based|ctra|rigidity|bending|torsion|hounsfield|finite\s+element|fe[am])\b", re.I)
_FE_TOKENS = re.compile(r"\b(fe|fea|finite\s+element|micromotion|boundary\s+conditions|nonlinear|yield|ultimate|fracture\s+load)\b", re.I)
_ANATOMY_OR_HISTORY = re.compile(
    r"(?:\bhistory\b.*\b(femur|hip|bone)\b|\bwhat\s+is\s+the\s+(femur|hip)\b|\banatomy\b.*\b(hip|femur)\b)",
    re.I
)
_PAPERS_INTENT = re.compile(r"\b(key\s+papers|suggest\s+papers|landmark|seminal|important|top\s+papers)\b", re.I)
CITE_TRIGGER = re.compile(
    r"\b(?:and\s+(?:cite|citations?|references?|studies?|papers?))\.?\s*$",
    re.IGNORECASE,
)


# ================== PUBMED & RETRIEVAL ==================
def fetch_pubmed_chunks(query_or_pmid: str, max_papers: int = 3) -> List[Dict[str, Any]]:
    retries = 1
    chunks: List[Dict[str, Any]] = []

    def _efetch(pmid: str):
        try:
            xml = _get_with_backoff(
                "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi",
                _ncbi_params({"db":"pubmed","id":pmid})
            )
            tree = ET.fromstring(xml)
            paras = [a.text for a in tree.findall(".//AbstractText") if a is not None and a.text]
            if paras:
                text = "\n".join(paras)
                chunks.append({"text": text, "source": "pubmed", "pmid": pmid})
        except Exception:
            pass

    if query_or_pmid.isdigit():
        _efetch(query_or_pmid); return chunks

    pmids: List[str] = []
    for attempt in range(retries + 1):
        try:
            xml = _get_with_backoff(
                "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi",
                _ncbi_params({"db":"pubmed","term":query_or_pmid, "retmax":max_papers})
            )
            root = ET.fromstring(xml)
            pmids = [e.text for e in root.findall(".//Id") if e is not None and e.text]
            break
        except Exception:
            if attempt == retries: return []
            time.sleep(0.5 * (2 ** attempt))

    for pmid in pmids[:max_papers]:
        _efetch(pmid)
    return chunks

STOPWORDS = set("the a an of and for with without to on in by from into over under how what why where when is are was were be been being this that these those it its as about".split())

def _parse_year(y: str) -> int:
    try:
        return int(re.findall(r"\d{4}", y or "")[0])
    except Exception:
        return 0

def _is_msk_paper(title: str, journal: str, year: str = "") -> bool:
    t = f"{title or ''} {journal or ''}".lower()
    must_body   = any(k in t for k in ["femur","femoral","hip","proximal femur","long bone","tibia","humerus"])
    must_method = any(k in t for k in ["ct","qct","finite element","structural rigidity","ctra","rigidity","bending","torsion"])
    if not (must_body and must_method): return False
    year_i = _parse_year(year)
    if year_i and not (2000 <= year_i <= 2025): return False
    return True

@lru_cache(maxsize=4096)
def fetch_pubmed_citations(query: str, max_results: int = 5) -> List[str]:
    try:
        xml = _get_with_backoff(
            "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi",
            _ncbi_params({"db":"pubmed","term":query, "retmax":max_results})
        )
        root = ET.fromstring(xml)
        pmids = [elem.text for elem in root.findall(".//Id") if elem is not None and elem.text]
        if not pmids: return []
    except Exception:
        return []

    try:
        xml = _get_with_backoff(
            "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esummary.fcgi",
            _ncbi_params({"db":"pubmed","id":",".join(pmids)})
        )
        summary_root = ET.fromstring(xml)
    except Exception:
        return []

    citations: List[str] = []
    for docsum in summary_root.findall(".//DocSum"):
        pmid = docsum.findtext("Id", default="")
        title = journal = year = doi = ""
        authors: List[str] = []
        for item in docsum.findall("Item"):
            name = item.get("Name", "")
            if name == "Title": title = item.text or ""
            elif name == "FullJournalName": journal = item.text or ""
            elif name == "PubDate": year = (item.text or "").split()[0]
            elif name == "AuthorList":
                for au in item.findall("Item"):
                    if au.text: authors.append(au.text)
            elif name == "ArticleIds":
                for sub in item.findall("Item"):
                    if sub.get("Name") == "doi":
                        doi = sub.text or ""
        if not _is_msk_paper(title, journal, year): continue
        first_author = authors[0] if authors else ""
        auth_str = f"{first_author} et al." if first_author else ""
        parts = [p for p in [auth_str, title, journal, year] if p]
        cit = ", ".join(parts).strip().rstrip(",")
        if pmid: cit += f"; PMID:{pmid}"
        if doi:  cit += f" DOI:{doi}"
        if cit: citations.append(cit)
    return citations[:max_results]

def _compact_terms(q: str) -> str:
    words = re.findall(r"[A-Za-z0-9\-]+", q.lower())
    keep = [w for w in words if w not in STOPWORDS and len(w) > 2]
    return " ".join(keep)[:200]

def _biomechish(q: str) -> bool:
    return bool(re.search(r"\b(femur|femoral|hip|bone|qct|ctra|rigidity|bending|torsion|elastic modulus|finite\s+element|fea)\b", q, re.I))
def _is_fe_override(q: str) -> bool:
    return bool(_FE_TOKENS.search(q))

_CONTRA_NO_EFFECT = re.compile(r"\b(no\s+significant\s+difference|no\s+effect|not\s+significant)\b", re.I)
_CONTRA_CHANGE    = re.compile(r"\b(increase[ds]?|decrease[ds]?|higher|lower|greater|reduced?)\b", re.I)
def _has_conflict(text: str) -> bool:
    return bool(_CONTRA_NO_EFFECT.search(text) and _CONTRA_CHANGE.search(text))

HARDCODED_CITS = {
    "EA": [
        "Morgan EF et al., Mechanical properties of cortical bone…, J Biomech, 2003; PMID:12547357",
        "Turner CH, Burr DB., Experimental techniques for bone mechanics, Bone, 1993; PMID:8252072"
    ],
    "EI": [
        "Courtney AC et al., Age-related reductions in the strength of the femur…, J Bone Miner Res, 1995; PMID:7584933",
        "Bell KL et al., Regional Heterogeneity of the Proximal Femur…, Bone, 1999; PMID:10574202"
    ],
    "GJ": [
        "Cowin SC., Bone Mechanics Handbook (torsion of bone cross-sections), CRC Press, 2001.",
        "Vollmer M et al., Long bone torsion testing methods, J Biomech, 1987; PMID:3670157"
    ]
}

def _fallback_cits_for(term: str) -> List[str]:
    return HARDCODED_CITS.get(term.upper(), [])

def detect_lab(q: str) -> str:
    ql = q.lower()
    if "freedman" in ql: return "freedman"
    if "alboro" in ql or "alborno" in ql: return "alboro"
    return "nazarian"

def build_lab_query(core_q: str, lab: str = "nazarian") -> str:
    topics = [
        "femur","femoral neck","hip","proximal femur",
        "CT","QCT","micro-CT","rigidity","CTRA","structural rigidity",
        "bending","torsional","axial","failure load","modulus","Hounsfield"
    ]
    ta = " OR ".join(f'"{t}"[Title/Abstract]' for t in topics)
    if lab == "freedman":
        author = '("Freedman BA"[Author] OR "Freedman"[Author])'
    elif lab == "alboro":
        author = '("Alboro"[Author] OR "Alborno"[Author])'
    else:
        author = '("Nazarian A"[Author] OR "Ara Nazarian"[Full Author Name])'
    date = '("2000"[Date - Publication] : "3000"[Date - Publication])'
    return f"{author} AND ({ta}) AND {date}"

def retrieve_context(query: str, top_k: int = 10) -> List[Dict[str, Any]]:
    q = query.strip()
    if _ANATOMY_OR_HISTORY.search(q) and not _biomechish(q):
        wiki = wiki_summary_allow(q, sentences=4)
        if wiki:
            dlog("WIKI", "Wikipedia biomechanics fallback hit")
            return [{"text": wiki, "source": "wikipedia"}]

    pm = re.search(r"pmid[:\s]*(\d+)", q, re.IGNORECASE)
    if pm:
        dlog("PMID", f"PMID inline {pm.group(1)}")
        return fetch_pubmed_chunks(pm.group(1), max_papers=1)

    if not (_CT_RIGIDITY_TOKENS.search(q) or _is_fe_override(q)):
        dlog("FALLBACK", "No CT/rigidity/FE tokens → try PubMed/Wiki")
        results = fetch_pubmed_chunks(q)
        if results:
            dlog("PUBMED", "PubMed search hit")
            return results
        wiki = wiki_summary_allow(q, sentences=3)
        if wiki:
            dlog("WIKI", "Wikipedia fallback hit")
            return [{"text": wiki, "source": "wikipedia"}]
        dlog("RETRIEVAL", "No results found")
        return []

    # FAISS path
    q_emb = embed_model.encode([q], convert_to_numpy=True).astype("float32")
    if _IS_IP:
        faiss.normalize_L2(q_emb)
    D, I = index.search(q_emb, top_k)
    results: List[Dict[str, Any]] = []
    for dist, idx_ in zip(D[0], I[0]):
        if idx_ < 0: continue
        item = all_chunks[idx_].copy()
        item["score"] = float(dist)
        item["text"]  = _hydrate_text(item)
        if not item["text"]: continue
        results.append(item)

    if results:
        scored: List[Dict[str, Any]] = []
        for it in results:
            rel = _rel_score(it.get("text", ""), it.get("title", ""), REL_CFG)
            rel = _boost_by_author(it.get("pmid"), rel, REL_CFG)
            rel = _mesh_boost(it.get("pmid"), rel, REL_CFG)
            it["_rel"] = rel
            scored.append(it)
        results = sorted(scored, key=lambda x: (x.get("_rel", 0), x.get("score", 0)), reverse=True)
        min_rel = int(REL_CFG.get("min_rel_to_use_faiss", 3))
        positives = [
            r for r in results
            if r.get("_rel", 0) >= min_rel and _MSK_MUST.search((r.get("title","")+" "+r.get("text","")))
        ]
        positives = _dedup_by_pmid_or_title(positives)
        if positives:
            dlog("FAISS", f"hit={len(positives)} (top rel={positives[0].get('_rel')} score={positives[0].get('score'):.3f})")
            return positives[:top_k]
        else:
            dlog("FALLBACK", "FAISS results off-topic → PubMed fallback")

    results = fetch_pubmed_chunks(q)
    if results:
        dlog("PUBMED", "PubMed search hit")
        return results

    wiki = wiki_summary_allow(q, sentences=3)
    if wiki:
        dlog("WIKI", "Wikipedia fallback hit")
        return [{"text": wiki, "source": "wikipedia"}]

    dlog("RETRIEVAL", "No results at all")
    return []

def build_prompt(chunks: List[Dict[str, Any]], question: str) -> str:
    header = (
        "You are Askstein (orthopedic biomechanics). Use ONLY the [Context] to answer. "
        "If the context is insufficient, say 'I don’t know based on the provided context.' "
        "Stay within musculoskeletal CT-based rigidity (EA, EI, GJ), femur/hip, CTRA/QCT, or FE modeling of these. "
        "Do not discuss cardiology, neurology, or unrelated domains."
    )
    cleaned = []
    per_chunk_chars = 900  # lower prompt length = lower KV memory
    for c in chunks:
        t = _to_text(c)
        if t: cleaned.append(t[:per_chunk_chars])
    context = "\n\n".join(cleaned)
    return f"{header}\n\n[Context]:\n{context}\n\n[Question]:\n{question}\n"

def _decode_generated(out_ids, in_len: int) -> str:
    gen = out_ids[0][in_len:]
    return tokenizer_lm.decode(gen, skip_special_tokens=True).lstrip(". \n").strip()

def _gen_once(inputs, args) -> Any:
    with torch.inference_mode():
        return model_lm.generate(**inputs, **args, use_cache=True)

def _generate(inputs, grounded: bool):
    args = GEN_ARGS_GROUNDED if grounded else GEN_ARGS_FALLBACK
    in_len = inputs["input_ids"].shape[-1]

    # First attempt
    try:
        out = _gen_once(inputs, args)
    except torch.cuda.OutOfMemoryError:
        try:
            torch.cuda.empty_cache()
        except Exception:
            pass
        small_args = dict(args)
        small_args["max_new_tokens"] = min(256, args.get("max_new_tokens", 256))
        # disable cache to save VRAM
        with torch.inference_mode():
            out = model_lm.generate(**inputs, **small_args, use_cache=False)

    if not AUTO_CONTINUE:
        return out

    # Auto-continue if we hit cap without EOS
    steps = 0
    while steps < AUTO_CONT_MAX_STEPS:
        seq = out[0]
        ended_with_eos = (seq[-1].item() == tokenizer_lm.eos_token_id)
        hit_cap = (seq.shape[0] - in_len) >= args["max_new_tokens"]
        if ended_with_eos or not hit_cap:
            break

        cont_inputs = {
            "input_ids": seq.unsqueeze(0),
            "attention_mask": torch.ones_like(seq).unsqueeze(0),
        }
        cont_inputs = {k: v.to(device) for k, v in cont_inputs.items()}
        cont_args = dict(args)
        cont_args["max_new_tokens"] = AUTO_CONT_NEW_TOKENS

        try:
            with torch.inference_mode():
                out = model_lm.generate(**cont_inputs, **cont_args, use_cache=True)
        except torch.cuda.OutOfMemoryError:
            try:
                torch.cuda.empty_cache()
            except Exception:
                pass
            with torch.inference_mode():
                out = model_lm.generate(**cont_inputs, **cont_args, use_cache=False)

        steps += 1

    return out

# ================== ANSWER SYNTHESIS ==================
def _synthesize_answer(chunks: List[Dict[str, Any]], question: str) -> str:
    prompt = build_prompt(chunks, question)
    inputs = tokenizer_lm(prompt, return_tensors="pt").to(device)
    in_len = inputs["input_ids"].shape[-1]
    out = _generate(inputs, grounded=True)
    answer = _decode_generated(out, in_len)
    return _post_clean(answer)

def _answer_from_chunks(chunks: List[Dict[str, Any]], question: str) -> str:
    joined = " ".join(_to_text(c) for c in chunks if _to_text(c))
    if _has_conflict(joined):
        dlog("SYNTH", "Conflict detected → summarize")
        return _synthesize_answer(chunks, question)
    return _synthesize_answer(chunks, question)

@lru_cache(maxsize=None)
def direct_llm_fallback(question: str) -> str:
    sys_prompt = (
        "You are Askstein (orthopedic biomechanics). If you lack enough domain context, say you don’t know. "
        "Avoid discussing non-musculoskeletal systems (cardiology, neurology). Do NOT invent references."
    )
    llm_prompt = f"{sys_prompt}\n\nQuestion: {question}\nAnswer:"
    inputs = tokenizer_lm(llm_prompt, return_tensors="pt").to(device)
    out = _generate(inputs, grounded=False)
    in_len = inputs["input_ids"].shape[-1]
    ans = _post_clean(_decode_generated(out, in_len))
    ans = re.sub(r"(?is)(^|\n)\s*references?:.*$", "", ans).strip()
    return "[LLM fallback — ungrounded]\n\n" + ans

# ================== PUBLIC API ==================
def ask(question: str) -> str:
    q = (question or "").strip()

    # --- PMID short-circuit ---------------------------------------------------
    m = re.search(r"pmid[:\s]*(\d+)", q, re.IGNORECASE)
    if m:
        pmid = m.group(1)
        chunks = fetch_pubmed_chunks(pmid, max_papers=1)
        return "\n".join(c.get("text", "") for c in chunks) or "Sorry, no abstract found."

    # --- Normalize query & detect 'and cite' / 'key papers' intent ------------
    wants_cite = bool(CITE_TRIGGER.search(q) or _PAPERS_INTENT.search(q))
    core_q = CITE_TRIGGER.sub("", q).strip().rstrip(".")
    core_q = _PAPERS_INTENT.sub("", core_q).strip()

    # --- Deterministic definitions/workflows (used in both paths) -------------
    det_text = deterministic_definitions_text(core_q)

    # ==========================================================================
    # Citation path
    # ==========================================================================
    if wants_cite:
        if det_text:
            explanation = det_text
            lq = core_q.lower()
            used_term = None
            if ("torsion" in lq) or ("gj" in lq):
                used_term = "GJ"
                pm_query = (
                    '(torsion[TiAb] OR "polar moment"[TiAb] OR GJ[TiAb]) AND '
                    '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND '
                    '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND '
                    '("2000"[DP] : "2025"[DP])'
                )
            elif ("bending" in lq) or ("ei" in lq):
                used_term = "EI"
                pm_query = (
                    '(bending[TiAb] OR "second moment"[TiAb] OR EI[TiAb]) AND '
                    '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND '
                    '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND '
                    '("2000"[DP] : "2025"[DP])'
                )
            else:
                used_term = "EA"
                pm_query = (
                    '("axial rigidity"[TiAb] OR EA[TiAb] OR "axial stiffness"[TiAb]) AND '
                    '("Bone and Bones"[MeSH] OR Femur[TiAb]) AND '
                    '("Finite Element Analysis"[MeSH] OR QCT[TiAb] OR CT[TiAb]) AND '
                    '("2000"[DP] : "2025"[DP])'
                )
            citations = fetch_pubmed_citations(pm_query, max_results=5)
            if not citations and used_term:
                citations = _fallback_cits_for(used_term)
        else:
            # Grounded explanation + progressively broader citation search
            explanation = _answer_from_chunks(retrieve_context(core_q, top_k=5), core_q)

            # Broad, de-biased search first
            compact = _compact_terms(core_q)
            pm_query = (
                f'({compact}) AND ("Bone and Bones"[MeSH] OR Femur[TiAb] OR Hip[TiAb] '
                f'OR Rigidity[TiAb] OR "Tomography, X-Ray Computed"[MeSH] OR "Finite Element Analysis"[MeSH]) '
                f'NOT (heart[TiAb] OR cardiac[TiAb] OR brain[TiAb] OR skull[TiAb] OR EGFR[TiAb]) '
                f'AND ("2000"[DP] : "2025"[DP])'
            )
            citations = fetch_pubmed_citations(pm_query, max_results=5)

            # If empty, bias toward lab-authored queries that are likely relevant
            if not citations:
                lab = detect_lab(core_q)
                pm_query = build_lab_query(core_q, lab=lab)
                citations = fetch_pubmed_citations(pm_query, max_results=5)

        resp = _post_clean(explanation)
        if citations:
            resp += "\n\nCitations:\n" + "\n".join(citations)
        else:
            resp += f'\n\nSorry, no relevant citations found for “{core_q}.”'
        return _ensure_min_answer(resp)

    # ==========================================================================
    # Non-citation path
    # ==========================================================================
    if det_text:
        dlog("ASK", "Deterministic definition/workflow fired")
        return det_text

    # If the query doesn't clearly hit MSK/FE tokens, try retrieval then fallback
    if not (_MSK_MUST.search(q) or _is_fe_override(q)):
        chunks = retrieve_context(q, top_k=5)
        if chunks:
            answer = _answer_from_chunks(chunks, q)
            return _ensure_min_answer(_post_clean(answer)) or direct_llm_fallback(q)
        return direct_llm_fallback(q)

    # Clear MSK/FE intent → do the normal grounded path
    chunks = retrieve_context(q, top_k=5)
    if not chunks:
        return direct_llm_fallback(q)

    answer = _answer_from_chunks(chunks, q)
    return _ensure_min_answer(_post_clean(answer)) or direct_llm_fallback(q)



def deterministic_definitions_text(core_q: str) -> Optional[str]:
    q_lower = core_q.lower()
    if "define axial rigidity" in q_lower or "what is axial rigidity" in q_lower:
        return ("Axial rigidity (EA) is Σ(Eᵢ·dAᵢ) across a CT slice; units: N. "
                "Modulus E per voxel comes from a density–modulus calibration; areas dAᵢ are voxel areas.")
    if "define bending rigidity" in q_lower or "what is bending rigidity" in q_lower:
        return ("Bending rigidity (EI) is Σ(Eᵢ·dAᵢ·yᵢ²) about a given axis; units: N·mm². "
                "yᵢ is distance to the neutral axis; computed slice-by-slice from QCT.")
    if ("define torsional rigidity" in q_lower) or ("what is torsional rigidity" in q_lower) or ("define gj" in q_lower):
        return ("Torsional rigidity (GJ) = shear modulus G times polar moment J. "
                "In QCT, J ≈ Σ(dAᵢ·rᵢ²) about the centroid; G ≈ E/(2(1+ν)).")
    if "qct" in q_lower and ("torsional" in q_lower or "gj" in q_lower):
        return ("From QCT, torsional rigidity is estimated as GJ, where J ≈ Σ(dAᵢ·rᵢ²) about the slice centroid and "
                "G = E/(2(1+ν)) from the voxel E map (ν≈0.3). Compute per-slice and report the minimum.")
    if re.search(r"\b(outline|steps|workflow|protocol)\b.*\b(ct|qct).*(rigidity|ea|ei|gj)", q_lower):
        return (
            "CT-based structural rigidity (CTRA/QCT) workflow:\n"
            "1) Acquire QCT (≤1 mm; density phantom).\n"
            "2) Preprocess & segment bone.\n"
            "3) HU→ρ; ρ→E calibration.\n"
            "4) Cross-sections along neck axis.\n"
            "5) EA, EI_x/EI_y, GJ (G≈E/(2(1+ν))).\n"
            "6) Extract minima & validate vs FEA/mech tests."
        )
    if re.search(r"\b(modulus)\b.*\brigidity\b|\bdefine\s+modulus\b", q_lower):
        return ("Elastic modulus (E) is a material property (Pa). "
                "Rigidity is structural (EA, EI, GJ). Modulus ≠ rigidity.")
    return None

# ================== UI: NAME GATE + PER-ANSWER FEEDBACK ==================
def _now_iso():
    return datetime.datetime.utcnow().replace(tzinfo=datetime.timezone.utc).isoformat()

def init_session():
    return {
        "session_id": str(uuid.uuid4()),
        "first_name": "",
        "last_name": "",
        "last_q": "",
        "last_a": "",
    }

def enter_app(first_name, last_name, state):
    first_name = (first_name or "").strip()
    last_name  = (last_name  or "").strip()
    if not first_name or not last_name:
        return gr.update(visible=True), gr.update(visible=False), state, "Please enter both first and last name."
    state["first_name"] = first_name
    state["last_name"]  = last_name
    return gr.update(visible=False), gr.update(visible=True), state, f"Welcome, {first_name}! You can start chatting."

from huggingface_hub import HfApi

def _push_file_to_dataset(local_path: str, repo_path: str) -> None:
    """
    Upload a local file to a *dataset* repo so the Space doesn't rebuild.
    Requires: HF_WRITE_TOKEN, LOGS_DATASET_ID
    """
    if not os.path.exists(local_path):
        dlog("UPLOAD", f"Skip: {local_path} does not exist"); return
    if not LOGS_DATASET_ID:
        dlog("UPLOAD", "Skip: LOGS_DATASET_ID not set"); return
    if not HF_WRITE_TOKEN:
        dlog("UPLOAD", "Skip: HF_WRITE_TOKEN not set"); return
    try:
        api = HfApi(token=HF_WRITE_TOKEN)
        api.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=repo_path,      # keep short names at the root
            repo_id=LOGS_DATASET_ID,
            repo_type="dataset",         # <- key change: dataset, not space
            commit_message=f"Update {repo_path}",
        )
        dlog("UPLOAD", f"Uploaded {repo_path} to dataset {LOGS_DATASET_ID}")
    except Exception as e:
        dlog("UPLOAD", f"Upload failed: {e}")


def _push_feedback_to_hub() -> None:
    """
    Back-compat name, but uploads to a *dataset* repo so the Space won't rebuild.
    Requires: PUSH_FEEDBACK=1, HF_WRITE_TOKEN, LOGS_DATASET_ID, and _push_file_to_dataset().
    """
    if not PUSH_FEEDBACK:
        return
    try:
        # Use the rotated filename if you enabled daily rotation
        remote_name = os.path.basename(FEEDBACK_PATH)
        _push_file_to_dataset(FEEDBACK_PATH, remote_name)
        dlog("UPLOAD", f"Feedback uploaded: {remote_name}")
    except Exception as e:
        dlog("UPLOAD", f"Feedback upload failed: {e}")


def _log_turn(state: Dict[str, Any], question: str, answer: str):
    """
    Append locally first, then (optionally) upload to a *dataset* repo.
    This avoids commits to the Space repo and prevents rebuilds after each turn.
    """
    rec = {
        "timestamp_utc": _now_iso(),
        "session_id": state.get("session_id", ""),
        "first_name":  state.get("first_name", ""),
        "last_name":   state.get("last_name", ""),
        "question":    question,
        "answer":      answer,
    }

    # Local append (ensure dir exists)
    try:
        os.makedirs(os.path.dirname(TRANSCRIPT_PATH) or ".", exist_ok=True)
        with open(TRANSCRIPT_PATH, "a", encoding="utf-8") as f:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")
    except Exception as e:
        dlog("LOG", f"Failed to write transcript locally: {e}")

    # Optional remote upload (dataset repo -> no Space rebuild)
    if PUSH_TRANSCRIPTS:
        try:
            remote_name = os.path.basename(TRANSCRIPT_PATH)
            _push_file_to_dataset(TRANSCRIPT_PATH, remote_name)
            dlog("UPLOAD", f"Transcript uploaded: {remote_name}")
        except Exception as e:
            dlog("UPLOAD", f"Transcript upload failed: {e}")


def save_feedback(rating, comment, state):
    if rating is None:
        return "Please select a rating (1–5).", gr.update(visible=True)
    row = {
        "timestamp_utc": _now_iso(),
        "session_id": state["session_id"],
        "first_name": state["first_name"],
        "last_name": state["last_name"],
        "question": state["last_q"],
        "answer": state["last_a"],
        "rating": rating,
        "comment": (comment or "").strip(),
    }
    header = ["timestamp_utc","session_id","first_name","last_name","question","answer","rating","comment"]
    try:
        file_exists = os.path.exists(FEEDBACK_PATH)
        with open(FEEDBACK_PATH, "a", newline="", encoding="utf-8") as f:
            w = csv.DictWriter(f, fieldnames=header)
            if not file_exists:
                w.writeheader()
            w.writerow(row)
        _push_feedback_to_hub()
        return "Thanks for the feedback! ✅", gr.update(visible=False)
    except Exception as e:
        return f"Failed to save feedback: {e}", gr.update(visible=True)

def predict(message, chat_history, state):
    msg = (message or "").strip()
    if not msg:
        return chat_history, "", gr.update(visible=False), None, "", state

    try:
        answer = ask(msg)
    except Exception as e:
        answer = f"Sorry — something went wrong: {e!r}"

    chat_history = (chat_history or []) + [(msg, answer)]
    state["last_q"] = msg
    state["last_a"] = answer

    try:
        _log_turn(state, msg, answer)
    except Exception:
        pass

    # free a bit of VRAM between turns
    try:
        torch.cuda.empty_cache()
    except Exception:
        pass

    return (
        chat_history,
        "",                          # clear input
        gr.update(visible=True),     # show feedback pane
        gr.update(value=None),       # reset rating
        gr.update(value=""),         # reset comment
        state
    )

# ================== UI ==================
with gr.Blocks(theme="soft") as demo:
    gr.Markdown("# Askstein — Orthopedic Biomechanics Chat (CT/QCT Rigidity, FE)")
    gr.Markdown("Grounded answers (FAISS + PubMed). Please enter your name to continue.")

    state = gr.State(init_session())

    # ---- Gate ----
    gate = gr.Group(visible=True)
    with gate:
        with gr.Row():
            first_tb = gr.Textbox(label="First name", placeholder="e.g., Shubh", scale=1)
            last_tb  = gr.Textbox(label="Last name",  placeholder="e.g., Laiwala", scale=1)
        enter_btn = gr.Button("Enter", variant="primary")
        gate_msg  = gr.Markdown("", elem_classes=["text-sm"])

    # ---- App (hidden until gate passes) ----
    app = gr.Group(visible=False)
    with app:
        chat = gr.Chatbot(label="Askstein", height=430, type="tuples")
        with gr.Row():
            user_in = gr.Textbox(placeholder="Ask a biomechanics question…", scale=5)
            send_btn = gr.Button("Send", variant="primary", scale=1)
        clear_btn = gr.Button("Clear chat")

        feedback_grp = gr.Group(visible=False)
        with feedback_grp:
            gr.Markdown("### How helpful was this answer?")
            rating = gr.Radio(choices=[1,2,3,4,5], label="Rating (1=poor, 5=great)")
            comment = gr.Textbox(label="Optional comment", placeholder="What was good or missing?")
            submit_fb = gr.Button("Submit feedback")
            fb_status = gr.Markdown("")

    # ---- Wiring (must stay inside Blocks) ----
    enter_btn.click(
        fn=enter_app,
        inputs=[first_tb, last_tb, state],
        outputs=[gate, app, state, gate_msg],
    )

    send_btn.click(
        fn=predict,
        inputs=[user_in, chat, state],
        outputs=[chat, user_in, feedback_grp, rating, comment, state],
        concurrency_limit=1,  # serialize LLM calls
    )

    user_in.submit(
        fn=predict,
        inputs=[user_in, chat, state],
        outputs=[chat, user_in, feedback_grp, rating, comment, state],
        concurrency_limit=1,  # serialize LLM calls
    )

    clear_btn.click(
        fn=lambda: ([], "", gr.update(visible=False), None, "", init_session()),
        inputs=None,
        outputs=[chat, user_in, feedback_grp, rating, comment, state],
        concurrency_limit=4,
    )

    submit_fb.click(
        fn=save_feedback,
        inputs=[rating, comment, state],
        outputs=[fb_status, feedback_grp],
        concurrency_limit=4,
    )

# ================== QUEUE & LAUNCH ==================
demo.queue(max_size=64)           # no deprecated concurrency_count
demo.launch(max_threads=8)