First_Pass.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from __future__ import annotations

"""
Askstein – Hybrid RAG (FAISS + PubMed), T4-small optimized (v2.4 FINAL)

Key points:
• T4-small friendly: device_map="auto", bounded max_memory (INT keys), OFFLOAD_DIR.
• One-time LoRA→base merge with graceful fallback if the adapter has unknown fields
  (e.g., 'corda_config' saved with a newer PEFT). If merge fails, we continue with base.
• QUANTIZE env: "4bit" (default), "8bit", or "none" for the merged weights.
• FAISS + PubMed + Wikipedia routing; deterministic EA/EI/GJ snippets; “…and cite”.
"""

# ==== Early env hygiene =======================================================
import os
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

# ==== Imports =================================================================
import re, json, time, sys, shutil, tempfile
from typing import List, Dict, Any, Optional
from functools import lru_cache
from xml.etree import ElementTree as ET

import numpy as np
import faiss
import requests

from sentence_transformers import SentenceTransformer
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# Wikipedia (enabled per-call)
import wikipedia
from wikipedia.exceptions import DisambiguationError, PageError

# ==== Small utilities =========================================================
def _env(name: str, default: str = "") -> str:
    v = os.getenv(name)
    return v if v is not None else default

def _pick(*candidates: str) -> str:
    here = os.path.dirname(os.path.abspath(__file__))
    for c in candidates:
        p = c if os.path.isabs(c) else os.path.join(here, c)
        if os.path.exists(p):
            return p
    return candidates[0]

class LOG:
    DEBUG = _env("DEBUG", "1").lower() not in ("0", "false", "no")
    @staticmethod
    def p(tag: str, msg: str):
        if LOG.DEBUG:
            print(f"[{tag}] {msg}")

# ==== Paths & Config ==========================================================
FAISS_PATH       = _env("FAISS_PATH", _pick("index.faiss", "faiss/index.faiss"))
META_PATH        = _env("META_PATH",  _pick("index_meta.filtered.json",
                                            "index_meta.filtered.jsonl",
                                            "faiss/index_meta.filtered.jsonl"))
REL_CONFIG_PATH  = _env("REL_CONFIG_PATH", _pick("relevance_config.json", "faiss/relevance_config.json"))

EMBED_MODEL_NAME = _env("EMBED_MODEL_NAME", "pritamdeka/S-PubMedBERT-MS-MARCO")

BASE_MODEL       = _env("BASE_MODEL", "mistralai/Mistral-7B-Instruct-v0.2")
ADAPTER_PATH     = _env("ADAPTER_PATH", _pick("lora_adapter", "adapters/mistral7b_fp16_lora"))
MERGED_MODEL_DIR = _env("MERGED_MODEL_DIR", _pick("merged-model", "/home/user/app/merged-model"))
FORCE_REMERGE    = _env("FORCE_REMERGE", "0") == "1"

OFFLOAD_DIR      = _env("OFFLOAD_DIR", _pick("offload", "/home/user/app/offload", "/tmp/offload"))
os.makedirs(OFFLOAD_DIR, exist_ok=True)
os.makedirs(MERGED_MODEL_DIR, exist_ok=True)

# Quantization: "4bit" (T4 default), "8bit", or "none"
QUANTIZE         = _env("QUANTIZE", "4bit").lower()

# ==== T4-friendly limits & toggles ===========================================
ALLOW_WIKIPEDIA = False
MAX_NEW_TOKENS_GROUNDED = 384
MAX_NEW_TOKENS_FALLBACK = 256
MIN_USEFUL_CHARS = 260
PROMPT_BUDGET_TOKENS = 6400
FE_TRIM_WORDS = 230

torch.manual_seed(42)
if torch.cuda.is_available():
    torch.backends.cuda.matmul.allow_tf32 = True  # perf on T4

