Add the PROTAC release
f12d3b3 verified 7 days ago
50 kB
#!/usr/bin/env python3
"""
batch_report.py
Evaluate one or more generation batches and produce a summary chart.
Usage:
    # Single batch (summary CSV from defog_guided.py):
    python batch_report.py ../generation_logs/defog_local_edit_20260424_003754_summary.csv
    # Multiple batches — one bar-group per batch:
    python batch_report.py ../generation_logs/*_summary.csv
    # Include a named baseline (plain SMILES .txt file, one per line):
    python batch_report.py ../generation_logs/*_summary.csv --baseline DeLinker:delinker_linkers.txt
    # Save chart to file instead of displaying:
    python batch_report.py ../generation_logs/*_summary.csv --out report.png
"""

import argparse
import subprocess
import sys

# ── Edit these to run without command-line arguments ─────────────────────────
# Set BATCH_FOLDER to a folder path to use all *_summary.csv files in it,
# or set BATCH_CSVS to a list of specific CSV paths.
# These are ignored when you pass paths on the command line.
BATCH_FOLDER: str | None = None          # e.g. "../generation_logs"
BATCH_CSVS:   list[str]  = [
    "../generation_logs/defog_local_edit_20260425_170213_summary.csv",  # COOH best (old)
    "../generation_logs/defog_local_edit_20260428_005755_summary.csv",  # COOH rerun
    "../generation_logs/defog_local_edit_20260425_173602_summary.csv",  # methyl best
]

# Human-readable labels for DeFoG batch runs (keyed by timestamp "YYYYMMDD_HHMMSS").
# Any timestamp not listed falls back to the timestamp string itself.
BATCH_LABELS: dict[str, str] = {
    "20260425_170213": "COOH Guided\n(orig)",
    "20260428_005755": "COOH Guided\n(rerun)",
    "20260425_173602": "Methyl Guided",
}
CHART_OUT:    str | None = "batch_report.png"  # saved next to this script

# ── Named baselines (plain SMILES .txt files, one SMILES per line) ────────────
# Add external model baselines here as {"Label": "/path/to/smiles.txt"}.
# Each baseline will appear as a bar-group alongside the DeFoG batches.
# Example:
#   BASELINES = {"DeLinker": "../generation_logs/delinker_linkers.txt"}
BASELINES: dict[str, str] = {
    # "DeLinker": "../generation_logs/delinker_linkers.txt",
}

# ── Full-PROTAC baselines (models that output assembled PROTACs, not linkers) ─
# Evaluated directly without assembly — SA/QED/MW/logP/RotBonds computed on the
# full molecule as-is. Use for models like DRlinker that generate whole PROTACs.
# Example:
#   PROTAC_BASELINES = {"DRlinker": "../generation_logs/drlinker_SyntaLinkerN_ms39_....txt"}
# Per-(protac-module) baseline dispatch. The active set is selected by
# --protac-module at runtime in main(). Each entry is {label: path}; empty
# strings are skipped. Add new targets here when their baseline outputs land.
# PROTAC baselines (LinkerNet/DiffPROTACs/DiffLinker outputs) are NOT included
# in this distribution. Add paths here if you want to include external baseline
# comparisons; otherwise the report compares only against itself.
PROTAC_BASELINES_BY_MODULE: dict[str, dict[str, str]] = {}
PROTAC_BASELINES: dict[str, str] = {}
# ─────────────────────────────────────────────────────────────────────────────
import tempfile
from pathlib import Path

import importlib
import importlib.util
import os
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from rdkit import Chem
from rdkit import RDConfig
from rdkit.Chem import Descriptors, QED as rdkit_qed


EVAL_SCRIPT = Path(__file__).parent / "evaluate_linkers.py"
PYTHON = sys.executable  # same interpreter (must be moca env)
LINKER_CSV = Path(__file__).parent.parent / "data" / "linker.csv"

_REPO_ROOT_DEFAULT = Path(__file__).parent.parent
_SRC_DEFAULT       = _REPO_ROOT_DEFAULT / "src"

# Fallback substructure patterns — overridden at startup by _load_protac_module()
_WARHEAD_PAT = Chem.MolFromSmarts("c1cc2ncncc2cc1")
_E3_PAT      = Chem.MolFromSmarts("N1CC(O)CC1")

# Active protac module name (set by main() via --protac-module)
_PROTAC_MODULE = "ErlotinibVHL"

# ── PROTACAnalyzer (lazy-loaded) ──────────────────────────────────────────────
_protac_analyzer_state: dict | None = None


def _get_protac_analyzer():
    """Lazy-load PROTACAnalyzer using the active _PROTAC_MODULE.
    Returns (analyzer, wh_q, e3_q, wh_reactive_comp, e3_reactive_comp).
    All are None on load failure.
    """
    global _protac_analyzer_state, _WARHEAD_PAT, _E3_PAT
    if _protac_analyzer_state is not None:
        s = _protac_analyzer_state
        return s["analyzer"], s["wh_q"], s["e3_q"], s["wh_reactive_comp"], s["e3_reactive_comp"]
    try:
        sys.path.insert(0, str(_SRC_DEFAULT / "classifier"))
        sys.path.insert(0, str(_SRC_DEFAULT / "protac"))
        import importlib as _il
        from protac_analyzer import PROTACAnalyzer
        _pmod = _il.import_module(_PROTAC_MODULE)
        WARHEAD_MOL    = _pmod.WARHEAD_MOL
        E3_MOL         = _pmod.E3_MOL
        WARHEAD_SMILES = _pmod.WARHEAD_SMILES
        E3_SMILES      = _pmod.E3_SMILES
        wh_rxn_smarts  = _pmod.WH_REACTIVE_SMARTS
        e3_rxn_smarts  = _pmod.E3_REACTIVE_SMARTS
        # Update fallback patterns from module if provided
        if hasattr(_pmod, 'WARHEAD_PAT') and _pmod.WARHEAD_PAT is not None:
            _WARHEAD_PAT = _pmod.WARHEAD_PAT
        if hasattr(_pmod, 'E3_PAT') and _pmod.E3_PAT is not None:
            _E3_PAT = _pmod.E3_PAT
        _COOH_OR_CHO = "[$([CX3](=[OX1])[OX2H1]),$([CX3H1](=[OX1]))]"
        analyzer = PROTACAnalyzer(
            warhead=WARHEAD_MOL,
            e3=E3_MOL,
            warhead_attachment_smarts=_COOH_OR_CHO,
            e3_attachment_smarts=_COOH_OR_CHO,
        )
        _wh_pat  = Chem.MolFromSmarts(wh_rxn_smarts)
        _e3_pat  = Chem.MolFromSmarts(e3_rxn_smarts)
        _protac_analyzer_state = {
            "analyzer":         analyzer,
            "wh_q":             Chem.MolFromSmiles(WARHEAD_SMILES),
            "e3_q":             Chem.MolFromSmiles(E3_SMILES),
            "wh_reactive_comp": {m[0] for m in WARHEAD_MOL.GetSubstructMatches(_wh_pat)},
            "e3_reactive_comp": {m[0] for m in E3_MOL.GetSubstructMatches(_e3_pat)},
        }
        print(f"[PROTACAnalyzer] loaded for {_PROTAC_MODULE} validity scoring")
    except Exception as e:
        print(f"[PROTACAnalyzer] could not load ({e}); falling back to substructure check",
              file=sys.stderr)
        _protac_analyzer_state = {
            "analyzer": None, "wh_q": None, "e3_q": None,
            "wh_reactive_comp": None, "e3_reactive_comp": None,
        }
    s = _protac_analyzer_state
    return s["analyzer"], s["wh_q"], s["e3_q"], s["wh_reactive_comp"], s["e3_reactive_comp"]


