duocthu/ingestion/ingestion/entities/catalog.py

from __future__ import annotations

import argparse
import json
import re
import unicodedata
from difflib import SequenceMatcher
from pathlib import Path

import fitz

from ingestion.validation.back_index import parse_back_index_see_aliases

WORD_RE = re.compile(r"\w+", re.UNICODE)
PAREN_RE = re.compile(r"\(([^()]*)\)")


def normalize_name(text: str) -> str:
    decomposed = unicodedata.normalize("NFKD", text.casefold()).replace("đ", "d")
    plain = "".join(char for char in decomposed if not unicodedata.combining(char))
    return " ".join(WORD_RE.findall(plain))


def _canonical_aliases(drug_name: str) -> set[str]:
    aliases = {drug_name.strip()}
    without_parentheses = PAREN_RE.sub("", drug_name).strip()
    if without_parentheses:
        aliases.add(without_parentheses)
    aliases.update(
        value.strip() for value in PAREN_RE.findall(drug_name) if value.strip()
    )
    return aliases


def _trade_aliases(monograph: dict) -> set[str]:
    section = monograph.get("sections", {}).get("ten_thuong_mai")
    if not section:
        return set()
    text = section.get("text", "")
    return {
        value.strip().strip(".")
        for value in re.split(r"[,;\n]", text)
        if value.strip().strip(".")
    }


def _read_monographs(path: Path) -> list[dict]:
    with path.open(encoding="utf-8") as handle:
        return [json.loads(line) for line in handle if line.strip()]


def build_entities(monographs_path: Path, pdf_path: Path) -> dict:
    monographs = _read_monographs(monographs_path)
    entities: dict[str, dict] = {}
    lookup: list[tuple[str, str, tuple[int, int]]] = []
    for monograph in monographs:
        drug_id = monograph["drug_id"]
        aliases = _canonical_aliases(monograph["drug_name"])
        aliases.update(_trade_aliases(monograph))
        aliases.add(drug_id.replace("_", " "))
        entities[drug_id] = {
            "drug_id": drug_id,
            "canonical_name": monograph["drug_name"],
            "aliases": aliases,
            "atc_codes": sorted(set(monograph.get("atc_codes", []))),
            "source_page_range": monograph.get("source_page_range"),
        }
        page_range = tuple(monograph["source_page_range"])
        lookup.extend(
            (normalize_name(alias), drug_id, page_range) for alias in _canonical_aliases(
                monograph["drug_name"],
            ) if normalize_name(alias)
        )

    unresolved = []
    ambiguous = []
    with fitz.open(pdf_path) as doc:
        index_aliases = parse_back_index_see_aliases(doc)
    for relation in index_aliases:
        target = normalize_name(relation.target)
        candidates = {
            drug_id for alias, drug_id, page_range in lookup
            if page_range[0] <= relation.printed_page - 1 <= page_range[1]
            and (
                target == alias
                or target.startswith(f"{alias} ")
                or target.endswith(f" {alias}")
                or f" {alias} " in target
            )
        }
        if not candidates:
            fuzzy = sorted(
                (
                    SequenceMatcher(None, target, alias).ratio(),
                    drug_id,
                )
                for alias, drug_id, page_range in lookup
                if page_range[0] <= relation.printed_page - 1 <= page_range[1]
            )
            if fuzzy and fuzzy[-1][0] >= 0.9:
                runner_up = fuzzy[-2][0] if len(fuzzy) > 1 else 0.0
                if fuzzy[-1][0] - runner_up >= 0.05:
                    candidates = {fuzzy[-1][1]}
        if len(candidates) == 1:
            entities[next(iter(candidates))]["aliases"].add(relation.alias)
        elif candidates:
            ambiguous.append(relation.alias)
        else:
            unresolved.append(relation.alias)

    output_entities = []
    for entity in entities.values():
        output_entities.append({
            **entity,
            "aliases": sorted(entity["aliases"], key=lambda value: value.casefold()),
        })
    return {
        "schema_version": 1,
        "entities": sorted(output_entities, key=lambda item: item["drug_id"]),
        "stats": {
            "entity_count": len(output_entities),
            "back_index_see_relations": len(index_aliases),
            "back_index_aliases_mapped": len(index_aliases) - len(unresolved) - len(ambiguous),
            "back_index_aliases_unresolved": len(unresolved),
            "back_index_aliases_ambiguous": len(ambiguous),
            "trade_name_sections": sum(
                "ten_thuong_mai" in row.get("sections", {}) for row in monographs
            ),
            "total_aliases": sum(len(row["aliases"]) for row in output_entities),
        },
        "unresolved_back_index_aliases": sorted(unresolved),
        "ambiguous_back_index_aliases": sorted(ambiguous),
    }


def write_entities(payload: dict, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--monographs", type=Path, required=True)
    parser.add_argument("--pdf", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    payload = build_entities(args.monographs, args.pdf)
    write_entities(payload, args.output)
    print(json.dumps(payload["stats"], ensure_ascii=False))


if __name__ == "__main__":
    main()