# mcp/nlp.py
import spacy
import scispacy
from scispacy.linking import EntityLinker

# Load a powerful biomedical model + UMLS linker
@spacy.util.cache_dir("~/.cache/scispacy")
def load_model():
    nlp = spacy.load("en_core_sci_scibert")
    # Resolve abbreviations then link to UMLS
    linker = EntityLinker(name="umls", resolve_abbreviations=True, threshold=0.75)
    nlp.add_pipe(linker)
    return nlp

nlp = load_model()

def extract_umls_concepts(text: str):
    """
    Returns a list of {cui, concept_name, score, semantic_types}.
    """
    doc = nlp(text)
    concepts = []
    for ent in doc.ents:
        for cui, score in ent._.umls_ents:
            meta = nlp.get_pipe("scispacy_linker").kb.cui_to_entity[cui]
            concepts.append({
                "cui": cui,
                "name": meta.canonical_name,
                "score": float(score),
                "types": meta.types  # list of semantic type strings
            })
    # Deduplicate by CUI, keep highest score
    seen = {}
    for c in concepts:
        prev = seen.get(c["cui"])
        if not prev or c["score"] > prev["score"]:
            seen[c["cui"]] = c
    return list(seen.values())