def _extract_linker_from_protac(protac_mol, warhead_q, e3_q):
    """Break assembled PROTAC into warhead / E3 / linker.
    Returns (linker_mol, n_interface_bonds, wh_match, e3_match) or (None, 0, None, None).
    n_interface_bonds == 2 means the linker cleanly bridges both fragments.
    wh_match / e3_match are the substructure match tuples needed for reactive-site checks.
    """
    wh_match = protac_mol.GetSubstructMatch(warhead_q)
    e3_match = protac_mol.GetSubstructMatch(e3_q)
    if not wh_match or not e3_match:
        return None, 0, None, None
    lk_set = set(range(protac_mol.GetNumAtoms())) - set(wh_match) - set(e3_match)
    if not lk_set:
        return None, 0, wh_match, e3_match
    bond_ids = [
        b.GetIdx() for b in protac_mol.GetBonds()
        if (b.GetBeginAtomIdx() in lk_set) != (b.GetEndAtomIdx() in lk_set)
    ]
    if not bond_ids:
        return None, 0, wh_match, e3_match
    try:
        break_mol = Chem.FragmentOnBonds(protac_mol, bond_ids, addDummies=False)
        for frag in Chem.GetMolFrags(break_mol, asMols=True):
            if frag.GetNumHeavyAtoms() == len(lk_set):
                return frag, len(bond_ids), wh_match, e3_match
    except Exception:
        pass
    return None, 0, wh_match, e3_match


def _check_reactive_attachment(protac_mol, wh_match, e3_match,
                                wh_reactive_comp, e3_reactive_comp):
    """Verify each linker–fragment interface bond attaches at the correct reactive site.
    Maps wh_reactive_comp / e3_reactive_comp (indices in free component mols) into
    PROTAC atom indices via the substructure matches, then checks that every interface
    bond's fragment-side atom is one of those reactive atoms.
    """
    lk_set      = set(range(protac_mol.GetNumAtoms())) - set(wh_match) - set(e3_match)
    wh_set      = set(wh_match)
    e3_set      = set(e3_match)
    wh_reactive = {wh_match[i] for i in wh_reactive_comp}
    e3_reactive = {e3_match[i] for i in e3_reactive_comp}

    wh_ok = e3_ok = False
    for b in protac_mol.GetBonds():
        s, e = b.GetBeginAtomIdx(), b.GetEndAtomIdx()
        if (s in lk_set) == (e in lk_set):
            continue
        frag_atom = s if s not in lk_set else e
        if frag_atom in wh_set:
            wh_ok = wh_ok or (frag_atom in wh_reactive)
        elif frag_atom in e3_set:
            e3_ok = e3_ok or (frag_atom in e3_reactive)
    return wh_ok and e3_ok

_pk_scorer_cache = None


def _get_pk_scorer():
    global _pk_scorer_cache
    if _pk_scorer_cache is not None:
        return _pk_scorer_cache
    try:
        sys.path.insert(0, str(_SRC_DEFAULT / "classifier"))
        sys.path.insert(0, str(_SRC_DEFAULT / "protac"))
        import importlib as _il
        from protac_classifier import EGFRPROPKScorer
        _pmod = _il.import_module(_PROTAC_MODULE)
        _Builder = _pmod.Builder
        _MODEL_DIR = str(_SRC_DEFAULT / "propk" / "models")
        _ATOM_DECODER = ["C", "N", "O", "F", "P", "S", "Cl", "Br", "I"]
        scorer = EGFRPROPKScorer(_Builder(), _MODEL_DIR, _ATOM_DECODER)
        _pk_scorer_cache = scorer
        print("[EGFRPROPKScorer] loaded")
    except Exception as e:
        print(f"[EGFRPROPKScorer] could not load ({e})", file=sys.stderr)
        _pk_scorer_cache = False
    return _pk_scorer_cache


def _compute_pk_scores(df: pd.DataFrame) -> pd.DataFrame:
    """Append EGFR-PROPK columns (T1/2, CL, Vss, composite) for the assembled
    PROTAC. Only meaningful for ErlotinibVHL; caller gates on protac_module."""
    scorer = _get_pk_scorer()
    pk_t_half, pk_cl, pk_vss, pk_score = [], [], [], []
    for _, row in df.iterrows():
        smi = row.get("protac_smiles")
        if scorer and smi and str(smi) not in ("", "nan", "None"):
            try:
                mol = Chem.MolFromSmiles(str(smi))
                if mol is not None:
                    pk = scorer.predict_pk(mol)
                    pk_t_half.append(pk["t_half"])
                    pk_cl.append(pk["cl"])
                    pk_vss.append(pk["vss"])
                    pk_score.append(pk["score"])
                    continue
            except Exception:
                pass
        pk_t_half.append(float("nan"))
        pk_cl.append(float("nan"))
        pk_vss.append(float("nan"))
        pk_score.append(float("nan"))
    df = df.copy()
    df["pk_t_half"] = pk_t_half
    df["pk_cl"]     = pk_cl
    df["pk_vss"]    = pk_vss
    df["pk_score"]  = pk_score
    return df


_training_smiles_cache: set | None = None


def _load_training_smiles() -> set:
    global _training_smiles_cache
    if _training_smiles_cache is None:
        df = pd.read_csv(LINKER_CSV)
        smiles = set()
        for smi in df["Smiles"].dropna():
            mol = Chem.MolFromSmiles(smi)
            if mol:
                smiles.add(Chem.MolToSmiles(mol))
        _training_smiles_cache = smiles
        print(f"Loaded {len(smiles)} training SMILES from linker.csv")
    return _training_smiles_cache


def _canonicalize(smiles_series) -> list:
    out = []
    for smi in smiles_series:
        mol = Chem.MolFromSmiles(smi)
        out.append(Chem.MolToSmiles(mol) if mol else None)
    return [s for s in out if s]


# ── helpers ──────────────────────────────────────────────────────────────────

# Lazy-loaded evaluate_linkers module — load once, reuse across all calls.
# Eliminates the ~30 sec subprocess startup × N calls = several minutes.
_eval_module_cache = None

