#!/usr/bin/env python3
"""
evaluate_linkers.py

Evaluate generated PROTAC linkers. The warhead/E3 pair is selected by
--protac-module (defaults to ErlotinibVHL = EGFR-VHL). Supported modules:
    ErlotinibVHL  — EGFR (gefitinib-piperazine) + VHL (VH032)
    BRD9VHL       — BRD9 (BI-7273-piperazine)   + VHL (VH032)
    KRASVHL       — KRAS (ACBI3-diazepane)      + VHL (VH032)
    BRD4VHL, ARCRBN — other registered modules (no DegradeMaster scoring)

For each generated SMILES:
  1. Skip chemically invalid molecules (all further steps on valid only)
  2. Run PROTACAnalyzer hard constraints (attachment sites, atom count,
     ring closure, SA score, assembled PROTAC MW)
  3. Attempt PROTAC assembly (warhead + linker + E3)
  4. Compute linker properties: SA score, QED, MW, logP, TPSA, HBD/HBA, rotbonds
  5. Compute assembled PROTAC PK properties
  6. Run DegradeMaster GCN model (ABILiLab/DegradeMaster) on the assembled
     PROTAC to predict degradation probability. Target/E3 pocket graphs and
     30-d feature vectors are dispatched per --protac-module:
       ErlotinibVHL → EGFR L858R pocket (2ITY), egfr_target_feature.npy
       BRD9VHL      → BRD9 pocket (case_study_1), brd9_target_feature.npy
       KRASVHL      → KRAS WT pocket (P01116 + ACBI3), krasvhl_target_feature.npy
     E3 pocket = VHL (P40337) for all three; E3 feature = vhl_e3_feature_egfr.npy
     (verified bit-identical across case_study_1 and case_study_2 entries).
  7. Print per-molecule table + aggregate summary

Usage:
    python evaluate_linkers.py path/to/final_smiles.txt
    python evaluate_linkers.py path/to/final_smiles.txt --out results.csv
    python evaluate_linkers.py path/to/final_smiles.txt --sa-max 5.0
    python evaluate_linkers.py path/to/final_smiles.txt --protac-module BRD9VHL
"""

import argparse
import importlib.util
import os
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from rdkit import Chem, RDConfig, RDLogger
from rdkit.Chem import Descriptors, MACCSkeys, QED, rdMolDescriptors

RDLogger.DisableLog("rdApp.*")

# Resolve src/ paths.
_REPO_ROOT = Path(__file__).parent.parent
_SRC = _REPO_ROOT / "src"
sys.path.insert(0, str(_SRC / "classifier"))
sys.path.insert(0, str(_SRC / "protac"))
_DM_ROOT = _SRC / "degrademaster"

# SA scorer from RDKit contrib
_sa_path = os.path.join(RDConfig.RDContribDir, "SA_Score", "sascorer.py")
_spec = importlib.util.spec_from_file_location("sascorer", _sa_path)
sascorer = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(sascorer)


# ---------------------------------------------------------------------------
# DegradeMaster integration
# ---------------------------------------------------------------------------

# Periodic table in order (118 elements) — must match DegradeMaster ATOMS list
_DM_ATOMS = [
    'H',  'He',
    'Li', 'Be', 'B',  'C',  'N',  'O',  'F',  'Ne',
    'Na', 'Mg', 'Al', 'Si', 'P',  'S',  'Cl', 'Ar',
    'K',  'Ca', 'Sc', 'Ti', 'V',  'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu', 'Zn',
    'Ga', 'Ge', 'As', 'Se', 'Br', 'Kr',
    'Rb', 'Sr', 'Y',  'Zr', 'Nb', 'Mo', 'Tc', 'Ru', 'Rh', 'Pd', 'Ag', 'Cd',
    'In', 'Sn', 'Sb', 'Te', 'I',  'Xe',
    'Cs', 'Ba', 'La', 'Ce', 'Pr', 'Nd', 'Pm', 'Sm', 'Eu', 'Gd', 'Tb', 'Dy',
    'Ho', 'Er', 'Tm', 'Yb', 'Lu',
    'Hf', 'Ta', 'W',  'Re', 'Os', 'Ir', 'Pt', 'Au', 'Hg', 'Tl', 'Pb', 'Bi',
    'Po', 'At', 'Rn',
    'Fr', 'Ra', 'Ac', 'Th', 'Pa', 'U',  'Np', 'Pu', 'Am', 'Cm', 'Bk',
    'Cf', 'Es', 'Fm', 'Md', 'No', 'Lr',
    'Rf', 'Db', 'Sg', 'Bh', 'Hs', 'Mt', 'Ds', 'Rg', 'Cn', 'Nh', 'Fl', 'Mc',
    'Lv', 'Ts', 'Og',
]