# ==== 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": _env("NCBI_EMAIL", ""),
    "ncbi_tool":  _env("NCBI_TOOL", "askstein"),
    "ncbi_apikey": _env("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 = json.load(f)
            if isinstance(user, dict):
                cfg.update(user)
    except Exception as e:
        LOG.p("rel-config", f"using defaults ({e})")
    return cfg

REL_CFG = load_rel_config(REL_CONFIG_PATH)
print("Loaded relevance config keys:", list(REL_CFG.keys()))
if LOG.DEBUG:
    print(f"[config] NCBI email set? {'yes' if REL_CFG.get('ncbi_email') else 'no'}")
    print(f"[config] NCBI api_key set? {'yes' if REL_CFG.get('ncbi_apikey') else 'no'}")

# ==== HTTP utils (session + backoff + circuit breaker) =======================
class _Http:
    session = requests.Session()
    session.headers.update({"User-Agent": "Askstein/1.0 (+https://hf.co/spaces)"} )

_EUTILS_DOWN_UNTIL = 0.0
_EUTILS_COOLDOWN   = 60.0

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

def _get_with_backoff(url: str, params: Dict[str, Any], tries: int = 3, base_sleep: float = 0.6, timeout: int = 10) -> str:
    global _EUTILS_DOWN_UNTIL
    if "eutils" in url and time.time() < _EUTILS_DOWN_UNTIL:
        raise RuntimeError("EUtils circuit breaker active")
    last_err = None
    for i in range(tries):
        try:
            if "eutils" in url and not REL_CFG.get("ncbi_apikey"):
                time.sleep(0.35)  # polite rate without key
            r = _Http.session.get(url, params=params, timeout=timeout)
            r.raise_for_status()
            return r.text
        except Exception as e:
            last_err = e
            if i == tries - 1:
                if "eutils" in url:
                    _EUTILS_DOWN_UNTIL = time.time() + _EUTILS_COOLDOWN
                raise
            time.sleep(base_sleep * (2 ** i))
    raise last_err if last_err else RuntimeError("Unknown HTTP error")

# ==== Wikipedia helpers =======================================================
_SANITIZE = re.compile(r"```.*?```|<\s*script[^>]*>.*?<\s*/\s*script\s*>", re.I | re.S)

def wiki_summary_allow(query: str, sentences: int = 3) -> 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.I)
        q = re.sub(r'\s+(located|location)\s*\?*$', '', q, flags=re.I).strip('?').strip()
        return wikipedia.summary(q, sentences=sentences)
    except (DisambiguationError, PageError, Exception):
        return None
    finally:
        globals()["ALLOW_WIKIPEDIA"] = prev

def wiki_summary_strong(query: str, sentences: int = 4) -> Optional[str]:
    try:
        results = wikipedia.search(query, results=5)
        for title in results:
            try:
                page = wikipedia.page(title, auto_suggest=False)
                text = (page.summary or "").strip()
                if not text:
                    continue
                if len(text) < 600 and page.content:
                    first_sec = page.content.split("\n\n")[1:2]
                    if first_sec:
                        text = text + "\n\n" + first_sec[0][:600]
                return _SANITIZE.sub("", text)
            except (DisambiguationError, PageError):
                continue
    except Exception:
        pass
    return None

# ==== Load FAISS + metadata + embedder =======================================
for pth in (FAISS_PATH, META_PATH):
    if not os.path.exists(pth):
        raise FileNotFoundError(f"Missing required file: {pth}")

print("Loading FAISS index…")
index = faiss.read_index(FAISS_PATH)
print("FAISS ntotal (rows):", index.ntotal)

print("Loading metadata…")
all_chunks: List[Dict[str, Any]] = []
with open(META_PATH, "r", encoding="utf-8") as f:
    if META_PATH.endswith(".json"):
        try:
            data = json.load(f)
            if isinstance(data, list):
                all_chunks.extend(data)
        except Exception:
            pass
    else:
        for line in f:
            try:
                all_chunks.append(json.loads(line))
            except Exception:
                pass
print(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}). Rebuild or fix META_PATH.")

print("Loading embedding model…", EMBED_MODEL_NAME)
embed_model = SentenceTransformer(EMBED_MODEL_NAME)

# Dim check + normalize for IP
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__

# ==== LLM (tokenizer + quant/merge cache) ====================================
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype  = torch.float16 if device == "cuda" else torch.float32
print("Loading LLM on", device)

tokenizer_lm = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=False)
if tokenizer_lm.pad_token_id is None:
    tokenizer_lm.pad_token = tokenizer_lm.eos_token