def _get_eval_module():
    global _eval_module_cache
    if _eval_module_cache is None:
        spec = importlib.util.spec_from_file_location("evaluate_linkers", EVAL_SCRIPT)
        mod  = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(mod)
        _eval_module_cache = mod
    return _eval_module_cache


def _call_evaluate_linkers(valid_smiles: list[str], label: str = "") -> pd.DataFrame:
    """Run evaluate_linkers.evaluate() in-process (no subprocess overhead).
    Forwards the active _PROTAC_MODULE so the right warhead/E3 builder and
    DegradMaster pocket+features are loaded.
    """
    ev = _get_eval_module()
    df = ev.evaluate(list(valid_smiles), sa_max=6.0, verbose=False,
                     protac_module=_PROTAC_MODULE)
    return df


def _std_valid_count(smiles_iter) -> int:
    """Count RDKit-parseable SMILES (no connectivity requirement).
    Disconnected fragments still count as parseable; the connectivity check
    falls under AsmVal (a disconnected linker can never canonically attach).
    """
    n = 0
    for smi in smiles_iter:
        if not smi:
            continue
        if Chem.MolFromSmiles(str(smi)) is not None:
            n += 1
    return n


def run_evaluate(summary_csv: Path) -> pd.DataFrame:
    """Extract valid SMILES from a generation summary CSV, run evaluate_linkers,
    return the per-molecule results DataFrame with uniqueness/novelty attached."""
    gen_df = pd.read_csv(summary_csv)
    valid_smiles = list(gen_df[gen_df["valid"] == 1]["smiles"].dropna())
    n_total      = len(gen_df)
    n_valid      = len(valid_smiles)
    n_std_valid  = _std_valid_count(gen_df["smiles"].fillna(""))

    train_smiles = _load_training_smiles()
    canon = _canonicalize(valid_smiles)
    unique_canon = set(canon)
    n_unique   = len(unique_canon)
    n_in_train = len(unique_canon & train_smiles)
    uniqueness = 100 * n_unique / n_valid if n_valid else 0.0
    novelty    = 100 * (n_unique - n_in_train) / n_unique if n_unique else 0.0

    df = _call_evaluate_linkers(valid_smiles, label=str(summary_csv))
    df["_n_total"]      = n_total
    df["_n_valid"]      = n_valid
    df["_n_std_valid"]  = n_std_valid
    df["_n_unique"]     = n_unique
    df["_uniqueness"]   = uniqueness
    df["_novelty"]      = novelty
    return df


def _read_smiles_file(smiles_path: Path):
    """Read a SMILES text file. Returns (smiles_list, n_total_override).
    Lines starting with '#' are treated as comments and skipped, EXCEPT that
    a line of the form '# n_total: N' encodes the true generation count
    (e.g. including invalid SMILES that were dropped before saving).
    """
    lines = Path(smiles_path).read_text().splitlines()
    n_total_override = None
    smiles = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        if line.startswith("#"):
            if line.lower().startswith("# n_total:"):
                try:
                    n_total_override = int(line.split(":", 1)[1].strip())
                except ValueError:
                    pass
        else:
            smiles.append(line)
    return smiles, n_total_override


def _diversity_stats(valid_smiles: list[str]):
    """Return (n_unique, uniqueness%, novelty%) over the canonical set."""
    n_valid = len(valid_smiles)
    train_smiles = _load_training_smiles()
    canon = _canonicalize(valid_smiles)
    unique_canon = set(canon)
    n_unique   = len(unique_canon)
    n_in_train = len(unique_canon & train_smiles)
    uniqueness = 100 * n_unique / n_valid       if n_valid  else 0.0
    novelty    = 100 * (n_unique - n_in_train) / n_unique if n_unique else 0.0
    return n_unique, uniqueness, novelty


def run_evaluate_smiles_file(smiles_path: Path) -> pd.DataFrame:
    """Load SMILES from a plain text file (one per line), run evaluate_linkers,
    and return the same per-molecule DataFrame as run_evaluate."""
    raw, n_total_override = _read_smiles_file(smiles_path)
    valid_smiles = [s for s in raw if Chem.MolFromSmiles(s)]
    n_total      = n_total_override if n_total_override is not None else len(raw)
    n_valid      = len(valid_smiles)
    n_std_valid  = _std_valid_count(raw)
    n_unique, uniqueness, novelty = _diversity_stats(valid_smiles)

    df = _call_evaluate_linkers(valid_smiles, label=str(smiles_path))
    df["_n_total"]      = n_total
    df["_n_valid"]      = n_valid
    df["_n_std_valid"]  = n_std_valid
    df["_n_unique"]     = n_unique
    df["_uniqueness"]   = uniqueness
    df["_novelty"]      = novelty
    return df


def _load_sascorer():
    _sa_path = os.path.join(RDConfig.RDContribDir, "SA_Score", "sascorer.py")
    spec = importlib.util.spec_from_file_location("sascorer", _sa_path)
    mod  = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


_sascorer = None

def _sa_score(mol):
    global _sascorer
    if _sascorer is None:
        _sascorer = _load_sascorer()
    return _sascorer.calculateScore(mol)


_ev_linkers_mod = None

def _get_run_degrademaster():
    """Lazily import _run_degrademaster from evaluate_linkers.py."""
    global _ev_linkers_mod
    if _ev_linkers_mod is None:
        try:
            spec = importlib.util.spec_from_file_location("evaluate_linkers", EVAL_SCRIPT)
            mod  = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(mod)
            _ev_linkers_mod = mod
        except Exception as e:
            print(f"[DegradeMaster] could not load evaluate_linkers: {e}", file=sys.stderr)
            _ev_linkers_mod = False
    if not _ev_linkers_mod:
        return None
    return getattr(_ev_linkers_mod, "_run_degrademaster", None)


