audit.py

#!/usr/bin/env python3
"""
MESS-Microbe Coverage Audit
===========================

Classifies each record into functional classes, checks every applicability
rule, and reports gaps. Produces both a per-record breakdown and a
database-wide summary.

Usage:
    python3 audit.py [--rules PATH] [--data PATH] [--out PATH]
"""

from __future__ import annotations

import argparse
import json
import sys
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any


# ----------------------------------------------------------------------------
# Path utilities
# ----------------------------------------------------------------------------

def get_path(record: dict, dotted: str):
    """Walk a dotted JSON path and return value or None."""
    cur: Any = record
    for part in dotted.split("."):
        if isinstance(cur, dict) and part in cur:
            cur = cur[part]
        else:
            return None
    return cur


def is_populated(value) -> bool:
    """A field counts as populated if it is non-None, non-empty array, non-empty dict, non-empty string."""
    if value is None:
        return False
    if isinstance(value, (list, dict, str)) and len(value) == 0:
        return False
    return True


# ----------------------------------------------------------------------------
# Functional class classification
# ----------------------------------------------------------------------------

def classify(record: dict) -> list[str]:
    """Assign a record to one or more functional classes per the rules.
    Classification is deliberately permissive — a record can be in multiple
    classes (e.g., G. sulfurreducens is electrogen + bidirectional_electroactive).
    """
    roles = set(record.get("mes_roles", []))
    is_ea = bool(record.get("is_electroactive"))
    eet_dir = record.get("eet_direction")
    long_range = (record.get("eet_mechanism") or {}).get("long_range_transport")

    classes: list[str] = []

    if is_ea and eet_dir in ("outward", "bidirectional") and (
        roles & {"electrogen", "bidirectional_eet"}
    ):
        classes.append("electrogen")

    if is_ea and eet_dir == "inward" and (
        roles & {"electrotroph", "acetogen", "homoacetogen",
                 "methanogen_acetoclastic", "methanogen_hydrogenotrophic", "methanogen_methylotrophic"}
    ) and "bidirectional_eet" not in roles:
        classes.append("electrotroph")

    if is_ea and eet_dir == "bidirectional":
        classes.append("bidirectional_electroactive")

    if "cable_bacterium" in roles or long_range is True:
        classes.append("cable_bacterium")

    if is_ea and (roles & {"acetogen", "homoacetogen"}):
        classes.append("acetogen_mes")

    if not is_ea and (roles & {"acetogen", "homoacetogen"}):
        classes.append("acetogen_non_mes")

    if (roles & {"methanogen_acetoclastic", "methanogen_hydrogenotrophic", "methanogen_methylotrophic"}) and "diet_partner_acceptor" in roles:
        classes.append("methanogen_diet")

    if (roles & {"methanogen_acetoclastic", "methanogen_hydrogenotrophic", "methanogen_methylotrophic"}) and "electrotroph" in roles:
        classes.append("methanogen_electrotroph")

    if roles & {"primary_fermenter", "secondary_fermenter"}:
        classes.append("fermenter")

    if (roles & {"syntrophic_h2_producer", "syntrophic_formate_producer"}) and not is_ea:
        classes.append("syntroph")

    if "sulfate_reducer" in roles:
        classes.append("sulfate_reducer")

    if is_ea and (
        roles & {"photoautotroph_oxygenic", "photoautotroph_anoxygenic", "photoheterotroph", "microalga"}
    ):
        classes.append("phototroph_electroactive")

    if not is_ea and not (roles & {
        "acetogen", "homoacetogen", "syntrophic_h2_producer", "syntrophic_formate_producer"
    }):
        # catches denitrifiers, hydrolyzers, contaminant degraders, etc., when not electroactive
        if not classes:
            classes.append("non_electroactive_community")

    if not classes:
        classes.append("unclassified")

    return classes


# ----------------------------------------------------------------------------
# Field applicability lookup
# ----------------------------------------------------------------------------