def _bnb_config() -> Optional[BitsAndBytesConfig]:
    if QUANTIZE == "4bit":
        return BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=torch.float16,
        )
    if QUANTIZE == "8bit":
        return BitsAndBytesConfig(load_in_8bit=True)
    return None

def _merged_present(path: str) -> bool:
    try:
        names = os.listdir(path)
        return any(n.endswith(".safetensors") for n in (names or []))
    except Exception:
        return False

def _max_memory_mapping():
    if torch.cuda.is_available():
        n = torch.cuda.device_count()
        mem = {i: "12GiB" for i in range(n)}  # INT keys required by accelerate>=0.30
        mem["cpu"] = "24GiB"
        return mem
    return None

def _safe_try_load_peft(base_model) -> Optional[AutoModelForCausalLM]:
    """
    Try to attach/merge the LoRA adapter. If the adapter config contains unknown fields
    (e.g., 'corda_config' from a newer PEFT), catch and return None to fall back.
    """
    if not os.path.exists(ADAPTER_PATH):
        LOG.p("PEFT", f"No adapter at '{ADAPTER_PATH}'.")
        return None
    try:
        peft_model = PeftModel.from_pretrained(
            base_model,
            ADAPTER_PATH,
            device_map="auto" if torch.cuda.is_available() else None,
            offload_folder=OFFLOAD_DIR,
        )
        merged = peft_model.merge_and_unload()
        try:
            merged.to(dtype=torch.float16)
        except Exception:
            pass
        LOG.p("MERGE", "LoRA merge successful.")
        return merged
    except TypeError as te:
        # Typical case: LoraConfig.__init__ got unexpected keyword arg 'corda_config'
        LOG.p("PEFT", f"Adapter incompatible with current PEFT: {te}. Using BASE MODEL only.")
        return None
    except Exception as e:
        LOG.p("PEFT", f"Failed to load adapter ({e}). Using BASE MODEL only.")
        return None

def _load_merged_or_merge() -> AutoModelForCausalLM:
    # 1) Use pre-merged weights if present and not forcing remerge
    if (not FORCE_REMERGE) and _merged_present(MERGED_MODEL_DIR):
        LOG.p("LOAD", f"Loading merged model from {MERGED_MODEL_DIR} (quant={QUANTIZE})")
        return AutoModelForCausalLM.from_pretrained(
            MERGED_MODEL_DIR,
            torch_dtype=(dtype if QUANTIZE == "none" else None),
            device_map="auto" if torch.cuda.is_available() else None,
            low_cpu_mem_usage=True,
            max_memory=_max_memory_mapping(),
            quantization_config=_bnb_config(),
        )

    # 2) Merge path: load base (no quant), try to attach & merge LoRA, save; on failure, save base.
    LOG.p("MERGE", "Merging LoRA into base weights (one-time)…")
    base = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        torch_dtype=dtype,
        device_map="auto" if torch.cuda.is_available() else None,
        low_cpu_mem_usage=True,
        max_memory=_max_memory_mapping(),
        offload_folder=OFFLOAD_DIR,
    )

    merged = _safe_try_load_peft(base)
    to_save = merged if merged is not None else base

    # Persist for faster cold starts (and to allow quantized reloads)
    tokenizer_lm.save_pretrained(MERGED_MODEL_DIR)
    to_save.save_pretrained(MERGED_MODEL_DIR, safe_serialization=True)
    LOG.p("MERGE", f"Saved {'merged' if merged is not None else 'base'} model to {MERGED_MODEL_DIR}")
    return to_save

model_lm = _load_merged_or_merge()
model_lm.eval()

GEN_ARGS_GROUNDED = dict(
    max_new_tokens=MAX_NEW_TOKENS_GROUNDED,
    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,
    do_sample=False,
    num_beams=1,
    no_repeat_ngram_size=3,
    repetition_penalty=1.05,
    eos_token_id=tokenizer_lm.eos_token_id,
)

def _generate(inputs, grounded: bool):
    args = GEN_ARGS_GROUNDED if grounded else GEN_ARGS_FALLBACK
    with torch.inference_mode():
        return model_lm.generate(**inputs, **args)

# ==== Text helpers ============================================================
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 _split_sents(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_sents(t)
    seen = set(); kept = []
    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)

