CupOFMoCA-PROTAC/src/classifier/protac_classifier.py

"""
PROTAC classifiers and guidance scorers for the DeFoG guidance loop.

All three classes accept DeFoG's molecule_list format and share _graph_to_mol.
They are called together on the same batch in each guidance loop iteration:

    PROTACClassifier      → 0/1 hard validity mask
    EGFRPROPKScorer       → continuous PK guidance (T½, CL, Vss)
    PROTACRDKitScorer     → continuous property guidance (QED, logP, TPSA, Fsp3, HBD, rotbonds)

DeFoG molecule_list format:  list of [atom_types, edge_types]
    atom_types : 1D int tensor, length n  — index into atom_decoder
    edge_types : 2D int tensor [n, n]     — 0=none 1=single 2=double 3=triple 4=aromatic ≥5=virtual
"""

from __future__ import annotations

import os
import pickle
from typing import Optional

import numpy as np
import torch
from rdkit import Chem, RDLogger
from rdkit.Chem import rdchem

from protac_analyzer import PROTACAnalyzer

RDLogger.DisableLog("rdApp.*")

# ---------------------------------------------------------------------------
# Shared graph → Mol conversion
# ---------------------------------------------------------------------------

_BOND_TYPES = [
    None,
    rdchem.BondType.SINGLE,
    rdchem.BondType.DOUBLE,
    rdchem.BondType.TRIPLE,
    rdchem.BondType.AROMATIC,
]


def _graph_to_mol(atom_types: torch.Tensor, edge_types: torch.Tensor,
                  atom_decoder: list[str]) -> Chem.RWMol:
    """Convert DeFoG (atom_types, edge_types) tensors to a Chem.RWMol."""
    mol = Chem.RWMol()
    for idx in atom_types:
        mol.AddAtom(Chem.Atom(atom_decoder[idx.item()]))

    edge_types = torch.triu(edge_types)
    edge_types[edge_types >= 5] = 0
    for i, j in torch.nonzero(edge_types).tolist():
        if i != j:
            mol.AddBond(i, j, _BOND_TYPES[edge_types[i, j].item()])
    return mol


def prebuild_mols(molecule_list: list, atom_decoder: list[str]) -> list:
    """Convert + sanitize each (atom_types, edge_types) once for shared scorer use.
    Returns list of sanitized RWMol or None if conversion/sanitization failed."""
    mols = []
    for atom_types, edge_types in molecule_list:
        try:
            m = _graph_to_mol(atom_types, edge_types, atom_decoder)
            Chem.SanitizeMol(m)
            mols.append(m)
        except Exception:
            mols.append(None)
    return mols


# ---------------------------------------------------------------------------
# Hard validity classifier
# ---------------------------------------------------------------------------