def run_evaluate_full_protac_file(smiles_path: Path) -> pd.DataFrame:
    """Evaluate full assembled PROTAC SMILES (e.g. from LinkerNet) directly.
    Validity gate (all must pass):
      1. RDKit-parseable SMILES
      2. Warhead + E3 substructure match found (can extract linker)
      3. Exactly 2 linker–fragment interface bonds
      4. Each interface bond attaches at the correct reactive site:
           warhead: piperazine NH  (_NH_SMARTS)
           E3:      tert-leu NH2   (_NH2_SMARTS)
      5. PROTACAnalyzer hard constraints on the extracted linker (all except
         attachment_sites, which is already verified by checks 3–4 above)
    SA and QED are computed on the extracted linker for a fair comparison with DeFoG.
    MW/logP/TPSA/RotBonds and DegradeMaster are computed on the full PROTAC.
    """
    raw, n_total_override = _read_smiles_file(smiles_path)
    n_valid_all = len(raw)
    n_total     = n_total_override if n_total_override is not None else n_valid_all
    n_std_valid = _std_valid_count(raw)

    dm_fn = _get_run_degrademaster()
    analyzer, wh_q, e3_q, wh_reactive_comp, e3_reactive_comp = _get_protac_analyzer()

    valid_rows, valid_smiles = [], []
    for smi in raw:
        mol = Chem.MolFromSmiles(smi)
        if mol is None:
            continue
        if len(Chem.GetMolFrags(mol)) > 1:
            continue  # disconnected — counts as invalid toward n_total

        # Extract linker fragment from assembled PROTAC
        linker_mol, n_iface, wh_match, e3_match = (
            _extract_linker_from_protac(mol, wh_q, e3_q)
            if wh_q is not None else (None, 0, None, None)
        )

        # Validity: PROTACAnalyzer structural checks + reactive-site attachment check.
        # attachment_sites is replaced by the more precise reactive-site check.
        if (analyzer is not None and linker_mol is not None and n_iface >= 2
                and wh_reactive_comp is not None):
            reactive_ok = _check_reactive_attachment(
                mol, wh_match, e3_match, wh_reactive_comp, e3_reactive_comp)
            checks = analyzer._run_checks(linker_mol)
            checks.pop("attachment_sites")
            is_assembled = int(reactive_ok and all(checks.values()))
        else:
            is_assembled = int(
                mol.HasSubstructMatch(_WARHEAD_PAT) and mol.HasSubstructMatch(_E3_PAT)
            )

        # SA and QED on extracted linker (matches DeFoG metric); fall back to full mol
        metric_mol = linker_mol if (linker_mol is not None and is_assembled) else mol
        try:
            sa   = round(_sa_score(metric_mol), 2)
            qed  = round(rdkit_qed.qed(metric_mol), 3)
            mw   = round(Descriptors.MolWt(mol), 1)
            logp = round(Descriptors.MolLogP(mol), 2)
            tpsa = round(Descriptors.TPSA(mol), 1)
            rotb = Descriptors.NumRotatableBonds(mol)
        except Exception:
            continue
        canon = Chem.MolToSmiles(mol)
        # DM-score every parseable+connected PROTAC, regardless of assembly_ok.
        # We want DM scores for non-canonically-assembled outputs too (e.g.,
        # DiffPROTACs places the linker at a non-reactive site → assembly_ok=False
        # but the SMILES is still a valid molecule the DM model can score).
        # is_assembled stays as a separate signal for the validity-gate metric.
        degrad = dm_fn(canon, _PROTAC_MODULE) if dm_fn else None
        valid_smiles.append(canon)
        # Descriptors (DM/SA/QED/MW/logP/TPSA/RotB) and the PK-window %s are gated
        # by is_assembled — we only report functional-PROTAC stats. Non-assembled
        # outputs still count toward Gen/StdVal/AsmVal but their descriptors are
        # NaN'd out so window-% denominators reflect only assembled molecules.
        nan = float("nan")
        valid_rows.append({
            "smiles":          canon,
            "assembly_ok":     is_assembled,
            "protac_smiles":   canon if is_assembled else None,
            "linker_sa":       sa   if is_assembled else nan,
            "linker_qed":      qed  if is_assembled else nan,
            "protac_mw":       mw   if is_assembled else nan,
            "protac_logp":     logp if is_assembled else nan,
            "protac_tpsa":     tpsa if is_assembled else nan,
            "protac_rotbonds": rotb if is_assembled else nan,
            "degrad_prob":     (degrad if degrad is not None else nan) if is_assembled else nan,
        })

    # n_valid = those passing ALL checks (reactive site + PROTACAnalyzer structural).
    # This matches DeFoG's n_valid semantics (PROTACAnalyzer valid during generation).
    n_valid = sum(1 for row in valid_rows if row["assembly_ok"])
    assembled_smiles = [row["smiles"] for row in valid_rows if row["assembly_ok"]]
    n_unique, uniqueness, novelty = _diversity_stats(assembled_smiles)

    if not valid_rows:
        df = pd.DataFrame(columns=[
            "smiles", "assembly_ok", "protac_smiles", "linker_sa", "linker_qed",
            "protac_mw", "protac_logp", "protac_tpsa", "protac_rotbonds", "degrad_prob",
        ])
    else:
        df = pd.DataFrame(valid_rows)
        df["_n_total"]    = n_total
        df["_n_valid"]    = n_valid
        df["_n_unique"]   = n_unique
        df["_uniqueness"] = uniqueness
        df["_novelty"]    = novelty
    df.attrs.update(_n_total=n_total, _n_valid=n_valid, _n_std_valid=n_std_valid,
                    _n_unique=n_unique, _uniqueness=uniqueness,
                    _novelty=novelty)
    return df


def batch_label(csv_path: Path) -> str:
    """Turn 'defog_local_edit_20260424_003754_summary.csv' → BATCH_LABELS entry or timestamp."""
    stem = csv_path.stem
    parts = stem.split("_")
    for i, p in enumerate(parts):
        if len(p) == 8 and p.isdigit():
            ts = "_".join(parts[i:i+2])
            return BATCH_LABELS.get(ts, ts)
    return BATCH_LABELS.get(stem, stem)


# ── Diversity metrics ────────────────────────────────────────────────────────
# (1) Intra-set: mean pairwise Tanimoto over generated outputs (lower = more diverse)
# (2) Training-set: max Tanimoto from each generated output to PROTAC-DB target+E3
#     training entries (lower = less memorization). 1 - this = "novelty distance."
_PROTACDB_BY_TARGET_CACHE: dict = {}
_TARGET_FILTER_BY_MODULE = {
    "ErlotinibVHL": ("EGFR", "VHL"),
    "BRD9VHL":      ("BRD9", "VHL"),
    "BRD4VHL":      ("BRD4", "VHL"),
    "KRASVHL":      ("KRAS", "VHL"),
    "ARCRBN":       ("AR-V7", "CRBN"),
    "ARCRBNv2":     ("AR-V7", "CRBN"),
    "BRD4CRBN":     ("BRD4", "CRBN"),
    "CDK6CRBN":     ("CDK6", "CRBN"),
}

def _load_protacdb_target_smiles(protac_module: str) -> list:
    """Cache and return PROTAC-DB SMILES filtered by the active target+E3 pair."""
    if protac_module in _PROTACDB_BY_TARGET_CACHE:
        return _PROTACDB_BY_TARGET_CACHE[protac_module]
    filt = _TARGET_FILTER_BY_MODULE.get(protac_module)
    if filt is None:
        _PROTACDB_BY_TARGET_CACHE[protac_module] = []
        return []
    target, e3 = filt
    name_json = Path(__file__).parent.parent / "src" / "degrademaster" / "data" / "PROTAC" / "name.json"
    if not name_json.exists():
        _PROTACDB_BY_TARGET_CACHE[protac_module] = []
        return []
    import json as _json
    with open(name_json) as f:
        name = _json.load(f)
    smis = [n["smiles"] for n in name.values()
            if target in n.get("tar_path", "") and e3 in n.get("e3_ligase_path", "")
            and "smiles" in n]
    _PROTACDB_BY_TARGET_CACHE[protac_module] = smis
    return smis


