Blame for src/mapping/codemapper.py

"""One-round cross-vocabulary CUI expansion, with caching and paced sessions."""
from __future__ import annotations

from dataclasses import replace
try:
    import fcntl
except ImportError:  # Keep the disabled optional stage importable on Windows.
    fcntl = None
import hashlib
from http.cookiejar import CookieJar
import json
import logging
import os
from pathlib import Path
import time
from urllib.error import HTTPError, URLError
from urllib.parse import urlencode
from urllib.request import HTTPCookieProcessor, Request, build_opener

import pandas as pd

from src.data.preprocessing import standardize_code
from src.llm.classify import ClassifiedCandidate
from src.utils.config import CodeMapperConfig


class CodeMapperError(RuntimeError):
    """Sanitized failure: never include credentials, cookies, or server bodies."""


def normalize_code_versions(response,vocabulary):
    """Select the pinned release from the current all-version API response."""
    _,separator,version=vocabulary.partition('@')
    if not separator:raise CodeMapperError('CodeMapper requires a pinned vocabulary version')
    result={}
    for code,value in response.items():
        if isinstance(value,list):
            matching=[item for item in value if isinstance(item,dict) and str(item.get('version'))==version]
            if len(matching)>1:raise CodeMapperError('Ambiguous CodeMapper code version response')
            if matching:result[code]=matching[0]
        elif isinstance(value,dict):
            result[code]=value  # Older cached/API responses were already release-specific.
        else:raise CodeMapperError('Unexpected CodeMapper code version response')
    return result


def normalize_concept_versions(response,targets):
    """Project series-level codes to requested releases, never to latest by default."""
    versions={target.partition('@')[0]:target.partition('@')[2] for target in targets}
    result=[]
    for concept in response:
        sources=[]
        for source in concept.get('sourceConcepts',[]):
            if 'codingSystem' in source:
                if source['codingSystem'] in targets:sources.append(source)
                continue
            series=source.get('vocId');version=versions.get(series)
            if version is None:continue
            matching=[item for item in source.get('versions',[]) if str(item.get('version'))==version]
            if matching:
                sources.append({**source,'codingSystem':series+'@'+version,
                                'preferredTerm':matching[0].get('term',source.get('preferredTerm'))})
        result.append({**concept,'sourceConcepts':sources})
    return result