class PROTACClassifier:
    """
    Hard validity scorer for PROTAC linker graphs produced by DeFoG.
    Returns 0/1 tensor — 1 only when:
      - PROTACAnalyzer structural checks pass (atom count, ring closure, MW, SA, repr.)
      - builder.assemble() succeeds (actual RDKit reaction with warhead + E3)

    Using assembly as the attachment check means guidance is steered by real
    chemistry — NH2, reactive NH, or any N that actually reacts — rather than
    a SMARTS proxy.

    Parameters
    ----------
    analyzer : PROTACAnalyzer
        Structural checks only (attachment_sites check is bypassed here).
    builder : PROTAC_Builder
        Used for the assembly-based attachment check.
    atom_decoder : list[str]
        Maps atom-type index → element symbol.
    """

    def __init__(self, analyzer: PROTACAnalyzer, builder, atom_decoder: list[str]):
        self.analyzer = analyzer
        self.builder = builder
        self.atom_decoder = atom_decoder
        # SMARTS for site counting. Requiring an alpha carbon on the carbonyl
        # rules out carbamic acids (N-C(=O)OH), which can't form stable amides.
        self._cooh_pat  = Chem.MolFromSmarts("[#6][CX3](=[OX1])[OX2H1]")
        self._cho_pat   = Chem.MolFromSmarts("[#6][CX3H1](=[OX1])")
        self._methyl_pat = Chem.MolFromSmarts("[CX4H3]")
        self._amine_pat  = Chem.MolFromSmarts("[NX3;!H0;!$(NC=O);!$(NS=O);!$(N=O)]")
        self._nh2_pat    = Chem.MolFromSmarts("[NX3H2;D1;!$(NC=O);!$(NS=O);!$(N=O)]")
        # Hard rejection patterns.
        self._alerts = [
            Chem.MolFromSmarts('[NX3][NX3]'),                # hydrazide
            Chem.MolFromSmarts('[#7]=[#7]'),                 # azo
            Chem.MolFromSmarts('[N+][O-]'),                  # N-oxide
            Chem.MolFromSmarts('[#7]-[OX2]'),                # N-O single bond
            Chem.MolFromSmarts('[NX3]C(=[OX1])[OX2H1]'),     # carbamate
        ]

    def get_scores(self, molecule_list: list,
                   fix_termini=False, mols=None) -> torch.Tensor:
        """Binary validity (0/1) over a list of molecules.

        Returns 1 if the mol passes (a) sanitize + SMILES round-trip,
        (b) no liability SMARTS, and (c) >= 2 reactive termini for the chosen
        fix_termini mode (cooh / methyl / amine / amino_acid). Returns 0 otherwise.
        Pass `mols` to skip graph_to_mol + sanitize.
        """
        methyl_mode     = (fix_termini == 'methyl')
        amine_mode      = (fix_termini == 'amine')
        amino_acid_mode = (fix_termini == 'amino_acid')
        if mols is None:
            mols = prebuild_mols(molecule_list, self.atom_decoder)
        scores = []
        for mol in mols:
            if mol is None:
                scores.append(0.0)
                continue
            try:
                # sanitize + SMILES round-trip
                try:
                    Chem.SanitizeMol(mol)
                    smi = Chem.MolToSmiles(mol)
                    if not smi or Chem.MolFromSmiles(smi) is None:
                        scores.append(0.0)
                        continue
                except Exception:
                    scores.append(0.0)
                    continue
                # liability filter
                if any(mol.HasSubstructMatch(p) for p in self._alerts):
                    scores.append(0.0)
                    continue
                # reactive termini count
                if amine_mode:
                    n_sites = len(mol.GetSubstructMatches(self._nh2_pat))
                elif amino_acid_mode:
                    n_sites = (len(mol.GetSubstructMatches(self._nh2_pat)) +
                               len(mol.GetSubstructMatches(self._cooh_pat)))
                else:
                    n_sites = (len(mol.GetSubstructMatches(self._cooh_pat)) +
                               len(mol.GetSubstructMatches(self._cho_pat)))
                    if methyl_mode:
                        n_sites += len(mol.GetSubstructMatches(self._methyl_pat))
                scores.append(1.0 if n_sites >= 2 else 0.0)
            except Exception:
                scores.append(0.0)
        return torch.tensor(scores, dtype=torch.float32)

    def __call__(self, molecule_list: list, mols=None) -> torch.Tensor:
        return self.get_scores(molecule_list, fix_termini=getattr(self, 'fix_termini', False), mols=mols)


# ---------------------------------------------------------------------------
# EGFR-PROPK soft PK guidance scorer
# ---------------------------------------------------------------------------

