"""
DeFoG guided PROTAC linker generation.

Local-edit loop: each step picks one graph position (atom or edge pair), queries DeFoG
for its predicted clean distribution, tries all candidate types, and commits the best.
Improvements persist: every step builds on all prior commits.

Objectives (fast_scoring layout):
  PROTACClassifier is the hard validity gate (structural checks + alert mask +
  >=2 SMARTS sites). It drives the validity mask and is NOT a term in the
  weighted sum, so it takes no weight.

  Soft scorers, in --weights positional order (paper Eq. 6):
    [0] PROTACLinkerSAScorer           — S1 synthetic accessibility in [0,1]    w=4.0
    [1] PROTACLinkerChemScorer         — S2 logP/MW/RotB/ether chemotype mean   w=2.0
    [2] PROTACLinkerConnectivityScorer — S3 connectivity (always runs)          w=4.0

  PreassembledPROTACScorer (Morgan-FP linker quality) was dropped after the
  weight ablation (paper Appendix J.5) showed that raising its weight lowered
  both assembly validity and DegradeMaster score.

Aggregated score = Σ weights[k] · soft_scores[k]. Candidates rejected by the
validity gate are masked to zero probability before exponential tilting.
"""

import argparse
import os
import sys
import json
import csv
from pathlib import Path
import random
from datetime import datetime
from enum import Enum
from typing import Optional, Tuple, Dict

import numpy as np
import torch
import torch.nn.functional as F
from tqdm import tqdm

# Add DeFoG to path so its modules resolve at import time.
# DeFoG's internal code does `from flow_matching import ...` (needs src/defog/src)
# and also `from src import utils` (needs src/defog/, so `src` resolves to defog's src).
_DEFOG_ROOT = os.path.normpath(
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "../src/defog")
)
_DEFOG_SRC  = os.path.join(_DEFOG_ROOT, "src")
for _p in (_DEFOG_ROOT, _DEFOG_SRC):
    if _p not in sys.path:
        sys.path.insert(0, _p)

# Add PROTAC sources (classifier + builders) to path.
# _REPO_SRC itself lets us do `from protac.<TARGET> import ...`,
# and the subdirs let us do `from protac_classifier import ...`, etc.
_REPO_SRC = os.path.normpath(
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "../src")
)
if _REPO_SRC not in sys.path:
    sys.path.insert(0, _REPO_SRC)
for _sub in ("classifier", "protac"):
    _p = os.path.join(_REPO_SRC, _sub)
    if _p not in sys.path:
        sys.path.insert(0, _p)

from flow_matching import flow_matching_utils  # noqa: E402
from utils import PlaceHolder                  # noqa: E402


class GuidanceMethod(Enum):
    BEST_OF_K        = "bok"
    EXPONENTIAL_TILT = "exp_tilt"
    HYBRID           = "hybrid"


# ---------------------------------------------------------------------------
# Generation tracker  (same structure as cupofmoca_guided.py)
# ---------------------------------------------------------------------------