def _trim_words(text: str, max_words: int = FE_TRIM_WORDS) -> str:
    words = text.split()
    if len(words) <= max_words:
        return text
    return " ".join(words[:max_words]).rstrip(",;:") + "…"

# ==== 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))
    if any(n in wl for n in names):
        return rel_base + int(cfg.get("author_weight", 3))
    return 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|boundary\s+conditions|nonlinear|yield|fracture\s+load|micromotion)\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)

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 _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 _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()
    body_ok = any(k in t for k in ["femur","femoral","femoral neck","proximal femur","hip"])
    method_ok = any(k in t for k in ["qct","quantitative computed tomography"," ct "," ct-",
                                     "finite element","fea","structural rigidity","rigidity","bending","torsion"])
    if any(k in t for k in ["humerus","shoulder","humeral"]) and not body_ok:
        return False
    if not (body_ok and method_ok):
        return False
    y = _parse_year(year)
    if y and not (2000 <= y <= 2025):
        return False
    return True

# ==== PubMed & citations ======================================================
def fetch_pubmed_chunks(query_or_pmid: str, max_papers: int = 3) -> List[Dict[str, Any]]:
    retries = 2
    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

@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]

# ==== PMC helpers =============================================================
@lru_cache(maxsize=4096)
def _pmid_to_pmcid(pmid: str) -> Optional[str]:
    try:
        xml = _get_with_backoff(
            "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/elink.fcgi",
            _ncbi_params({"dbfrom":"pubmed","db":"pmc","id":pmid})
        )
        root = ET.fromstring(xml)
        for link in root.findall(".//LinkSetDb/Link/Id"):
            if link.text:
                return link.text.strip()
    except Exception:
        pass
    return None

def fetch_pmc_paras(pmid: str, max_paras: int = 2) -> List[str]:
    pmc = _pmid_to_pmcid(str(pmid))
    if not pmc:
        return []
    try:
        xml = _get_with_backoff(
            "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi",
            _ncbi_params({"db":"pmc","id":pmc})
        )
        root = ET.fromstring(xml)
        paras = []
        for sec in root.findall(".//body//sec"):
            for p in sec.findall("p"):
                if p.text and len(p.text.strip()) > 200:
                    paras.append(p.text.strip())
                if len(paras) >= max_paras:
                    break
            if len(paras) >= max_paras:
                break
        return paras
    except Exception:
        return []

# ==== FE/Anatomy routing helpers =============================================
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|fracture|healing|cortical)\b", q, re.I))

def _is_fe_override(q: str) -> bool:
    return bool(_FE_TOKENS.search(q))

# ==== Conflict detector =======================================================
_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))

# ==== Canonical formulas injection ===========================================
CANON_PATTERNS = {
    r"\b(axial (rigidity|stiffness)|\bea\b)\b": ("Axial Rigidity (EA)", "EA = Σ_i (E_i · dA_i)"),
    r"\b(bending (rigidity|stiffness)|\bei\b)\b": ("Bending Rigidity (EI)", "EI = Σ_i (E_i · dA_i · y_i^2)"),
    r"\b(torsional (rigidity|stiffness)|\bgj\b)\b": ("Torsional Rigidity (GJ)", "GJ = G·J;  G = E/(2(1+ν)),  J = Σ_i (dA_i · r_i^2)"),
}
def _maybe_inject_formula(q_lower: str, chunks: List[Dict[str, Any]]) -> bool:
    for pat, (label, text) in CANON_PATTERNS.items():
        if re.search(pat, q_lower):
            chunks.insert(0, {"text": f"{label}:\n{text}", "source": "injected"})
            return True
    return False

# ==== “…and cite” curated fallbacks ==========================================
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(), [])

# ==== Lab detection (lightweight) ============================================
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}"