class EGFRPROPKScorer:
    """
    Soft PK guidance using fine-tuned CatBoost models from Zhang et al. 2026
    (EGFR-PROPK). Predicts T½, CL, Vss on the assembled PROTAC.

    score = sigmoid(t½_raw) − sigmoid(CL_raw)

    GIN features (300 dims) are zero-padded — DGL not available in this env.
    The remaining 2563 structural features still provide useful guidance.

    Parameters
    ----------
    builder : PROTAC_Builder
        Used to assemble linker + warhead + E3 into a full PROTAC.
    model_dir : str
        Directory containing t0.5_10_finetune.pkl, CL_1_finetune.pkl,
        Vss_1_finetune.pkl.
    atom_decoder : list[str]
        Maps atom-type index → element symbol.
    half_life_weight, clearance_weight : float
        Relative weights for the composite score.
    """

    # Feature vector layout matching EGFR-PROPK training (maplight_gnn.py)
    _N_MORGAN = 1024
    _N_AVALON = 1024
    _N_ERG    = 315
    _N_GIN    = 300   # zero-padded — DGL unavailable
    _RDKIT_DESC_NAMES = [
        "BalabanJ", "BertzCT", "Chi0", "Chi0n", "Chi0v", "Chi1",
        "Chi1n", "Chi1v", "Chi2n", "Chi2v", "Chi3n", "Chi3v", "Chi4n", "Chi4v",
        "EState_VSA1", "EState_VSA10", "EState_VSA11", "EState_VSA2", "EState_VSA3",
        "EState_VSA4", "EState_VSA5", "EState_VSA6", "EState_VSA7", "EState_VSA8",
        "EState_VSA9", "ExactMolWt", "FpDensityMorgan1", "FpDensityMorgan2",
        "FpDensityMorgan3", "FractionCSP3", "HallKierAlpha", "HeavyAtomCount",
        "HeavyAtomMolWt", "Ipc", "Kappa1", "Kappa2", "Kappa3", "LabuteASA",
        "MaxAbsEStateIndex", "MaxAbsPartialCharge", "MaxEStateIndex", "MaxPartialCharge",
        "MinAbsEStateIndex", "MinAbsPartialCharge", "MinEStateIndex", "MinPartialCharge",
        "MolLogP", "MolMR", "MolWt", "NHOHCount", "NOCount", "NumAliphaticCarbocycles",
        "NumAliphaticHeterocycles", "NumAliphaticRings", "NumAromaticCarbocycles",
        "NumAromaticHeterocycles", "NumAromaticRings", "NumHAcceptors", "NumHDonors",
        "NumHeteroatoms", "NumRadicalElectrons", "NumRotatableBonds",
        "NumSaturatedCarbocycles", "NumSaturatedHeterocycles", "NumSaturatedRings",
        "NumValenceElectrons", "PEOE_VSA1", "PEOE_VSA10", "PEOE_VSA11", "PEOE_VSA12",
        "PEOE_VSA13", "PEOE_VSA14", "PEOE_VSA2", "PEOE_VSA3", "PEOE_VSA4", "PEOE_VSA5",
        "PEOE_VSA6", "PEOE_VSA7", "PEOE_VSA8", "PEOE_VSA9", "RingCount", "SMR_VSA1",
        "SMR_VSA10", "SMR_VSA2", "SMR_VSA3", "SMR_VSA4", "SMR_VSA5", "SMR_VSA6",
        "SMR_VSA7", "SMR_VSA8", "SMR_VSA9", "SlogP_VSA1", "SlogP_VSA10", "SlogP_VSA11",
        "SlogP_VSA12", "SlogP_VSA2", "SlogP_VSA3", "SlogP_VSA4", "SlogP_VSA5",
        "SlogP_VSA6", "SlogP_VSA7", "SlogP_VSA8", "SlogP_VSA9", "TPSA",
        "VSA_EState1", "VSA_EState10", "VSA_EState2", "VSA_EState3", "VSA_EState4",
        "VSA_EState5", "VSA_EState6", "VSA_EState7", "VSA_EState8", "VSA_EState9",
        "fr_Al_COO", "fr_Al_OH", "fr_Al_OH_noTert", "fr_ArN", "fr_Ar_COO", "fr_Ar_N",
        "fr_Ar_NH", "fr_Ar_OH", "fr_COO", "fr_COO2", "fr_C_O", "fr_C_O_noCOO",
        "fr_C_S", "fr_HOCCN", "fr_Imine", "fr_NH0", "fr_NH1", "fr_NH2", "fr_N_O",
        "fr_Ndealkylation1", "fr_Ndealkylation2", "fr_Nhpyrrole", "fr_SH",
        "fr_aldehyde", "fr_alkyl_carbamate", "fr_alkyl_halide", "fr_allylic_oxid",
        "fr_amide", "fr_amidine", "fr_aniline", "fr_aryl_methyl", "fr_azide", "fr_azo",
        "fr_barbitur", "fr_benzene", "fr_benzodiazepine", "fr_bicyclic", "fr_diazo",
        "fr_dihydropyridine", "fr_epoxide", "fr_ester", "fr_ether", "fr_furan",
        "fr_guanido", "fr_halogen", "fr_hdrzine", "fr_hdrzone", "fr_imidazole",
        "fr_imide", "fr_isocyan", "fr_isothiocyan", "fr_ketone", "fr_ketone_Topliss",
        "fr_lactam", "fr_lactone", "fr_methoxy", "fr_morpholine", "fr_nitrile",
        "fr_nitro", "fr_nitro_arom", "fr_nitro_arom_nonortho", "fr_nitroso",
        "fr_oxazole", "fr_oxime", "fr_para_hydroxylation", "fr_phenol",
        "fr_phenol_noOrthoHbond", "fr_phos_acid", "fr_phos_ester", "fr_piperdine",
        "fr_piperzine", "fr_priamide", "fr_prisulfonamd", "fr_pyridine", "fr_quatN",
        "fr_sulfide", "fr_sulfonamd", "fr_sulfone", "fr_term_acetylene", "fr_tetrazole",
        "fr_thiazole", "fr_thiocyan", "fr_thiophene", "fr_unbrch_alkane", "fr_urea", "qed",
    ]

    def __init__(self, builder, model_dir: str, atom_decoder: list[str],
                 half_life_weight: float = 1.0, clearance_weight: float = 1.0):
        from rdkit.Chem import DataStructs
        from rdkit.Chem.rdMolDescriptors import GetHashedMorganFingerprint
        from rdkit.Avalon.pyAvalonTools import GetAvalonCountFP
        from rdkit.Chem import rdReducedGraphs
        from rdkit.ML.Descriptors.MoleculeDescriptors import MolecularDescriptorCalculator

        self.builder = builder
        self.atom_decoder = atom_decoder
        self.half_life_weight = half_life_weight
        self.clearance_weight = clearance_weight

        self._DataStructs = DataStructs
        self._GetHashedMorganFP = GetHashedMorganFingerprint
        self._GetAvalonFP = GetAvalonCountFP
        self._GetErGFP = rdReducedGraphs.GetErGFingerprint
        self._desc_calc = MolecularDescriptorCalculator(self._RDKIT_DESC_NAMES)

        def _load(name):
            with open(os.path.join(model_dir, name), "rb") as f:
                return pickle.load(f)

        self._model_t_half = _load("t0.5_10_finetune.pkl")
        self._model_cl     = _load("CL_1_finetune.pkl")
        self._model_vss    = _load("Vss_1_finetune.pkl")

    def _assemble(self, linker_mol: Chem.Mol) -> Optional[Chem.Mol]:
        amide_only = getattr(self, 'fix_termini', None) != 'methyl'
        smi = Chem.MolToSmiles(linker_mol)
        if smi == getattr(self, '_asm_smi', None):
            return self._asm_mol
        try:
            result = self.builder.assemble(linker_mol, return_all=False, amide_only=amide_only)
            smiles = result if isinstance(result, str) else result.get("protac_smiles")
            mol = Chem.MolFromSmiles(smiles) if smiles else None
        except Exception:
            mol = None
        self._asm_smi, self._asm_mol = smi, mol
        return mol

    def _featurize(self, mol: Chem.Mol) -> np.ndarray:
        def _fp_arr(fp):
            a = np.zeros((0,), dtype=np.int8)
            self._DataStructs.ConvertToNumpyArray(fp, a)
            return a

        vec = np.concatenate([
            _fp_arr(self._GetHashedMorganFP(mol, nBits=self._N_MORGAN, radius=2)),
            _fp_arr(self._GetAvalonFP(mol, nBits=self._N_AVALON)),
            np.array(self._GetErGFP(mol)),
            np.array(self._desc_calc.CalcDescriptors(mol)),
            np.zeros(self._N_GIN, dtype=np.float32),
        ])
        return np.nan_to_num(vec, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float32)

    def _score_protac(self, protac_mol: Chem.Mol) -> float:
        X = self._featurize(protac_mol).reshape(1, -1)
        t_half = float(self._model_t_half.predict(X)[0])
        cl     = float(self._model_cl.predict(X)[0])
        sigmoid = lambda x: 1.0 / (1.0 + np.exp(-x))
        return self.half_life_weight * sigmoid(t_half) - self.clearance_weight * sigmoid(cl)

    def predict_pk(self, protac_mol: Chem.Mol) -> dict:
        """Return raw EGFR-PROPK predictions for the assembled PROTAC.

        Keys: t_half (hours, log-scale), cl (clearance), vss (volume of distribution),
        plus the composite score (sigmoid(t_half) − sigmoid(cl)). Use this for
        post-hoc validation / reporting; use _score_protac for guidance scalars.
        """
        X = self._featurize(protac_mol).reshape(1, -1)
        t_half = float(self._model_t_half.predict(X)[0])
        cl     = float(self._model_cl.predict(X)[0])
        vss    = float(self._model_vss.predict(X)[0])
        sigmoid = lambda x: 1.0 / (1.0 + np.exp(-x))
        return {
            "t_half":  t_half,
            "cl":      cl,
            "vss":     vss,
            "score":   self.half_life_weight * sigmoid(t_half) - self.clearance_weight * sigmoid(cl),
        }

    def get_scores(self, molecule_list: list) -> torch.Tensor:
        scores = []
        for atom_types, edge_types in molecule_list:
            try:
                linker = _graph_to_mol(atom_types, edge_types, self.atom_decoder)
                Chem.SanitizeMol(linker)
                protac = self._assemble(linker)
                scores.append(self._score_protac(protac) if protac is not None else 0.0)
            except Exception:
                scores.append(0.0)
        return torch.tensor(scores, dtype=torch.float32)

    def __call__(self, molecule_list: list) -> torch.Tensor:
        return self.get_scores(molecule_list)