_DM_BOND_MAP = {
    Chem.BondType.SINGLE:   1,
    Chem.BondType.DOUBLE:   2,
    Chem.BondType.TRIPLE:   3,
    Chem.BondType.AROMATIC: 4,
}

# SMARTS for amide bonds — matched before H addition, converted to type 5
# to match mol2 file convention used during DegradeMaster training.
_AMIDE_PAT = Chem.MolFromSmarts('[NX3;H0,H1,H2][CX3](=[OX1])')

_dm_cache: dict = {}   # lazy-loaded singleton keyed by str(_DM_ROOT)


def _dm_smiles_to_graph(smi: str):
    """Convert PROTAC SMILES to a PyG Data graph matching DegradeMaster's
    training format: explicit H atoms, amide bonds as type 5."""
    from torch_geometric.data import Data
    mol = Chem.MolFromSmiles(smi)
    if mol is None:
        return None
    # Identify amide bonds (N-C=O) before H addition (atom indices stable here)
    amide_bonds = {
        (min(m[0], m[1]), max(m[0], m[1]))
        for m in mol.GetSubstructMatches(_AMIDE_PAT)
    }
    mol = Chem.AddHs(mol)
    atoms = [
        _DM_ATOMS.index(a.GetSymbol()) if a.GetSymbol() in _DM_ATOMS else len(_DM_ATOMS)
        for a in mol.GetAtoms()
    ]
    e1, e2, ea = [], [], []
    for bond in mol.GetBonds():
        i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        bt = _DM_BOND_MAP.get(bond.GetBondType(), 1)
        if (min(i, j), max(i, j)) in amide_bonds:
            bt = 5  # amide bond type as in mol2 training data
        e1.append(i); e2.append(j); ea.append(bt)
    return Data(
        x=torch.tensor(atoms),
        edge_index=torch.tensor([e1 + e2, e2 + e1]),
        edge_attr=torch.tensor(ea + ea),
    )


def _dm_extract_graph(data, slices, idx: int):
    """Extract single graph entry from an InMemoryDataset collated tensor."""
    from torch_geometric.data import Data
    xs, xe   = slices['x'][idx].item(),          slices['x'][idx + 1].item()
    es, ee   = slices['edge_index'][idx].item(),  slices['edge_index'][idx + 1].item()
    eas, eae = slices['edge_attr'][idx].item(),   slices['edge_attr'][idx + 1].item()
    return Data(
        x=data.x[xs:xe],
        edge_index=data.edge_index[:, es:ee],
        edge_attr=data.edge_attr[eas:eae],
    )


