03_convert_html_to_csv.py

#!/usr/bin/env python3
"""Convert EUCAST MIC HTML pages into normalized CSV outputs."""

from __future__ import annotations

import csv
import re
from pathlib import Path
from typing import Dict, Iterable, List, Sequence, Tuple

from bs4 import BeautifulSoup  # type: ignore[import]
from rpy2.robjects import StrVector, r  # type: ignore[import]
from rpy2.robjects.packages import importr  # type: ignore[import]

SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent

SPECIES_CSV = SCRIPT_DIR / "species_eucast_with_amr.csv"
HTML_DIR = SCRIPT_DIR / "eucast_mic_html"
ANTIBIOTICS_REFERENCE = PROJECT_ROOT / "antibiotics.csv"

OUTPUT_COMBOS = SCRIPT_DIR / "mic_combinations.csv"
OUTPUT_VALUES = SCRIPT_DIR / "mic_values.csv"

SUMMARY_HEADERS = {
    "Distributions": "distribution_count",
    "Observations": "observation_count",
    "(T)ECOFF": "ecoff_value",
    "Confidence interval": "confidence_interval",
}

INVALID_ECOFFS = {"ID", "IE"}

SEPARATOR_PATTERN = re.compile(r"[‐‑‒–—-]")
PAREN_PATTERN = re.compile(r"\([^)]*\)")

AMR = importr("AMR")


def normalize_text(text: str) -> str:
    return " ".join(text.split())


def load_species_map(path: Path) -> Dict[str, Dict[str, str]]:
    with path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        mapping: Dict[str, Dict[str, str]] = {}
        for row in reader:
            mapping[row["species_id"]] = {
                "species_name": row["species_name"],
                "amr_mo": row["amr_mo"],
                "phenotype": row.get("phenotype", ""),
            }
    return mapping


def normalize_variants(value: str) -> List[str]:
    """Return canonical matching keys (plain + variant without parentheses)."""

    def _normalize(text: str) -> str:
        text = SEPARATOR_PATTERN.sub("/", text.lower())
        text = re.sub(r"\s*/\s*", "/", text)
        text = re.sub(r"\s+", " ", text)
        return text.strip()

    cleaned = value.strip()
    if not cleaned:
        return []

    variants = {_normalize(cleaned)}
    stripped = PAREN_PATTERN.sub(" ", cleaned)
    stripped = re.sub(r"\s+", " ", stripped).strip()
    if stripped and stripped != cleaned:
        variants.add(_normalize(stripped))

    return list(variants)


def load_antibiotic_reference(path: Path) -> Dict[str, str]:
    lookup: Dict[str, str] = {}
    with path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            amr_code = row.get("amr_code", "").strip()
            if not amr_code:
                continue
            for column in (
                "full_name_en",
                "full_name_de",
                "short_name_en",
                "short_name_de",
            ):
                value = row.get(column, "")
                if value:
                    for variant in normalize_variants(value):
                        lookup.setdefault(variant, amr_code)
            for column in ("synonyms_en", "synonyms_de"):
                value = row.get(column, "")
                if value:
                    for synonym in value.split(";"):
                        synonym = synonym.strip()
                        if synonym:
                            for variant in normalize_variants(synonym):
                                lookup.setdefault(variant, amr_code)
    return lookup


def resolve_antibiotic_code(
    name: str, reference: Dict[str, str]
) -> Tuple[str, str]:
    for variant in normalize_variants(name):
        if variant in reference:
            return reference[variant], "reference"

    mos = AMR.as_ab(StrVector([name]))
    code = r["as.character"](mos)[0]
    if code not in ("NA", "NULL", "UNKNOWN"):
        return code, "AMR::as_ab"

    raise ValueError(
        f"Unable to resolve AMR code for antibiotic '{name}'. "
        "Add it to antibiotics.csv or adjust the resolver."
    )


def classify_ecoff(value: str) -> Tuple[str, str]:
    value = value.strip()
    if not value:
        return "", "missing"
    if value in INVALID_ECOFFS:
        return "", "invalid"
    if value == "-":
        return "", "less_than_three"
    if value.startswith("((") and value.endswith("))"):
        return value[2:-2].strip(), "forced_ecoff"
    if value.startswith("(") and value.endswith(")"):
        return value[1:-1].strip(), "tentative_ecoff"
    return value, "value"


def parse_confidence_interval(raw: str) -> Tuple[str, str]:
    raw = raw.strip()
    if not raw:
        return "", ""
    match = re.match(r"^\s*([0-9]+(?:\.[0-9]+)?)\s*-\s*([0-9]+(?:\.[0-9]+)?)\s*$", raw)
    if not match:
        return "", ""
    return match.group(1), match.group(2)


def to_int(text: str) -> int:
    text = text.strip()
    return int(text) if text.isdigit() else 0


