graph/builder.py

"""
NetworkX DiGraph construction from extracted entities + trials.
Upserts nodes and edges (merges PMIDs, recomputes confidence average).
Pre-populates with ALS seed entities from config.
"""
from __future__ import annotations

import json
from pathlib import Path

import networkx as nx

from config import (
    DERIVED_SEEDS_PATH,
    ENTITIES_PATH,
    KG_MIN_EDGE_CONFIDENCE,
    MANUAL_SEEDS_PATH,
    TRIALS_PATH,
)
from extraction.normalizer import normalize_entity
from logging_config import get_logger
from models import PaperExtractionResult

_logger = get_logger("graph.builder")


def build_graph(
    entities_path: Path = ENTITIES_PATH,
    trials_path: Path = TRIALS_PATH,
) -> nx.DiGraph:
    """Build the ALS knowledge graph. Returns a populated DiGraph."""
    G: nx.DiGraph = nx.DiGraph()

    _add_seed_entities(G)
    _logger.info(f"Seeded graph: {G.number_of_nodes()} seed nodes")

    if entities_path.exists():
        n_papers = _add_extracted_entities(G, entities_path)
        _logger.info(
            f"After extraction: {G.number_of_nodes()} nodes, "
            f"{G.number_of_edges()} edges from {n_papers} papers"
        )
    else:
        _logger.warning(f"Entities file not found: {entities_path} — skipping NER enrichment")

    if trials_path.exists():
        n_trials = _add_trials(G, trials_path)
        _logger.info(f"Added {n_trials} trial nodes")
    else:
        _logger.warning(f"Trials file not found: {trials_path}")

    return G


def _add_seed_entities(G: nx.DiGraph) -> None:
    type_map = {
        "genes": "Gene",
        "proteins": "Protein",
        "compounds": "Compound",
        "mechanisms": "Mechanism",
        "phenotypes": "Phenotype",
    }

    manual = json.loads(MANUAL_SEEDS_PATH.read_text())

    derived: dict[str, list[str]] = {}
    if DERIVED_SEEDS_PATH.exists():
        derived = json.loads(DERIVED_SEEDS_PATH.read_text())

    for category, entity_type in type_map.items():
        seen: set[str] = set()
        for name in manual.get(category, []):
            canonical_id = normalize_entity(name, entity_type)
            seen.add(canonical_id)
            _upsert_node(G, canonical_id, {
                "type": entity_type,
                "display_name": name,
                "paper_count": 0,
                "evidence_pmids": [],
                "is_seed": True,
            })
        for name in derived.get(category, []):
            canonical_id = normalize_entity(name, entity_type)
            if canonical_id in seen:
                continue
            seen.add(canonical_id)
            _upsert_node(G, canonical_id, {
                "type": entity_type,
                "display_name": name,
                "paper_count": 0,
                "evidence_pmids": [],
                "is_seed": False,
                "is_derived_seed": True,
            })