class CodeMapperClient:
    def __init__(self, config: CodeMapperConfig):
        self.config = config
        self.cache_dir = Path(config.cache_dir)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.opener = build_opener(HTTPCookieProcessor(CookieJar()))
        self.logged_in = False
        self.calls: list[dict] = []
        self.cache_hits = 0

    def _pace(self):
        if fcntl is None:
            raise CodeMapperError("Shared CodeMapper pacing currently requires macOS or Linux")
        # Share the limiter across client instances/processes using this cache.
        with (self.cache_dir / "rate.lock").open("a+") as handle:
            fcntl.flock(handle, fcntl.LOCK_EX)
            handle.seek(0)
            previous = float(handle.read() or 0)
            delay = previous + 60 / self.config.requests_per_minute - time.time()
            if delay > 0:
                time.sleep(delay)
            handle.seek(0)
            handle.truncate()
            handle.write(str(time.time()))
            handle.flush()

    def _request(self, path, method="POST", form=None):
        fields = [(key, item) for key, value in (form or {}).items()
                  for item in (value if isinstance(value, list) else [value])]
        request = Request(self.config.base_url.rstrip("/") + path,
                          data=urlencode(fields).encode() if form is not None else None,
                          method=method, headers={"Accept": "application/json"})
        for attempt in range(self.config.max_retries):
            self._pace()
            logging.getLogger(__name__).info('CodeMapper %s %s (attempt %s)',method,path,attempt+1)
            started = time.perf_counter()
            try:
                with self.opener.open(request, timeout=self.config.timeout_seconds) as response:
                    raw = response.read()
                    status = response.status
                self.calls.append({"path": path, "status": status,
                                   "seconds": round(time.perf_counter() - started, 4),
                                   "bytes": len(raw)})
                return json.loads(raw) if raw else None
            except HTTPError as exc:
                self.calls.append({"path": path, "status": exc.code,
                                   "seconds": round(time.perf_counter() - started, 4)})
                if exc.code not in {429, 500, 502, 503, 504}:
                    raise CodeMapperError(f"CodeMapper {path}: HTTP {exc.code}") from None
                try:
                    retry_after = float(exc.headers.get("Retry-After", 0))
                except ValueError:
                    retry_after = 0
                if attempt + 1 < self.config.max_retries:
                    time.sleep(max(retry_after, 2 ** attempt))
            except (URLError, TimeoutError, OSError):
                if attempt + 1 < self.config.max_retries:
                    time.sleep(2 ** attempt)
            except (ValueError, UnicodeDecodeError):
                raise CodeMapperError(f"CodeMapper {path}: invalid JSON") from None
        raise CodeMapperError(f"CodeMapper {path}: failed after retries")

    def _login(self):
        if self.logged_in:
            return
        username, password = os.getenv("CODEMAPPER_USERNAME"), os.getenv("CODEMAPPER_PASSWORD")
        if not username or not password:
            raise CodeMapperError("Set CODEMAPPER_USERNAME and CODEMAPPER_PASSWORD in .env")
        response = self._request("/session", form={"username": username, "password": password})
        if not isinstance(response, dict) or response.get("success") is not True:
            raise CodeMapperError("CodeMapper login unsuccessful")
        self.logged_in = True

    def close(self):
        if self.logged_in:
            try:
                self._request("/session", "DELETE")
            finally:
                self.logged_in = False

    def _cache_path(self, kind, identity):
        key = hashlib.sha256(json.dumps([self.config.base_url, kind, identity],
                                       sort_keys=True).encode()).hexdigest()
        return self.cache_dir / f"{key}.json"

    def _load(self, kind, identity):
        path = self._cache_path(kind, identity)
        if path.exists():
            self.cache_hits += 1
            return json.loads(path.read_text())
        return None

    def _save(self, kind, identity, value):
        path = self._cache_path(kind, identity)
        temporary = path.with_suffix(f".{os.getpid()}.tmp")
        temporary.write_text(json.dumps(value))
        temporary.replace(path)

    def cuis_for_codes(self, vocabulary, codes):
        result, missing = {}, []
        for code in sorted(set(codes)):
            cached = self._load("code_cuis", [vocabulary, code])
            if cached is None:
                missing.append(code)
            elif cached:
                result[code] = cached
        for start in range(0, len(missing), self.config.batch_size):
            batch = missing[start:start + self.config.batch_size]
            self._login()
            response = self._request("/legacy/code-mapper/cuis-for-codes",
                                     form={"vocId": vocabulary.partition('@')[0], "codes": batch})
            if not isinstance(response, dict):
                raise CodeMapperError("Unexpected code-to-CUI response")
            response=normalize_code_versions(response,vocabulary)
            for code in batch:
                value = response.get(code, {})
                self._save("code_cuis", [vocabulary, code], value)
                if value:
                    result[code] = value
        return result

    def concepts(self, cuis, targets):
        result, missing = [], []
        targets = sorted(set(targets))
        for cui in sorted(set(cuis)):
            cached = self._load("concept", [targets, cui])
            if cached is None:
                missing.append(cui)
            else:
                result.extend(cached)
        for start in range(0, len(missing), self.config.batch_size):
            batch = missing[start:start + self.config.batch_size]
            self._login()
            response = self._request("/legacy/code-mapper/umls-concepts",
                                     form={"cuis": batch, "vocIds": sorted({target.partition('@')[0] for target in targets})})
            if not isinstance(response, list):
                raise CodeMapperError("Unexpected concept response")
            response=normalize_concept_versions(response,targets)
            for cui in batch:
                values = [c for c in response if c.get("cui") == cui]
                self._save("concept", [targets, cui], values)
                result.extend(values)
        return result