# ---------------------------------------------------------------------------
# RDKit property soft guidance scorer
# ---------------------------------------------------------------------------

class PROTACRDKitScorer:
    """
    Soft property guidance using RDKit descriptors on the assembled PROTAC.
    Returns weighted mean of per-property scores, each in [0, 1].

    All windows are derived from PROTAC-DB (n=6,110 unique assembled PROTACs):
      flat top  = [mean − 1σ, mean + 1σ]  (score = 1)
      outer     = [mean − 2σ, mean + 2σ]  (score ramps 0 → 1)

    Parameters
    ----------
    builder : PROTAC_Builder
        Used to assemble linker + warhead + E3 into a full PROTAC.
    atom_decoder : list[str]
        Maps atom-type index → element symbol.
    weights : dict[str, float] | None
        Per-property weight. Keys: qed, logp, tpsa, fsp3, hbd, rotbonds.
    active : set[str] | None
        Subset of properties to score. Defaults to all.
    """

    # Trapezoidal: (lo, lo_ramp, hi_ramp, hi)
    # Score = 0 outside [lo, hi], linearly ramps to 1 in shoulders, flat 1 in [lo_ramp, hi_ramp].
    _PROPERTY_WINDOWS = {
        "qed":      (0.001, 0.015, 0.145, 0.210),
        "logp":     (0.7,   3.1,   7.9,   10.3),
        "tpsa":     (119.0, 165.0, 259.0, 306.0),
        "fsp3":     (0.24,  0.33,  0.51,  0.60),
        "hbd":      (0.7,   2.4,   6.0,   7.8),
        "rotbonds": (5.0,   12.0,  25.0,  32.0),
    }

    def __init__(self, builder, atom_decoder: list[str],
                 weights: dict | None = None, active: set | None = None):
        from rdkit.Chem import Descriptors, QED, rdMolDescriptors

        self.builder = builder
        self.atom_decoder = atom_decoder
        self.active = active if active is not None else set(self._PROPERTY_WINDOWS)
        self.weights = {k: (weights or {}).get(k, 1.0) for k in self.active}

        self._prop_fns = {
            "qed":      QED.qed,
            "logp":     Descriptors.MolLogP,
            "tpsa":     Descriptors.TPSA,
            "fsp3":     rdMolDescriptors.CalcFractionCSP3,
            "hbd":      lambda m: float(Descriptors.NumHDonors(m)),
            "rotbonds": lambda m: float(Descriptors.NumRotatableBonds(m)),
        }

    @staticmethod
    def _trapezoid(v, lo, lo_ramp, hi_ramp, hi) -> float:
        if v <= lo or v >= hi:
            return 0.0
        if v <= lo_ramp:
            return (v - lo) / (lo_ramp - lo)
        if v >= hi_ramp:
            return (hi - v) / (hi - hi_ramp)
        return 1.0

    def _score_property(self, name: str, mol: Chem.Mol) -> float:
        return self._trapezoid(self._prop_fns[name](mol), *self._PROPERTY_WINDOWS[name])

    def _assemble(self, linker_mol: Chem.Mol) -> Optional[Chem.Mol]:
        amide_only = getattr(self, 'fix_termini', None) != 'methyl'
        smi = Chem.MolToSmiles(linker_mol)
        if smi == getattr(self, '_asm_smi', None):
            return self._asm_mol
        try:
            result = self.builder.assemble(linker_mol, return_all=False, amide_only=amide_only)
            smiles = result if isinstance(result, str) else result.get("protac_smiles")
            mol = Chem.MolFromSmiles(smiles) if smiles else None
        except Exception:
            mol = None
        self._asm_smi, self._asm_mol = smi, mol
        return mol

    def _score_protac(self, protac_mol: Chem.Mol) -> float:
        per_prop = {k: self._score_property(k, protac_mol) for k in self.active}
        total_w = sum(self.weights[k] for k in per_prop)
        return sum(self.weights[k] * v for k, v in per_prop.items()) / total_w if total_w else 0.0

    def get_scores(self, molecule_list: list) -> torch.Tensor:
        scores = []
        for atom_types, edge_types in molecule_list:
            try:
                linker = _graph_to_mol(atom_types, edge_types, self.atom_decoder)
                Chem.SanitizeMol(linker)
                protac = self._assemble(linker)
                scores.append(self._score_protac(protac) if protac is not None else 0.0)
            except Exception:
                scores.append(0.0)
        return torch.tensor(scores, dtype=torch.float32)

    def __call__(self, molecule_list: list) -> torch.Tensor:
        return self.get_scores(molecule_list)