def applicability_for(field_path: str, classes: list[str], rules: dict) -> str:
    """Look up applicability level (REQUIRED/EXPECTED/OPTIONAL/NA) for a given
    field across the record's functional classes. Strictest wins:
      REQUIRED > EXPECTED > OPTIONAL > NA
    Exception: if ALL applicable classes say NA, the field is NA.
    """
    rule = rules.get(field_path)
    if rule is None:
        return "OPTIONAL"

    levels = []
    for cls in classes:
        lvl = rule.get(cls)
        if lvl is None:
            continue
        levels.append(lvl)
    if not levels:
        levels = [rule.get("_default", "OPTIONAL")]

    # If every applicable class says NA, the field is NA
    if all(lvl == "NA" for lvl in levels):
        return "NA"

    # Strictness ordering (excluding NA which we just handled)
    order = {"REQUIRED": 3, "EXPECTED": 2, "OPTIONAL": 1, "NA": 0}
    non_na = [lvl for lvl in levels if lvl != "NA"]
    return max(non_na, key=lambda x: order[x])


# ----------------------------------------------------------------------------
# Audit logic
# ----------------------------------------------------------------------------

def audit_record(record: dict, rules: dict) -> dict:
    classes = classify(record)
    field_rules = rules["applicability_rules"]

    results = {
        "id": record.get("id"),
        "scientific_name": record.get("scientific_name"),
        "classes": classes,
        "missing_required": [],
        "missing_expected": [],
        "populated_optional": [],
        "applicable_required": 0,
        "applicable_expected": 0,
        "populated_required": 0,
        "populated_expected": 0,
    }

    for field_path, _rule in field_rules.items():
        if field_path.startswith("_"):
            continue
        applic = applicability_for(field_path, classes, field_rules)
        value = get_path(record, field_path)
        populated = is_populated(value)

        if applic == "REQUIRED":
            results["applicable_required"] += 1
            if populated:
                results["populated_required"] += 1
            else:
                results["missing_required"].append(field_path)
        elif applic == "EXPECTED":
            results["applicable_expected"] += 1
            if populated:
                results["populated_expected"] += 1
            else:
                results["missing_expected"].append(field_path)
        elif applic == "OPTIONAL" and populated:
            results["populated_optional"].append(field_path)

    # compute completeness percentages
    req_pct = (
        100 * results["populated_required"] / results["applicable_required"]
        if results["applicable_required"] > 0 else 100
    )
    exp_pct = (
        100 * results["populated_expected"] / results["applicable_expected"]
        if results["applicable_expected"] > 0 else 100
    )
    results["required_pct"] = round(req_pct, 1)
    results["expected_pct"] = round(exp_pct, 1)

    # derive achieved completeness tier
    thresholds = rules["completeness_thresholds"]
    achieved = "stub"
    for tier in ("partial", "comprehensive"):
        t = thresholds[tier]
        if req_pct >= t["required_pct"] and exp_pct >= t["expected_pct"]:
            achieved = tier
    results["achieved_completeness"] = achieved
    results["declared_completeness"] = (record.get("data_quality") or {}).get("completeness")
    results["completeness_match"] = (
        achieved == results["declared_completeness"]
        or (results["declared_completeness"] == "expert_reviewed" and achieved == "comprehensive")
    )

    return results


def audit_database(data_path: Path, rules_path: Path) -> dict:
    rules = json.loads(rules_path.read_text())
    data = json.loads(data_path.read_text())
    records = data["microbes"]

    record_audits = [audit_record(r, rules) for r in records]

    # Database-wide rollup
    class_counts = Counter()
    for ra in record_audits:
        for cls in ra["classes"]:
            class_counts[cls] += 1

    completeness_audit = Counter()
    for ra in record_audits:
        completeness_audit[(ra["declared_completeness"], ra["achieved_completeness"])] += 1

    declared_match = sum(1 for ra in record_audits if ra["completeness_match"])

    # Aggregate gaps across DB
    field_gap_counts = defaultdict(lambda: {"required_missing": 0, "expected_missing": 0})
    for ra in record_audits:
        for f in ra["missing_required"]:
            field_gap_counts[f]["required_missing"] += 1
        for f in ra["missing_expected"]:
            field_gap_counts[f]["expected_missing"] += 1

    return {
        "n_records": len(records),
        "class_counts": dict(class_counts),
        "completeness_audit": {f"{k[0]}_declared_{k[1]}_achieved": v for k, v in completeness_audit.items()},
        "completeness_match_count": declared_match,
        "completeness_match_pct": round(100 * declared_match / len(records), 1),
        "field_gap_counts": dict(field_gap_counts),
        "record_audits": record_audits,
    }


# ----------------------------------------------------------------------------
# Markdown report
# ----------------------------------------------------------------------------