def _compute_diversity(df: pd.DataFrame, protac_module: str) -> dict:
    """Compute (intra-set diversity, training-set similarity) for the assembled
    PROTAC SMILES in df. Uses Morgan fingerprints (r=2, 2048 bits, Tanimoto).
    For DeFoG runs we use df['protac_smiles']; for full-PROTAC baselines
    df['smiles'] holds the assembled PROTAC directly."""
    from rdkit.Chem import AllChem
    from rdkit.DataStructs import TanimotoSimilarity

    # Prefer protac_smiles (DeFoG runs); fall back to smiles (full-PROTAC baselines
    # like LinkerNet/DiffPROTACs/DiffLinker, where protac_smiles may be all NaN
    # if the analyzer's validity gate rejected the non-canonical chemistry).
    smis = []
    if "protac_smiles" in df.columns:
        smis = [s for s in df["protac_smiles"].dropna() if s]
    if not smis and "smiles" in df.columns:
        smis = [s for s in df["smiles"].dropna() if s]
    fps = []
    for s in smis:
        m = Chem.MolFromSmiles(s)
        if m is None: continue
        fps.append(AllChem.GetMorganFingerprintAsBitVect(m, 2, 2048))

    out = {
        "intra_diversity_mean":   float("nan"),
        "intra_diversity_max":    float("nan"),
        "train_sim_mean":         float("nan"),
        "train_sim_pct_ge_0_8":   float("nan"),
        "n_diversity_scored":     len(fps),
    }
    if len(fps) < 2:
        return out

    # 1. Intra-set pairwise Tanimoto (mean + max)
    pairs = []
    for i in range(len(fps)):
        for j in range(i + 1, len(fps)):
            pairs.append(TanimotoSimilarity(fps[i], fps[j]))
    out["intra_diversity_mean"] = float(np.mean(pairs))
    out["intra_diversity_max"]  = float(max(pairs))

    # 2. Max Tanimoto to PROTAC-DB target+E3 training entries.
    # Cache training fingerprints across baseline calls within the same target run
    # (saves ~100ms × 4 baselines for the FP rebuild).
    global _PROTACDB_FP_CACHE
    if '_PROTACDB_FP_CACHE' not in globals():
        _PROTACDB_FP_CACHE = {}
    if protac_module in _PROTACDB_FP_CACHE:
        train_fps = _PROTACDB_FP_CACHE[protac_module]
    else:
        train_smis = _load_protacdb_target_smiles(protac_module)
        train_fps = []
        for s in train_smis:
            m = Chem.MolFromSmiles(s)
            if m is None: continue
            train_fps.append(AllChem.GetMorganFingerprintAsBitVect(m, 2, 2048))
        _PROTACDB_FP_CACHE[protac_module] = train_fps
    if train_fps:
        max_sims = [max(TanimotoSimilarity(q, t) for t in train_fps) for q in fps]
        out["train_sim_mean"]       = float(np.mean(max_sims))
        out["train_sim_pct_ge_0_8"] = 100.0 * sum(1 for s in max_sims if s >= 0.8) / len(max_sims)
    return out


def compute_stats(df: pd.DataFrame) -> dict:
    def _meta(col):
        if len(df) > 0 and col in df.columns:
            return df[col].iloc[0]
        return df.attrs.get(col, 0)

    n_total     = _meta("_n_total")
    n_valid     = _meta("_n_valid")
    n_std_valid = _meta("_n_std_valid")
    assembled   = df["assembly_ok"].sum() if "assembly_ok" in df.columns else 0

    # Scores ≥ 0.99 are almost certainly OOD extrapolation artefacts (N-O bonds,
    # exotic heterocycles absent from DegradMaster training data).  Exclude them
    # from the mean/std so the summary reflects calibrated predictions only.
    # (0.99 chosen because real PROTAC-DB actives top out around 0.95 across
    # tested target+E3 pairs; anything ≥0.99 is over the realistic ceiling.)
    _DEGRAD_OOD_THRESHOLD = 0.99
    _scored_raw = df["degrad_prob"].dropna() if "degrad_prob" in df.columns else pd.Series([], dtype=float)
    scored      = _scored_raw[_scored_raw < _DEGRAD_OOD_THRESHOLD]
    n_ood       = int((_scored_raw >= _DEGRAD_OOD_THRESHOLD).sum())
    logp_in = df["protac_logp"].dropna()     if "protac_logp"     in df.columns else pd.Series([], dtype=float)
    mw_in   = df["protac_mw"].dropna()       if "protac_mw"       in df.columns else pd.Series([], dtype=float)
    rot_in  = df["protac_rotbonds"].dropna() if "protac_rotbonds"  in df.columns else pd.Series([], dtype=float)

    # Windows match protac_classifier.py _PROPERTY_WINDOWS outer bounds (mean ± 2σ, PROTAC-DB).
    # Ranges = PROTAC DB mean ± 1σ (logP: 5.0±2.4, MW: 940±168, RotB: 18.5±6.4).
    _LOGP_LO, _LOGP_HI   = 2.5,  7.4
    _MW_LO,   _MW_HI     = 771,  1108
    _ROT_LO,  _ROT_HI    = 12,   25

    # All window/threshold percentages denominate over n_total so that
    # molecules that failed assembly or scoring count as failures, not absences.
    return {
        "std_valid_rate":     100 * n_std_valid / n_total if n_total else 0,
        "asm_valid_rate":     100 * n_valid / n_total if n_total else 0,
        "assembly_rate":      100 * assembled / n_valid if n_valid else 0,
        # SA score / QED: means over valid rows only (NaN rows auto-excluded)
        "sa_mean":            df["linker_sa"].dropna().mean()  if "linker_sa"  in df.columns else float("nan"),
        "sa_std":             df["linker_sa"].dropna().std()   if "linker_sa"  in df.columns else float("nan"),
        "qed_mean":           df["linker_qed"].dropna().mean() if "linker_qed" in df.columns else float("nan"),
        "qed_std":            df["linker_qed"].dropna().std()  if "linker_qed" in df.columns else float("nan"),
        # DegradeMaster: thresholds over count of scored (excludes OOD-flagged + NaN)
        "degrad_mean":        scored.mean() if len(scored) else float("nan"),
        "degrad_std":         scored.std()  if len(scored) else float("nan"),
        "degrad_ge03":        100 * (scored >= 0.3).sum() / len(scored) if len(scored) else 0,
        "degrad_ge05":        100 * (scored >= 0.5).sum() / len(scored) if len(scored) else 0,
        "degrad_ge07":        100 * (scored >= 0.7).sum() / len(scored) if len(scored) else 0,
        "degrad_n_ood":       n_ood,
        # PK windows: % over the count of non-NaN values (not n_valid which gates by
        # canonical-attachment validity — would give >100% for baselines whose metrics
        # are computed even when non-canonically attached).
        "logp_in_window_pct": 100 * logp_in.between(_LOGP_LO, _LOGP_HI).sum() / len(logp_in) if len(logp_in) else 0,
        "logp_mean":          logp_in.mean() if len(logp_in) else float("nan"),
        "mw_in_window_pct":   100 * mw_in.between(_MW_LO, _MW_HI).sum()        / len(mw_in)   if len(mw_in)   else 0,
        "mw_mean":            mw_in.mean() if len(mw_in) else float("nan"),
        "rot_in_window_pct":  100 * rot_in.between(_ROT_LO, _ROT_HI).sum()     / len(rot_in)  if len(rot_in)  else 0,
        "n_scored":           len(scored),
        "n_total":            n_total,
        "n_valid":            n_valid,
        # EGFR-PROPK (Zhang et al. 2026) post-hoc PK predictions; only
        # populated for ErlotinibVHL, NaN-filled otherwise.
        "pk_t_half_mean":     df["pk_t_half"].dropna().mean() if "pk_t_half" in df.columns else float("nan"),
        "pk_t_half_std":      df["pk_t_half"].dropna().std()  if "pk_t_half" in df.columns else float("nan"),
        "pk_cl_mean":         df["pk_cl"].dropna().mean()     if "pk_cl"     in df.columns else float("nan"),
        "pk_cl_std":          df["pk_cl"].dropna().std()      if "pk_cl"     in df.columns else float("nan"),
        "pk_vss_mean":        df["pk_vss"].dropna().mean()    if "pk_vss"    in df.columns else float("nan"),
        "pk_vss_std":         df["pk_vss"].dropna().std()     if "pk_vss"    in df.columns else float("nan"),
        "pk_score_mean":      df["pk_score"].dropna().mean()  if "pk_score"  in df.columns else float("nan"),
        "pk_n":               df["pk_t_half"].dropna().shape[0] if "pk_t_half" in df.columns else 0,
        # Diversity (all over n_valid, computed before dedup)
        "uniqueness":         _meta("_uniqueness"),
        "novelty":            _meta("_novelty"),
        # Structural diversity (Morgan-FP Tanimoto):
        #   intra_diversity_mean: mean pairwise Tanimoto in generated set
        #     (LOWER = more chemically diverse outputs)
        #   train_sim_mean:       mean (max sim to PROTAC-DB target+E3 training)
        #     (LOWER = less memorization of training data; 1 - x ≈ novelty distance)
        #   train_sim_pct_ge_0_8: % of outputs that are >0.80 Tanimoto-similar
        #     to a training entry — direct memorization indicator
        **_compute_diversity(df, _PROTAC_MODULE),
    }