def _add_extracted_entities(G: nx.DiGraph, entities_path: Path) -> int:
    n_papers = 0
    with open(entities_path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            raw = json.loads(line)
            pmid = raw["pmid"]
            n_papers += 1

            # Add / update entity nodes
            for ent in raw.get("entities", []):
                canonical_id = ent.get("canonical_id", "")
                if not canonical_id:
                    continue
                if G.has_node(canonical_id):
                    G.nodes[canonical_id]["paper_count"] += 1
                    G.nodes[canonical_id]["evidence_pmids"].append(pmid)
                    # Update confidence as running average
                    cur = G.nodes[canonical_id].get("confidence", 0.7)
                    G.nodes[canonical_id]["confidence"] = (cur + ent.get("confidence", 0.7)) / 2
                else:
                    _upsert_node(G, canonical_id, {
                        "type": ent.get("type", "Unknown"),
                        "display_name": ent.get("name", canonical_id),
                        "paper_count": 1,
                        "evidence_pmids": [pmid],
                        "confidence": ent.get("confidence", 0.7),
                        "is_seed": False,
                    })

            # Add / update relationship edges
            for rel in raw.get("relationships", []):
                source = rel.get("source", "")
                target = rel.get("target", "")
                rel_type = rel.get("relation_type", "")
                if not source or not target or not rel_type:
                    continue

                # Ensure both endpoints exist as nodes
                for node_id in (source, target):
                    if not G.has_node(node_id):
                        _upsert_node(G, node_id, {
                            "type": "Unknown",
                            "display_name": node_id.split(":", 1)[-1],
                            "paper_count": 0,
                            "evidence_pmids": [],
                            "is_seed": False,
                        })

                conf = rel.get("confidence", 0.7)
                if G.has_edge(source, target):
                    edge = G[source][target]
                    if rel_type not in edge.get("relation_types", []):
                        edge.setdefault("relation_types", [edge.get("relation_type", rel_type)])
                        edge["relation_types"].append(rel_type)
                    if pmid not in edge["evidence_pmids"]:
                        edge["evidence_pmids"].append(pmid)
                    # Running average confidence
                    edge["confidence"] = (edge["confidence"] + conf) / 2
                else:
                    G.add_edge(source, target, **{
                        "relation_type": rel_type,
                        "relation_types": [rel_type],
                        "evidence_pmids": [pmid],
                        "confidence": conf,
                        "evidence_text": rel.get("evidence_text", ""),
                    })

    return n_papers


def _add_trials(G: nx.DiGraph, trials_path: Path) -> int:
    n_trials = 0
    with open(trials_path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            trial = json.loads(line)
            nct_id = trial.get("nct_id", "")
            if not nct_id:
                continue

            node_id = f"trial:{nct_id}"
            _upsert_node(G, node_id, {
                "type": "ClinicalTrial",
                "display_name": trial.get("title", nct_id)[:120],
                "nct_id": nct_id,
                "phase": trial.get("phase", ""),
                "status": trial.get("status", ""),
                "url": trial.get("url", f"https://clinicaltrials.gov/study/{nct_id}"),
                "paper_count": 0,
                "evidence_pmids": [],
            })

            # Link trial to known target entities
            for target_name in trial.get("target_entities", []):
                # Try gene, compound, protein
                matched = False
                for etype in ("Gene", "Compound", "Protein", "Mechanism"):
                    candidate_id = normalize_entity(target_name, etype)
                    if G.has_node(candidate_id):
                        if not G.has_edge(node_id, candidate_id):
                            G.add_edge(node_id, candidate_id, **{
                                "relation_type": "TESTED_IN",
                                "relation_types": ["TESTED_IN"],
                                "evidence_pmids": [],
                                "confidence": 1.0,
                                "evidence_text": "",
                            })
                        matched = True
                        break
                if not matched:
                    # Novel entity from trial with no paper evidence — canonicalize as
                    # compound (trial targets are drugs) but type stays Unknown since
                    # we can't confirm mechanism/class without literature support.
                    fallback_id = normalize_entity(target_name, "Compound")
                    _upsert_node(G, fallback_id, {
                        "type": "Unknown",
                        "display_name": target_name,
                        "paper_count": 0,
                        "evidence_pmids": [],
                        "confidence": 0.0,
                        "is_seed": False,
                        "is_trial_derived": True,
                    })
                    G.add_edge(node_id, fallback_id, **{
                        "relation_type": "TESTED_IN",
                        "relation_types": ["TESTED_IN"],
                        "evidence_pmids": [],
                        "confidence": 1.0,
                        "evidence_text": "",
                    })
                    _logger.debug(f"Trial {nct_id}: created trial-derived node {fallback_id!r} (type=Unknown)")

            n_trials += 1
    return n_trials


def _upsert_node(G: nx.DiGraph, node_id: str, attrs: dict) -> None:
    if G.has_node(node_id):
        G.nodes[node_id].update(attrs)
    else:
        G.add_node(node_id, **attrs)