def render_markdown_report(audit: dict) -> str:
    lines: list[str] = []
    lines.append("# MESS-Microbe Coverage Audit Report\n")
    lines.append(f"**Records audited:** {audit['n_records']}\n")
    lines.append(f"**Declared = Achieved completeness:** "
                 f"{audit['completeness_match_count']} / {audit['n_records']} "
                 f"({audit['completeness_match_pct']}%)\n")

    lines.append("\n## Functional class distribution\n")
    lines.append("| Class | Records |")
    lines.append("| --- | ---: |")
    for cls, n in sorted(audit["class_counts"].items(), key=lambda x: -x[1]):
        lines.append(f"| {cls} | {n} |")

    lines.append("\n## Per-record audit\n")
    lines.append("| ID | Classes | Declared | Achieved | Match | Req % | Exp % | Top missing REQUIRED |")
    lines.append("| --- | --- | --- | --- | :---: | ---: | ---: | --- |")
    for ra in audit["record_audits"]:
        cls_str = ", ".join(ra["classes"])
        match = "✓" if ra["completeness_match"] else "✗"
        miss_r = ", ".join(ra["missing_required"][:3]) or "—"
        if len(ra["missing_required"]) > 3:
            miss_r += f" (+{len(ra['missing_required']) - 3} more)"
        lines.append(
            f"| `{ra['id']}` | {cls_str} | {ra['declared_completeness']} | "
            f"{ra['achieved_completeness']} | {match} | "
            f"{ra['required_pct']} | {ra['expected_pct']} | {miss_r} |"
        )

    # Database-wide gap report
    lines.append("\n## Database-wide field gap report\n")
    lines.append("Fields ranked by total missing across applicable records.\n")
    gap_rows = []
    for field, counts in audit["field_gap_counts"].items():
        total = counts["required_missing"] + counts["expected_missing"]
        gap_rows.append((total, field, counts["required_missing"], counts["expected_missing"]))
    gap_rows.sort(reverse=True)

    if gap_rows:
        lines.append("| Field | Missing as REQUIRED | Missing as EXPECTED | Total |")
        lines.append("| --- | ---: | ---: | ---: |")
        for total, field, req, exp in gap_rows[:30]:
            lines.append(f"| `{field}` | {req} | {exp} | {total} |")
    else:
        lines.append("_No gaps detected._\n")

    # Per-record detail blocks
    lines.append("\n## Per-record detail\n")
    for ra in audit["record_audits"]:
        lines.append(f"### `{ra['id']}` — {ra['scientific_name']}\n")
        lines.append(f"- **Classes:** {', '.join(ra['classes'])}")
        lines.append(f"- **Declared completeness:** {ra['declared_completeness']}; **achieved:** {ra['achieved_completeness']}")
        lines.append(f"- **Required coverage:** {ra['populated_required']}/{ra['applicable_required']} ({ra['required_pct']}%)")
        lines.append(f"- **Expected coverage:** {ra['populated_expected']}/{ra['applicable_expected']} ({ra['expected_pct']}%)")
        if ra["missing_required"]:
            lines.append("- **Missing REQUIRED:** " + ", ".join(f"`{f}`" for f in ra["missing_required"]))
        if ra["missing_expected"]:
            lines.append("- **Missing EXPECTED:** " + ", ".join(f"`{f}`" for f in ra["missing_expected"][:10]))
            if len(ra["missing_expected"]) > 10:
                lines.append(f"  (+{len(ra['missing_expected']) - 10} more)")
        lines.append("")

    return "\n".join(lines)


# ----------------------------------------------------------------------------
# CLI
# ----------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--rules", default="./applicability_rules.json")
    parser.add_argument("--data", default="../data/mess-microbes.seed.json")
    parser.add_argument("--out", default="../reports/coverage_audit.md")
    args = parser.parse_args()

    base = Path(__file__).parent
    rules_path = (base / args.rules).resolve()
    data_path = (base / args.data).resolve()
    out_path = (base / args.out).resolve()

    if not rules_path.exists():
        print(f"ERROR: rules file not found: {rules_path}", file=sys.stderr); sys.exit(1)
    if not data_path.exists():
        print(f"ERROR: data file not found: {data_path}", file=sys.stderr); sys.exit(1)

    audit = audit_database(data_path, rules_path)
    report = render_markdown_report(audit)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(report)

    print(f"Records audited: {audit['n_records']}")
    print(f"Declared=Achieved: {audit['completeness_match_count']}/{audit['n_records']} ({audit['completeness_match_pct']}%)")
    print(f"Report written to: {out_path}")


if __name__ == "__main__":
    main()