# ── plotting ──────────────────────────────────────────────────────────────────

def make_chart(labels: list[str], stats: list[dict], out_path: str | None):
    n = len(labels)
    x = np.arange(n)
    colors = plt.cm.tab10(np.linspace(0, 0.7, n))

    fig, axes = plt.subplots(3, 3, figsize=(14, 12))
    fig.suptitle("PROTAC Linker Batch Report", fontsize=14, fontweight="bold")

    def bar(ax, values, errs=None, title="", ylabel="", note="", ideal_dir="", color_bars=True):
        bars = ax.bar(x, values, width=max(0.3, 0.6), color=colors if color_bars else "steelblue",
                      yerr=errs, capsize=4, error_kw={"linewidth": 1.2})
        ax.set_title(title, fontsize=10, fontweight="bold")
        ax.set_ylabel(ylabel, fontsize=8)
        ax.set_xticks(x)
        ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=7)
        ax.yaxis.set_tick_params(labelsize=8)
        if note:
            ax.text(0.02, 0.97, note, transform=ax.transAxes,
                    fontsize=7, va="top", color="#555555",
                    bbox=dict(boxstyle="round,pad=0.2", fc="white", alpha=0.7))
        if ideal_dir:
            ax.text(0.98, 0.97, ideal_dir, transform=ax.transAxes,
                    fontsize=8, va="top", ha="right", color="#1a6e1a" if "higher" in ideal_dir else "#8b0000",
                    fontweight="bold")
        for bar_obj, val in zip(bars, values):
            if not np.isnan(val):
                ax.text(bar_obj.get_x() + bar_obj.get_width() / 2,
                        bar_obj.get_height() + (max(values) * 0.01 if max(values) > 0 else 0.01),
                        f"{val:.2f}", ha="center", va="bottom", fontsize=7)

    # 1. Standard Valid vs Assembly Valid
    # Full-PROTAC baselines (LinkerNet/DiffPROTACs/DiffLinker) are already assembled;
    # the strict assembly check falls through to a lax substructure-presence test,
    # so reporting their rate is misleading. Suppress with NaN.
    ax = axes[0, 0]
    bw = 0.35
    _vrates = [
        ([s["std_valid_rate"] for s in stats], "#64B5F6", "Standard Valid\n(RDKit-parseable)"),
        ([float("nan") if s.get("is_full_protac_baseline") else s["asm_valid_rate"]
          for s in stats], "#1565C0", "Assembly Valid\n(all structural checks)"),
    ]
    for i, (vals, vc, vl) in enumerate(_vrates):
        _offs = (i - 0.5) * bw
        bars = ax.bar(x + _offs, vals, width=bw, color=vc, alpha=0.88, label=vl)
        for b, v in zip(bars, vals):
            if not np.isnan(v):
                ax.text(b.get_x() + b.get_width() / 2,
                        b.get_height() + (max(vals) * 0.01 if max(vals) > 0 else 0.5),
                        f"{v:.0f}%", ha="center", va="bottom", fontsize=6)
    ax.set_title("Validity Rate (%)", fontsize=10, fontweight="bold")
    ax.set_ylabel("% of generated samples", fontsize=8)
    ax.set_xticks(x); ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=7)
    ax.yaxis.set_tick_params(labelsize=8)
    ax.set_ylim(0, 120)
    ax.legend(fontsize=7)
    ax.text(0.98, 0.97, "↑ higher better", transform=ax.transAxes,
            fontsize=8, va="top", ha="right", color="#1a6e1a", fontweight="bold")

    # 2. SA score (lower = better)
    bar(axes[0, 1],
        [s["sa_mean"] for s in stats],
        errs=[s["sa_std"] for s in stats],
        title="Linker SA Score (mean ± std)",
        ylabel="SA score",
        note="1=easy, 10=hard to synthesise",
        ideal_dir="↓ lower better")
    axes[0, 1].set_ylim(0, max(s["sa_mean"] + s["sa_std"] for s in stats) * 1.2)
    axes[0, 1].axhline(6, color="red", linestyle="--", linewidth=0.8, label="hard cutoff (6)")
    axes[0, 1].legend(fontsize=7)

    # 3. QED (higher = better)
    bar(axes[0, 2],
        [s["qed_mean"] for s in stats],
        errs=[s["qed_std"] for s in stats],
        title="Linker QED (mean ± std)",
        ylabel="QED (0–1)",
        note="drug-likeness score",
        ideal_dir="↑ higher better")
    axes[0, 2].set_ylim(0, 1.0)

    # 4. DegradeMaster prob (higher = better)
    bar(axes[1, 0],
        [s["degrad_mean"] for s in stats],
        errs=[s["degrad_std"] for s in stats],
        title="DegradeMaster Prob (mean ± std)",
        ylabel="degradation probability (0–1)",
        note="GCN; EGFR L858R + VHL\nactive ref mean ≈ 0.37",
        ideal_dir="↑ higher better")
    axes[1, 0].set_ylim(0, 1.0)
    axes[1, 0].axhline(0.3, color="orange", linestyle="--", linewidth=0.8, label="≥0.3 active")
    axes[1, 0].axhline(0.5, color="red",    linestyle="--", linewidth=0.8, label="≥0.5 confident")
    axes[1, 0].legend(fontsize=7)

    # 5. % degrad_prob >= threshold (stacked-style grouped bar)
    ax = axes[1, 1]
    bw = 0.25
    offsets = [-bw, 0, bw]
    thresh_labels = ["≥0.3", "≥0.5", "≥0.7"]
    thresh_colors = ["#f0a500", "#e05c00", "#b00000"]
    thresh_keys   = ["degrad_ge03", "degrad_ge05", "degrad_ge07"]
    for i, (k, tc, tl) in enumerate(zip(thresh_keys, thresh_colors, thresh_labels)):
        vals = [s[k] for s in stats]
        bars = ax.bar(x + offsets[i], vals, width=bw, color=tc, alpha=0.85, label=tl)
        for b, v in zip(bars, vals):
            ax.text(b.get_x() + b.get_width()/2, b.get_height() + 0.5,
                    f"{v:.0f}%", ha="center", va="bottom", fontsize=6)
    ax.set_title("% Assembled PROTACs ≥ Degrad Threshold", fontsize=10, fontweight="bold")
    ax.set_ylabel("% of scored compounds")
    ax.set_xticks(x); ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=7)
    ax.set_ylim(0, 115)
    ax.legend(fontsize=7, title="threshold", title_fontsize=7)
    ax.text(0.98, 0.97, "↑ higher better", transform=ax.transAxes,
            fontsize=8, va="top", ha="right", color="#1a6e1a", fontweight="bold")

    # 6. PK window pass rates
    ax = axes[1, 2]
    bw = 0.25
    offsets = [-bw, 0, bw]
    pk_labels  = ["logP [2.5,7.4]", "MW [771,1108]", "RotBonds [12,25]"]
    pk_keys    = ["logp_in_window_pct", "mw_in_window_pct", "rot_in_window_pct"]
    pk_colors  = ["#2196F3", "#4CAF50", "#9C27B0"]
    for i, (k, pc, pl) in enumerate(zip(pk_keys, pk_colors, pk_labels)):
        vals = [s[k] for s in stats]
        bars = ax.bar(x + offsets[i], vals, width=bw, color=pc, alpha=0.85, label=pl)
        for b, v in zip(bars, vals):
            ax.text(b.get_x() + b.get_width()/2, b.get_height() + 0.5,
                    f"{v:.0f}%", ha="center", va="bottom", fontsize=6)
    ax.set_title("PROTAC PK Window Pass Rates", fontsize=10, fontweight="bold")
    ax.set_ylabel("% of assembled PROTACs in window")
    ax.set_xticks(x); ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=7)
    ax.set_ylim(0, 115)
    ax.legend(fontsize=7)
    ax.text(0.98, 0.97, "↑ higher better", transform=ax.transAxes,
            fontsize=8, va="top", ha="right", color="#1a6e1a", fontweight="bold")

    # 7. Uniqueness / Novelty
    ax = axes[2, 0]
    bw = 0.3
    offsets = [-bw/2, bw/2]
    div_labels  = ["Uniqueness", "Novelty"]
    div_keys    = ["uniqueness", "novelty"]
    div_colors  = ["#00897B", "#1E88E5"]
    div_notes   = ["↑ no duplicates", "↑ not in train set"]
    for i, (k, dc, dl) in enumerate(zip(div_keys, div_colors, div_labels)):
        vals = [s[k] for s in stats]
        bars = ax.bar(x + offsets[i], vals, width=bw, color=dc, alpha=0.85, label=f"{dl} ({div_notes[i]})")
        for b, v in zip(bars, vals):
            ax.text(b.get_x() + b.get_width()/2, b.get_height() + 0.5,
                    f"{v:.0f}%", ha="center", va="bottom", fontsize=6)
    ax.set_title("Diversity Metrics (% of unique valid)", fontsize=10, fontweight="bold")
    ax.set_ylabel("% of unique valid linkers")
    ax.set_xticks(x); ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=7)
    ax.set_ylim(0, 115)
    ax.legend(fontsize=7)

    axes[2, 1].set_visible(False)
    axes[2, 2].set_visible(False)

    plt.tight_layout()

    if out_path:
        plt.savefig(out_path, dpi=150, bbox_inches="tight")
        print(f"Chart saved to {out_path}")
    else:
        plt.show()