# ---------------------------------------------------------------------------
# Connectivity soft guidance scorer
# ---------------------------------------------------------------------------

class PROTACLinkerConnectivityScorer:
    """
    Soft signal pushing the linker toward a single connected molecule with both
    anchors in the same fragment.  Returns score in [0, 1]:

      1.0                                 — connected (single fragment)
      frag_frac / num_fragments           — disconnected with anchors in same fragment
      frag_frac / num_fragments / 2       — disconnected with anchors in different fragments
      0.0                                 — sanitization fails

    Scaling by 1/num_fragments penalizes "almost-connected" (95%-largest-fragment +
    1 stray atom) much more harshly than the older bare frag_frac, which previously
    rewarded such candidates at ~0.95 and gave guidance no incentive to merge the stray.
    With this scaling, 2-fragment cases peak at 0.5 even when the largest fragment is
    99% of atoms; only true single-fragment graphs reach 1.0. Anchor halving stacks
    on top to favour bridging the warhead/E3 anchors specifically.
    """

    def __init__(self, atom_decoder: list[str]):
        self.atom_decoder = atom_decoder

    def get_scores(self, molecule_list: list, mols=None) -> torch.Tensor:
        anchor_nodes = getattr(self, 'anchor_nodes', None)
        if mols is None:
            mols = prebuild_mols(molecule_list, self.atom_decoder)
        scores = []
        for mol in mols:
            if mol is None:
                scores.append(0.0)
                continue
            try:
                frags = Chem.GetMolFrags(mol)
                num_frags = len(frags)
                if num_frags == 1:
                    scores.append(1.0)
                    continue
                frag_frac = max(len(f) for f in frags) / mol.GetNumAtoms()
                # Harsher: scale by 1/num_fragments so any extra fragment caps the
                # score below 0.5, regardless of how small the stray fragment is.
                score = frag_frac / num_frags
                if anchor_nodes is not None and len(anchor_nodes) >= 2 and mol.GetNumAtoms() > max(anchor_nodes):
                    anchor_frag = next((f for f in frags if anchor_nodes[0] in f), None)
                    if anchor_frag is None or any(a not in anchor_frag for a in anchor_nodes[1:]):
                        score *= 0.5
                scores.append(score)
            except Exception:
                scores.append(0.0)
        return torch.tensor(scores, dtype=torch.float32)

    def __call__(self, molecule_list: list, mols=None) -> torch.Tensor:
        return self.get_scores(molecule_list, mols=mols)