# Per-target DegradeMaster wiring. Each entry maps a protac-module name to
# the target pocket graph + 30-d feature vector. E3 pocket and E3 feature are
# shared (VHL VH032 across all three modules).
#   pocket_graph_file — torch.save'd Data object, atoms in periodic-table index
#   target_feature_file — np.save'd (30,) float64 vector
#   notes — human-readable description for the run summary
_DM_TARGETS = {
    "ErlotinibVHL": {
        "pocket_graph_file":   "egfr_pocket_graph.pt",
        "target_feature_file": "egfr_target_feature.npy",
        "e3_pocket_source":    ("processed_slot", "ligase_pocket.pt", 0),  # VHL slot 0
        "e3_feature_file":     "vhl_e3_feature_egfr.npy",
        "notes": "EGFR L858R pocket (PDB 2ITY); E3=VHL; training-set features",
    },
    "BRD9VHL": {
        "pocket_graph_file":   "brd9_pocket_graph.pt",
        "target_feature_file": "brd9_target_feature.npy",
        "e3_pocket_source":    ("processed_slot", "ligase_pocket.pt", 0),  # VHL
        "e3_feature_file":     "vhl_e3_feature_egfr.npy",
        "notes": "BRD9 pocket (case_study_1, BRD9_Q9H8M2); E3=VHL; "
                 "target_feature_case1[0]",
    },
    "ARCRBNv2": {
        "pocket_graph_file":   "arcrbn_pocket_graph.pt",      # AR LBD (1E3G + R18)
        "target_feature_file": "arcrbn_target_feature.npy",   # PROTAC-DB AR-V7+CRBN row
        "e3_pocket_source":    ("file", "crbn_pocket_graph.pt"),  # CRBN (4TZ4 + LVY)
        "e3_feature_file":     "crbn_e3_feature.npy",         # PROTAC-DB CRBN feature
        "notes": "AR LBD pocket (1E3G + R1881 ligand); E3=CRBN (4TZ4 + lenalidomide); "
                 "AR/CRBN features from PROTAC-DB AR-V7+CRBN row",
    },
    "BRD4CRBN": {
        "pocket_graph_file":   "brd4_pocket_graph.pt",        # BRD4 (BRD4_O60885 + smina)
        "target_feature_file": "brd4crbn_target_feature.npy", # PROTAC-DB BRD4+CRBN row
        "e3_pocket_source":    ("file", "crbn_pocket_graph.pt"),  # CRBN (4TZ4 + LVY)
        "e3_feature_file":     "crbn_e3_feature.npy",         # PROTAC-DB CRBN feature
        "notes": "BRD4 pocket (BRD4_O60885 + smina-docked warhead); E3=CRBN (4TZ4 + "
                 "lenalidomide); BRD4/CRBN features from PROTAC-DB BRD4+CRBN row",
    },
    "CDK6CRBN": {
        "pocket_graph_file":   "cdk6crbn_pocket_graph.pt",    # CDK6 (5L2I + LQQ palbociclib)
        "target_feature_file": "cdk6crbn_target_feature.npy", # PROTAC-DB CDK6+CRBN row
        "e3_pocket_source":    ("file", "crbn_pocket_graph.pt"),  # CRBN (4TZ4 + LVY)
        "e3_feature_file":     "crbn_e3_feature.npy",         # PROTAC-DB CRBN feature
        "notes": "CDK6 pocket (5L2I + LQQ palbociclib-class ligand); E3=CRBN (4TZ4 + "
                 "lenalidomide); CDK6/CRBN features from PROTAC-DB CDK6+CRBN row. "
                 "Uses cooh (diacid) linker chemistry — same as BRD9-VHL.",
    },
    # BRD4VHL — pocket+features are wired (brd4_pocket_graph.pt,
    # brd4_target_feature.npy) but the published case-study checkpoint
    # saturates at ≥0.99 on 44% of generated linkers (vs 4% on real
    # PROTAC-DB BRD4+VHL actives), so DM scores on generated PROTACs
    # are OOD-overconfident and not trustworthy. Re-enable here if a
    # checkpoint trained on broader data becomes available.
    #
    # KRASVHL — same: krasvhl_pocket_graph.pt and krasvhl_target_feature.npy
    # exist, but the case-study checkpoint collapses to exactly 0.000 on
    # all KRAS features (verified across all 18 case_study_2 entries).
    #
    # ARCRBN (v1) — original, OOD-broken. Use ARCRBNv2 instead.
}