def expand_narrow_matches(classified: list[ClassifiedCandidate], corpus: pd.DataFrame,
                          vocabularies: list[str], config: CodeMapperConfig,
                          *, client=None):
    """Promote/add direct cross-vocabulary matches from a frozen Narrow seed set.

    Only codes in the loaded/preprocessed corpus are eligible. Same-vocabulary
    links and descendants are excluded. No new additions become seeds this round.
    Mapping confidence is 0 (unassessed), not an invented model probability.
    """
    seeds = [c for c in classified if c.label == "Narrow" and c.vocabulary in vocabularies]
    versions = {v: f"{v}@{config.vocabulary_versions[v]}" for v in vocabularies}
    reverse = {version: voc for voc, version in versions.items()}
    by_cui = {}
    client = client or CodeMapperClient(config)
    started = time.perf_counter()
    seed_counts, unmapped = {}, []
    try:
        for voc in vocabularies:
            group = [c for c in seeds if c.vocabulary == voc]
            if not group:
                continue
            seed_counts[voc] = len(group)
            response = client.cuis_for_codes(versions[voc], [c.code for c in group])
            for seed in group:
                cuis = response.get(seed.code, {}).get("cuis", [])
                if not cuis:
                    unmapped.append({"code": seed.code, "vocabulary": voc})
                for cui in cuis:
                    by_cui.setdefault(cui, []).append(seed)
        concepts = client.concepts(by_cui, list(versions.values())) if by_cui else []
    finally:
        client.close()
    available = {(standardize_code(c, v), v): (c, d)
                 for c, v, d in corpus[["code", "vocabulary", "description"]].itertuples(index=False, name=None)}
    links, unavailable = {}, {}
    for concept in concepts:
        for source in concept.get("sourceConcepts", []):
            voc = reverse.get(source["codingSystem"])
            if voc is None:
                continue
            key = (standardize_code(source["id"], voc), voc)
            provenance = [{"seed_code": s.code, "seed_vocabulary": s.vocabulary,
                           "seed_label": s.label, "cui": concept["cui"],
                           "source_version": versions[s.vocabulary],
                           "target_version": source["codingSystem"]}
                          for s in by_cui.get(concept["cui"], []) if s.vocabulary != voc]
            if not provenance:
                continue
            target = links if key in available else unavailable
            target.setdefault(key, []).extend(provenance)
    merged = list(classified)
    positions = {(standardize_code(c.code, c.vocabulary), c.vocabulary): i
                 for i, c in enumerate(merged)}
    added, promoted, already_narrow = [], [], 0
    for key, provenance in sorted(links.items()):
        index = positions.get(key)
        previous = merged[index] if index is not None else None
        if previous is not None and previous.label == "Narrow":
            already_narrow += 1
            continue
        code, description = available[key]
        item = ClassifiedCandidate(code, description, key[1], "Narrow", 0.0,
                                   "Included by direct cross-vocabulary CodeMapper CUI mapping; "
                                   "not independently classified for EDF scope.",
                                   decision_origin="codemapper", mapping_provenance=provenance,
                                   prior_label=previous.label if previous else None,
                                   prior_confidence=previous.confidence if previous else None)
        if previous is None:
            merged.append(item)
            added.append({"code": code, "vocabulary": key[1]})
        else:
            merged[index] = replace(item, retrieval_score=previous.retrieval_score)
            promoted.append({"code": code, "vocabulary": key[1], "prior_label": previous.label})
    info = {"enabled": True, "rounds": 1, "descendants_used": False,
            "direct_links": [{"code": available[key][0], "vocabulary": key[1],
                              "provenance": provenance} for key, provenance in sorted(links.items())],
            "seeds": len(seeds), "seed_counts": seed_counts, "unique_cuis": len(by_cui),
            "unmapped_seeds": unmapped, "new_codes": added, "promoted_codes": promoted,
            "already_narrow": already_narrow,
            "unavailable_codes": [{"code": c, "vocabulary": v, "provenance": p}
                                  for (c, v), p in sorted(unavailable.items())],
            "calls": client.calls, "cache_hits": client.cache_hits,
            "seconds": round(time.perf_counter() - started, 3)}
    return merged, info