# ---------------------------------------------------------------------------
# SA score soft guidance scorer  (scores the linker, not the assembled PROTAC)
# ---------------------------------------------------------------------------

class PROTACLinkerSAScorer:
    """
    Soft guidance scorer based on synthetic accessibility of the linker itself.

    Maps SA score [1, sa_threshold] → [1.0, 0.0] linearly (soft gradient).
    SA > sa_threshold hard-zeros to 0.0, acting as both a guidance penalty and
    a hard validity gate. Scoring the linker (not the assembled PROTAC) keeps
    the signal focused on what the generator actually controls.
    """

    def __init__(self, atom_decoder: list[str], sa_threshold: float = 6.0):
        import sys
        from rdkit.Chem import RDConfig
        sys.path.insert(0, os.path.join(RDConfig.RDContribDir, "SA_Score"))
        import sascorer as _sascorer
        self._sascorer = _sascorer
        self.atom_decoder = atom_decoder
        self.sa_threshold = sa_threshold

    def get_scores(self, molecule_list: list, mols=None) -> torch.Tensor:
        if mols is None:
            mols = prebuild_mols(molecule_list, self.atom_decoder)
        scores = []
        for mol in mols:
            if mol is None:
                scores.append(0.0)
                continue
            try:
                sa = self._sascorer.calculateScore(mol)
                scores.append(max(0.0, (self.sa_threshold - sa) / (self.sa_threshold - 1.0)))
            except Exception:
                scores.append(0.0)
        return torch.tensor(scores, dtype=torch.float32)

    def __call__(self, molecule_list: list, mols=None) -> torch.Tensor:
        return self.get_scores(molecule_list, mols=mols)