def find_html_path(species_id: str) -> Path | None:
    matches = sorted(HTML_DIR.glob(f"{species_id}_*.html"))
    if matches:
        return matches[0]
    return None


def extract_headers(table: BeautifulSoup) -> Tuple[List[str], List[Tuple[int, str]]]:
    header_cells = table.thead.find_all("th") if table.thead else []
    headers = [normalize_text(cell.get_text()) for cell in header_cells]
    dilution_columns: List[Tuple[int, str]] = []
    for idx, label in enumerate(headers):
        if idx == 0 or label in SUMMARY_HEADERS or not label:
            continue
        dilution_columns.append((idx, label))
    return headers, dilution_columns


def parse_table_rows(
    table: BeautifulSoup,
    headers: List[str],
    dilution_columns: List[Tuple[int, str]],
) -> Iterable[List[str]]:
    for tbody in table.find_all("tbody"):
        for tr in tbody.find_all("tr"):
            cells = tr.find_all("td")
            if len(cells) != len(headers):
                continue
            yield [normalize_text(cell.get_text()) for cell in cells]


def main() -> None:
    species_map = load_species_map(SPECIES_CSV)
    antibiotic_lookup = load_antibiotic_reference(ANTIBIOTICS_REFERENCE)

    combo_rows: List[Dict[str, str]] = []
    value_rows: List[Dict[str, str]] = []
    combo_id = 1

    for species_id in sorted(species_map.keys(), key=lambda x: int(x)):
        metadata = species_map[species_id]
        html_path = find_html_path(species_id)
        if html_path is None:
            print(f"[WARN] Missing HTML for species {species_id} ({metadata['species_name']})")
            continue

        soup = BeautifulSoup(html_path.read_text(encoding="utf-8"), "html.parser")
        table = soup.find("table", id="search-results-table")
        if not table:
            print(f"[WARN] No MIC table in {html_path}")
            continue

        headers, dilution_columns = extract_headers(table)
        try:
            distribution_idx = headers.index("Distributions")
            observation_idx = headers.index("Observations")
            ecoff_idx = headers.index("(T)ECOFF")
            confidence_idx = headers.index("Confidence interval")
        except ValueError as exc:
            print(f"[WARN] Missing expected column in {html_path}: {exc}")
            continue

        for row in parse_table_rows(table, headers, dilution_columns):
            antibiotic_name = row[0]
            if not antibiotic_name:
                continue

            ecoff_text = row[ecoff_idx]
            if ecoff_text in INVALID_ECOFFS:
                continue

            ecoff_value, ecoff_annotation = classify_ecoff(ecoff_text)
            antibiotic_code, source = resolve_antibiotic_code(antibiotic_name, antibiotic_lookup)
            if source == "AMR::as_ab":
                print(f"[INFO] Antibiotic '{antibiotic_name}' resolved via AMR::as_ab → {antibiotic_code}")

            distribution_count = to_int(row[distribution_idx])
            observation_count = to_int(row[observation_idx])
            confidence_lower, confidence_upper = parse_confidence_interval(row[confidence_idx])

            combo_rows.append(
                {
                    "combo_id": str(combo_id),
                    "species_amr_mo": metadata["amr_mo"],
                    "species_name": metadata["species_name"],
                    "phenotype": metadata.get("phenotype", ""),
                    "antibiotic_name": antibiotic_name,
                    "antibiotic_amr_code": antibiotic_code,
                    "distribution_count": str(distribution_count),
                    "observation_count": str(observation_count),
                    "ecoff_value": ecoff_value,
                    "ecoff_annotation": ecoff_annotation,
                    "confidence_lower": confidence_lower,
                    "confidence_upper": confidence_upper,
                }
            )

            for idx, label in dilution_columns:
                count = to_int(row[idx])
                value_rows.append(
                    {
                        "combo_id": str(combo_id),
                        "dilution_mg_l": label,
                        "count": str(count),
                    }
                )

            combo_id += 1

    write_csv(
        OUTPUT_COMBOS,
        combo_rows,
        [
            "combo_id",
            "species_amr_mo",
            "species_name",
            "phenotype",
            "antibiotic_name",
            "antibiotic_amr_code",
            "distribution_count",
            "observation_count",
            "ecoff_value",
            "ecoff_annotation",
            "confidence_lower",
            "confidence_upper",
        ],
    )
    write_csv(
        OUTPUT_VALUES,
        value_rows,
        ["combo_id", "dilution_mg_l", "count"],
    )
    print(
        f"Wrote {len(combo_rows)} combinations to {OUTPUT_COMBOS} "
        f"and {len(value_rows)} MIC rows to {OUTPUT_VALUES}"
    )


def write_csv(path: Path, rows: Sequence[Dict[str, str]], fieldnames: List[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


if __name__ == "__main__":
    main()