def _load_degrademaster(protac_module: str = "ErlotinibVHL"):
    """Load DegradeMaster model and per-target pocket data (lazy, cached).

    Returns (model, e3_graph, target_graph, e3_feat, tar_feat, notes)
    or None if the target is unsupported / data files are missing.
    """
    key = (str(_DM_ROOT), protac_module)
    if key in _dm_cache:
        return _dm_cache[key]

    if not _DM_ROOT.exists():
        return None

    target_cfg = _DM_TARGETS.get(protac_module)
    if target_cfg is None:
        # Unsupported module — emit once, cache the negative result.
        print(f"[DegradeMaster] no DegradeMaster wiring for protac module "
              f"'{protac_module}' — skipping model", file=sys.stderr)
        _dm_cache[key] = None
        return None

    # Import DegradeMaster model module
    dm_model_path = _DM_ROOT / "model.py"
    try:
        _spec2 = importlib.util.spec_from_file_location("dm_model", dm_model_path)
        dm_model_mod = importlib.util.module_from_spec(_spec2)
        _spec2.loader.exec_module(dm_model_mod)
    except Exception as exc:
        print(f"[DegradeMaster] failed to load model.py: {exc} — skipping model", file=sys.stderr)
        _dm_cache[key] = None
        return None
    GraphConv   = dm_model_mod.GraphConv
    ProtacModel = dm_model_mod.ProtacModel

    # Pocket graphs and features:
    #   Target pocket     — pre-built per-target .pt under data/case_study/
    #   Target feature    — per-target .npy (30-d), built from DegradeMaster's
    #                       case_study features (one row per target).
    #   E3 pocket         — dispatched per target_cfg["e3_pocket_source"]:
    #                       ("processed_slot", file.pt, idx) → slot from a
    #                                          processed/ collated InMemoryDataset
    #                       ("file", filename.pt) → standalone .pt file
    #   E3 feature vec    — per-target file (vhl_e3_feature_egfr.npy is the
    #                       canonical VHL feature; CRBN targets use a different
    #                       .npy extracted from PROTAC-DB).
    cs_proc = _DM_ROOT / "data" / "case_study" / "processed"
    cs_data = _DM_ROOT / "data" / "case_study"
    try:
        # E3 pocket dispatch
        e3_src = target_cfg["e3_pocket_source"]
        if e3_src[0] == "processed_slot":
            _, fname, idx = e3_src
            lp_data, lp_slices = torch.load(cs_proc / fname, weights_only=False)
            e3_graph = _dm_extract_graph(lp_data, lp_slices, idx)
        elif e3_src[0] == "file":
            _, fname = e3_src
            e3_graph = torch.load(cs_data / fname, weights_only=False)
        else:
            raise ValueError(f"unknown e3_pocket_source kind: {e3_src[0]}")

        target_graph = torch.load(cs_data / target_cfg["pocket_graph_file"],
                                   weights_only=False)
        tar_feat = np.load(cs_data / target_cfg["target_feature_file"], allow_pickle=True)
        e3_feat  = np.load(cs_data / target_cfg["e3_feature_file"],     allow_pickle=True)
    except FileNotFoundError as exc:
        print(f"[DegradeMaster] missing data file for {protac_module}: {exc} "
              f"— skipping model", file=sys.stderr)
        _dm_cache[key] = None
        return None

    # Build GCN model (checkpoint/1000 uses GCN with hidden_size=128)
    prot_gcn = GraphConv(num_embeddings=118, graph_dim=167, hidden_size=128)
    lig_gcn  = GraphConv(num_embeddings=118, graph_dim=30,  hidden_size=128)
    tar_gcn  = GraphConv(num_embeddings=118, graph_dim=30,  hidden_size=128)
    model    = ProtacModel(prot_gcn, lig_gcn, tar_gcn, 128)

    ckpt_path = _DM_ROOT / "checkpoint" / "1000" / "model_for_case_study.pth"
    ckpt = torch.load(ckpt_path, map_location="cpu")
    model.load_state_dict(ckpt["OurModel"])
    model.eval()

    result = (model, e3_graph, target_graph, e3_feat, tar_feat, target_cfg["notes"])
    _dm_cache[key] = result
    return result


def _run_degrademaster(protac_smi: str,
                       protac_module: str = "ErlotinibVHL") -> "float | None":
    """Return DegradeMaster degradation probability (0–1) for an assembled
    PROTAC SMILES, or None if the model is unavailable or SMILES is invalid."""
    dm = _load_degrademaster(protac_module)
    if dm is None:
        return None

    from torch_geometric.data import Batch

    model, e3_graph, target_graph, e3_feat, tar_feat, _ = dm

    pg = _dm_smiles_to_graph(protac_smi)
    if pg is None:
        return None

    maccs = np.array(MACCSkeys.GenMACCSKeys(Chem.MolFromSmiles(protac_smi)), dtype=np.float32)
    feat  = torch.tensor(
        np.concatenate([maccs, tar_feat, e3_feat]).astype(np.float32)
    ).unsqueeze(0)   # (1, 227)

    with torch.no_grad():
        logits, _ = model(
            Batch.from_data_list([pg]),
            Batch.from_data_list([e3_graph]),
            Batch.from_data_list([target_graph]),
            feat,
        )
        prob = torch.softmax(logits, dim=-1)
    return float(prob[0, 1].item())