# ==== Retrieval ===============================================================
def retrieve_context(query: str, top_k: int = 10) -> List[Dict[str, Any]]:
    q = query.strip()

    if _ANATOMY_OR_HISTORY.search(q):
        wiki = wiki_summary_allow(q, sentences=4) or wiki_summary_strong(q, sentences=4)
        if wiki:
            LOG.p("WIKI", "Anatomy/History → Wikipedia")
            return [{"text": wiki, "source": "wikipedia"}]

    pm = re.search(r"pmid[:\s]*(\d+)", q, re.I)
    if pm:
        LOG.p("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)):
        LOG.p("FALLBACK", "No CT/FE tokens → robust PubMed/Wiki")
        qx = q.lower()
        compact = _compact_terms(qx)
        passes = [
            f'"{qx}"[Title/Abstract]',
            f'({compact}) AND (hip[TiAb] OR femur[TiAb] OR femoral[TiAb] OR tibia[TiAb] OR "long bone"[TiAb] '
            f'OR fracture[TiAb] OR healing[TiAb] OR cortical[TiAb] OR trabecular[TiAb] OR mouse[TiAb] OR murine[TiAb]) '
            f'AND ("2000"[DP] : "3000"[DP])',
            build_lab_query(qx, lab=detect_lab(qx)),
            f'({compact}) AND ("Bone and Bones"[MeSH] OR Femur[MeSH] OR Tibia[MeSH] OR '
            f'"Fractures, Bone"[MeSH] OR "Wound Healing"[MeSH] OR "Tomography, X-Ray Computed"[MeSH] OR '
            f'"Finite Element Analysis"[MeSH]) AND ("2000"[DP] : "3000"[DP])',
        ]
        if re.search(r"\b(how|why|impact|effect|influence)\b", qx):
            passes.append(f'({compact}) AND review[ptyp] AND ("2010"[DP] : "3000"[DP])')

        seen_pmids, fetched = set(), []
        for term in passes:
            for c in fetch_pubmed_chunks(term, max_papers=20):
                pmid = str(c.get("pmid") or "")
                if pmid and pmid in seen_pmids:
                    continue
                seen_pmids.add(pmid); fetched.append(c)
            if len(fetched) >= 20:
                break

        for it in fetched:
            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
        fetched.sort(key=lambda x: x.get("_rel", 0), reverse=True)

        if fetched:
            LOG.p("PUBMED", f"Robust PubMed hit: {len(fetched)}")
            enriched = []
            for r in fetched[:8]:
                enriched.append(r)
                pmid = r.get("pmid")
                if pmid:
                    try:
                        paras = fetch_pmc_paras(str(pmid), max_paras=1)
                    except Exception:
                        paras = []
                    for p in paras:
                        enriched.append({"text": p, "source": "pmc", "pmid": pmid})
            return enriched[:top_k]

        wiki = wiki_summary_strong(qx, sentences=4)
        if wiki:
            LOG.p("WIKI", "Wikipedia strong fallback hit")
            return [{"text": wiki, "source": "wikipedia"}]
        LOG.p("RETRIEVAL", "No results found in robust fallback")
        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)
        t = _to_text(item)
        if not t:
            pmid = str(item.get("pmid") or "")
            if pmid.isdigit():
                abs_chunks = fetch_pubmed_chunks(pmid, max_papers=1)
                if abs_chunks:
                    t = abs_chunks[0].get("text","")
        if not t:
            continue
        item["text"] = t
        results.append(item)

    if results:
        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
        results = sorted(results, 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","")))
        ]
        seen = set(); deduped: List[Dict[str, Any]] = []
        for r in positives:
            key = str(r.get("pmid") or "").strip() \
                  or (r.get("title") or "").strip().lower()[:120] \
                  or (r.get("text") or "").strip().lower()[:200]
            if key in seen: continue
            seen.add(key); deduped.append(r)
        if deduped:
            LOG.p("FAISS", f"FAISS hit={len(deduped)} (top rel={deduped[0].get('_rel')} score={deduped[0].get('score'):.3f})")
            return deduped[:top_k]
        else:
            LOG.p("FALLBACK", "FAISS results off-topic → PubMed fallback")

    results = fetch_pubmed_chunks(q)
    if results:
        LOG.p("PUBMED", "PubMed search hit")
        return results

    if _biomechish(q):
        wiki = wiki_summary_allow(q, sentences=3)
        if wiki:
            LOG.p("WIKI", "Wikipedia biomechanics fallback")
            return [{"text": wiki, "source": "wikipedia"}]

    LOG.p("RETRIEVAL", "No results found at all")
    return []