# ── main ──────────────────────────────────────────────────────────────────────

def _print_stats(s: dict):
    print(f"  Generated / Std Valid / Asm Valid : {s['n_total']} / {round(s['std_valid_rate']*s['n_total']/100)} / {s['n_valid']}")
    print(f"  Standard valid rate  (↑)          : {s['std_valid_rate']:.1f}%  (RDKit-parseable SMILES)")
    print(f"  Assembly valid rate  (↑)          : {s['asm_valid_rate']:.1f}%  (canonical attachment + analyzer pass)")
    print(f"  SA score (↓ lower better)          : {s['sa_mean']:.2f} ± {s['sa_std']:.2f}")
    print(f"  QED      (↑ higher better)         : {s['qed_mean']:.3f} ± {s['qed_std']:.3f}")
    ood_note = f"  ({s['degrad_n_ood']} OOD-flagged ≥0.99 excluded)" if s.get('degrad_n_ood') else ""
    print(f"  DegradMaster prob (↑ higher)       : {s['degrad_mean']:.3f} ± {s['degrad_std']:.3f}{ood_note}")
    print(f"    ≥ 0.3: {s['degrad_ge03']:.1f}%   ≥ 0.5: {s['degrad_ge05']:.1f}%   ≥ 0.7: {s['degrad_ge07']:.1f}%")
    print(f"  PROTAC logP in [2.5,7.4]  (↑)     : {s['logp_in_window_pct']:.1f}%  (mean {s['logp_mean']:.2f})")
    print(f"  PROTAC MW   in [771,1108] (↑)     : {s['mw_in_window_pct']:.1f}%  (mean {s['mw_mean']:.0f} Da)")
    print(f"  PROTAC RotB in [12,25]    (↑)     : {s['rot_in_window_pct']:.1f}%")
    # EGFR-PROPK predictions (only populated for EGFR-VHL target)
    if s.get('pk_n', 0) > 0:
        print(f"  EGFR-PROPK T½ (↑)                 : {s['pk_t_half_mean']:.2f} ± {s['pk_t_half_std']:.2f}  (n={s['pk_n']}, raw)")
        print(f"  EGFR-PROPK CL (↓)                 : {s['pk_cl_mean']:.2f} ± {s['pk_cl_std']:.2f}  (raw)")
        print(f"  EGFR-PROPK Vss                    : {s['pk_vss_mean']:.2f} ± {s['pk_vss_std']:.2f}  (raw)")
        print(f"  EGFR-PROPK score (σ(T½)−σ(CL)) (↑): {s['pk_score_mean']:.3f}")
    print(f"  Uniqueness   (↑ higher better)     : {s['uniqueness']:.1f}%")
    print(f"  Novelty      (↑ higher better)     : {s['novelty']:.1f}%  (not in training set)")
    # Structural diversity (Morgan FP / Tanimoto)
    nd = s.get('n_diversity_scored', 0)
    if nd >= 2:
        intra = s['intra_diversity_mean']
        train_sim = s['train_sim_mean']
        memo_pct = s['train_sim_pct_ge_0_8']
        print(f"  Intra-set diversity (↓ lower=more diverse): mean Tanimoto = {intra:.3f}  (n={nd} pairs, max={s['intra_diversity_max']:.3f})")
        if not np.isnan(train_sim):
            print(f"  Memorization (↓ lower=less memorized) : mean max-Tanimoto to PROTAC-DB target+E3 = {train_sim:.3f}")
            print(f"    {memo_pct:.1f}% of outputs are >0.80 Tanimoto-similar to a training entry (memorized)")