class GenerationTracker:

    def __init__(self, output_dir="./generation_logs"):
        self.output_dir = output_dir
        os.makedirs(output_dir, exist_ok=True)
        self.run_id = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.step_logs = []
        self.final_results = []

    def log_step(self, step: int, scores, detailed_results=None,
                 improved_indices=None, guidance_info: Optional[Dict] = None):
        step_info = {
            "step":             step,
            "scores":           scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else scores,
            "num_valid":        int((scores > 0).sum()) if hasattr(scores, "__iter__") else 0,
            "improved_indices": improved_indices or [],
            "guidance_info":    guidance_info or {},
        }
        self.step_logs.append(step_info)

    def log_final(self, molecule_list, detailed_results):
        for i, (_, result) in enumerate(zip(molecule_list, detailed_results)):
            self.final_results.append({
                "sample_idx":  i,
                "valid":       result.get("valid", 0),
                "smiles":      result.get("smiles", ""),
                "checks":      result.get("checks", {}),
                "properties":  result.get("properties", {}),
            })

    def save(self, filename_prefix="defog_guided"):
        step_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
        with open(step_path, "w") as f:
            json.dump(self.step_logs, f, indent=2, default=str)

        final_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_final.json")
        with open(final_path, "w") as f:
            json.dump(self.final_results, f, indent=2, default=str)

        csv_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_summary.csv")
        with open(csv_path, "w", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(["sample_idx", "valid", "smiles", "mw", "logp", "tpsa", "hbd", "rotbonds"])
            for r in self.final_results:
                props = r.get("properties", {})
                writer.writerow([
                    r["sample_idx"], r["valid"], r["smiles"],
                    props.get("mw", ""), props.get("xlogp", ""),
                    props.get("polararea", ""), props.get("hbonddonor", ""),
                    props.get("rotbonds", ""),
                ])

        print(f"\n=== Generation logs saved ===")
        print(f"  Step logs:  {step_path}")
        print(f"  Final:      {final_path}")
        print(f"  CSV:        {csv_path}")
        return step_path, final_path, csv_path

    def print_step_summary(self, step, scores, detailed_results=None, guidance_method=None):
        scores_list = scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else list(scores)
        # PROTACClassifier returns binary 0/1: 1 = passed all PROTACAnalyzer structural
        # checks + alert mask + has ≥2 SMARTS attachment sites; 0 = failed any.
        # Final assembly (actual RDKit reaction trial) is checked separately at the end.
        num_valid = sum(1 for s in scores_list if s > 0)
        print(f"\n--- Step {step} ({guidance_method or 'unknown'}) --- classifier-valid: {num_valid}/{len(scores_list)}")

    def print_final_summary(self):
        total     = len(self.final_results)
        num_valid = sum(1 for r in self.final_results if r.get("valid", 0) > 0)
        print("\n" + "=" * 60)
        print("FINAL GENERATION SUMMARY")
        print("=" * 60)
        print(f"Total molecules: {total}   Valid: {num_valid}")
        for r in self.final_results:
            tag = "✓" if r.get("valid") == 1 else "✗"
            smi = r.get("smiles", "") or "INVALID (sanitization failed)"
            print(f"  [{tag}] {r['sample_idx']:3d}  {smi[:60]}")


# ---------------------------------------------------------------------------
# DeFoGGenerator  (mirrors CupOfMoCAGenerator)
# ---------------------------------------------------------------------------

class DeFoGGenerator:
    """
    DeFoG-based guided PROTAC linker generator.

    Parameters
    ----------
    model : GraphDiscreteFlowModel
        Pre-loaded DeFoG model (pytorch_lightning module).
    atom_decoder : list[str]
        Maps atom-type index → element symbol (from dataset_infos.atom_decoder).
    objectives : list
        Soft scorers used in the weighted sum S(z). Final entry is treated as
        the connectivity scorer (always runs even when validity_fn = 0).
    validity_fn : callable | None
        Binary 0/1 hard gate (PROTACClassifier). Output drives the validity mask
        in exp_tilt sampling. Does not contribute to S(z).
    args : argparse.Namespace
    tracker : GenerationTracker | None
    """

    def __init__(self, model, atom_decoder: list, objectives: list, args,
                 validity_fn=None, tracker=None):
        self.model       = model
        self.atom_decoder = atom_decoder
        self.objectives  = objectives
        self.validity_fn = validity_fn
        self.args        = args
        self.num_objectives = len(objectives)
        self.tracker     = tracker
        self.device      = next(model.parameters()).device
        self.guidance_method = GuidanceMethod(args.guidance_method)

        # Vocab mask: restrict to atom types seen in PROTAC linker training data.
        # Zeroes out I, P, Cl, Br from pred_X before guidance — model can never pick them.
        _LINKER_ATOMS = {"C", "N", "O", "S", "F"}
        self.vocab_mask = torch.tensor(
            [1.0 if sym in _LINKER_ATOMS else 0.0 for sym in atom_decoder],
            device=self.device,
        )  # [d_X]

        self.sample_T       = getattr(args, 'num_steps', None) or model.sample_T
        self.node_dist      = model.node_dist
        self.time_distorter = model.time_distorter


    # -------------------------------------------------------------------------
    # Initial graph generation (x0 from prior μ₀)
    # -------------------------------------------------------------------------

    def _generate_single_x0(self) -> Tuple[torch.Tensor, torch.Tensor,
                                            torch.Tensor, torch.Tensor]:
        """
        Sample initial noisy graphs from the limit distribution.
        Returns X, E, y, node_mask — all on self.device.
        """
        batch_size = self.args.num_samples

        # Sample number of nodes per graph
        if self.args.num_nodes is not None:
            n_nodes = self.args.num_nodes * torch.ones(
                batch_size, device=self.device, dtype=torch.int
            )
        else:
            n_nodes = self.node_dist.sample_n(batch_size, self.device)
            # Clamp to PROTAC linker atom range
            n_nodes = n_nodes.clamp(self.args.min_nodes, self.args.max_nodes)

        n_max = n_nodes.max().item()
        arange    = torch.arange(n_max, device=self.device).unsqueeze(0).expand(batch_size, -1)
        node_mask = arange < n_nodes.unsqueeze(1)

        z_T = flow_matching_utils.sample_discrete_feature_noise(
            limit_dist=self.model.limit_dist, node_mask=node_mask
        )
        X, E, y = z_T.X, z_T.E, z_T.y
        return X, E, y, node_mask, n_nodes

    # Per-type terminus geometry configs.
    # 'cooh'   — two COOH groups at nodes 0-5; anchor carbons at 0 and 3
    # 'methyl' — CH3(0)–CH2(1)–[generated]–CH2(2)–CH3(3) scaffold;
    #            terminal CH3s (0, 3) bond only to their inner CH2 anchor
    _TERMINUS_CONFIGS = {
        'cooh': {
            'n_terminus_nodes': 6,
            'min_total_nodes':  8,    # 6 terminus + at least 2 free bridge atoms
            'anchor_nodes':     (0, 3),
            'fixed_atom_pos':   frozenset({0, 1, 2, 3, 4, 5}),
            'fixed_o_nodes':    frozenset({1, 2, 4, 5}),
            'fixed_edge_pos':   frozenset({
                (0, 1), (1, 0), (0, 2), (2, 0),   # terminus-1: C=O, C-OH
                (3, 4), (4, 3), (3, 5), (5, 3),   # terminus-2: C=O, C-OH
                (0, 3), (3, 0),                    # no direct C–C between anchor atoms
            }),
        },
        'methyl': {
            'n_terminus_nodes': 4,
            'min_total_nodes':  6,    # 4 terminus + at least 2 free bridge atoms
            'anchor_nodes':     (1, 2),
            'fixed_atom_pos':   frozenset({0, 1, 2, 3}),
            'fixed_o_nodes':    frozenset({0, 3}),   # terminal CH3s — bond only to their CH2
            'fixed_edge_pos':   frozenset({
                (0, 1), (1, 0),   # CH3–CH2 bond at warhead end
                (2, 3), (3, 2),   # CH2–CH3 bond at E3 end
                (1, 2), (2, 1),   # no direct anchor–anchor bond (force bridge through free atoms)
            }),
        },
        'amine': {
            'n_terminus_nodes': 2,
            'min_total_nodes':  4,    # 2 terminus N + at least 2 free bridge atoms
            'anchor_nodes':     (0, 1),
            'fixed_atom_pos':   frozenset({0, 1}),
            'fixed_o_nodes':    frozenset(),          # no O nodes to isolate
            'fixed_edge_pos':   frozenset({
                (0, 1), (1, 0),   # no direct N–N bond between the two terminus atoms
            }),
        },
        'amino_acid': {
            # Asymmetric: N(0) on warhead end, COOH(1-3) on VHL end.
            # Used for BRD4-VHL and any PROTAC where warhead provides COOH
            # and E3 provides NH2 (so linker must bridge with amine + acid termini).
            'n_terminus_nodes': 4,
            'min_total_nodes':  6,    # 4 terminus + at least 2 free bridge atoms
            'anchor_nodes':     (0, 1),
            'fixed_atom_pos':   frozenset({0, 1, 2, 3}),
            'fixed_o_nodes':    frozenset({2, 3}),    # O atoms of the COOH group
            'fixed_edge_pos':   frozenset({
                (1, 2), (2, 1),   # C=O (double, carbonyl)
                (1, 3), (3, 1),   # C-OH (single, hydroxyl)
                (0, 1), (1, 0),   # no direct N–C(cooh) bond between anchors
            }),
        },
    }

    @staticmethod
    def _limit_anchor_bonds(E: torch.Tensor, i: int, n: int,
                            anchor: int, blocked: set, max_extra: int = 1) -> None:
        """Cap the number of bonds from `anchor` to free atoms (in-place on E).

        For atoms outside the pinned `blocked` set, count bonds and zero out all
        but the first `max_extra` of them. Keeps the lowest-indexed bond
        (deterministic). Prevents the anchor from being absorbed into a ring
        (e.g., the pinned NH2 nitrogen ending up as part of a pyrazole) which
        would break the terminal-NH2 / single-bridge-bond invariant required for
        amide-bond assembly with the warhead/E3.
        """
        free_bonds = [nb for nb in range(n)
                      if nb != anchor and nb not in blocked
                      and E[i, anchor, nb].argmax().item() > 0]
        if len(free_bonds) > max_extra:
            for nb in free_bonds[max_extra:]:
                E[i, anchor, nb] = 0.0;  E[i, anchor, nb, 0] = 1.0
                E[i, nb, anchor] = 0.0;  E[i, nb, anchor, 0] = 1.0

    def _fix_termini(self, X: torch.Tensor, E: torch.Tensor,
                     n_nodes: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Hard-pin terminus groups after every step so guidance can never revert them.

        'cooh'       — nodes 0-5: two COOH groups (carbonyl-C + C=O + C-OH each).
                       Anchor carbons (0, 3) connect freely to bridge atoms (6+).
        'methyl'     — nodes 0-3: CH3(0)–CH2(1) ... CH2(2)–CH3(3) scaffold.
                       Terminal CH3s (0, 3) locked to only bond their inner CH2 (1 or 2).
        'amine'      — nodes 0-1: two N atoms (nucleophiles for warhead/E3 COOH).
        'amino_acid' — nodes 0-3: N(0) warhead-side amine + COOH(1-3) VHL-side acid.
                       Asymmetric — use for BRD4-VHL where warhead provides COOH and
                       E3 (VHL) provides NH2, so the linker bridges with NH2 and COOH.

        Each anchor atom is also capped to at most ONE bond to free atoms (via
        _limit_anchor_bonds) so the model can't absorb the terminus into a ring.
        """
        terminus_type = getattr(self.args, 'fix_termini', None)
        if not terminus_type:
            return X, E
        cfg = self._TERMINUS_CONFIGS[terminus_type]
        try:
            C_IDX = self.atom_decoder.index("C")
            O_IDX = self.atom_decoder.index("O") if terminus_type == 'cooh' else None
        except ValueError:
            return X, E

        for i in range(X.shape[0]):
            n = n_nodes[i].item()
            if n < cfg['n_terminus_nodes']:
                continue

            if terminus_type == 'cooh':
                # Atom types: C, O, O, C, O, O
                for pos, atype in [(0, C_IDX), (1, O_IDX), (2, O_IDX),
                                   (3, C_IDX), (4, O_IDX), (5, O_IDX)]:
                    X[i, pos] = 0.0
                    X[i, pos, atype] = 1.0
                # Internal bonds: C=O (double=2), C-OH (single=1)
                for r, c, btype in [
                    (0, 1, 2), (1, 0, 2), (0, 2, 1), (2, 0, 1),
                    (3, 4, 2), (4, 3, 2), (3, 5, 1), (5, 3, 1),
                ]:
                    E[i, r, c] = 0.0
                    E[i, r, c, btype] = 1.0
                # O atoms bond only to their own carbonyl C
                for o_node, c_node in [(1, 0), (2, 0), (4, 3), (5, 3)]:
                    for other in range(n):
                        if other == c_node:
                            continue
                        E[i, o_node, other] = 0.0;  E[i, o_node, other, 0] = 1.0
                        E[i, other, o_node] = 0.0;  E[i, other, o_node, 0] = 1.0
                # Block direct C–C bond between anchor carbons
                E[i, 0, 3] = 0.0;  E[i, 0, 3, 0] = 1.0
                E[i, 3, 0] = 0.0;  E[i, 3, 0, 0] = 1.0
                # Cap each anchor C to ≤1 bond to free atoms (single bridge attachment).
                # Blocked = self + the COOH O atoms + the other anchor C.
                self._limit_anchor_bonds(E, i, n, 0, blocked={0, 1, 2, 3}, max_extra=1)
                self._limit_anchor_bonds(E, i, n, 3, blocked={0, 3, 4, 5}, max_extra=1)

            elif terminus_type == 'methyl':
                # Scaffold: CH3(0)–CH2(1)–[bridge]–CH2(2)–CH3(3)
                # Terminal CH3s (0, 3) bond only to their inner CH2 anchor.
                for pos in (0, 1, 2, 3):
                    X[i, pos] = 0.0
                    X[i, pos, C_IDX] = 1.0
                # Force CH3–CH2 single bonds at each end
                for r, c in [(0, 1), (1, 0), (2, 3), (3, 2)]:
                    E[i, r, c] = 0.0;  E[i, r, c, 1] = 1.0
                # Terminal CH3 atoms (0, 3) bond ONLY to their inner CH2
                for terminal, inner in [(0, 1), (3, 2)]:
                    for other in range(n):
                        if other == inner:
                            continue
                        E[i, terminal, other] = 0.0;  E[i, terminal, other, 0] = 1.0
                        E[i, other, terminal] = 0.0;  E[i, other, terminal, 0] = 1.0
                # Block direct anchor–anchor bond (force bridge through free atoms)
                E[i, 1, 2] = 0.0;  E[i, 1, 2, 0] = 1.0
                E[i, 2, 1] = 0.0;  E[i, 2, 1, 0] = 1.0
                # Inner CH2 anchors (1, 2) bond ONCE to the bridge; everything else free.
                # (The terminal CH3s 0, 3 are already pinned to their inner CH2 only.)
                self._limit_anchor_bonds(E, i, n, 1, blocked={0, 1, 2}, max_extra=1)
                self._limit_anchor_bonds(E, i, n, 2, blocked={1, 2, 3}, max_extra=1)

            elif terminus_type == 'amine':
                # Two terminal N atoms at nodes 0 and 1.
                # Each N is the nucleophilic attachment point; warhead/E3 provide COOH.
                try:
                    N_IDX = self.atom_decoder.index("N")
                except ValueError:
                    continue
                for pos in (0, 1):
                    X[i, pos] = 0.0
                    X[i, pos, N_IDX] = 1.0
                # Block direct N–N bond between the two terminus atoms
                E[i, 0, 1] = 0.0;  E[i, 0, 1, 0] = 1.0
                E[i, 1, 0] = 0.0;  E[i, 1, 0, 0] = 1.0
                # Cap each anchor N to ≤1 bond to free atoms — keeps it as primary NH2,
                # not absorbed into a ring or turned into a secondary amine.
                self._limit_anchor_bonds(E, i, n, 0, blocked={0, 1}, max_extra=1)
                self._limit_anchor_bonds(E, i, n, 1, blocked={0, 1}, max_extra=1)

            elif terminus_type == 'amino_acid':
                # Node 0: N (warhead-side amine — nucleophile for warhead COOH)
                # Nodes 1,2,3: COOH group (VHL-side — electrophile for VHL NH2)
                #   node 1 = carbonyl C, node 2 = C=O oxygen, node 3 = C-OH oxygen
                try:
                    N_IDX = self.atom_decoder.index("N")
                    O_IDX = self.atom_decoder.index("O")
                except ValueError:
                    continue
                X[i, 0] = 0.0;  X[i, 0, N_IDX] = 1.0
                X[i, 1] = 0.0;  X[i, 1, C_IDX] = 1.0
                for pos in (2, 3):
                    X[i, pos] = 0.0;  X[i, pos, O_IDX] = 1.0
                # COOH internal bonds: C=O (double, type 2) and C-OH (single, type 1)
                for r, c, btype in [(1, 2, 2), (2, 1, 2), (1, 3, 1), (3, 1, 1)]:
                    E[i, r, c] = 0.0;  E[i, r, c, btype] = 1.0
                # O atoms bond only to their COOH carbon (node 1)
                for o_node in (2, 3):
                    for other in range(n):
                        if other == 1:
                            continue
                        E[i, o_node, other] = 0.0;  E[i, o_node, other, 0] = 1.0
                        E[i, other, o_node] = 0.0;  E[i, other, o_node, 0] = 1.0
                # Block direct N(0)–C(1) bond between the two anchor atoms
                E[i, 0, 1] = 0.0;  E[i, 0, 1, 0] = 1.0
                E[i, 1, 0] = 0.0;  E[i, 1, 0, 0] = 1.0
                # Cap each anchor to ≤1 bond to free atoms:
                # - N(0) capped → primary NH2 (not absorbed into a ring)
                # - C(1) capped → carboxyl C with single bridge bond (not 2 substituents)
                self._limit_anchor_bonds(E, i, n, 0, blocked={0, 1, 2, 3}, max_extra=1)
                self._limit_anchor_bonds(E, i, n, 1, blocked={0, 1, 2, 3}, max_extra=1)

        return X, E

    def _check_validity(self, X: torch.Tensor, E: torch.Tensor,
                        n_nodes: torch.Tensor) -> torch.Tensor:
        """Run the hard validity gate (PROTACClassifier)."""
        molecule_list = self._to_molecule_list(X, E, n_nodes)
        if self.validity_fn is None:
            return torch.ones(len(molecule_list), dtype=torch.bool, device=self.device)
        return self.validity_fn(molecule_list) > 0

    def _to_molecule_list(self, X: torch.Tensor, E: torch.Tensor,
                          n_nodes: torch.Tensor) -> list:
        """Convert batch tensors to DeFoG molecule_list format.

        X may be one-hot [B, N, d_X] or index [B, N]; same for E.
        _graph_to_mol expects index tensors, so collapse if needed.
        """
        mol_list = []
        for i in range(X.shape[0]):
            n = n_nodes[i].item()
            atom_types = X[i, :n]
            edge_types = E[i, :n, :n]
            if atom_types.dim() == 2:      # one-hot → indices
                atom_types = atom_types.argmax(-1)
            if edge_types.dim() == 3:      # one-hot → indices
                edge_types = edge_types.argmax(-1)
            mol_list.append([atom_types.cpu(), edge_types.cpu()])
        return mol_list

    def generate_x0(self) -> Tuple[torch.Tensor, ...]:
        """Generate initial graphs, optionally enforcing validity."""
        print("Generating initial graphs (x0 from limit distribution)...")
        X, E, y, node_mask, n_nodes = self._generate_single_x0()

        if self.args.require_valid_x0:
            print("Enforcing initial validity constraint...")
            for retry in range(self.args.max_x0_retries):
                validity = self._check_validity(X, E, n_nodes)
                num_valid   = validity.sum().item()
                num_invalid = self.args.num_samples - num_valid

                if num_invalid == 0:
                    print(f"✓ All {self.args.num_samples} samples valid!")
                    break

                print(f"  Retry {retry + 1}/{self.args.max_x0_retries}: "
                      f"{num_valid}/{self.args.num_samples} valid")

                invalid_idx = torch.where(~validity)[0]
                orig = self.args.num_samples
                self.args.num_samples = num_invalid
                X_new, E_new, _, nm_new, nn_new = self._generate_single_x0()
                self.args.num_samples = orig

                X[invalid_idx]         = X_new
                E[invalid_idx]         = E_new
                node_mask[invalid_idx] = nm_new
                n_nodes[invalid_idx]   = nn_new

        X, E = self._fix_termini(X, E, n_nodes)
        return X, E, y, node_mask, n_nodes

    # -------------------------------------------------------------------------
    # Score helpers  (identical to cupofmoca_guided)
    # -------------------------------------------------------------------------

    def _get_scores(self, molecule_list: list) -> torch.Tensor:
        """S(z) — shape [num_objectives + 1, batch_size].

        Row layout:
          [0]    validity_fn (PROTACClassifier) — binary 0/1, hard gate only
                 (not in the weighted sum; consumed by validity_mask in exp_tilt)
          [1:-1] soft scorers — fast / sa / chem (skipped if validity = 0 for all)
          [-1]   connectivity — always runs even when validity = 0

        Returning the validity row as row 0 keeps a stable interface with
        callers that read `raw[0]` for the validity mask.
        """
        from protac_classifier import prebuild_mols as _prebuild
        mols = _prebuild(molecule_list, self.atom_decoder)

        batch = len(molecule_list)
        if self.validity_fn is not None:
            validity = self.validity_fn(molecule_list, mols=mols).to(self.device)
        else:
            validity = torch.ones(batch, device=self.device)

        if self.num_objectives == 0:
            return validity.unsqueeze(0)

        # Connectivity (last) always runs.
        conn_scores = self.objectives[-1](molecule_list, mols=mols).to(self.device)

        n_mid = self.num_objectives - 1  # objectives[:-1]
        if n_mid > 0 and (validity == 0).all():
            zeros = torch.zeros(n_mid, batch, device=self.device)
            return torch.stack([validity, *[zeros[i] for i in range(n_mid)], conn_scores], dim=0)

        mid_scores = [obj(molecule_list, mols=mols).to(self.device) for obj in self.objectives[:-1]]
        return torch.stack([validity, *mid_scores, conn_scores], dim=0)

    def _compute_aggregated_score(self, scores: torch.Tensor,
                                  weights: torch.Tensor) -> torch.Tensor:
        """
        scores  : [1 + num_objectives, batch]  (row 0 is validity, ignored here)
        weights : [num_objectives, 1]

        Weighted sum over the soft scorers only. The validity row from _get_scores
        is consumed separately by the validity_mask in exp_tilt.
        """
        soft = scores[1:]
        if self.args.score_aggregation == "weighted_min":
            return torch.min(weights * soft, dim=0).values
        return (weights * soft).sum(dim=0)

    # -------------------------------------------------------------------------
    # Guidance methods — one position per step, same structure as cupofmoca_guided.py
    # The position is one atom OR one edge pair; the method selects the type to place there.
    # -------------------------------------------------------------------------

    def _best_of_k_guidance(
        self,
        x_idx: torch.Tensor, e_idx: torch.Tensor,
        ptype: str, pos,
        pred_X_i: torch.Tensor, pred_E_i: torch.Tensor,
        weights: torch.Tensor, current_score: float, n: int,
    ) -> Tuple[torch.Tensor, torch.Tensor, bool]:
        """Try every candidate type at one position; commit the highest-scoring one.

        All scorers are run on every candidate (with validity-first short-circuit
        inside _get_scores: invalid candidates skip the soft scorers). Commit the
        new type only if it strictly improves over the current score.
        """
        n_atom = len(self.atom_decoder)
        n_bond = pred_E_i.shape[-1]

        if ptype == "A":
            if pos >= n:
                return x_idx, e_idx, False
            best_score = current_score
            best_type = x_idx[pos].item()
            for atype in range(n_atom):
                if self.vocab_mask[atype].item() == 0.0:
                    continue
                trial_x = x_idx.clone(); trial_x[pos] = atype
                s = self._compute_aggregated_score(
                    self._get_scores([[trial_x, e_idx]]), weights
                ).item()
                if s > best_score:
                    best_score = s; best_type = atype
            if best_type != x_idx[pos].item():
                x_idx = x_idx.clone(); x_idx[pos] = best_type
                return x_idx, e_idx, True

        elif ptype == "E":
            r, c = pos
            if r >= n or c >= n:
                return x_idx, e_idx, False
            best_score = current_score
            best_type = e_idx[r, c].item()
            for btype in range(n_bond):
                trial_e = e_idx.clone(); trial_e[r, c] = btype; trial_e[c, r] = btype
                s = self._compute_aggregated_score(
                    self._get_scores([[x_idx, trial_e]]), weights
                ).item()
                if s > best_score:
                    best_score = s; best_type = btype
            if best_type != e_idx[r, c].item():
                e_idx = e_idx.clone(); e_idx[r, c] = best_type; e_idx[c, r] = best_type
                return x_idx, e_idx, True

        return x_idx, e_idx, False

    def _exponential_tilt_guidance(
        self,
        x_idx: torch.Tensor, e_idx: torch.Tensor,
        ptype: str, pos,
        pred_X_i: torch.Tensor, pred_E_i: torch.Tensor,
        beta: float, weights: torch.Tensor, current_score: float, n: int,
    ) -> Tuple[torch.Tensor, torch.Tensor, bool]:
        """
        Exponential tilting: p'(c) ∝ p(c) · exp(β·S(c)) · 1_valid(c).

        For atoms: tries top_k types from model distribution.
        For edges: tries all bond types.
        Optional acceptance criterion: reject if selected score < current_score.
        """
        n_atom = len(self.atom_decoder)
        n_bond = pred_E_i.shape[-1]

        if ptype == "A":
            if pos >= n:
                return x_idx, e_idx, False
            top_k = min(getattr(self.args, 'top_k', n_atom), n_atom)
            topk_out   = pred_X_i[pos, :n_atom].topk(top_k)
            candidates = topk_out.indices.cpu()
            base_probs = topk_out.values.cpu()

            agg_list, clf_list = [], []
            for atype in candidates.tolist():
                tx = x_idx.clone(); tx[pos] = atype
                raw = self._get_scores([[tx, e_idx]])
                agg_list.append(self._compute_aggregated_score(raw, weights).item())
                clf_list.append(raw[0, 0].item())
            agg_scores   = torch.tensor(agg_list)
            clf_scores   = torch.tensor(clf_list)
            current_type = x_idx[pos].item()

        elif ptype == "E":
            r, c = pos
            if r >= n or c >= n:
                return x_idx, e_idx, False
            candidates = torch.arange(n_bond)
            base_probs = pred_E_i[r, c, :n_bond].cpu()

            agg_list, clf_list = [], []
            for btype in candidates.tolist():
                te = e_idx.clone(); te[r, c] = btype; te[c, r] = btype
                raw = self._get_scores([[x_idx, te]])
                agg_list.append(self._compute_aggregated_score(raw, weights).item())
                clf_list.append(raw[0, 0].item())
            agg_scores   = torch.tensor(agg_list)
            clf_scores   = torch.tensor(clf_list)
            current_type = e_idx[r, c].item()

        else:
            return x_idx, e_idx, False

        # Hard gate: zero out any candidate type whose PROTACClassifier (objectives[0])
        # rejects (liability or insufficient termini). The soft scorers (connectivity,
        # SA, ChemScorer, FP-LogReg) provide the climbable gradient for the rest.
        validity_mask = (clf_scores > 0).float()

        log_base   = torch.log(base_probs.clamp(min=1e-10))
        log_tilted = log_base + beta * agg_scores + torch.log(validity_mask + 1e-10)
        log_max    = log_tilted.max()

        if torch.isinf(log_max) or torch.isnan(log_max):
            final_probs = base_probs.clamp(min=1e-10)
        else:
            tilted = torch.exp(log_tilted - log_max)
            psum   = tilted.sum()
            final_probs = tilted / psum if psum > 1e-10 else base_probs.clamp(min=1e-10)

        final_probs = final_probs.clamp(min=1e-10)
        final_probs = final_probs / final_probs.sum()

        selected_idx   = torch.multinomial(final_probs, 1).item()
        selected_type  = candidates[selected_idx].item()
        selected_score = agg_scores[selected_idx].item()

        if self.args.use_acceptance and selected_score < current_score:
            return x_idx, e_idx, False

        if selected_type == current_type:
            return x_idx, e_idx, False

        # All scorers already ran in the candidate loop above (no two-stage),
        # so the selected_score is already the full aggregated score.
        if ptype == "A":
            x_idx = x_idx.clone(); x_idx[pos] = selected_type
            return x_idx, e_idx, True
        else:
            r, c = pos
            e_idx = e_idx.clone(); e_idx[r, c] = selected_type; e_idx[c, r] = selected_type
            return x_idx, e_idx, True

    def _hybrid_guidance(
        self,
        x_idx: torch.Tensor, e_idx: torch.Tensor,
        ptype: str, pos,
        pred_X_i: torch.Tensor, pred_E_i: torch.Tensor,
        beta: float, weights: torch.Tensor, current_score: float, progress: float, n: int,
    ) -> Tuple[torch.Tensor, torch.Tensor, bool]:
        """BoK (exploration) for steps < 50%; exp_tilt (refinement) for steps ≥ 50%."""
        if progress < 0.5:
            return self._best_of_k_guidance(
                x_idx, e_idx, ptype, pos, pred_X_i, pred_E_i, weights, current_score, n)
        else:
            return self._exponential_tilt_guidance(
                x_idx, e_idx, ptype, pos, pred_X_i, pred_E_i, beta, weights, current_score, n)

    # -------------------------------------------------------------------------
    # Main generation loop — local-edit style (Algorithm 2 on graphs)
    # -------------------------------------------------------------------------

    def generate(self) -> list:
        """
        Run local-edit guided generation.

        Each step picks one position (atom or edge pair, rotating through all positions),
        queries the model, tries all candidate types, and commits the best-scoring one.
        Improvements persist across all steps.
        """
        X, E, y, node_mask, n_nodes = self.generate_x0()

        batch_size   = self.args.num_samples
        n_atom_types = len(self.atom_decoder)
        n_bond_types = E.shape[-1]   # read from model tensor, not hardcoded

        # weights apply to the soft scorers only (len = num_objectives).
        if self.args.weights is None:
            weights = torch.tensor(
                [1.0 / self.num_objectives] * self.num_objectives,
                device=self.device,
            ).view(-1, 1)
        else:
            raw_w = list(self.args.weights)
            # Backward compat: if too many weights are passed (e.g. the
            # old 5-vector form with classifier+fast leading), strip the
            # extra leading entries so the trailing positions still line up
            # with the connectivity scorer.
            while len(raw_w) > self.num_objectives:
                raw_w = raw_w[1:]
            if len(raw_w) != self.num_objectives:
                raise ValueError(
                    f"len(weights)={len(raw_w)} must equal num_objectives="
                    f"{self.num_objectives} (got {self.args.weights})"
                )
            weights = torch.tensor(raw_w, device=self.device).view(-1, 1)

        # Position schedule: atoms first, then upper-triangle edge pairs.
        # Uses n_max so every atom and edge in the largest graph gets visited.
        n_max = n_nodes.max().item()
        atom_positions = [("A", p) for p in range(n_max)]
        edge_positions = [("E", (r, c)) for r in range(n_max) for c in range(r + 1, n_max)]
        all_positions  = edge_positions + atom_positions
        period         = len(all_positions)

        warmup_steps  = min(getattr(self.args, 'warmup_steps', 0), self.sample_T - 1)
        guided_steps  = self.sample_T - warmup_steps

        _gm_desc = {
            "bok":      "exhaustive atom/bond search (model logits unused)",
            "exp_tilt": "model-prior × exp(β·score) sampling",
            "hybrid":   "bok → exp_tilt at 50% progress",
        }.get(self.guidance_method.value, "")
        print(f"\nGuidance:  {self.guidance_method.value}  ({_gm_desc})")
        print(f"Weights:   {weights.squeeze().tolist()}")
        print(f"β range:   [{self.args.beta_min}, {self.args.beta_max}]")
        print(f"Schedule:  {n_max} atoms + {n_max*(n_max-1)//2} edges = {period} positions, "
              f"{guided_steps} guided steps (positions sampled per-step by model disagreement)")
        if warmup_steps:
            print(f"Warmup:    {warmup_steps} DeFoG denoising steps "
                  f"(t=0→{warmup_steps/self.sample_T:.2f}) before guidance")

        with torch.no_grad():
            # ------------------------------------------------------------------
            # Warmup: run DeFoG's own denoising kernel to get structured starting
            # structures.  Guidance operates on noise; warmup replaces those first
            # steps with proper flow-matching transitions so guided loop starts from
            # real chemistry rather than the limit distribution.
            # ------------------------------------------------------------------
            if warmup_steps:
                for t_int in tqdm(range(warmup_steps), desc="DeFoG warmup"):
                    t_array = t_int * torch.ones((batch_size, 1), device=self.device).float()
                    t_norm  = t_array / self.sample_T
                    if "absorb" in self.model.cfg.model.transition and t_int == 0:
                        t_norm = t_norm + 1e-6
                    s_array = t_array + 1
                    s_norm  = s_array / self.sample_T
                    t_norm  = self.time_distorter.sample_ft(t_norm, self.model.cfg.sample.time_distortion)
                    s_norm  = self.time_distorter.sample_ft(s_norm, self.model.cfg.sample.time_distortion)
                    sampled_s, _ = self.model.sample_p_zs_given_zt(t_norm, s_norm, X, E, y, node_mask)
                    X, E, y = sampled_s.X, sampled_s.E, sampled_s.y
                    X, E = self._fix_termini(X, E, n_nodes)  # re-pin after each warmup step

                warmup_mols   = self._to_molecule_list(X, E, n_nodes)
                current_scores = self._get_scores(warmup_mols)
                warmup_agg     = self._compute_aggregated_score(current_scores, weights)
                print(f"\nPost-warmup scores: {warmup_agg.cpu().numpy()}")

            for t_int in tqdm(range(warmup_steps, self.sample_T), desc="DeFoG local-edit"):

                progress = t_int / max(self.sample_T - 1, 1)
                beta = self.args.beta_min + (self.args.beta_max - self.args.beta_min) * progress

                t_array = t_int * torch.ones((batch_size, 1), device=self.device).float()
                t_norm  = t_array / self.sample_T

                if "absorb" in self.model.cfg.model.transition and t_norm.min() < 1e-6:
                    t_norm = t_norm + 1e-6

                t_norm = self.time_distorter.sample_ft(t_norm, self.model.cfg.sample.time_distortion)

                noisy_data = {"X_t": X, "E_t": E, "y_t": y, "t": t_norm, "node_mask": node_mask}
                extra_data = self.model.compute_extra_data(noisy_data)
                pred       = self.model.forward(noisy_data, extra_data, node_mask)
                pred_X     = F.softmax(pred.X, dim=-1)   # [B, n_max, d_X]
                pred_E     = F.softmax(pred.E, dim=-1)   # [B, n_max, n_max, d_E]

                # Mask disallowed atom types and renormalize
                pred_X = pred_X * self.vocab_mask
                pred_X = pred_X / pred_X.sum(dim=-1, keepdim=True).clamp(min=1e-10)

                # Score cache: only mols changed in the previous step's commit
                # have stale cache entries. Use cached `current_scores` directly
                # (initialized post-warmup; refreshed for changed indices below).
                current_agg = self._compute_aggregated_score(current_scores, weights)

                if t_int == warmup_steps:
                    print(f"\nInitial scores: {current_scores.cpu().numpy()}")

                # ── Position sampling weighted by model disagreement ────────────
                # Per-position disagreement = 1 − p_θ(current_type | context).
                # High value ⇒ model wants to change this slot ⇒ worth visiting.
                # Pinned-terminus positions get weight 0 so they're never sampled.
                _ttype = getattr(self.args, 'fix_termini', None)
                _cfg = self._TERMINUS_CONFIGS[_ttype] if _ttype else None

                # Atom disagreement: gather p_θ(current_type) per atom slot, average over batch.
                _x_cur = X.argmax(-1)                                                     # [B, n_max]
                _atom_p_cur = pred_X.gather(-1, _x_cur.unsqueeze(-1)).squeeze(-1)         # [B, n_max]
                _atom_dis = (1.0 - _atom_p_cur).mean(0)                                   # [n_max]
                # Edge disagreement: same idea over edge slots.
                _e_cur = E.argmax(-1)                                                     # [B, n_max, n_max]
                _edge_p_cur = pred_E.gather(-1, _e_cur.unsqueeze(-1)).squeeze(-1)         # [B, n_max, n_max]
                _edge_dis = (1.0 - _edge_p_cur).mean(0)                                   # [n_max, n_max]

                _pos_weights = torch.zeros(period, device=self.device)
                for _idx, (_pt, _ps) in enumerate(all_positions):
                    if _pt == "A":
                        if _cfg and _ps in _cfg['fixed_atom_pos']:
                            continue
                        _pos_weights[_idx] = _atom_dis[_ps]
                    else:
                        _r, _c = _ps
                        if _cfg and (_r in _cfg['fixed_o_nodes'] or _c in _cfg['fixed_o_nodes']
                                     or _ps in _cfg['fixed_edge_pos']):
                            continue
                        _pos_weights[_idx] = _edge_dis[_r, _c]

                # Fall back to uniform if model is fully confident everywhere.
                if _pos_weights.sum() < 1e-10:
                    _pos_weights = torch.ones_like(_pos_weights)
                    if _cfg:
                        for _idx, (_pt, _ps) in enumerate(all_positions):
                            if _pt == "A" and _ps in _cfg['fixed_atom_pos']:
                                _pos_weights[_idx] = 0.0
                            elif _pt == "E":
                                _r, _c = _ps
                                if (_r in _cfg['fixed_o_nodes'] or _c in _cfg['fixed_o_nodes']
                                        or _ps in _cfg['fixed_edge_pos']):
                                    _pos_weights[_idx] = 0.0

                # Position sampling mode:
                #   "weighted" (default) — multinomial sample over all eligible positions,
                #                          weighted by disagreement (long-tail exploration)
                #   "topk"               — zero out all but top-K most-disagreed positions
                #                          before sampling (focused exploitation)
                _pos_mode = getattr(self.args, 'position_sampling', 'weighted')
                if _pos_mode == 'topk':
                    _k = min(getattr(self.args, 'top_positions', 50), period)
                    _topk = torch.topk(_pos_weights, _k)
                    _mask = torch.zeros_like(_pos_weights)
                    _mask[_topk.indices] = _pos_weights[_topk.indices]
                    _pos_weights = _mask
                    if _pos_weights.sum() < 1e-10:
                        _pos_weights = torch.ones_like(_pos_weights)

                _sampled_idx = torch.multinomial(_pos_weights, 1).item()
                ptype, pos = all_positions[_sampled_idx]

                new_X = X.clone()
                new_E = E.clone()

                improved_indices = []
                for i in range(batch_size):
                    n = n_nodes[i].item()
                    # Current graph as index tensors, clamped to real atom/bond types
                    x_idx = X[i, :n].argmax(-1).cpu().clamp(0, n_atom_types - 1)
                    e_idx = E[i, :n, :n].argmax(-1).cpu().clamp(0, n_bond_types - 1)
                    current_score_i = current_agg[i].item()

                    pred_X_i = pred_X[i]
                    pred_E_i = pred_E[i]

                    if self.guidance_method == GuidanceMethod.BEST_OF_K:
                        x_idx, e_idx, changed = self._best_of_k_guidance(
                            x_idx, e_idx, ptype, pos, pred_X_i, pred_E_i,
                            weights, current_score_i, n)
                    elif self.guidance_method == GuidanceMethod.EXPONENTIAL_TILT:
                        x_idx, e_idx, changed = self._exponential_tilt_guidance(
                            x_idx, e_idx, ptype, pos, pred_X_i, pred_E_i,
                            beta, weights, current_score_i, n)
                    elif self.guidance_method == GuidanceMethod.HYBRID:
                        x_idx, e_idx, changed = self._hybrid_guidance(
                            x_idx, e_idx, ptype, pos, pred_X_i, pred_E_i,
                            beta, weights, current_score_i, progress, n)

                    if changed:
                        improved_indices.append(i)

                    new_X[i, :n] = F.one_hot(x_idx.long().to(self.device), X.shape[-1]).float()
                    new_E[i, :n, :n] = F.one_hot(e_idx.long().to(self.device), E.shape[-1]).float()

                X, E = new_X, new_E
                X, E = self._fix_termini(X, E, n_nodes)  # prevent any reversion

                if improved_indices:
                    # Only re-score the molecules that actually changed. Stitch the
                    # new column scores into `current_scores` (the cross-step cache).
                    post_mols = self._to_molecule_list(X, E, n_nodes)
                    changed_mols = [post_mols[i] for i in improved_indices]
                    changed_scores = self._get_scores(changed_mols)  # [num_obj, len(changed)]
                    for j, i in enumerate(improved_indices):
                        current_scores[:, i] = changed_scores[:, j]
                    new_agg = self._compute_aggregated_score(current_scores, weights)
                    score_improved = [i for i in improved_indices
                                      if new_agg[i].item() > current_agg[i].item()]
                    if score_improved:
                        print(f"\nStep {t_int} [{ptype} pos={pos}]: improved {score_improved}")
                        for i in score_improved:
                            print(f"  [{i}] {current_agg[i]:.4f} → {new_agg[i]:.4f}")

                if self.tracker:
                    self.tracker.log_step(
                        t_int, current_scores[0], improved_indices,
                        guidance_info={"ptype": ptype, "pos": str(pos)}
                    )
                    if t_int % 10 == 0 or t_int == self.sample_T - 1:
                        self.tracker.print_step_summary(
                            t_int, current_scores[0], guidance_method="local_edit"
                        )

        # Strip virtual classes, collapse one-hot → indices
        X, E, y = self.model.noise_dist.ignore_virtual_classes(X, E, y)
        final = PlaceHolder(X=X, E=E, y=y).mask(node_mask, collapse=True)
        return self._to_molecule_list(final.X, final.E, n_nodes)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main(args):
    # PROTAC source dirs already on sys.path (set at module load above).
    # Re-assert here for safety in case main() is called from a wrapper.
    _src_root = os.path.normpath(
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "../src")
    )
    if _src_root not in sys.path:
        sys.path.insert(0, _src_root)
    for _sub in ("classifier", "protac"):
        _p = os.path.join(_src_root, _sub)
        if _p not in sys.path:
            sys.path.insert(0, _p)

    from rdkit import Chem
    from rdkit.Chem import Descriptors, rdMolDescriptors
    from protac_analyzer import PROTACAnalyzer
    from protac_classifier import PROTACClassifier, EGFRPROPKScorer, PROTACRDKitScorer, PROTACLinkerSAScorer, PROTACLinkerConnectivityScorer, _graph_to_mol
    _pmod_name = getattr(args, 'protac_module', 'ErlotinibVHL')
    import importlib as _il
    _pmod = _il.import_module(f"protac.{_pmod_name}")
    WARHEAD_MOL    = _pmod.WARHEAD_MOL
    E3_MOL         = _pmod.E3_MOL
    _ProtacBuilder = _pmod.Builder

    class PreassembledPROTACScorer:
        """
        Fast linker scorer: no PROTAC assembly at scoring time.

        Pre-computes warhead (EGFR erlotinib) + E3 ligase (VHL) contributions to
        Scores linker quality using a logistic regression classifier trained to
        distinguish real PROTAC linkers (COOH-terminated, from PROTAC-DB) from
        generated junk (N-COOH, N-F, S-F, disconnected, etc.).

        Returns score in [0, 2]:
          0                 — disconnected or unsanitizable linker
          max_frag/n_atoms  — soft connectivity gradient for disconnected mols
          1 + clf_prob      — [1, 2] for connected mols; clf_prob ∈ [0,1] is
                              the classifier's probability of being a real linker

        Use as objectives[0] with sa_scorer as objectives[1].
        Full assembly is still run after generation for final validity reporting.
        """
        _ATOM_MIN, _ATOM_MAX = 3, 20

        # Hard liability filters (same set as ChemScorer / PROTACClassifier).
        _LIABILITY_PATS = [
            Chem.MolFromSmarts('[NX3][NX3]'),                # hydrazide
            Chem.MolFromSmarts('[#7]=[#7]'),                 # azo
            Chem.MolFromSmarts('[N+][O-]'),                  # N-oxide
            Chem.MolFromSmarts('[#7]-[OX2]'),                # N-O
            Chem.MolFromSmarts('[NX3]C(=[OX1])[OX2H1]'),     # carbamate
        ]

        def __init__(self, atom_decoder, clf_path):
            import pickle
            from rdkit.Chem import AllChem as _AllChem
            self.atom_decoder = atom_decoder
            self._allchem = _AllChem
            with open(clf_path, 'rb') as f:
                saved = pickle.load(f)
            self._clf = saved['clf']
            print(f"[PreassembledPROTACScorer] loaded linker quality classifier "
                  f"({len(saved['pos_smiles'])} pos / {len(saved['neg_smiles'])} neg)")

        def _fingerprint(self, mol):
            import numpy as np
            fp = self._allchem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=1024)
            return np.array(list(fp), dtype=np.float32).reshape(1, -1)

        def __call__(self, molecule_list, mols=None):
            from protac_classifier import prebuild_mols as _prebuild
            if mols is None:
                mols = _prebuild(molecule_list, self.atom_decoder)
            scores = []
            for mol in mols:
                if mol is None:
                    scores.append(0.0)
                    continue
                frag_mols = Chem.GetMolFrags(mol, asMols=True)
                if len(frag_mols) > 1:
                    largest = max(frag_mols, key=lambda f: f.GetNumHeavyAtoms())
                    n_heavy_largest = largest.GetNumHeavyAtoms()
                    if n_heavy_largest < self._ATOM_MIN:
                        scores.append(0.0)
                        continue
                    if any(largest.HasSubstructMatch(p) for p in self._LIABILITY_PATS):
                        scores.append(0.0)
                        continue
                    frac = n_heavy_largest / mol.GetNumHeavyAtoms()
                    try:
                        prob = float(self._clf.predict_proba(self._fingerprint(largest))[0, 1])
                    except Exception:
                        prob = 0.0
                    scores.append(prob * frac)
                    continue
                n_heavy = mol.GetNumHeavyAtoms()
                if not (self._ATOM_MIN <= n_heavy <= self._ATOM_MAX):
                    scores.append(0.0)
                    continue
                if any(mol.HasSubstructMatch(p) for p in self._LIABILITY_PATS):
                    scores.append(0.0)
                    continue
                try:
                    prob = float(self._clf.predict_proba(self._fingerprint(mol))[0, 1])
                except Exception:
                    prob = 0.0
                scores.append(prob)
            return torch.tensor(scores, dtype=torch.float32)

    class PROTACLinkerChemScorer:
        """
        Scores linker chemistry toward assembled PROTAC target ranges (no assembly needed).
        MW and RotBonds windows are derived by subtracting the fixed warhead+E3 offset from
        the assembled PROTAC target ranges [771-1108 Da, 12-25 RotBonds].
        logP scored directly on linker (training mean 0.48±1.37).
        Hard zero for MC liabilities: hydrazines (N–N), azo (N=N), N-oxides,
        any N–O single bond (hydroxylamines, alkoxyamines, oxime ethers), and
        carbamates (N directly bonded to a COOH carbonyl C).
        """
        # Assembled PROTAC target ranges (from PROTAC-DB)
        _PROTAC_MW_LO,   _PROTAC_MW_HI   = 771.0, 1108.0
        _PROTAC_ROTB_LO, _PROTAC_ROTB_HI = 12.0,  25.0

        _ETHER_O    = Chem.MolFromSmarts('[OX2;!$(O=C)]')
        _HYDRAZINE  = Chem.MolFromSmarts('[NX3][NX3]')             # N–N single (incl. hydrazine, hydrazide)
        _AZO        = Chem.MolFromSmarts('[#7]=[#7]')              # N=N (azo, diazo)
        _NOXIDE     = Chem.MolFromSmarts('[N+][O-]')
        _NO_BOND    = Chem.MolFromSmarts('[#7]-[OX2]')             # any N–O single bond (catches hydroxylamines, alkoxyamines, oxime ethers — sp2 + sp3)
        _CARBAMATE  = Chem.MolFromSmarts('[NX3]C(=[OX1])[OX2H1]')  # N directly bonded to COOH carbonyl C — kills cooh-mode "diacid" candidates that grew an N next to a pinned carboxyl. Does NOT match α-amino acids (N–CH–COOH) where N is on the α-C, not the carbonyl.

        @staticmethod
        def _trapezoid(v, lo, lo_ramp, hi_ramp, hi):
            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 __init__(self, atom_decoder, warhead_mol=None, e3_mol=None):
            from rdkit.Chem import Descriptors as _Desc
            self.atom_decoder = atom_decoder
            self._logp = _Desc.MolLogP
            self._mw   = _Desc.MolWt
            self._rotb = _Desc.NumRotatableBonds

            # Compute fixed warhead+E3 offset once (subtract 2×18 Da for two bond formations)
            if warhead_mol is not None and e3_mol is not None:
                fixed_mw   = _Desc.MolWt(warhead_mol) + _Desc.MolWt(e3_mol) - 36.0
                fixed_rotb = float(_Desc.NumRotatableBonds(warhead_mol) +
                                   _Desc.NumRotatableBonds(e3_mol))
            else:
                fixed_mw, fixed_rotb = 0.0, 0.0

            # Derive linker-level target windows; disable guidance if window collapses
            _MW_MIN_RANGE   = 50.0   # Da  — skip MW guidance if flat-top narrower than this
            _ROTB_MIN_RANGE = 3.0    # bonds — skip RotB guidance if flat-top narrower than this
            mw_lo   = max(0.0, self._PROTAC_MW_LO   - fixed_mw)
            mw_hi   = max(0.0, self._PROTAC_MW_HI   - fixed_mw)
            rotb_lo = max(0.0, self._PROTAC_ROTB_LO - fixed_rotb)
            rotb_hi = max(0.0, self._PROTAC_ROTB_HI - fixed_rotb)
            shoulder = 0.2
            mw_span   = mw_hi   - mw_lo
            rotb_span = rotb_hi - rotb_lo
            self._mw_active   = mw_span   >= _MW_MIN_RANGE
            self._rotb_active = rotb_span >= _ROTB_MIN_RANGE
            if self._mw_active:
                self._mw_window = (mw_lo - mw_span*shoulder, mw_lo, mw_hi, mw_hi + mw_span*shoulder)
            if self._rotb_active:
                self._rotb_window = (rotb_lo - rotb_span*shoulder, rotb_lo, rotb_hi, rotb_hi + rotb_span*shoulder)
            print(f"[ChemScorer] linker MW target:   "
                  f"[{mw_lo:.0f}, {mw_hi:.0f}] Da (fixed {fixed_mw:.0f} Da)"
                  f"{'' if self._mw_active else '  ← SKIPPED (window too narrow)'}")
            print(f"[ChemScorer] linker RotB target: "
                  f"[{rotb_lo:.0f}, {rotb_hi:.0f}] (fixed {fixed_rotb:.0f})"
                  f"{'' if self._rotb_active else '  ← SKIPPED (window too narrow)'}")

        def __call__(self, molecule_list, mols=None):
            from protac_classifier import prebuild_mols as _prebuild
            if mols is None:
                mols = _prebuild(molecule_list, self.atom_decoder)
            scores = []
            for mol in mols:
                if mol is None:
                    scores.append(0.0); continue
                if len(Chem.GetMolFrags(mol)) > 1:
                    scores.append(0.0); continue
                if (mol.HasSubstructMatch(self._HYDRAZINE) or
                        mol.HasSubstructMatch(self._AZO) or
                        mol.HasSubstructMatch(self._NOXIDE) or
                        mol.HasSubstructMatch(self._NO_BOND) or
                        mol.HasSubstructMatch(self._CARBAMATE)):
                    scores.append(0.0); continue
                n_ether = len(mol.GetSubstructMatches(self._ETHER_O))
                ether_score = min(n_ether / 4.0, 1.0)
                try:
                    logp = self._logp(mol)
                    mw   = self._mw(mol)
                    rotb = float(self._rotb(mol))
                except Exception:
                    logp, mw, rotb = 5.0, 500.0, 20.0
                logp_score = self._trapezoid(logp, -3.0, -1.0, 2.0, 5.0)
                mw_score   = self._trapezoid(mw,   *self._mw_window) if self._mw_active   else 1.0
                rotb_score = self._trapezoid(rotb, *self._rotb_window) if self._rotb_active else 1.0
                scores.append(0.25 * ether_score + 0.25 * logp_score +
                               0.25 * mw_score   + 0.25 * rotb_score)
            return torch.tensor(scores, dtype=torch.float32)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # Load DeFoG model
    # load_from_checkpoint fails because train_metrics/sampling_metrics are not saved in hparams.
    # Instead: pull cfg + dataset_infos directly from the checkpoint, pass dummy metrics.
    print(f"Loading DeFoG checkpoint: {args.checkpoint}")
    from graph_discrete_flow_model import GraphDiscreteFlowModel
    from analysis.visualization import MolecularVisualization

    class _DummyMetrics:
        def __call__(self, *a, **kw): return {}
        def reset(self): pass
        def log_epoch_metrics(self, *a, **kw): pass

    raw_ckpt = torch.load(args.checkpoint, map_location='cpu', weights_only=False)
    hp = raw_ckpt['hyper_parameters']
    model = GraphDiscreteFlowModel(
        cfg=hp['cfg'],
        dataset_infos=hp['dataset_infos'],
        train_metrics=_DummyMetrics(),
        sampling_metrics=_DummyMetrics(),
        visualization_tools=hp.get('visualization_tools') or MolecularVisualization(
            remove_h=True, dataset_infos=hp['dataset_infos']
        ),
        extra_features=hp.get('extra_features') or _DummyMetrics(),
        domain_features=hp.get('domain_features') or _DummyMetrics(),
        test_labels=hp.get('test_labels'),
    )
    model.load_state_dict(raw_ckpt['state_dict'], strict=False)
    model.to(device)
    model.eval()
    print("Model loaded.")

    _KNOWN_DECODERS = {
        "moses":   ["C", "N", "S", "O", "F", "Cl", "Br", "H"],
        "zinc":    ["C", "N", "O", "F", "P", "S", "Cl", "Br", "I"],
        "qm9":     ["C", "N", "O", "F"],
        "guacamol": ["C", "N", "O", "F", "P", "S", "Cl", "Br", "I"],
    }
    if args.dataset is not None:
        atom_decoder = _KNOWN_DECODERS[args.dataset]
        print(f"Atom decoder (from --dataset {args.dataset}): {atom_decoder}")
    else:
        atom_decoder = model.dataset_info.atom_decoder
        print(f"Atom decoder (from checkpoint): {atom_decoder}")

    # Build objectives
    # Guidance SMARTS: linker must have at least one reactive attachment per end.
    #   COOH  = [CX3](=[OX1])[OX2H1]   — direct amide bond
    #   CHO   = [CX3H1](=[OX1])         — post-assembly aldehyde, restored to COOH
    #   CH3   = [CX4H3]                 — terminal methyl for C-N alkylation
    # When fix_termini=methyl the CH3 groups are pre-pinned so they match immediately,
    # making the attachment score trivially 2 and letting other objectives steer.
    if args.fix_termini == 'methyl':
        _ATTACH_SMARTS = "[$([CX3;!$([CX3][N,n])](=[OX1])[OX2H1]),$([CX3H1](=[OX1])),$([CX4H3])]"
    elif args.fix_termini == 'amine':
        # Linker presents N termini; warhead/E3 provide COOH electrophiles.
        _ATTACH_SMARTS = "[NX3;!H0;!$(NC=O);!$(NS=O);!$(N=O)]"
    elif args.fix_termini == 'amino_acid':
        # Linker presents one NH2 (warhead end) + one COOH (VHL end).
        # Match either — assembly score requires ≥2 hits (one of each type).
        _ATTACH_SMARTS = "[$([NX3H2]),$([CX3](=[OX1])[OX2H1])]"
    else:
        _ATTACH_SMARTS = "[$([CX3;!$([CX3][N,n])](=[OX1])[OX2H1]),$([CX3H1](=[OX1]))]"
    analyzer    = PROTACAnalyzer(WARHEAD_MOL, E3_MOL,
                                 _ATTACH_SMARTS, _ATTACH_SMARTS,
                                 sa_score_max=5.0)
    builder     = _ProtacBuilder()
    classifier  = PROTACClassifier(analyzer, builder, atom_decoder)
    classifier.fix_termini  = args.fix_termini
    conn_scorer = PROTACLinkerConnectivityScorer(atom_decoder)
    if args.fix_termini:
        # Anchor nodes for the connectivity scorer: halves its soft score when the
        # two termini end up in different fragments, so guidance is steered toward
        # bridging the warhead/E3 anchors specifically.
        conn_scorer.anchor_nodes = DeFoGGenerator._TERMINUS_CONFIGS[args.fix_termini]['anchor_nodes']
    sa_scorer   = PROTACLinkerSAScorer(atom_decoder, sa_threshold=4.5)

    if getattr(args, 'fast_scoring', False):
        if args.fix_termini not in ('cooh', 'amine', 'amino_acid'):
            raise ValueError("--fast_scoring requires --fix_termini cooh, amine, or amino_acid")
        # Classifier dispatch priority:
        #   1. Explicit --clf_path on CLI (override)
        #   2. protac_linker_clf_<protac_module>.pkl  (target-specific, if present)
        #   3. protac_linker_clf_<termini>.pkl  (termini-specific fallback)
        #   4. protac_linker_clf.pkl  (generic default)
        _termini_clf = {
            'amine':      "protac_linker_clf_amine.pkl",
            'amino_acid': "protac_linker_clf_amino_acid.pkl",
        }.get(args.fix_termini, "protac_linker_clf.pkl")
        _clf_dir     = Path(__file__).resolve().parents[1] / "src" / "classifier"
        _module_clf  = f"protac_linker_clf_{args.protac_module}.pkl"
        _module_path = _clf_dir / _module_clf
        if getattr(args, 'clf_path', None):
            _clf_path = args.clf_path
        elif _module_path.exists():
            _clf_path = str(_module_path)
            print(f"[fast_scoring] using target-specific classifier: {_module_clf}")
        else:
            _clf_path = str(_clf_dir / _termini_clf)
            print(f"[fast_scoring] using termini classifier: {_termini_clf}")
        chem_scorer = PROTACLinkerChemScorer(atom_decoder, WARHEAD_MOL, E3_MOL)
        # classifier is the hard validity gate, not in the soft sum.
        # PreassembledPROTACScorer was dropped after the ablation showed it
        # hurts AsmVal/DM when weighted up.
        OBJECTIVE_FUNCTIONS = [sa_scorer, chem_scorer, conn_scorer]
        print("Scoring mode: FAST (SA + linker chem + connectivity; PROTACClassifier as hard gate)")
    else:
        if not getattr(args, 'model_dir', None):
            raise ValueError("--model_dir is required for full scoring mode (or use --fast_scoring)")
        pk_scorer   = EGFRPROPKScorer(builder, args.model_dir, atom_decoder)
        pk_scorer.fix_termini   = args.fix_termini
        prop_scorer = PROTACRDKitScorer(builder, atom_decoder)
        prop_scorer.fix_termini = args.fix_termini
        OBJECTIVE_FUNCTIONS = [pk_scorer, prop_scorer, sa_scorer, conn_scorer]
        print("Scoring mode: FULL (PROPK + RDKit + SA + connectivity; PROTACClassifier as hard gate)")

    tracker   = GenerationTracker(output_dir=args.output_dir)
    generator = DeFoGGenerator(model, atom_decoder, OBJECTIVE_FUNCTIONS, args,
                                validity_fn=classifier, tracker=tracker)

    all_molecules = []

    for batch_idx in range(args.num_batches):
        print(f"\n{'=' * 60}")
        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
        print(f"{'=' * 60}")

        molecule_list = generator.generate()
        final_scores  = generator._get_scores(molecule_list)

        print(f"\n--- Batch {batch_idx + 1} Results ---")
        print(f"Validity scores: {final_scores[0].cpu().numpy()}")

        all_molecules.extend(molecule_list)

        linker_mols = []
        for atom_types, edge_types in molecule_list:
            try:
                mol = _graph_to_mol(atom_types, edge_types, atom_decoder)
                Chem.SanitizeMol(mol)
                linker_mols.append(mol)
            except Exception:
                linker_mols.append(Chem.RWMol())
        detailed = analyzer.get_detailed_scores(linker_mols)
        tracker.log_final(molecule_list, detailed)

    tracker.print_final_summary()
    tracker.save(filename_prefix="defog_local_edit")

    # Strict final validation — attempt actual assembly via builder (NH2 or reactive NH)
    # This is the ground truth: if builder.assemble() succeeds, the linker is genuinely
    # assembly-ready regardless of which N type was used.
    strict_pass = 0
    print("\n=== Strict assembly-ready check (actual RDKit assembly attempt) ===")
    with open(args.output_file, "w") as f:
        for atom_types, edge_types in all_molecules:
            try:
                mol = _graph_to_mol(atom_types, edge_types, atom_decoder)
                Chem.SanitizeMol(mol)
                smi = Chem.MolToSmiles(mol)
                # Reject disconnected molecules — a valid linker must be a single fragment
                if "." in smi:
                    f.write(f"DISCONNECTED {smi}\n")
                    continue
                try:
                    amide_only_strict = (args.fix_termini == 'cooh')
                    builder.assemble(mol, return_all=False, amide_only=amide_only_strict)
                    tag = "AssemblyReady"
                    strict_pass += 1
                except Exception:
                    tag = "Valid"
                f.write(f"{tag} {smi}\n")
            except Exception:
                f.write("INVALID\n")
    print(f"Assembly-ready (strict): {strict_pass}/{len(all_molecules)}")

    print(f"\nSMILES saved to: {args.output_file}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="DeFoG Guided PROTAC Linker Generation")

    # Model
    parser.add_argument("--checkpoint", type=str, required=True,
                        help="Path to DeFoG model checkpoint (.ckpt)")
    parser.add_argument("--model_dir", type=str, default=None,
                        help="Path to EGFR-PROPK models directory (required unless --fast_scoring)")
    parser.add_argument("--clf_path", type=str, default=None,
                        help="Path to protac_linker_clf.pkl (default: smiles/protac_linker_clf.pkl)")
    parser.add_argument("--protac_module", type=str, default="ErlotinibVHL",
                        help="Builder module name under MoCA/smiles/protac/ "
                             "(e.g. ErlotinibVHL, ARCRBN). Must export WARHEAD_MOL, E3_MOL, Builder.")
    parser.add_argument("--dataset", type=str, default=None,
                        choices=["moses", "zinc", "qm9", "guacamol"],
                        help="Override atom decoder (default: read from checkpoint)")

    # Generation
    parser.add_argument("--num_samples",  type=int, default=10)
    parser.add_argument("--num_batches",  type=int, default=5)
    parser.add_argument("--num_steps",    type=int, default=None,
                        help="Override denoising steps (default: read from checkpoint, typically 1000)")
    parser.add_argument("--num_nodes",    type=int, default=None,
                        help="Fixed number of linker atoms (None = sample from node_dist)")
    parser.add_argument("--min_nodes",    type=int, default=10)
    parser.add_argument("--max_nodes",    type=int, default=25)

    # Warmup: run DeFoG's own denoising kernel for this many steps before guided local edits.
    # Replaces the first N steps of the loop with proper flow-matching transitions so
    # guidance starts from structured molecules rather than pure limit-distribution noise.
    # Good default: 40–60% of --num_steps (e.g., 400 for a 1000-step run).
    parser.add_argument("--warmup_steps", type=int, default=0,
                        help="DeFoG denoising steps before guided local-edit loop (default: 0 = disabled)")

    # Guidance — one graph position per step, mirrors CupOfMoca's per-token approach
    parser.add_argument("--guidance_method", type=str, default="exp_tilt",
                        choices=["bok", "exp_tilt", "hybrid"],
                        help="exp_tilt (default): use DeFoG model logits as base prob, tilt by score (preserves "
                             "learned chemistry like ring closures). bok: pure score-driven brute force "
                             "(ignores model — empirically kills ring formation). hybrid: bok→exp_tilt.")
    parser.add_argument("--beta_min",        type=float, default=1.0,
                        help="Initial guidance strength β (annealed from beta_min → beta_max)")
    parser.add_argument("--beta_max",        type=float, default=20.0,
                        help="Final guidance strength β")
    parser.add_argument("--top_k",           type=int,   default=9,
                        help="Atom-type candidates for exp_tilt (default: all types)")
    parser.add_argument("--use_acceptance",  action="store_true",
                        help="Reject exp_tilt proposals that lower the score (mirrors CupOfMoca)")
    parser.add_argument("--score_aggregation", type=str, default="weighted_sum",
                        choices=["weighted_sum", "weighted_min"])
    parser.add_argument("--weights",         type=float, nargs="+", default=None,
                        help="Objective weights e.g. 1.0 0.5 0.5")

    # Position-sampling strategy for the per-step position pick
    parser.add_argument("--position_sampling", type=str, default="weighted",
                        choices=["weighted", "topk"],
                        help="weighted (default): multinomial across all eligible positions "
                             "weighted by 1−p_θ(current_type) — long-tail exploration. "
                             "topk: zero out all but the top-K most-disagreed positions, then "
                             "sample within those — focused exploitation.")
    parser.add_argument("--top_positions", type=int, default=50,
                        help="K for --position_sampling topk (default 50). "
                             "Smaller = more focused; larger = more exploration. "
                             "Ignored when --position_sampling weighted.")

    # Terminus-conditioned generation
    # Pre-pins reactive terminus groups so the attachment score is satisfied immediately
    # and guidance concentrates on connecting the bridge and optimising other objectives.
    #   cooh   — two COOH groups at nodes 0-5 (amide-bond assembly path)
    #   methyl — CH3(0)–CH2(1)…CH2(2)–CH3(3) scaffold (C-N alkylation path, 77% of PROTAC-DB)
    parser.add_argument("--fix_termini", default=None,
                        choices=["cooh", "methyl", "amine", "amino_acid"],
                        help="Pin terminus groups: 'cooh' (nodes 0-5, amide path), "
                             "'methyl' (nodes 0-3, C-N alkylation path), "
                             "'amine' (nodes 0-1, N termini — warhead/E3 provide COOH), "
                             "'amino_acid' (nodes 0-3, asymmetric N+COOH — for BRD4-VHL)")

    # Initial generation
    parser.add_argument("--require_valid_x0", action="store_true")
    parser.add_argument("--max_x0_retries",   type=int, default=10)

    # Fast scoring: skip full PROTAC assembly at every step.
    # Pre-computes warhead + E3 ligase property contributions (MW, logP, RotB) once;
    # scores each candidate linker by estimating full PROTAC properties via offsets.
    # Eliminates PROTACAnalyzer assembly calls (~10-50x faster per scoring step).
    # Only valid with --fix_termini cooh (COOH termini guaranteed, so assembly check is moot).
    # Reverts to full assembly for final validity check after generation.
    parser.add_argument("--fast_scoring", action="store_true",
                        help="Use PreassembledPROTACScorer: skip assembly, estimate PROTAC "
                             "properties via warhead+E3 offsets. Supports cooh, amine, and "
                             "amino_acid (each has its own dedicated classifier).")

    # Output
    parser.add_argument("--output_file", type=str, default="./protac_linkers.smi")
    parser.add_argument("--output_dir",  type=str, default="./generation_logs")
    parser.add_argument("--seed",        type=int, default=None)

    args = parser.parse_args()

    if args.fix_termini:
        cfg = DeFoGGenerator._TERMINUS_CONFIGS[args.fix_termini]
        if args.min_nodes < cfg['min_total_nodes']:
            args.min_nodes = cfg['min_total_nodes']
            print(f"--fix_termini {args.fix_termini} requires at least {cfg['min_total_nodes']} nodes; "
                  f"setting --min_nodes {cfg['min_total_nodes']}")

    if args.seed is not None:
        random.seed(args.seed)
        np.random.seed(args.seed)
        torch.manual_seed(args.seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed_all(args.seed)

    main(args)