# ---------------------------------------------------------------------------
# Property helpers
# ---------------------------------------------------------------------------

def linker_props(mol: Chem.Mol) -> dict:
    return {
        "linker_mw":       round(Descriptors.MolWt(mol), 1),
        "linker_logp":     round(Descriptors.MolLogP(mol), 2),
        "linker_tpsa":     round(Descriptors.TPSA(mol), 1),
        "linker_hbd":      Descriptors.NumHDonors(mol),
        "linker_hba":      Descriptors.NumHAcceptors(mol),
        "linker_rotbonds": Descriptors.NumRotatableBonds(mol),
        "linker_rings":    rdMolDescriptors.CalcNumRings(mol),
        "linker_hac":      mol.GetNumHeavyAtoms(),
        "linker_sa":       round(sascorer.calculateScore(mol), 2),
        "linker_qed":      round(QED.qed(mol), 3),
    }


def protac_props(mol: Chem.Mol) -> dict:
    mw   = Descriptors.MolWt(mol)
    logp = Descriptors.MolLogP(mol)
    hbd  = Descriptors.NumHDonors(mol)
    hba  = Descriptors.NumHAcceptors(mol)
    return {
        "protac_mw":       round(mw, 1),
        "protac_logp":     round(logp, 2),
        "protac_tpsa":     round(Descriptors.TPSA(mol), 1),
        "protac_hbd":      hbd,
        "protac_hba":      hba,
        "protac_rotbonds": Descriptors.NumRotatableBonds(mol),
        "protac_hac":      mol.GetNumHeavyAtoms(),
        # Lipinski Ro5 violations (PROTACs routinely violate, but track anyway)
        "ro5_violations":  int(mw > 500) + int(logp > 5) + int(hbd > 5) + int(hba > 10),
    }


# ---------------------------------------------------------------------------
# Main evaluation
# ---------------------------------------------------------------------------