def main():
    parser = argparse.ArgumentParser(description="Batch report for generated PROTAC linkers")
    parser.add_argument("csvs", nargs="*", help="Generation summary CSV file(s) (*_summary.csv)")
    parser.add_argument("--out", default=None, help="Save chart to this path (PNG/PDF). Default: display.")
    parser.add_argument(
        "--baseline", action="append", default=[], metavar="LABEL:PATH",
        help="Named baseline from a plain linker SMILES .txt file. "
             "Format: 'DeLinker:/path/to/smiles.txt'. Repeatable.",
    )
    parser.add_argument(
        "--protac-baseline", action="append", default=[], metavar="LABEL:PATH",
        dest="protac_baseline",
        help="Named baseline from a full assembled PROTAC SMILES .txt file "
             "(e.g. DRlinker output). Evaluated directly without assembly. "
             "Format: 'DRlinker:/path/to/smiles.txt'. Repeatable.",
    )
    parser.add_argument(
        "--protac-module", default=None, dest="protac_module",
        help="Builder module name under MoCA/smiles/protac/ "
             "(e.g. ErlotinibVHL, ARCRBN). Overrides _PROTAC_MODULE constant.",
    )
    args = parser.parse_args()

    # Apply protac module selection before any lazy-loads
    if args.protac_module:
        global _PROTAC_MODULE
        _PROTAC_MODULE = args.protac_module

    # Resolve which DeFoG batch CSVs to process (CLI takes priority over in-file constants)
    if args.csvs:
        csv_paths = [Path(p) for p in args.csvs]
    elif BATCH_FOLDER:
        csv_paths = sorted(Path(BATCH_FOLDER).glob("*_summary.csv"))
    elif BATCH_CSVS:
        csv_paths = [Path(p) for p in BATCH_CSVS]
    else:
        csv_paths = []

    # Merge BASELINES constant with --baseline CLI args (CLI appends on top)
    baselines: dict[str, str] = dict(BASELINES)
    for spec in args.baseline:
        if ":" not in spec:
            parser.error(f"--baseline must be LABEL:PATH, got: {spec!r}")
        lbl, path = spec.split(":", 1)
        baselines[lbl] = path

    # Pick per-target baseline set based on the active protac-module.
    # Falls back to the legacy AR-CRBN set if the module name isn't registered.
    _module_baselines = PROTAC_BASELINES_BY_MODULE.get(_PROTAC_MODULE, PROTAC_BASELINES)
    protac_baselines: dict[str, str] = dict(_module_baselines)
    for spec in args.protac_baseline:
        if ":" not in spec:
            parser.error(f"--protac-baseline must be LABEL:PATH, got: {spec!r}")
        lbl, path = spec.split(":", 1)
        protac_baselines[lbl] = path

    if not csv_paths and not baselines and not protac_baselines:
        parser.error(
            "Provide CSV paths on the command line, set BATCH_FOLDER / BATCH_CSVS, "
            "or add entries to BASELINES / PROTAC_BASELINES at the top of the file."
        )

    out_path = args.out or CHART_OUT

    labels, all_stats = [], []

    for csv_path in sorted(csv_paths):
        csv_path = Path(csv_path)
        label = batch_label(csv_path)
        print(f"\n{'='*60}")
        print(f"Batch: {label}  ({csv_path.name})")
        print(f"{'='*60}")

        df = run_evaluate(csv_path)
        if _PROTAC_MODULE == "ErlotinibVHL":
            df = _compute_pk_scores(df)
        s = compute_stats(df)
        labels.append(label)
        all_stats.append(s)
        _print_stats(s)

    for label, smiles_path in baselines.items():
        p = Path(smiles_path)
        print(f"\n{'='*60}")
        print(f"Baseline: {label}  ({p.name})")
        print(f"{'='*60}")

        df = run_evaluate_smiles_file(p)
        if _PROTAC_MODULE == "ErlotinibVHL":
            df = _compute_pk_scores(df)
        s = compute_stats(df)
        labels.append(label)
        all_stats.append(s)
        _print_stats(s)

    for label, smiles_path in protac_baselines.items():
        p = Path(smiles_path)
        print(f"\n{'='*60}")
        print(f"PROTAC Baseline: {label}  ({p.name})")
        print(f"{'='*60}")

        df = run_evaluate_full_protac_file(p)
        if _PROTAC_MODULE == "ErlotinibVHL":
            df = _compute_pk_scores(df)
        s = compute_stats(df)
        # Full-PROTAC baselines are already assembled; the strict assembly check
        # silently falls back to a lax substructure-presence test, so the rate is
        # not directly comparable to MoCA's. Suppress it from the report.
        s["is_full_protac_baseline"] = True
        labels.append(label)
        all_stats.append(s)
        _print_stats(s)

    make_chart(labels, all_stats, out_path)


if __name__ == "__main__":
    main()