# ==== Prompting & Generation ==================================================
def build_prompt(chunks: List[Dict[str, Any]], question: str, prompt_budget=PROMPT_BUDGET_TOKENS) -> 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."
    )
    question_block = f"[Question]:\n{question}\n"
    header_ids = tokenizer_lm(header, return_tensors="pt").input_ids[0]
    q_ids = tokenizer_lm(question_block, return_tensors="pt").input_ids[0]
    remaining = max(256, prompt_budget - len(header_ids) - len(q_ids))

    ctx_texts, used = [], 0
    for c in chunks:
        t = _to_text(c)
        if not t: continue
        ids = tokenizer_lm(t, return_tensors="pt").input_ids[0]
        if used + len(ids) > remaining: break
        used += len(ids); ctx_texts.append(t)

    context = "\n\n".join(ctx_texts)
    return f"{header}\n\n[Context]:\n{context}\n\n{question_block}"

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 _synthesize_answer(chunks: List[Dict[str, Any]], question: str) -> str:
    LOG.p("SYNTH", "Multi-chunk synthesis pass")
    prompt = build_prompt(chunks, question)
    inputs = tokenizer_lm(prompt, return_tensors="pt")
    if torch.cuda.is_available():
        inputs = {k: v.to("cuda") for k, v in inputs.items()}
    out = _generate(inputs, grounded=True)
    in_len = inputs["input_ids"].shape[-1]
    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):
        LOG.p("SYNTH", "Conflict detected (no-diff vs change) → summarize")
        return _synthesize_answer(chunks, question)
    return _synthesize_answer(chunks, question)

# ==== Deterministic biomech definitions ======================================
def deterministic_definitions_text(core_q: str) -> Optional[str]:
    q = core_q.lower()
    if "define axial rigidity" in q or "what is axial rigidity" in q:
        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 or "what is bending rigidity" in q:
        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 or "what is torsional rigidity" in q or "define gj" in q:
        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 and ("torsional" in q or "gj" in q):
        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 along the shaft/neck and report minima "
                "or location-specific values. Note: this is an engineering approximation; full torsion may require FEA.")
    if re.search(r"\b(outline|steps|workflow|protocol)\b.*\b(ct|qct).*(rigidity|ea|ei|gj)", q):
        return (
            "CT-based structural rigidity (CTRA/QCT) workflow:\n"
            "1) Acquire QCT of proximal femur (≤1 mm slice; in-phantom density calibration).\n"
            "2) Preprocess (bias/beam-hardening correction; resample to isotropic voxels).\n"
            "3) Segment bone → cortical vs trabecular (threshold + morphology cleanup).\n"
            "4) HU→ρ (mgHA/cm³) via phantom; ρ→E using calibrated density–modulus map.\n"
            "5) Define cross-sections along the femoral neck axis (every 1–2 mm).\n"
            "6) EA = Σ(Eᵢ·dAᵢ); EI_x/EI_y = Σ(Eᵢ·dAᵢ·yᵢ²/xᵢ²); GJ ≈ Σ(dAᵢ·rᵢ²)·G.\n"
            "7) Extract minima (e.g., min(EI)) as fracture-relevant metrics.\n"
            "8) Validate vs tests/subject-specific FEA; report units & axes.\n"
            "9) QC overlays, centroid alignment, axis consistency, unit checks.\n"
            "10) Output min/mean EA/EI/GJ with locations; compare across time/subjects."
        )
    if re.search(r"\b(modulus)\b.*\brigidity\b|\bdefine\s+modulus\b", q):
        return ("Elastic modulus (E) is a material property (stress/strain; Pa). "
                "Rigidity is a structural property: EA (axial), EI (bending), GJ (torsion). Modulus ≠ rigidity.")
    return None