def evaluate(smiles_list: list[str], sa_max: float = 6.0,
             verbose: bool = True,
             protac_module: str = "ErlotinibVHL",
             fix_termini: "str | None" = None) -> pd.DataFrame:
    from protac_analyzer import PROTACAnalyzer
    import importlib as _il
    _pmod = _il.import_module(protac_module)
    Builder     = _pmod.Builder
    warhead_mol = _pmod.WARHEAD_MOL
    e3_mol      = _pmod.E3_MOL

    # Linker-side attachment SMARTS — must match what the linker exposes after
    # generation. Mirrors the dispatch in defog_guided.py.
    #   cooh       : both ends are COOH (or post-assembly CHO) — diacid linker
    #   amine      : both ends are secondary/primary NH — diamine linker
    #   amino_acid : one end NH2, one end COOH — asymmetric linker
    #   None / default : cooh (backwards-compat)
    if fix_termini == 'amine':
        _WH_ATTACH = "[NX3;!H0;!$(NC=O);!$(NS=O);!$(N=O)]"
        _E3_ATTACH = "[NX3;!H0;!$(NC=O);!$(NS=O);!$(N=O)]"
    elif fix_termini == 'amino_acid':
        _WH_ATTACH = "[$([NX3H2]),$([CX3](=[OX1])[OX2H1])]"
        _E3_ATTACH = "[$([NX3H2]),$([CX3](=[OX1])[OX2H1])]"
    else:  # 'cooh' or None
        _COOH_OR_CHO = "[$([CX3](=[OX1])[OX2H1]),$([CX3H1](=[OX1]))]"
        _WH_ATTACH = _COOH_OR_CHO
        _E3_ATTACH = _COOH_OR_CHO
    analyzer = PROTACAnalyzer(
        warhead=warhead_mol,
        e3=e3_mol,
        warhead_attachment_smarts=_WH_ATTACH,
        e3_attachment_smarts=_E3_ATTACH,
        sa_score_max=sa_max,
    )
    builder = Builder()

    # Probe DegradeMaster availability once before the loop
    dm_loaded     = _load_degrademaster(protac_module)
    dm_available  = dm_loaded is not None
    dm_notes      = dm_loaded[5] if dm_available else ""

    rows = []
    n_invalid = 0

    for raw in smiles_list:
        smi = raw.split()[-1] if raw.startswith("Valid") else raw.strip()
        if not smi:
            continue

        mol = Chem.MolFromSmiles(smi)

        # Step 1 — chemical validity
        if mol is None:
            n_invalid += 1
            continue

        canon = Chem.MolToSmiles(mol)
        row = {"smiles": canon}

        # Step 2 — linker properties
        row.update(linker_props(mol))

        # Step 3 — hard constraints
        checks = analyzer._run_checks(mol)
        row["linker_valid"] = all(checks.values())
        row.update({f"check_{k}": v for k, v in checks.items()})

        # Step 4 — PROTAC assembly (attempt regardless of linker_valid,
        #           so we can still score partially valid linkers)
        protac_smi  = None
        protac_mol  = None
        assembly_ok = False
        try:
            protac_smi = builder.assemble(mol, return_all=False, amide_only=False)
            protac_mol = Chem.MolFromSmiles(protac_smi)
            assembly_ok = protac_mol is not None
        except Exception:
            pass

        row["assembly_ok"]   = assembly_ok
        row["protac_smiles"] = protac_smi or ""

        # Step 5 — assembled PROTAC PK properties
        if assembly_ok:
            row.update(protac_props(protac_mol))
            passed, val_checks = builder.validate_protac(protac_mol)
            row["protac_valid"] = passed
        else:
            for k in ["protac_mw", "protac_logp", "protac_tpsa", "protac_hbd",
                      "protac_hba", "protac_rotbonds", "protac_hac", "ro5_violations"]:
                row[k] = None
            row["protac_valid"] = False

        # Step 6 — DegradeMaster degradation probability
        if assembly_ok and dm_available:
            degrad_prob = _run_degrademaster(protac_smi, protac_module)
            row["degrad_prob"]  = round(degrad_prob, 4) if degrad_prob is not None else None
            row["degrad_notes"] = dm_notes if degrad_prob is not None else "DegradeMaster inference failed"
        else:
            row["degrad_prob"]  = None
            row["degrad_notes"] = "assembly failed" if not assembly_ok else "model unavailable"

        rows.append(row)

    df = pd.DataFrame(rows)

    if verbose:
        _print_summary(df, n_invalid, len(smiles_list), dm_available, dm_notes)

    return df


