scripts/audit_unifac_groups.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

from rdkit import Chem

sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from pino.thermo.unifac import DOUFSG_AVAILABLE_GROUPS, get_dortmund_subgroups


def audit_registry_thermodynamics(registry_path: str | Path) -> dict[str, Any]:
    """Scan a JSON-Lines registry and flag compounds that cannot be represented in DOUFSG."""
    print("=== Beginning Comprehensive Dortmund UNIFAC (DOUFSG) Group Audit ===")

    with Path(registry_path).open("r", encoding="utf-8") as f:
        compounds = [json.loads(line) for line in f]

    unifac_failures = []
    total = len(compounds)
    inspected = 0

    for c in compounds:
        cas = c.get("cas", "")
        smiles = c.get("smiles", "")
        name = c.get("name", "Unknown")

        if not smiles:
            unifac_failures.append({
                "cas": cas,
                "name": name,
                "smiles": smiles,
                "reason": "Missing SMILES string",
                "missing_subgroup_ids": [],
            })
            continue

        try:
            mol = Chem.MolFromSmiles(smiles)
            if mol is None:
                raise ValueError("Invalid SMILES structure")

            molecule_subgroups = get_dortmund_subgroups(smiles)
            inspected += 1

            # Check if all discovered structural fragments exist in the DOUFSG parameter set
            missing_in_doufsg = [
                int(gid) for gid in molecule_subgroups.keys()
                if int(gid) not in DOUFSG_AVAILABLE_GROUPS
            ]
            if missing_in_doufsg:
                unifac_failures.append({
                    "cas": cas,
                    "name": name,
                    "smiles": smiles,
                    "reason": f"Subgroup IDs {missing_in_doufsg} not present in active DOUFSG definitions",
                    "missing_subgroup_ids": missing_in_doufsg,
                })

            if not molecule_subgroups:
                # Either fragmentation failed or every subgroup was missing.
                # Re-check by trying a raw ugropy call to report the missing subgroup names.
                reason = "UNIFAC fragmentation returned no usable DOUFSG subgroups"
                try:
                    import ugropy
                    raw_groups = ugropy.Groups(smiles, identifier_type="smiles").dortmund.subgroups
                    unknown = [
                        name for name in raw_groups.keys()
                        if int(getattr(name, "id", 0)) not in DOUFSG_AVAILABLE_GROUPS
                    ]
                    if unknown:
                        reason = f"ugropy groups {unknown} have no matching DOUFSG parameter"
                except Exception:
                    pass
                unifac_failures.append({
                    "cas": cas,
                    "name": name,
                    "smiles": smiles,
                    "reason": reason,
                    "missing_subgroup_ids": [],
                })

        except Exception as e:
            unifac_failures.append({
                "cas": cas,
                "name": name,
                "smiles": smiles,
                "reason": f"UNIFAC fragmentation error: {str(e)}",
                "missing_subgroup_ids": [],
            })

    print(f"\nAudit complete. Inspected {total} materials.")
    print(f"Successfully derived DOUFSG subgroups for {inspected} materials.")

    if unifac_failures:
        print(f"⚠️ Found {len(unifac_failures)} thermodynamic integration blockers:")
        for fail in unifac_failures[:10]:
            print(f"  - CAS [{fail['cas']}] ({fail['name']}): {fail['reason']}")
        if len(unifac_failures) > 10:
            print(f"  ... and {len(unifac_failures) - 10} more")
    else:
        print("✅ SUCCESS: 100% of registry compounds have valid, un-fragmented DOUFSG definitions.")

    return {
        "total": total,
        "inspected": inspected,
        "failures": unifac_failures,
    }


def main() -> int:
    parser = argparse.ArgumentParser(description="Audit registry compounds against DOUFSG")
    parser.add_argument(
        "--registry",
        default="data/aroma_chemicals.jsonl",
        help="Path to the JSON-Lines registry file",
    )
    parser.add_argument(
        "--output",
        default="data/unifac_missing_groups.json",
        help="Path to write the failure report",
    )
    args = parser.parse_args()

    result = audit_registry_thermodynamics(args.registry)

    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(result, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"\nBlockers exported to {output_path}.")

    return 0 if not result["failures"] else 1


if __name__ == "__main__":
    raise SystemExit(main())