# ==== Orchestrator ============================================================
def ask(question: str) -> str:
    q = question.strip()

    m = re.search(r"pmid[:\s]*(\d+)", q, re.I)
    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."

    if _PAPERS_INTENT.search(q):
        core_q = re.sub(_PAPERS_INTENT, "", q, flags=re.I).strip() or "CT/QCT structural rigidity femur hip finite element"
        compact = _compact_terms(core_q)
        pm_query = (
            f'(({compact}) AND (hip[TiAb] OR femur[TiAb] OR femoral[TiAb])) AND '
            '("Finite Element Analysis"[MeSH Terms] OR finite element[TiAb] OR QCT[TiAb] OR CT[TiAb] OR rigidity[TiAb]) '
            'AND ("2000"[DP] : "2025"[DP])'
        )
        cits = fetch_pubmed_citations(pm_query, max_results=5)
        return "Recommended papers:\n" + "\n".join(f"- {c}" for c in cits) if cits else \
               "Sorry, I couldn’t find good MSK/rigidity papers for that query."

    comp = re.match(r"(.+?)\s+and\s+(?:cite|references?|studies?|papers?)", q, flags=re.I)
    if comp:
        core_q = comp.group(1).strip()
        det_text = deterministic_definitions_text(core_q)
        used_term = None
        if det_text:
            explanation = det_text
            lq = core_q.lower()
            if ("torsional" 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] OR "Long bone"[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) or _fallback_cits_for(used_term or "")
        else:
            chunks = retrieve_context(core_q, top_k=5)
            explanation = _answer_from_chunks(chunks, core_q) if chunks else "I don’t know based on the provided context."
            pm_query  = f'"{core_q}"[Title/Abstract]'
            citations = fetch_pubmed_citations(pm_query, max_results=5)
            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)
            if not citations:
                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)
        resp = (det_text or explanation)
        if citations:
            resp += "\n\nCitations:\n" + "\n".join(citations)
        else:
            resp += f"\n\nSorry, no relevant citations found for “{core_q}.”"
        ans = _ensure_min_answer(_post_clean(resp))
        if _is_fe_override(core_q):
            ans = _trim_words(ans, FE_TRIM_WORDS)
        return ans

    det_answer = deterministic_definitions_text(q)
    if det_answer:
        LOG.p("ASK", "Deterministic definition/workflow fired")
        return det_answer

    if not (_MSK_MUST.search(q) or _is_fe_override(q)):
        chunks = retrieve_context(q, top_k=5)
        if chunks:
            _maybe_inject_formula(q.lower(), chunks)
            ans = _answer_from_chunks(chunks, q)
            ans = _ensure_min_answer(_post_clean(ans))
            if _is_fe_override(q):
                ans = _trim_words(ans, FE_TRIM_WORDS)
            return ans
        sys_prompt = (
            "You are Askstein (orthopedic biomechanics). Prefer concise, factual answers. "
            "If you lack sufficient evidence, say so briefly and propose what studies/data would answer it. "
            "Avoid non-MSK domains."
        )
        llm_prompt = f"{sys_prompt}\n\nQ: {q}\nA:"
        inputs = tokenizer_lm(llm_prompt, return_tensors="pt")
        if torch.cuda.is_available():
            inputs = {k: v.to("cuda") for k, v in inputs.items()}
        out = _generate(inputs, grounded=False)
        in_len = inputs["input_ids"].shape[-1]
        ans = _post_clean(_decode_generated(out, in_len))
        return _ensure_min_answer(ans)

    chunks = retrieve_context(q, top_k=5)
    if not chunks:
        sys_prompt = (
            "You are Askstein (orthopedic biomechanics). Prefer concise, factual answers. "
            "If you lack sufficient evidence, say so briefly and propose what studies/data would answer it."
        )
        llm_prompt = f"{sys_prompt}\n\nQ: {q}\nA:"
        inputs = tokenizer_lm(llm_prompt, return_tensors="pt")
        if torch.cuda.is_available():
            inputs = {k: v.to("cuda") for k, v in inputs.items()}
        out = _generate(inputs, grounded=False)
        in_len = inputs["input_ids"].shape[-1]
        ans = _post_clean(_decode_generated(out, in_len))
        return _ensure_min_answer(ans)

    _maybe_inject_formula(q.lower(), chunks)
    ans = _answer_from_chunks(chunks, q)
    ans = _ensure_min_answer(_post_clean(ans))
    if _is_fe_override(q):
        ans = _trim_words(ans, FE_TRIM_WORDS)
    return ans

# ==== Minimal CLI =============================================================
if __name__ == "__main__":
    print("=== Askstein CLI === (type 'exit' to quit)")
    try:
        while True:
            try:
                q = input("You: ")
            except EOFError:
                break
            if not q:
                continue
            if q.lower() in ("exit","quit"):
                break
            try:
                out = ask(q)
                print("Askstein:", out, "\n")
            except Exception as e:
                print("[error]", repr(e))
    except KeyboardInterrupt:
        pass