def _print_summary(df: pd.DataFrame, n_invalid: int, n_total: int,
                   dm_available: bool = True, dm_notes: str = ""):
    n_valid_mol  = len(df)
    n_lk_valid   = df["linker_valid"].sum()
    n_assembled  = df["assembly_ok"].sum()
    n_pt_valid   = df["protac_valid"].sum() if "protac_valid" in df else 0

    print("=" * 65)
    print("PROTAC LINKER EVALUATION SUMMARY")
    print("=" * 65)
    print(f"  Total input SMILES         : {n_total}")
    print(f"  Chemically invalid (skipped): {n_invalid}")
    print(f"  Chemically valid linkers   : {n_valid_mol}  "
          f"({n_valid_mol/n_total*100:.1f}%)")
    print(f"  Pass all hard constraints  : {n_lk_valid}  "
          f"({n_lk_valid/n_total*100:.1f}%)")
    print(f"  Successfully assembled     : {n_assembled}  "
          f"({n_assembled/n_total*100:.1f}%)")
    print(f"  Assembled + PROTAC valid   : {n_pt_valid}  "
          f"({n_pt_valid/n_total*100:.1f}%)")

    print()
    print("── Linker properties (valid mols only) ──────────────────────")
    for col, label in [
        ("linker_sa",       "SA score     (lower=easier synth)"),
        ("linker_qed",      "QED          (higher=drug-like)  "),
        ("linker_mw",       "Linker MW    (Da)               "),
        ("linker_logp",     "Linker logP                     "),
        ("linker_tpsa",     "Linker TPSA  (Ų)               "),
        ("linker_rotbonds", "Linker RotBonds                 "),
    ]:
        if col in df.columns:
            vals = df[col].dropna()
            print(f"  {label}: mean={vals.mean():.2f}  "
                  f"std={vals.std():.2f}  "
                  f"[{vals.min():.2f}, {vals.max():.2f}]")

    if n_assembled > 0:
        assembled = df[df["assembly_ok"]]
        print()
        print("── Assembled PROTAC PK properties ───────────────────────────")
        for col, label, window in [
            ("protac_mw",       "PROTAC MW    (Da)  [700–1200]  ", (700, 1200)),
            ("protac_logp",     "PROTAC logP        [−2–7]      ", (-2, 7)),
            ("protac_tpsa",     "PROTAC TPSA  (Ų)             ", None),
            ("protac_hbd",      "PROTAC HBD                    ", None),
            ("protac_hba",      "PROTAC HBA                    ", None),
            ("protac_rotbonds", "PROTAC RotBonds   [10–30]      ", (10, 30)),
            ("ro5_violations",  "Ro5 violations                ", None),
        ]:
            if col in assembled.columns:
                vals = assembled[col].dropna()
                in_window = ""
                if window:
                    frac = ((vals >= window[0]) & (vals <= window[1])).mean()
                    in_window = f"  {frac*100:.0f}% in window"
                print(f"  {label}: mean={vals.mean():.2f}  "
                      f"std={vals.std():.2f}  "
                      f"[{vals.min():.2f}, {vals.max():.2f}]{in_window}")

        # DegradeMaster degradation probability
        if "degrad_prob" in assembled.columns:
            dp = assembled["degrad_prob"].dropna()
            if len(dp) > 0:
                print()
                if dm_available:
                    print("── DegradeMaster degradation probability (0–1) ──────────────")
                    print(f"   (GCN checkpoint/1000; {dm_notes})")
                else:
                    print("── Degradation probability (model unavailable) ───────────────")
                print(f"  Scored compounds           : {len(dp)}")
                print(f"  Mean prob                  : {dp.mean():.3f}")
                print(f"  Std                        : {dp.std():.3f}")
                print(f"  Range                      : [{dp.min():.3f}, {dp.max():.3f}]")
                thresholds = [0.3, 0.5, 0.7]
                for t in thresholds:
                    n_above = (dp >= t).sum()
                    print(f"  ≥ {t:.1f} (likely active)       : {n_above}  ({n_above/len(dp)*100:.1f}%)")

    print()
    print("── Hard constraint breakdown ─────────────────────────────────")
    check_cols = [c for c in df.columns if c.startswith("check_")]
    for col in check_cols:
        name = col.replace("check_", "")
        pct = df[col].mean() * 100
        print(f"  {name:<30}: {pct:5.1f}% pass")
    print("=" * 65)


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

def main():
    parser = argparse.ArgumentParser(description="Evaluate generated PROTAC linkers")
    parser.add_argument("smiles_file", help="Path to generated SMILES file (one per line)")
    parser.add_argument("--out", default=None, help="Save results CSV to this path")
    parser.add_argument("--sa-max", type=float, default=6.0,
                        help="Max SA score for linker validity (default 6.0)")
    parser.add_argument("--protac-module", dest="protac_module", default="ErlotinibVHL",
                        help="Name of the protac builder module (e.g. ErlotinibVHL, "
                             "BRD9VHL, KRASVHL, BRD4VHL, ARCRBN). Selects warhead/E3 and "
                             "DegradeMaster pocket+features. Default: ErlotinibVHL.")
    parser.add_argument("--quiet", action="store_true", help="Suppress printed summary")
    args = parser.parse_args()

    path = Path(args.smiles_file)
    if not path.exists():
        print(f"Error: file not found: {path}", file=sys.stderr)
        sys.exit(1)

    with open(path) as f:
        smiles = [line.strip() for line in f if line.strip()]

    print(f"Loaded {len(smiles)} SMILES from {path}\n")

    df = evaluate(smiles, sa_max=args.sa_max, verbose=not args.quiet,
                  protac_module=args.protac_module)

    if args.out:
        df.to_csv(args.out, index=False)
        print(f"\nResults saved to {args.out}")


if __name__ == "__main__":
    main()