# genesis/ontology.py
"""
Ontology Expansion for GENESIS-AI
Uses UMLS & BioPortal to find related terms, synonyms, and hierarchies.
"""

import os
import requests

UMLS_API_KEY = os.getenv("UMLS_API_KEY")
BIOPORTAL_API_KEY = os.getenv("BIOPORTAL_API_KEY")

def expand_with_bioportal(term):
    """Expand term using BioPortal ontology search."""
    try:
        url = f"https://data.bioontology.org/search?q={term}&apikey={BIOPORTAL_API_KEY}"
        r = requests.get(url)
        r.raise_for_status()
        data = r.json()
        return list({res["prefLabel"] for res in data.get("collection", []) if "prefLabel" in res})
    except Exception as e:
        return {"error": str(e)}

def expand_with_umls(term):
    """Expand term using UMLS API."""
    try:
        url = f"https://uts-ws.nlm.nih.gov/rest/search/current?string={term}&apiKey={UMLS_API_KEY}"
        r = requests.get(url)
        r.raise_for_status()
        data = r.json()
        return [res["name"] for res in data.get("result", {}).get("results", []) if "name" in res]
    except Exception as e:
        return {"error": str(e)}

def expand_terms_with_ontology(term):
    """Combine BioPortal and UMLS expansions."""
    bioportal_terms = expand_with_bioportal(term)
    umls_terms = expand_with_umls(term)
    results = set()
    if isinstance(bioportal_terms, list):
        results.update(bioportal_terms)
    if isinstance(umls_terms, list):
        results.update(umls_terms)
    return list(results) if results else [term]