Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,17 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.DS_Store
+.idea/
+.vscode/
+.env
+generation_logs/
+results/
+outputs/
+*.log
+
+# Hugging Face and PyTorch downloads should remain outside this repository.
+.cache/
+

CITATION.cff

@@ -0,0 +1,30 @@
+cff-version: 1.2.0
+message: "If you use this software, please cite the CupOFMoCA paper."
+title: "CupOFMoCA: Coupled Objective-Guided Discrete Flows for Molecular Conjugate Assembly"
+type: software
+authors:
+  - family-names: Zhang
+    given-names: Ruoxi Rosie
+  - family-names: Li
+    given-names: Ziang
+  - family-names: Gu
+    given-names: Jiatao
+  - family-names: Chatterjee
+    given-names: Pranam
+license: Apache-2.0
+repository-code: "https://huggingface.co/ChatterjeeLab/CupOFMoCA"
+preferred-citation:
+  type: conference-paper
+  title: "CupOFMoCA: Coupled Objective-Guided Discrete Flows for Molecular Conjugate Assembly"
+  authors:
+    - family-names: Zhang
+      given-names: Ruoxi Rosie
+    - family-names: Li
+      given-names: Ziang
+    - family-names: Gu
+      given-names: Jiatao
+    - family-names: Chatterjee
+      given-names: Pranam
+  collection-title: "Advances in Neural Information Processing Systems"
+  year: 2026
+

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,16 @@
+CupOFMoCA
+Copyright 2026 CupOFMoCA authors
+
+The code in this repository is licensed under Apache License 2.0. Model
+weights and third-party model downloads can be subject to their source model
+licenses; users are responsible for reviewing those terms.
+
+third_party/SMILESToPeptide contains the exact peptide-to-SMILES conversion
+helpers used in the submitted PDC experiments. The implementation is adapted
+from p2smi by Aaron Feller, distributed under the MIT License. Its license is
+included at third_party/SMILESToPeptide/LICENSE-p2smi.
+
+Runtime model dependencies include:
+  - facebook/esm2_t33_650M_UR50D
+  - aaronfeller/PeptideCLM-23M-all
+

README.md

@@ -1,3 +1,226 @@
 ---
-license: cc-by-nc-nd-4.0
+license: apache-2.0
+library_name: pytorch
+tags:
+  - chemistry
+  - drug-discovery
+  - molecular-generation
+  - peptide
+  - discrete-flow
+  - peptide-drug-conjugate
 ---
+
+# CupOFMoCA
+
+Code and PDC model artifacts for **CupOFMoCA: Coupled Objective-Guided Discrete Flows for
+Molecular Conjugate Assembly** (NeurIPS 2026).
+
+CupOFMoCA performs coupled, objective-guided discrete-flow generation while enforcing chemical
+assembly constraints. This release supports the paper's peptide-drug conjugate (PDC) experiments
+across SMILES, natural amino-acid (AA), and unnatural amino-acid (UAA) representations.
+
+## Model overview
+
+CupOFMoCA is a research pipeline for generating peptide-drug conjugate candidates. It couples a
+pretrained discrete-flow generator with property guidance and conjugation-validity checks during
+sampling. This repository contains three generator checkpoints and the property scorers used by the
+released PDC pipeline:
+
+| Checkpoint | Representation | Role |
+|---|---|---|
+| `SMILESReDi_base.ckpt` | peptide SMILES | pretrained SMILES discrete-flow generator |
+| `PepReDi_base.pt` | canonical amino acids | pretrained AA discrete-flow generator |
+| `UAAReDi_base.pt` | extended amino-acid vocabulary | pretrained UAA discrete-flow generator |
+
+The checkpoints are loaded by the repository's custom PyTorch code; this is not a Transformers
+pipeline and does not provide a hosted-inference widget.
+
+### Intended uses
+
+- Reproduce and inspect the paper's PDC generation experiments.
+- Generate computational PDC candidates for research and method development.
+- Extend or compare constrained, multi-objective molecular-generation methods.
+
+The release is not intended for clinical decisions, direct therapeutic use, or as a substitute for
+chemical synthesis, experimental validation, toxicology, or safety assessment.
+
+### Training data
+
+- **AA-PepReDi:** peptides from PepNN and BioLip2, together with 6--49-residue sequences from
+  PPIRef, using an 80:10:10 train/validation/test split.
+- **UAA-PepReDi:** 150,000 synthetic sequences constructed with `p2smi` across combinations of
+  non-canonical and D-amino-acid composition, using an 80:10:10 split.
+- **SMILESReDi:** the released pretrained SMILESReDi weights are used without additional
+  CupOFMoCA fine-tuning.
+
+Architecture, optimization, data construction, and scorer details are reported in Appendix E of
+the paper. Training datasets are not redistributed in this repository.
+
+### Evaluation
+
+The paper evaluates generated PDCs for assembly validity and predicted functional properties across
+ERBB2, GNRHR, and NRP1, using SMILES, AA, and UAA representations. The reported results and
+experimental protocol are in the paper; the submitted evaluation routines are included under
+`smiles/evaluation/`.
+
+### Limitations
+
+- Outputs and objective values are computational predictions, not experimentally validated
+  compounds or guarantees of efficacy, synthesizability, toxicity, or safety.
+- The released configurations cover three receptor targets and fixed target-specific conjugation
+  chemistries; performance outside this scope has not been established.
+- The release contains the paper's PDC stack only. Its PROTAC experiments require separate code and
+  artifacts.
+- Results can vary with hardware, CUDA/PyTorch versions, and revisions of models downloaded at
+  runtime.
+
+## Repository layout
+
+```text
+checkpoints/                 generator weights
+configs/                     paper settings and target sequences
+docs/                        model-artifact details and checksums
+scripts/run_experiment.py    unified experiment runner
+scripts/verify_release.py    release integrity checker
+smiles/
+  generation/               generation entry points
+  evaluation/               metrics and analyzers
+  models/                   ReDi model and data code
+  conjugation/              PDC assembly builders
+  scoring/                  property scorers
+  smiles_tokenizer/         tokenizer code and vocabulary
+third_party/                 peptide-to-SMILES conversion helper
+```
+
+The release currently covers the PDC implementation. The PROTAC experiments described in the paper
+use a separate code and artifact stack that is not included here.
+
+## 1. Environment
+
+A CUDA GPU is strongly recommended. The tested environment used Python 3.9, PyTorch 2.5.1, and
+CUDA 12.1.
+
+```bash
+conda env create -f environment.yml
+conda activate cupofmoca
+```
+
+Alternatively, install a PyTorch build appropriate for your CUDA driver and then install the pinned
+packages:
+
+```bash
+pip install -r requirements.txt
+```
+
+On first use, Transformers downloads `facebook/esm2_t33_650M_UR50D` and
+`aaronfeller/PeptideCLM-23M-all` from Hugging Face. Set `HF_HOME` to a persistent cache if needed.
+
+Verify the bundled artifacts and source paths before running an experiment:
+
+```bash
+python scripts/verify_release.py
+```
+
+## 2. Run PDC generation
+
+Run commands from the repository root. Preview a configured command without loading models:
+
+```bash
+python scripts/run_experiment.py \
+  --representation uaa \
+  --variant cupofmoca \
+  --method guided \
+  --target NRP1 \
+  --dry-run
+```
+
+Run the experiment:
+
+```bash
+python scripts/run_experiment.py \
+  --representation uaa \
+  --variant cupofmoca \
+  --method guided \
+  --target NRP1
+```
+
+Run a small smoke test:
+
+```bash
+python scripts/run_experiment.py \
+  --representation aa --variant baseline --method staged --target ERBB2 \
+  --num-candidates 1
+```
+
+Outputs are written under `outputs/` in separate run directories. Existing outputs are protected
+unless `--force` is supplied. Arguments after `--` are passed to the underlying generator:
+
+```bash
+python scripts/run_experiment.py \
+  --representation smiles --variant cupofmoca --method guided --target GNRHR \
+  -- --top_k 100
+```
+
+### Configurations
+
+| Target | Protein | Conjugation builder | Default SMILES length |
+|---|---|---|---:|
+| ERBB2 | receptor tyrosine-protein kinase erbB-2 | `SSCpt` | 80 |
+| GNRHR | gonadotropin-releasing hormone receptor | `AoaDau` | 80 |
+| NRP1 | neuropilin-1 | `DSPASN38` | 100 |
+
+AA and UAA runs use length 10. Shared settings include 128 optimization steps, top-p 1.0,
+objective weights `[1, 1]`, and acceptance sampling with beta 50 for CupOFMoCA. The machine-readable
+configuration is in `configs/paper.json`; target sequences are in `configs/targets.fasta`.
+
+### Direct entry points
+
+The full argument surfaces remain available under `smiles/generation/`:
+
+```text
+guided.py                   cupofmoca_guided.py
+guided_aa.py                cupofmoca_guided_aa.py
+guided_uaa.py               cupofmoca_guided_uaa.py
+staged.py                   cupofmoca_staged.py
+staged_aa.py                cupofmoca_staged_aa.py
+staged_uaa.py               cupofmoca_staged_uaa.py
+```
+
+Run one directly as a module, for example:
+
+```bash
+python -m smiles.generation.cupofmoca_guided_uaa --help
+```
+
+`smiles/evaluation/run_metrics.py` and `smiles/evaluation/run_metrics_uaa.py` contain the submitted evaluation routines.
+Use `python <file> --help` to inspect their arguments.
+
+## 3. Included checkpoints
+
+The release includes the SMILES, AA, and UAA generator weights together with the binding,
+hemolysis, non-fouling, and solubility scorers needed by the PDC pipeline. Artifact roles, sizes, and
+SHA-256 values are listed in [docs/ARTIFACTS.md](docs/ARTIFACTS.md).
+
+The AA-input affinity predictor is not included because the PDC generation configuration uses the
+bundled SMILES-input scorer (`--affinity_type smiles`).
+
+Large `*.pt` and `*.ckpt` files are configured for Git LFS. The repository is approximately 1.9 GB
+before environment and model-cache downloads.
+
+The checkpoint formats are loaded through PyTorch. As with any pickle-backed PyTorch artifact,
+only load weights obtained from a source you trust.
+
+## 4. Notes
+
+- The paper reports NVIDIA RTX A6000 48 GB GPUs. Lower-memory devices may require smaller batches.
+- Exact bitwise output can vary across CUDA, PyTorch, GPU, and downloaded model revisions.
+- Generated structures are computational hypotheses that require experimental validation. Property
+  models are predictors, not clinical or safety guarantees.
+- The peptide-to-SMILES helper used by the experiments is vendored under `third_party/` with its
+  upstream license notice.
+
+## 5. Citation and license
+
+Citation metadata is provided in `CITATION.cff`. The repository code is released under Apache-2.0;
+see `LICENSE` and `NOTICE`. Third-party components and downloaded model artifacts retain their
+respective terms.

checkpoints/PepReDi_base.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3b6437edee514bd7adb9aacea776f3dd97c59ebc7b4928b390eafdd87eaeb8c9
+size 344474053

checkpoints/SMILESReDi_base.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b2a6049e7ce6c2ef2ab0da0987ee14a5c7e3df737c5c32202631f5c230a81df
+size 1038182356

checkpoints/UAAReDi_base.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa9c25692f308e87befee15cecdd4a24d9e2270e039b12d073023390632b2a8d
+size 344824232

configs/paper.json

@@ -0,0 +1,20 @@
+{
+  "default_seeds": [42, 123, 456],
+  "num_candidates_per_seed": 100,
+  "num_samples_per_batch": 1,
+  "optimization_steps": 128,
+  "top_p": 1.0,
+  "top_k": 200,
+  "cascade_k": 8,
+  "beta_min": 50.0,
+  "beta_max": 50.0,
+  "weights": [1.0, 1.0],
+  "require_valid_x0": true,
+  "max_x0_retries": 10,
+  "targets": {
+    "ERBB2": {"builder": "SSCpt", "smiles_length": 80},
+    "GNRHR": {"builder": "AoaDau", "smiles_length": 80},
+    "NRP1": {"builder": "DSPASN38", "smiles_length": 100}
+  },
+  "representation_lengths": {"aa": 10, "uaa": 10}
+}

configs/targets.fasta

@@ -0,0 +1,6 @@
+>ERBB2
+MELAALCRWGLLLALLPPGAASTQVCTGTDMKLRLPASPETHLDMLRHLYQGCQVVQGNLELTYLPTNASLSFLQDIQEVQGYVLIAHNQVRQVPLQRLRIVRGTQLFEDNYALAVLDNGDPLNNTTPVTGASPGGLRELQLRSLTEILKGGVLIQRNPQLCYQDTILWKDIFHKNNQLALTLIDTNRSRACHPCSPMCKGSRCWGESSEDCQSLTRTVCAGGCARCKGPLPTDCCHEQCAAGCTGPKHSDCLACLHFNHSGICELHCPALVTYNTDTFESMPNPEGRYTFGASCVTACPYNYLSTDVGSCTLVCPLHNQEVTAEDGTQRCEKCSKPCARVCYGLGMEHLREVRAVTSANIQEFAGCKKIFGSLAFLPESFDGDPASNTAPLQPEQLQVFETLEEITGYLYISAWPDSLPDLSVFQNLQVIRGRILHNGAYSLTLQGLGISWLGLRSLRELGSGLALIHHNTHLCFVHTVPWDQLFRNPHQALLHTANRPEDECVGEGLACHQLCARGHCWGPGPTQCVNCSQFLRGQECVEECRVLQGLPREYVNARHCLPCHPECQPQNGSVTCFGPEADQCVACAHYKDPPFCVARCPSGVKPDLSYMPIWKFPDEEGACQPCPINCTHSCVDLDDKGCPAEQRASPLTSIISAVVGILLVVVLGVVFGILIKRRQQKIRKYTMRRLLQETELVEPLTPSGAMPNQAQMRILKETELRKVKVLGSGAFGTVYKGIWIPDGENVKIPVAIKVLRENTSPKANKEILDEAYVMAGVGSPYVSRLLGICLTSTVQLVTQLMPYGCLLDHVRENRGRLGSQDLLNWCMQIAKGMSYLEDVRLVHRDLAARNVLVKSPNHVKITDFGLARLLDIDETEYHADGGKVPIKWMALESILRRRFTHQSDVWSYGVTVWELMTFGAKPYDGIPAREIPDLLEKGERLPQPPICTIDVYMIMVKCWMIDSECRPRFRELVSEFSRMARDPQRFVVIQNEDLGPASPLDSTFYRSLLEDDDMGDLVDAEEYLVPQQGFFCPDPAPGAGGMVHHRHRSSSTRSGGGDLTLGLEPSEEEAPRSPLAPSEGAGSDVFDGDLGMGAAKGLQSLPTHDPSPLQRYSEDPTVPLPSETDGYVAPLTCSPQPEYVNQPDVRPQPPSPREGPLPAARPAGATLERPKTLSPGKNGVVKDVFAFGGAVENPEYLTPQGGAAPQPHPPPAFSPAFDNLYYWDQDPPERGAPPSTFKGTPTAENPEYLGLDVPV
+>GNRHR
+MANSASPEQNQNHCSAINNSIPLMQGNLPTLTLSGKIRVTVTFFLFLLSATFNASFLLKLQKWTQKKEKGKKLSRMKLLLKHLTLANLLETLIVMPLDGMWNITVQWYAGELLCKVLSYLKLFSMYAPAFMMVVISLDRSLAITRPLALKSNSKVGQSMVGLAWILSSVFAGPQLYIFRMIHLADSSGQTKVFSQCVTHCSFSQWWHQAFYNFFTFSCLFIIPLFIMLICNAKIIFTLTRVLHQDPHELQLNQSKNNIPRARLKTLKMTVAFATSFTVCWTPYYVLGIWYWFDPEMLNRLSDPVNHFFFLFAFLNPCFDPLIYGYFSL
+>NRP1
+MERGLPLLCAVLALVLAPAGAFRNDKCGDTIKIESPGYLTSPGYPHSYHPSEKCEWLIQAPDPYQRIMINFNPHFDLEDRDCKYDYVEVFDGENENGHFRGKFCGKIAPPPVVSSGPFLFIKFVSDYETHGAGFSIRYEIFKRGPECSQNYTTPSGVIKSPGFPEKYPNSLECTYIVFVPKMSEIILEFESFDLEPDSNPPGGMFCRYDRLEIWDGFPDVGPHIGRYCGQKTPGRIRSSSGILSMVFYTDSAIAKEGFSANYSVLQSSVSEDFKCMEALGMESGEIHSDQITASSQYSTNWSAERSRLNYPENGWTPGEDSYREWIQVDLGLLRFVTAVGTQGAISKETKKKYYVKTYKIDVSSNGEDWITIKEGNKPVLFQGNTNPTDVVVAVFPKPLITRFVRIKPATWETGISMRFEVYGCKITDYPCSGMLGMVSGLISDSQITSSNQGDRNWMPENIRLVTSRSGWALPPAPHSYINEWLQIDLGEEKIVRGIIIQGGKHRENKVFMRKFKIGYSNNGSDWKMIMDDSKRKAKSFEGNNNYDTPELRTFPALSTRFIRIYPERATHGGLGLRMELLGCEVEAPTAGPTTPNGNLVDECDDDQANCHSGTGDDFQLTGGTTVLATEKPTVIDSTIQSEFPTYGFNCEFGWGSHKTFCHWEHDNHVQLKWSVLTSKTGPIQDHTGDGNFIYSQADENQKGKVARLVSPVVYSQNSAHCMTFWYHMSGSHVGTLRVKLRYQKPEEYDQLVWMAIGHQGDHWKEGRVLLHKSLKLYQVIFEGEIGKGNLGGIAVDDISINNHISQEDCAKPADLDKKNPEIKIDETGSTPGYEGEGEGDKNISRKPGNVLKTLDPILITIIAMSALGVLLGAVCGVVLYCACWHNGMSERNLSALENYNFELVDGVKLKKDKLNTQSTYSEA

docs/ARTIFACTS.md

@@ -0,0 +1,32 @@
+# Model artifacts
+
+The release includes every local checkpoint needed by the submitted PDC generation code. Run
+`python scripts/verify_release.py` to verify byte counts, SHA-256 hashes, and private-path hygiene.
+The machine-readable manifest is in `docs/artifacts.json`.
+
+| Artifact | Purpose | Size | SHA-256 |
+|---|---|---:|---|
+| `checkpoints/SMILESReDi_base.ckpt` | SMILES generator | 1,038,182,356 B | `8b2a6049e7ce6c2ef2ab0da0987ee14a5c7e3df737c5c32202631f5c230a81df` |
+| `checkpoints/PepReDi_base.pt` | natural-AA generator | 344,474,053 B | `3b6437edee514bd7adb9aacea776f3dd97c59ebc7b4928b390eafdd87eaeb8c9` |
+| `checkpoints/UAAReDi_base.pt` | UAA generator | 344,824,232 B | `fa9c25692f308e87befee15cecdd4a24d9e2270e039b12d073023390632b2a8d` |
+| `smiles/scoring/checkpoints/binding-affinity.pt` | SMILES binding scorer | 132,489,414 B | `8daf2ce5ec84045e94bbe41b0a57669287ac937a2a51e97da1af36ae8dc5e16f` |
+| `smiles/scoring/checkpoints/hemolysis-xgboost.json` | hemolysis scorer | 1,932,561 B | `69552199fa7722369211aa0e33adc9850816fb816f555ce9e207c73e0543b0ee` |
+| `smiles/scoring/checkpoints/nonfouling-xgboost.json` | non-fouling scorer | 645,478 B | `08364879c255e236cf0a337505e99e36824f1a72090f5829d3f7870d4ca11644` |
+| `smiles/scoring/checkpoints/solubility-xgboost.json` | solubility scorer | 196,887 B | `3c6a534db85550c947d864aa5cf0b1a2862333bc50eb4b6d5053074cc7d66f02` |
+
+ESM-2 (`facebook/esm2_t33_650M_UR50D`) and PeptideCLM
+(`aaronfeller/PeptideCLM-23M-all`) are downloaded from Hugging Face on first use and therefore are
+not duplicated here.
+
+The optional AA-input binding predictor was not used for the reported PDC runs, which used
+`--affinity_type smiles`, and is intentionally excluded (it would add about 2.7 GB). The public
+wrapper passes the reported SMILES-input setting explicitly.
+
+When publishing with `hf upload`, the Hugging Face Hub client handles large-file storage. If this
+release is mirrored to a Git LFS remote:
+
+```bash
+git lfs install
+git lfs track "*.pt" "*.ckpt"
+git add .gitattributes checkpoints smiles/scoring/checkpoints
+```

docs/artifacts.json

@@ -0,0 +1,38 @@
+[
+  {
+    "path": "checkpoints/SMILESReDi_base.ckpt",
+    "bytes": 1038182356,
+    "sha256": "8b2a6049e7ce6c2ef2ab0da0987ee14a5c7e3df737c5c32202631f5c230a81df"
+  },
+  {
+    "path": "checkpoints/PepReDi_base.pt",
+    "bytes": 344474053,
+    "sha256": "3b6437edee514bd7adb9aacea776f3dd97c59ebc7b4928b390eafdd87eaeb8c9"
+  },
+  {
+    "path": "checkpoints/UAAReDi_base.pt",
+    "bytes": 344824232,
+    "sha256": "fa9c25692f308e87befee15cecdd4a24d9e2270e039b12d073023390632b2a8d"
+  },
+  {
+    "path": "smiles/scoring/checkpoints/binding-affinity.pt",
+    "bytes": 132489414,
+    "sha256": "8daf2ce5ec84045e94bbe41b0a57669287ac937a2a51e97da1af36ae8dc5e16f"
+  },
+  {
+    "path": "smiles/scoring/checkpoints/hemolysis-xgboost.json",
+    "bytes": 1932561,
+    "sha256": "69552199fa7722369211aa0e33adc9850816fb816f555ce9e207c73e0543b0ee"
+  },
+  {
+    "path": "smiles/scoring/checkpoints/nonfouling-xgboost.json",
+    "bytes": 645478,
+    "sha256": "08364879c255e236cf0a337505e99e36824f1a72090f5829d3f7870d4ca11644"
+  },
+  {
+    "path": "smiles/scoring/checkpoints/solubility-xgboost.json",
+    "bytes": 196887,
+    "sha256": "3c6a534db85550c947d864aa5cf0b1a2862333bc50eb4b6d5053074cc7d66f02"
+  }
+]
+

environment.yml

@@ -0,0 +1,28 @@
+name: cupofmoca
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+dependencies:
+  - python=3.9
+  - pip
+  - pytorch=2.5.1
+  - pytorch-cuda=12.1
+  - numpy=1.26.4
+  - scipy=1.13.1
+  - pandas=2.3.2
+  - scikit-learn=1.6.1
+  - xgboost=2.1.4
+  - matplotlib=3.9.4
+  - seaborn=0.13.2
+  - pillow=11.3.0
+  - rdkit=2025.03.6
+  - pip:
+      - lightning==2.5.5
+      - transformers==4.41.2
+      - datasets==4.0.0
+      - fair-esm==2.0.0
+      - SmilesPE==0.0.3
+      - easydict==1.13
+      - tqdm==4.67.1
+

requirements.txt

@@ -0,0 +1,18 @@
+torch==2.5.1
+lightning==2.5.5
+transformers==4.41.2
+datasets==4.0.0
+fair-esm==2.0.0
+rdkit==2025.3.6
+SmilesPE==0.0.3
+numpy==1.26.4
+scipy==1.13.1
+pandas==2.3.2
+scikit-learn==1.6.1
+xgboost==2.1.4
+matplotlib==3.9.4
+seaborn==0.13.2
+pillow==11.3.0
+easydict==1.13
+tqdm==4.67.1
+

scripts/run_experiment.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python3
+"""Run reproducible CupOFMoCA PDC experiments.
+
+This wrapper keeps the submitted generation programs intact while providing one
+portable interface, bundled checkpoint defaults, named targets, and collision-
+safe outputs.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shlex
+import subprocess
+import sys
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+CONFIG = json.loads((ROOT / "configs" / "paper.json").read_text())
+
+CHECKPOINTS = {
+    "smiles": ROOT / "checkpoints" / "SMILESReDi_base.ckpt",
+    "aa": ROOT / "checkpoints" / "PepReDi_base.pt",
+    "uaa": ROOT / "checkpoints" / "UAAReDi_base.pt",
+}
+
+
+def read_fasta(path: Path) -> dict[str, str]:
+    records: dict[str, str] = {}
+    name = None
+    chunks: list[str] = []
+    for raw in path.read_text().splitlines():
+        line = raw.strip()
+        if not line:
+            continue
+        if line.startswith(">"):
+            if name is not None:
+                records[name] = "".join(chunks)
+            name, chunks = line[1:].split()[0], []
+        elif name is None:
+            raise ValueError(f"Sequence before FASTA header in {path}")
+        else:
+            chunks.append(line)
+    if name is not None:
+        records[name] = "".join(chunks)
+    return records
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--representation", choices=("smiles", "aa", "uaa"), required=True)
+    parser.add_argument("--variant", choices=("baseline", "cupofmoca"), default="cupofmoca")
+    parser.add_argument("--method", choices=("guided", "staged"), default="guided")
+    parser.add_argument("--target", choices=tuple(CONFIG["targets"]), required=True)
+    parser.add_argument(
+        "--seeds", type=int, nargs="+", default=CONFIG["default_seeds"],
+        help="Random seeds for generation.",
+    )
+    parser.add_argument("--num-candidates", type=int, default=CONFIG["num_candidates_per_seed"])
+    parser.add_argument("--gen-len", type=int, default=None)
+    parser.add_argument("--checkpoint", type=Path, default=None)
+    parser.add_argument("--output-root", type=Path, default=ROOT / "outputs")
+    parser.add_argument("--python", default=sys.executable)
+    parser.add_argument("--dry-run", action="store_true")
+    parser.add_argument("--force", action="store_true", help="Allow an existing output to be replaced.")
+    parser.add_argument(
+        "extra_args", nargs=argparse.REMAINDER,
+        help="Arguments after '--' are appended to the underlying generator command.",
+    )
+    args = parser.parse_args()
+    if args.num_candidates < 1:
+        parser.error("--num-candidates must be positive")
+    return args
+
+
+def build_command(args: argparse.Namespace, seed: int, target_seq: str, output: Path) -> list[str]:
+    suffix = "" if args.representation == "smiles" else f"_{args.representation}"
+    prefix = "cupofmoca_" if args.variant == "cupofmoca" else ""
+    program = f"smiles.generation.{prefix}{args.method}{suffix}"
+    checkpoint = (args.checkpoint or CHECKPOINTS[args.representation]).resolve()
+    target_cfg = CONFIG["targets"][args.target]
+    gen_len = args.gen_len or (
+        target_cfg["smiles_length"]
+        if args.representation == "smiles"
+        else CONFIG["representation_lengths"][args.representation]
+    )
+
+    cmd = [
+        args.python, "-u", "-m", program,
+        "--checkpoint", str(checkpoint),
+        "--target", target_seq,
+        "--num_samples", str(CONFIG["num_samples_per_batch"]),
+        "--num_batches", str(args.num_candidates),
+        "--gen_len", str(gen_len),
+        "--optimization_steps", str(CONFIG["optimization_steps"]),
+        "--weights", *map(str, CONFIG["weights"]),
+        "--top_p", str(CONFIG["top_p"]),
+        "--seed", str(seed),
+        "--output_file", str(output),
+    ]
+
+    if args.representation in {"aa", "uaa"}:
+        cmd += ["--tokenizer_type", "amino"]
+    if args.representation == "aa":
+        # This is the scorer used by the PDC paper runs. The optional AA-input
+        # scorer is not needed and its multi-GB checkpoint is intentionally omitted.
+        cmd += ["--affinity_type", "smiles"]
+
+    if args.method == "guided":
+        cmd += [
+            "--builder_type", target_cfg["builder"],
+            "--require_valid_x0",
+            "--max_x0_retries", str(CONFIG["max_x0_retries"]),
+            "--output_dir", str(output.parent / "logs"),
+        ]
+
+    if args.variant == "cupofmoca":
+        cmd += [
+            "--use_acceptance",
+            "--beta_min", str(CONFIG["beta_min"]),
+            "--beta_max", str(CONFIG["beta_max"]),
+        ]
+        if args.method == "guided":
+            cmd += ["--guidance_method", "exp_tilt"]
+            if args.representation in {"aa", "uaa"}:
+                cmd += ["--cascade_k", str(CONFIG["cascade_k"])]
+        else:
+            cmd += ["--output_dir", str(output.parent / "logs")]
+
+    extra = list(args.extra_args)
+    if extra and extra[0] == "--":
+        extra = extra[1:]
+    return cmd + extra
+
+
+def main() -> int:
+    args = parse_args()
+    targets = read_fasta(ROOT / "configs" / "targets.fasta")
+    target_seq = targets[args.target]
+    failures = 0
+
+    for seed in args.seeds:
+        run_dir = (
+            args.output_root.resolve()
+            / args.target
+            / args.representation
+            / f"{args.variant}_{args.method}"
+            / f"seed_{seed}"
+        )
+        output = run_dir / "molecules.txt"
+        if output.exists() and not args.force and not args.dry_run:
+            print(f"Refusing to overwrite {output}; pass --force to replace it.", file=sys.stderr)
+            failures += 1
+            continue
+        if output.exists() and args.force:
+            output.unlink()
+        command = build_command(args, seed, target_seq, output)
+        display = " ".join(
+            f"<{args.target} sequence, {len(target_seq)} aa>" if item == target_seq else shlex.quote(item)
+            for item in command
+        )
+        print(display)
+        if args.dry_run:
+            continue
+        run_dir.mkdir(parents=True, exist_ok=True)
+        completed = subprocess.run(command, cwd=ROOT, check=False)
+        if completed.returncode:
+            failures += 1
+            print(f"Seed {seed} failed with exit code {completed.returncode}.", file=sys.stderr)
+
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+

scripts/verify_release.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+"""Verify bundled PDC artifacts and portability-sensitive source files."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import sys
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+MANIFEST = ROOT / "docs" / "artifacts.json"
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for block in iter(lambda: handle.read(8 * 1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def main() -> int:
+    entries = json.loads(MANIFEST.read_text())
+    failures = 0
+    for entry in entries:
+        path = ROOT / entry["path"]
+        if not path.is_file():
+            print(f"MISSING  {entry['path']}")
+            failures += 1
+            continue
+        size_ok = path.stat().st_size == entry["bytes"]
+        hash_ok = sha256(path) == entry["sha256"]
+        state = "OK" if size_ok and hash_ok else "FAILED"
+        print(f"{state:7}  {entry['path']}")
+        failures += int(not (size_ok and hash_ok))
+
+    private_roots = tuple("/" + stem for stem in ("scratch/", "Users/", "home/"))
+    for path in list((ROOT / "smiles").rglob("*.py")) + list((ROOT / "scripts").rglob("*.py")):
+        text = path.read_text(errors="replace")
+        if any(root in text for root in private_roots):
+            print(f"PRIVATE  {path.relative_to(ROOT)}")
+            failures += 1
+
+    print("Release verification passed." if not failures else f"Verification found {failures} problem(s).")
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+

smiles/__init__.py

@@ -0,0 +1 @@
+"""CupOFMoCA peptide-drug conjugate implementation."""

smiles/conjugation/AoaDau.py

@@ -0,0 +1,732 @@
+"""
+AoaDaunorubicin PDC Builder
+
+Conjugation chemistry: Lysine ε-amine + Aoa-Daunorubicin (amide bond formation)
+
+Reference PDC: PDC_00208 (GnRH-III-[8Lys(Dau=Aoa)])
+- Receptor: GNRHR
+- Linker: Aminooxyacetic acid (Aoa)
+- Payload: Daunorubicin
+
+Architecture:
+    Peptide──Lys──NH──C(=O)──CH₂──O──N═C──Daunorubicin
+                  │         │
+              amide bond   Aoa linker    oxime bond (pre-formed)
+"""
+
+from rdkit import Chem
+from rdkit.Chem import RWMol, AllChem
+
+from .pdc_builder import PDC_Builder
+
+
+# ============================================================================
+# Aoa-Daunorubicin PDC Builder
+# ============================================================================
+
+class AoaDaunorubicin_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for Aminooxyacetic acid (Aoa) - Daunorubicin conjugates
+    
+    Conjugation chemistry (amide bond formation):
+    - Peptide site: Lysine ε-amine (side chain -NH2)
+    - Linker-payload site: Carboxylic acid of Aoa (already conjugated to Dau via oxime)
+    - Reaction: Amide bond formation (Lys-NH2 + HOOC-Aoa-Dau → Lys-NH-CO-Aoa-Dau)
+    
+    Architecture:
+        Peptide──Lys──NH──C(=O)──CH₂──O──N═C(CH₃)──Daunorubicin
+                      └── amide ──┘    └── oxime (pre-formed) ──┘
+    
+    Reference: PDC_00208 (GnRH-III-[8Lys(Dau=Aoa)])
+    """
+    
+    # Default Aoa-Daunorubicin SMILES with free carboxylic acid
+    # Structure: HOOC-CH2-O-N=C(CH3)-Daunorubicin
+    # The oxime bond is already formed between Aoa and Dau ketone
+    DEFAULT_AOA_DAU_SMILES = (
+        "OC(=O)CO/N=C(\\C)[C@@]1(O)Cc2c(O)c3c(c(O)c2[C@H](OC2C[C@@H](N)[C@@H](O)[C@@H](C)O2)C1)"
+        "C(=O)c1c(OC)cccc1C3=O"
+    )
+    
+    def __init__(self, aoa_dau_smiles=None):
+        """
+        Initialize Aoa-Daunorubicin builder
+        
+        Args:
+            aoa_dau_smiles: Optional custom Aoa-Dau SMILES. If None, uses default.
+                           Must have free carboxylic acid for conjugation.
+        """
+        if aoa_dau_smiles is None:
+            aoa_dau_smiles = self.DEFAULT_AOA_DAU_SMILES
+        
+        super().__init__(aoa_dau_smiles, component_name="Aoa-Daunorubicin (linker-payload)")
+    
+    def _verify_reactive_group(self):
+        """Verify carboxylic acid group is present for amide bond formation"""
+        # Pattern: Carboxylic acid (COOH)
+        carboxylic_acid_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+        
+        if not self.linker_payload_mol.HasSubstructMatch(carboxylic_acid_pattern):
+            print(" Warning: No carboxylic acid group detected in Aoa-Dau")
+        else:
+            print(" Aoa-Daunorubicin loaded with carboxylic acid reactive group")
+        
+        # Also verify oxime bond is present (Aoa-Dau linkage)
+        oxime_pattern = Chem.MolFromSmarts('[#6]-[#8]-[#7]=[#6]')
+        if self.linker_payload_mol.HasSubstructMatch(oxime_pattern):
+            print(" Oxime bond (Aoa-Dau linkage) detected")
+        else:
+            print(" Warning: No oxime bond detected - check Aoa-Dau structure")
+    
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find all lysine ε-amine groups in peptide
+        
+        Returns list of dicts with site information including atom index
+        """
+        return self.find_peptide_lysines(peptide_smiles)
+    
+    def find_peptide_lysines(self, peptide_smiles):
+        """
+        Find all lysine residues with free ε-amine in peptide
+        
+        Lysine side chain: -CH2-CH2-CH2-CH2-NH2 (CCCCN)
+        
+        Returns:
+            List of dicts with 'atom_idx' (nitrogen index) and metadata
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+        
+        lysines = []
+        
+        # SMARTS pattern for Lysine ε-amine: 4 CH2 groups ending in NH2
+        # This specifically matches side chain, not N-terminus
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        matches = peptide_mol.GetSubstructMatches(lys_pattern)
+        
+        for match in matches:
+            # match indices: 0=CH2, 1=CH2, 2=CH2, 3=CH2, 4=NH2
+            nitrogen_idx = match[4]
+            
+            lysines.append({
+                'atom_idx': nitrogen_idx,
+                'side_chain_atoms': list(match),
+                'type': 'Lys_epsilon_amine'
+            })
+        
+        # Also count N-terminus for informational purposes
+        # N-terminus: NH2 attached to backbone carbon (CH for most AAs, CH2 for Gly)
+        n_terminus_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        n_term_matches = peptide_mol.GetSubstructMatches(n_terminus_pattern)
+        
+        if n_term_matches:
+            print(f"  (Also found {len(n_term_matches)} N-terminus α-amine(s) - excluded by default)")
+        
+        return lysines
+    
+    def count_all_nh2_sites(self, peptide_smiles):
+        """
+        Count all NH2 sites in peptide (both Lys ε-amines and N-terminus)
+        
+        Useful for understanding what sites are available.
+        
+        Returns:
+            Dict with counts: {'lys_epsilon': n, 'n_terminus': n, 'other': n, 'total': n}
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+        
+        counts = {'lys_epsilon': 0, 'n_terminus': 0, 'other': 0, 'total': 0}
+        
+        # Lys ε-amine: CH2-CH2-CH2-CH2-NH2 (4-carbon side chain)
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        counts['lys_epsilon'] = len(peptide_mol.GetSubstructMatches(lys_pattern))
+        
+        # N-terminus: NH2 attached to backbone carbon (either CH for most AAs, or CH2 for Gly)
+        # Pattern: NH2-C-C(=O) where C is part of backbone
+        n_term_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        counts['n_terminus'] = len(peptide_mol.GetSubstructMatches(n_term_pattern))
+        
+        # Total NH2 (any primary amine)
+        all_nh2_pattern = Chem.MolFromSmarts('[NH2]')
+        total_nh2 = len(peptide_mol.GetSubstructMatches(all_nh2_pattern))
+        
+        counts['other'] = total_nh2 - counts['lys_epsilon'] - counts['n_terminus']
+        counts['total'] = total_nh2
+        
+        return counts
+    
+    def _find_carboxylic_acid_atoms(self, mol):
+        """
+        Find the carboxylic acid carbon (the reactive site for amide formation)
+        
+        Returns: (carbon_idx, oxygen_idx) of the C(=O)OH group
+        """
+        # Pattern: C(=O)OH
+        cooh_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+        matches = mol.GetSubstructMatches(cooh_pattern)
+        
+        if not matches:
+            return None, None
+        
+        # Return first match: (carbon_idx, hydroxyl_oxygen_idx)
+        match = matches[0]
+        return match[0], match[1]
+    
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """
+        Perform amide bond conjugation between NH2 groups and Aoa carboxylic acid
+        
+        Args:
+            peptide_mol: RDKit Mol object
+            site_info: Not used (kept for compatibility)
+            exclude_n_terminus: If True, skip N-terminus conjugation
+        
+        Returns: List of all possible conjugate SMILES
+        """
+        return self._manual_conjugation(peptide_mol, site_info, exclude_n_terminus)
+    
+    def _manual_conjugation(self, peptide_mol, lysine_info, exclude_n_terminus=True):
+        """
+        Create ALL possible conjugates by forming amide bond at NH2 sites
+        
+        Chemistry:
+        - Any-NH2 + HOOC-Aoa-Dau → Any-NH-C(=O)-Aoa-Dau + H2O
+        
+        Args:
+            peptide_mol: RDKit Mol object of peptide
+            lysine_info: Not used (kept for interface compatibility)
+            exclude_n_terminus: If True, only conjugate at Lys ε-amines (4-carbon chain),
+                               not at N-terminus α-amine. Default: True
+        
+        Returns:
+            List of all possible conjugate SMILES
+        """
+        if exclude_n_terminus:
+            # SMARTS that matches only Lys ε-amine: requires 4 CH2 chain before NH2
+            # This includes the connection to the peptide backbone (*) to preserve the full peptide
+            # Pattern: *-CH2-CH2-CH2-CH2-NH2 (where * is the connection to the peptide)
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:5][CH2][CH2][CH2][CH2:4][NH2:1].[C:2](=[O:3])[OH]>>[*:5][CH2][CH2][CH2][CH2:4][NH:1][C:2](=[O:3])'
+            )
+        else:
+            # Match ANY primary amine (includes N-terminus)
+            # Include wildcard connection to preserve the rest of the molecule
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:4][NH2:1].[C:2](=[O:3])[OH]>>[*:4][NH:1][C:2](=[O:3])'
+            )
+        
+        # Run the reaction - returns all possible products
+        products = amide_rxn.RunReactants((peptide_mol, self.linker_payload_mol))
+        
+        if not products:
+            raise ValueError("Amide bond formation reaction failed - no products generated")
+        
+        # Collect ALL valid conjugates
+        conjugates = []
+        seen_smiles = set()  # Deduplicate
+        
+        for product_tuple in products:
+            try:
+                conjugate_mol = product_tuple[0]
+                Chem.SanitizeMol(conjugate_mol)
+                conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+                
+                # Deduplicate (same SMILES = same molecule)
+                if conjugate_smiles not in seen_smiles:
+                    seen_smiles.add(conjugate_smiles)
+                    conjugates.append(conjugate_smiles)
+            except Exception as e:
+                # Skip invalid products
+                continue
+        
+        if not conjugates:
+            raise ValueError("All products failed sanitization")
+        
+        return conjugates
+    
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Aoa-Daunorubicin to peptide via amide bond at NH2 groups
+        
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain NH2)
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines,
+                               not at N-terminus. Set to False to include N-terminus.
+            return_all: If True, return all possible conjugates. If False, return first only.
+        
+        Returns:
+            If return_all=True: List of all possible conjugate SMILES
+            If return_all=False: Single conjugate SMILES (first one)
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+        
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Lys ε-amine site(s)")
+        
+        # Get all conjugates
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+        
+        site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2 (including N-terminus)"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [{site_type}]")
+        
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+    
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides
+        
+        Args:
+            peptide_smiles_list: List of peptide SMILES strings
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates per peptide
+        
+        Returns:
+            Dict with:
+                'conjugates': List of lists (if return_all=True) or list of SMILES
+                'successful': List of successful peptide indices
+                'failed': List of (index, error) tuples
+                'total_conjugates': Total number of conjugates generated
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+        
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles, 
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+        
+        print(f"\n=== Summary ===")
+        print(f"Peptides: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, Failed: {len(results['failed'])}")
+        print(f"Total conjugates generated: {results['total_conjugates']}")
+        
+        return results
+    
+    def flatten_batch_results(self, results):
+        """
+        Flatten batch conjugation results into a single list of all conjugates
+        
+        Args:
+            results: Dict returned by batch_conjugate()
+        
+        Returns:
+            List of all conjugate SMILES (flattened)
+        """
+        all_conjugates = []
+        for conjugate_list in results['conjugates']:
+            if conjugate_list is not None:
+                if isinstance(conjugate_list, list):
+                    all_conjugates.extend(conjugate_list)
+                else:
+                    all_conjugates.append(conjugate_list)
+        return all_conjugates
+    
+    # =========================================================================
+    # Attachment Site Compatibility Checks (Lys-amine + Aoa-COOH specific)
+    # =========================================================================
+    
+    def has_aoa_amide_conjugation(self, mol):
+        """
+        Check if molecule has the Aoa-amide conjugation pattern
+        
+        Pattern: N-C(=O)-CH2-O-N=C (amide-CH2-O-oxime linkage)
+        This is the signature of successful Lys-Aoa conjugation.
+        
+        Returns:
+            bool: True if Aoa conjugation pattern is present
+        """
+        # Specific Aoa-conjugation pattern
+        # Looking for: N-C(=O)-C-O-N=C (amide-CH2-O-oxime linkage)
+        aoa_conjugation_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH2][OX2][NX2]=[#6]')
+        if aoa_conjugation_pattern:
+            return mol.HasSubstructMatch(aoa_conjugation_pattern)
+        return False
+    
+    def has_oxime_bond(self, mol):
+        """
+        Check if molecule has oxime bond (Aoa-Dau linkage)
+        
+        Pattern: C-O-N=C
+        This confirms the Aoa-Dau linker-payload is intact.
+        """
+        oxime_pattern = Chem.MolFromSmarts('[#6]-[#8]-[#7]=[#6]')
+        if oxime_pattern:
+            return mol.HasSubstructMatch(oxime_pattern)
+        return False
+    
+    def has_unreacted_aoa_carboxylic_acid(self, mol):
+        """
+        Check if molecule has unreacted Aoa carboxylic acid
+        
+        The specific Aoa pattern: HOOC-CH2-O-N=C
+        If this is present, the Aoa hasn't been conjugated to the peptide.
+        
+        Note: We specifically check for the Aoa-COOH pattern, not just any COOH,
+        because peptides naturally have carboxylic acids (C-terminus, Asp, Glu).
+        
+        Returns:
+            bool: True if unreacted Aoa-COOH is present (bad)
+        """
+        # Specific pattern: Free COOH attached to CH2-O-N=C (Aoa pattern)
+        unreacted_aoa_pattern = Chem.MolFromSmarts('[CX3](=[OX1])[OX2H1][CH2][OX2][NX2]=[#6]')
+        if unreacted_aoa_pattern:
+            return mol.HasSubstructMatch(unreacted_aoa_pattern)
+        return False
+    
+    def find_amide_conjugation_sites(self, mol):
+        """
+        Find amide bonds that are part of the Aoa conjugation
+        
+        Specifically looks for: N-C(=O)-CH2-O-N=C pattern
+        
+        Returns:
+            List of (nitrogen_idx, carbonyl_carbon_idx) tuples
+        """
+        aoa_conjugation_pattern = Chem.MolFromSmarts('[NX3:1][CX3:2](=[OX1])[CH2][OX2][NX2]=[#6]')
+        conjugation_sites = []
+        
+        if aoa_conjugation_pattern:
+            matches = mol.GetSubstructMatches(aoa_conjugation_pattern)
+            for match in matches:
+                nitrogen_idx = match[0]
+                carbonyl_c_idx = match[1]
+                conjugation_sites.append((nitrogen_idx, carbonyl_c_idx))
+        
+        return conjugation_sites
+    
+    def check_valency_at_conjugation_site(self, mol, nitrogen_idx):
+        """
+        Check valency at amide nitrogen conjugation site
+        
+        For Lys-Aoa conjugation, the nitrogen should be:
+        - Amide nitrogen (bonded to carbonyl carbon)
+        - Secondary amine (2 bonds to heavy atoms)
+        - No charge
+        
+        Args:
+            mol: RDKit Mol object
+            nitrogen_idx: Index of the nitrogen atom to check
+        
+        Returns:
+            bool: True if valency is valid
+        """
+        atom = mol.GetAtomWithIdx(nitrogen_idx)
+        
+        if atom.GetSymbol() != 'N':
+            return True  # Not nitrogen, skip
+        
+        degree = atom.GetDegree()  # Bonds to heavy atoms
+        total_hs = atom.GetTotalNumHs()
+        formal_charge = atom.GetFormalCharge()
+        
+        # Get neighbor information
+        neighbors = [mol.GetAtomWithIdx(n.GetIdx()) for n in atom.GetNeighbors()]
+        
+        # Count carbonyl carbons (C=O) among neighbors
+        carbonyl_neighbors = 0
+        for neighbor in neighbors:
+            if neighbor.GetSymbol() == 'C':
+                # Check if this carbon has a double bond to oxygen
+                for n2 in neighbor.GetNeighbors():
+                    if n2.GetSymbol() == 'O':
+                        bond = mol.GetBondBetweenAtoms(neighbor.GetIdx(), n2.GetIdx())
+                        if bond and bond.GetBondType() == Chem.BondType.DOUBLE:
+                            carbonyl_neighbors += 1
+                            break
+        
+        # ========== Valid Amide Nitrogen States ==========
+        
+        # 1. Secondary amide: N bonded to 2 heavy atoms (one is carbonyl C), 1 H
+        #    R-NH-C(=O)-R' 
+        if degree == 2 and total_hs == 1 and formal_charge == 0 and carbonyl_neighbors >= 1:
+            return True
+        
+        # 2. Tertiary amide: N bonded to 3 heavy atoms (one is carbonyl C), 0 H
+        #    R-N(R')-C(=O)-R'' 
+        if degree == 3 and total_hs == 0 and formal_charge == 0 and carbonyl_neighbors >= 1:
+            return True
+        
+        # 3. Aromatic N in peptide backbone is okay
+        if atom.GetIsAromatic():
+            return True
+        
+        # ========== Invalid States ==========
+        
+        # Quaternary nitrogen with positive charge (not typical for amide)
+        if degree >= 4 and formal_charge > 0:
+            return False
+        
+        # N with no carbonyl neighbor but expected to be amide
+        if carbonyl_neighbors == 0:
+            # This might be a free amine - need context to determine if that's okay
+            # For Aoa conjugation, we expect amide formation
+            pass  # Allow for now, other checks will catch issues
+        
+        return True
+    
+    def check_attachment_site_compatibility(self, mol):
+        """
+        Check if Lys-Aoa amide attachment site chemistry is valid
+        
+        For Aoa-Daunorubicin conjugation, validates:
+        1. Has Aoa-amide conjugation pattern (N-C(=O)-CH2-O-N=C)
+        2. Has oxime bond intact (Aoa-Dau linkage)
+        3. No unreacted Aoa-COOH (conjugation must have happened)
+        4. Valid valency at conjugation nitrogen sites
+        
+        Args:
+            mol: RDKit Mol object of the conjugate
+        
+        Returns:
+            bool: True if attachment site is valid, False otherwise
+        """
+        # ===== CRITICAL CHECK: Must have Aoa-amide conjugation pattern =====
+        if not self.has_aoa_amide_conjugation(mol):
+            return False
+        
+        # ===== Check oxime bond is intact =====
+        if not self.has_oxime_bond(mol):
+            return False
+        
+        # ===== Check no unreacted Aoa-COOH =====
+        # If the original Aoa carboxylic acid pattern is still present,
+        # conjugation didn't happen properly
+        if self.has_unreacted_aoa_carboxylic_acid(mol):
+            return False
+        
+        # ===== Check valency at conjugation sites =====
+        conjugation_sites = self.find_amide_conjugation_sites(mol)
+        
+        for nitrogen_idx, carbonyl_c_idx in conjugation_sites:
+            if not self.check_valency_at_conjugation_site(mol, nitrogen_idx):
+                return False
+        
+        return True
+    
+    def _basic_validation(self, mol):
+        """
+        Basic validation checks for Aoa-Daunorubicin conjugates
+        
+        Includes attachment site compatibility check with valency validation.
+        
+        Checks:
+        1. Attachment site compatibility (Aoa-amide pattern, oxime, valency)
+        2. Daunorubicin core present
+        3. Aminosugar present (daunosamine)
+        
+        Note: We do NOT check for absence of COOH because peptides naturally
+        have carboxylic acids (C-terminus, Asp, Glu residues).
+        """
+        checks = {}
+        
+        # Check 1: Has amide bond (general check)
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[#6]')
+        checks['has_amide_bond'] = mol.HasSubstructMatch(amide_pattern)
+        
+        # Check 2: Oxime bond present (Aoa-Dau linkage: C-O-N=C)
+        checks['has_oxime_bond'] = self.has_oxime_bond(mol)
+        
+        # Check 3: Specific Aoa-conjugation pattern
+        checks['has_aoa_conjugation'] = self.has_aoa_amide_conjugation(mol)
+        
+        # Check 4: No unreacted Aoa-COOH
+        checks['no_unreacted_aoa'] = not self.has_unreacted_aoa_carboxylic_acid(mol)
+        
+        # Check 5: Full attachment site compatibility (includes valency)
+        checks['attachment_site_valid'] = self.check_attachment_site_compatibility(mol)
+        
+        # Check 6: Daunorubicin core present (anthraquinone with specific substitution)
+        anthraquinone = Chem.MolFromSmarts('[#6]1[#6][#6][#6]2[#6](=[#8])[#6]3[#6][#6][#6][#6][#6]3[#6](=[#8])[#6]2[#6]1')
+        checks['has_daunorubicin_core'] = mol.HasSubstructMatch(anthraquinone)
+        
+        # Check 7: Has aminosugar (daunosamine - 6-membered ring with N)
+        aminosugar = Chem.MolFromSmarts('[OX2]1[#6][#6][#6]([NX3])[#6][#6]1')
+        checks['has_aminosugar'] = mol.HasSubstructMatch(aminosugar)
+        
+        # Find and report conjugation sites for diagnostics
+        conjugation_sites = self.find_amide_conjugation_sites(mol)
+        checks['conjugation_sites'] = conjugation_sites
+        
+        # Overall pass/fail - attachment_site_valid is the key check
+        checks['_passed'] = checks['attachment_site_valid']
+        
+        return checks
+    
+    def validate_conjugate(self, conjugate_smiles):
+        """
+        Validate that the conjugate is chemically reasonable
+        
+        Returns:
+            (passed: bool, checks: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+        
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+        
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+        
+        return passed, checks
+
+
+# ============================================================================
+# Main - Test Examples  
+# ============================================================================
+
+if __name__ == "__main__":
+    
+    print("=" * 70)
+    print("Aoa-Daunorubicin PDC Builder - Test")
+    print("=" * 70 + "\n")
+    
+    # Initialize builder
+    builder = AoaDaunorubicin_PDC_Builder()
+    
+    # Reference peptide: GnRH-III (PEP00168)
+    gnrh_iii_smiles = (
+        "[H]NCCCC[C@H](NC(=O)[C@H](Cc1cn([H])c2ccccc12)NC(=O)[C@H](CC(=O)O)"
+        "NC(=O)[C@H](Cc1cn([H])cn1)NC(=O)[C@H](CO)NC(=O)[C@H](Cc1cn([H])c2ccccc12)"
+        "NC(=O)[C@H](Cc1cn([H])cn1)NC(=O)[C@@H]1CCC(=O)N1)C(=O)N1CCC[C@H]1C(=O)NCC(N)=O"
+    )
+    
+    # Test peptides
+    test_peptides = [
+        # Simple test: Gly-Lys-Gly (has N-terminus NH2 + Lys ε-amine = 2 NH2 sites)
+        "NCC(=O)N[C@@H](CCCCN)C(=O)NCC(=O)O",
+        # Lys-Arg (has N-terminus NH2 + Lys ε-amine = 2 NH2 sites)  
+        "N[C@@H](CCCCN)C(=O)N[C@@H](CCCNC(=N)N)C(=O)O",
+        # GnRH-III
+        gnrh_iii_smiles,
+    ]
+    
+    print("\n" + "=" * 70)
+    print("Testing NH2 Site Detection")
+    print("=" * 70 + "\n")
+    
+    for i, pep in enumerate(test_peptides):
+        name = ["Gly-Lys-Gly", "Lys-Arg", "GnRH-III"][i]
+        counts = builder.count_all_nh2_sites(pep)
+        print(f"{name}:")
+        print(f"  Lys ε-amines: {counts['lys_epsilon']}")
+        print(f"  N-terminus:   {counts['n_terminus']}")
+        print(f"  Other NH2:    {counts['other']}")
+        print(f"  Total NH2:    {counts['total']}")
+        print()
+    
+    print("\n" + "=" * 70)
+    print("Testing: exclude_n_terminus=True (DEFAULT)")
+    print("=" * 70 + "\n")
+    
+    peptide = test_peptides[0]  # Gly-Lys-Gly
+    print(f"Input peptide (Gly-Lys-Gly): {peptide}")
+    
+    try:
+        conjugates = builder.conjugate_to_peptide(peptide, exclude_n_terminus=True)
+        print(f"\n SUCCESS!")
+        print(f"Number of conjugates: {len(conjugates)} (should be 1 - Lys only)")
+        
+        for i, conj in enumerate(conjugates):
+            valid, checks = builder.validate_conjugate(conj)
+            print(f"  Conjugate {i+1}: {'PASSED' if valid else 'FAILED'}")
+            print(conj)
+            
+    except Exception as e:
+        print(f"\n FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    print("\n" + "=" * 70)
+    print("Testing: exclude_n_terminus=False (include N-terminus)")
+    print("=" * 70 + "\n")
+    
+    print(f"Input peptide (Gly-Lys-Gly): {peptide}")
+    
+    try:
+        conjugates = builder.conjugate_to_peptide(peptide, exclude_n_terminus=False)
+        print(f"\n SUCCESS!")
+        print(f"Number of conjugates: {len(conjugates)} (should be 2 - Lys + N-terminus)")
+        
+        for i, conj in enumerate(conjugates):
+            valid, checks = builder.validate_conjugate(conj)
+            print(f"  Conjugate {i+1}: {'PASSED' if valid else 'FAILED'}")
+            print(conj)
+            
+    except Exception as e:
+        print(f"\n FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    print("\n" + "=" * 70)
+    print("Testing GnRH-III (has only 1 Lys at position 8)")
+    print("=" * 70 + "\n")
+    
+    print(f"Input peptide (GnRH-III): {gnrh_iii_smiles[:60]}...")
+    print("\nNote: GnRH-III sequence is Glp-His-Trp-Ser-His-Asp-Trp-Lys-Pro-Gly")
+    print("      Position 4 is Serine (not Lysine), so only 1 Lys at position 8")
+    
+    # First check the NH2 sites
+    counts = builder.count_all_nh2_sites(gnrh_iii_smiles)
+    print(f"\nNH2 site breakdown:")
+    print(f"  Lys ε-amines: {counts['lys_epsilon']} (expected: 1 for Lys8)")
+    print(f"  N-terminus:   {counts['n_terminus']} (expected: 0 - Glp is cyclic)")
+    print(f"  Other:        {counts['other']}")
+    
+    try:
+        conjugates = builder.conjugate_to_peptide(gnrh_iii_smiles, exclude_n_terminus=True)
+        print(f"\n SUCCESS!")
+        print(f"Number of conjugates: {len(conjugates)}")
+        
+        for i, conj in enumerate(conjugates):
+            valid, checks = builder.validate_conjugate(conj)
+            status = "PASSED" if valid else "FAILED"
+            print(f"  Conjugate {i+1}: {status} - {len(conj)} chars")
+            print(conj)
+            
+    except Exception as e:
+        print(f"\n FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    print("\n" + "=" * 70)
+    print("Batch Conjugation Test (exclude_n_terminus=True)")
+    print("=" * 70 + "\n")
+    
+    results = builder.batch_conjugate(test_peptides, exclude_n_terminus=True)
+    
+    print(f"\nResults breakdown:")
+    for i in results['successful']:
+        conjugates = results['conjugates'][i]
+        num_valid = sum(1 for c in conjugates if builder.validate_conjugate(c)[0])
+        name = ["Gly-Lys-Gly", "Lys-Arg", "GnRH-III"][i]
+        print(f"  {name}: {len(conjugates)} conjugate(s), {num_valid} valid")

smiles/conjugation/CptFactory.py

@@ -0,0 +1,312 @@
+"""
+Camptothecin PDC Factory - Drug-level builder with linker sampling
+
+Instead of specifying a particular linker-drug combo (e.g. SSCpt, DSPACpt),
+the user specifies only the drug (Camptothecin) and a linker is sampled
+according to ERBB2+Cpt entry frequencies from PDCdb.
+
+Linker frequencies (ERBB2 + Camptothecin):
+===========================================
+    Disulfide bond           (SSCpt)     - 7 entries
+    Maleimide-hexanoate      (MalHexCpt) - 5 entries
+    3-Disulfanylpropanoic acid (DSPACpt) - 5 entries
+    2-(Disulfanyl)ethyl carbonate (DSECCpt) - 5 entries
+                                    Total: 22 entries
+
+Sampling probabilities:
+    SSCpt:    7/22 ≈ 0.318
+    MalHexCpt: 5/22 ≈ 0.227
+    DSPACpt:  5/22 ≈ 0.227
+    DSECCpt:  5/22 ≈ 0.227
+
+Usage:
+======
+    # Drug-only mode - linker is sampled per peptide
+    factory = CamptothecinPDC_Factory()
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+
+    # Fixed linker mode - use a specific linker
+    factory = CamptothecinPDC_Factory(linker='SSCpt')
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+
+    # All linkers mode - generate conjugates with every linker
+    factory = CamptothecinPDC_Factory(linker='all')
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+"""
+
+import numpy as np
+import contextlib
+import io
+
+from .SSCpt import DisulfideCamptothecin_PDC_Builder
+from .DSPACpt import DisulfanylpropanoicAcidCpt_PDC_Builder
+from .MalHexCpt import MaleimideHexanoateCpt_PDC_Builder
+from .DSECCpt import DisulfanylethylCarbonateCpt_PDC_Builder
+
+
+# ============================================================================
+# Linker registry: name → (builder_class, PDCdb frequency, description)
+# ============================================================================
+
+ERBB2_CPT_LINKER_REGISTRY = {
+    'SSCpt': {
+        'class': DisulfideCamptothecin_PDC_Builder,
+        'frequency': 7,
+        'linker_id': 'LIN00005',
+        'description': 'Disulfide bond (ester + carbonate variants)',
+        'pdc_refs': ['PDC_00294', 'PDC_00295'],
+    },
+    'MalHexCpt': {
+        'class': MaleimideHexanoateCpt_PDC_Builder,
+        'frequency': 5,
+        'linker_id': 'LIN00031',
+        'description': '4-yl 6-(2,5-dioxo-2,5-dihydro-1H-pyrrol-1-yl)hexanoate',
+        'pdc_refs': ['PDC_00336'],
+    },
+    'DSPACpt': {
+        'class': DisulfanylpropanoicAcidCpt_PDC_Builder,
+        'frequency': 5,
+        'linker_id': 'LIN00012',
+        'description': '3-Disulfanylpropanoic acid',
+        'pdc_refs': ['PDC_00335'],
+    },
+    'DSECCpt': {
+        'class': DisulfanylethylCarbonateCpt_PDC_Builder,
+        'frequency': 5,
+        'linker_id': 'LIN00039',
+        'description': '2-(Disulfanyl)ethyl hydrogen carbonate',
+        'pdc_refs': ['PDC_00337'],
+    },
+}
+
+
+def _compute_linker_probabilities(registry=None):
+    """Compute sampling probabilities from frequency counts."""
+    if registry is None:
+        registry = ERBB2_CPT_LINKER_REGISTRY
+    names = list(registry.keys())
+    freqs = np.array([registry[n]['frequency'] for n in names], dtype=np.float64)
+    probs = freqs / freqs.sum()
+    return names, probs
+
+
+class CamptothecinPDC_Factory:
+    """
+    Drug-level PDC factory for Camptothecin conjugates.
+
+    Wraps all four ERBB2+Cpt linker builders and provides:
+    - 'sample' mode (default): randomly pick a linker per call according to
+      PDCdb entry frequencies
+    - 'all' mode: generate conjugates with every linker
+    - fixed linker mode: always use a specific linker (e.g. 'SSCpt')
+
+    The interface mirrors the individual builder's `conjugate_to_peptide` so it
+    can be used as a drop-in replacement in the Conjugation classifier.
+    """
+
+    def __init__(self, linker='sample', rng_seed=None):
+        """
+        Args:
+            linker: One of:
+                - 'sample' (default): sample linker per peptide according to
+                  ERBB2+Cpt frequency weights
+                - 'all': use every linker and pool conjugates
+                - A specific linker name ('SSCpt', 'DSPACpt', 'MalHexCpt', 'DSECCpt')
+            rng_seed: Optional seed for reproducible linker sampling
+        """
+        self.linker_mode = linker
+        self.registry = ERBB2_CPT_LINKER_REGISTRY
+        self.rng = np.random.default_rng(rng_seed)
+
+        # Pre-compute probabilities
+        self._linker_names, self._linker_probs = _compute_linker_probabilities(self.registry)
+
+        # Validate fixed linker
+        if linker not in ('sample', 'all') and linker not in self.registry:
+            raise ValueError(
+                f"Unknown linker '{linker}'. "
+                f"Choose from: 'sample', 'all', {list(self.registry.keys())}"
+            )
+
+        # Pre-instantiate builders (suppress verbose init output)
+        self._builders = {}
+        for name, entry in self.registry.items():
+            with contextlib.redirect_stdout(io.StringIO()):
+                self._builders[name] = entry['class']()
+
+        # Track which linker was last sampled (useful for logging)
+        self.last_sampled_linker = None
+
+    # ------------------------------------------------------------------
+    # Linker selection
+    # ------------------------------------------------------------------
+
+    def sample_linker(self):
+        """Sample a linker name according to frequency weights."""
+        idx = self.rng.choice(len(self._linker_names), p=self._linker_probs)
+        return self._linker_names[idx]
+
+    def get_builder(self, linker_name):
+        """Return the pre-instantiated builder for a given linker name."""
+        return self._builders[linker_name]
+
+    @property
+    def linker_probabilities(self):
+        """Return dict of linker name → sampling probability."""
+        return dict(zip(self._linker_names, self._linker_probs.tolist()))
+
+    # ------------------------------------------------------------------
+    # Conjugation interface (mirrors individual builders)
+    # ------------------------------------------------------------------
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Camptothecin to peptide, selecting linker by mode.
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Lys -NH2)
+            exclude_n_terminus: If True, only conjugate at Lys ε-amines
+            return_all: If True, return all positional conjugates
+
+        Returns:
+            List of conjugate SMILES.
+            In 'sample' mode: conjugates from the sampled linker.
+            In 'all' mode: conjugates from every linker (pooled).
+            In fixed mode: conjugates from the specified linker.
+        """
+        if self.linker_mode == 'all':
+            self.last_sampled_linker = 'all'
+            return self._conjugate_all_linkers(peptide_smiles, exclude_n_terminus, return_all)
+        elif self.linker_mode == 'sample':
+            chosen = self.sample_linker()
+            self.last_sampled_linker = chosen
+            return self._conjugate_with_linker(chosen, peptide_smiles, exclude_n_terminus, return_all)
+        else:
+            self.last_sampled_linker = self.linker_mode
+            return self._conjugate_with_linker(self.linker_mode, peptide_smiles, exclude_n_terminus, return_all)
+
+    def _conjugate_with_linker(self, linker_name, peptide_smiles, exclude_n_terminus, return_all):
+        """Run conjugation with a single, specific linker."""
+        builder = self._builders[linker_name]
+        with contextlib.redirect_stdout(io.StringIO()):
+            conjugates = builder.conjugate_to_peptide(
+                peptide_smiles,
+                exclude_n_terminus=exclude_n_terminus,
+                return_all=return_all,
+            )
+        # Normalize to list
+        if not isinstance(conjugates, list):
+            conjugates = [conjugates] if conjugates else []
+        return conjugates
+
+    def _conjugate_all_linkers(self, peptide_smiles, exclude_n_terminus, return_all):
+        """Generate conjugates with every linker and pool them."""
+        all_conjugates = []
+        seen = set()
+        for name in self._linker_names:
+            try:
+                conjs = self._conjugate_with_linker(name, peptide_smiles, exclude_n_terminus, return_all)
+                for c in conjs:
+                    if c not in seen:
+                        seen.add(c)
+                        all_conjugates.append(c)
+            except Exception:
+                continue
+        if not all_conjugates:
+            raise ValueError("No conjugates generated with any linker")
+        return all_conjugates
+
+    # ------------------------------------------------------------------
+    # Validation (delegates to the appropriate builder)
+    # ------------------------------------------------------------------
+
+    def validate_conjugate(self, conjugate_smiles, linker_name=None):
+        """
+        Validate a conjugate SMILES.
+
+        If linker_name is given, validate with that builder.
+        Otherwise, try all builders and return the best result.
+        """
+        if linker_name and linker_name in self._builders:
+            return self._builders[linker_name].validate_conjugate(conjugate_smiles)
+
+        # Try each builder; return first passing result
+        for name, builder in self._builders.items():
+            passed, checks = builder.validate_conjugate(conjugate_smiles)
+            if passed:
+                checks['matched_linker'] = name
+                return True, checks
+
+        # None passed - return last failure
+        return False, checks
+
+    def check_attachment_site_compatibility(self, conjugate_smiles, linker_name=None):
+        """Delegate to appropriate builder's attachment check."""
+        if linker_name and linker_name in self._builders:
+            return self._builders[linker_name].check_attachment_site_compatibility(conjugate_smiles)
+
+        for name, builder in self._builders.items():
+            passed, details = builder.check_attachment_site_compatibility(conjugate_smiles)
+            if passed:
+                details['matched_linker'] = name
+                return True, details
+        return False, details
+
+    # ------------------------------------------------------------------
+    # Batch interface
+    # ------------------------------------------------------------------
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides.
+
+        In 'sample' mode each peptide gets a freshly sampled linker.
+
+        Returns:
+            Dict with 'conjugates', 'successful', 'failed', 'total_conjugates',
+            and 'linkers_used' (list of linker names, one per peptide).
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0,
+            'linkers_used': [],
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all,
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+                # Track which linker was used (only meaningful in sample/fixed mode)
+                results['linkers_used'].append(
+                    self.linker_mode if self.linker_mode not in ('sample', 'all') else None
+                )
+            except Exception as e:
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+                results['linkers_used'].append(None)
+
+        return results
+
+    # ------------------------------------------------------------------
+    # Informational
+    # ------------------------------------------------------------------
+
+    def __repr__(self):
+        probs = self.linker_probabilities
+        lines = [f"CamptothecinPDC_Factory(mode='{self.linker_mode}')"]
+        for name, prob in probs.items():
+            desc = self.registry[name]['description']
+            lines.append(f"  {name:12s}  p={prob:.3f}  ({desc})")
+        return '\n'.join(lines)

smiles/conjugation/DSEBSN38.py

@@ -0,0 +1,377 @@
+"""
+4-(2-Disulfanylethoxy)-4-oxobutanoic Acid - SN-38 PDC Builder (DSEBSN38)
+
+Conjugation chemistry: Cysteine thiol + Linker thiol → Disulfide bond (S-S)
+
+Reference PDCs:
+- PDC_02100 (SN38(10)-4-CRGDK)
+- PDC_02103 (SN38(10)-4-(Ac)CGVQIRGDK)
+- PDC_02105 (SN38(10)-4-(Ac)CRGDK)
+- PDC_02108 (SN38(10)-4-CDRGDK)
+- PDC_02109 (SN38(10)-4-(Ac)CDRGDK)
+- PDC_02111 (SN38(20)-4-CRGDK)
+
+Receptor: NRP1 (Neuropilin-1)
+Linker: LIN00168 - 4-(2-Disulfanylethoxy)-4-oxobutanoic acid
+Payload: 7-Ethyl-10-hydroxycamptothecin (SN-38)
+Architecture: Monomer, Cys thiol conjugation (disulfide bond)
+
+Chemistry Note:
+===============
+LIN00168 provides a longer disulfide-ester linkage via succinic acid scaffold:
+    HOOC─CH₂─CH₂─C(═O)─O─CH₂─CH₂─S─SH
+
+            butanoic acid         ethoxy-disulfanyl
+
+Linker-to-Drug Connection (Ester at drug OH, thiol for peptide Cys):
+    SN-38(10-O)─C(═O)─CH₂─CH₂─C(═O)─O─CH₂─CH₂─SH
+
+Full PDC Architecture:
+    Peptide──Cys──S─S──CH₂─CH₂─O─C(═O)─CH₂─CH₂─C(═O)─O──SN-38
+                  └── disulfide ─┘       └── linker ──┘  └ ester ┘
+
+Usage:
+======
+    builder = DisulfanylethOxyButanoicAcidSN38_PDC_Builder()
+    conjugates = builder.conjugate_to_peptide(peptide_smiles)
+"""
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, Descriptors
+
+from .pdc_builder import PDC_Builder
+
+
+# ============================================================================
+# 4-(2-Disulfanylethoxy)-4-oxobutanoic Acid - SN-38 PDC Builder
+# ============================================================================
+
+class DisulfanylethOxyButanoicAcidSN38_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for 4-(2-Disulfanylethoxy)-4-oxobutanoic Acid + SN-38 conjugates
+
+    Conjugation chemistry (disulfide bond formation):
+    - Peptide site: Cysteine thiol (side chain -SH)
+    - Linker-payload site: Terminal thiol (-SH) on linker
+    - Reaction: Peptide-Cys-SH + HS-Linker-Drug → Peptide-Cys-S-S-Linker-Drug
+
+    Linker: LIN00168 - 4-(2-Disulfanylethoxy)-4-oxobutanoic acid
+        IUPAC: 4-(2-(disulfanyl)ethoxy)-4-oxobutanoic acid
+        Connectivity: HOOC-CH2-CH2-C(=O)-O-CH2-CH2-S-SH
+        Succinic acid monoester with 2-(disulfanyl)ethanol
+
+    Drug: 7-Ethyl-10-hydroxycamptothecin (SN-38)
+        Connection: Ester at C-10 hydroxyl (default) or C-20 hydroxyl
+
+    Reference: PDC_02100, PDC_02111 (NRP1)
+    """
+
+    # ========================================================================
+    # Linker-Drug SMILES Variants
+    # ========================================================================
+
+    # SN-38 with DSEB ester at C-10 (phenolic OH), free thiol
+    # Drug(10-O)-C(=O)-CH2-CH2-C(=O)-O-CH2-CH2-SH
+    POS10_LINKER_SN38_SMILES = (
+        "CCC1(O)C(=O)OCc2c1cc1n(c2=O)Cc2c(CC)c3cc(OC(=O)CCC(=O)OCCS)ccc3nc2-1"
+    )
+
+    # SN-38 with DSEB ester at C-20 (tertiary OH), free thiol
+    POS20_LINKER_SN38_SMILES = (
+        "CCC1(OC(=O)CCC(=O)OCCS)C(=O)OCc2c1cc1n(c2=O)Cc2c(CC)c3cc(O)ccc3nc2-1"
+    )
+
+    DEFAULT_LINKER_SN38_SMILES = POS10_LINKER_SN38_SMILES
+
+    DRUG_POSITIONS = {
+        '10': POS10_LINKER_SN38_SMILES,
+        '20': POS20_LINKER_SN38_SMILES,
+    }
+
+    def __init__(self, linker_drug_smiles=None, drug_position='all'):
+        """
+        Initialize DSEB-SN38 builder
+
+        Args:
+            linker_drug_smiles: Optional custom Linker-Drug SMILES.
+            drug_position: 'all' (default), '10', or '20'.
+        """
+        self.drug_position = drug_position
+
+        if linker_drug_smiles is not None:
+            self._variants = [linker_drug_smiles]
+            self._variant_names = ['custom']
+            super().__init__(
+                linker_drug_smiles,
+                component_name="Custom-DSEB-SN38 (NRP1)"
+            )
+        elif drug_position == 'all':
+            self._variants = [self.POS10_LINKER_SN38_SMILES, self.POS20_LINKER_SN38_SMILES]
+            self._variant_names = ['pos10', 'pos20']
+            super().__init__(
+                self.POS10_LINKER_SN38_SMILES,
+                component_name="DSEB-SN38 (LIN00168, all positions)"
+            )
+        else:
+            if drug_position not in self.DRUG_POSITIONS:
+                raise ValueError(f"Unknown drug_position '{drug_position}'. "
+                                 f"Choose from: 'all', '10', '20'")
+            smiles = self.DRUG_POSITIONS[drug_position]
+            self._variants = [smiles]
+            self._variant_names = [f'pos{drug_position}']
+            super().__init__(
+                smiles,
+                component_name=f"DSEB-SN38 (LIN00168, C-{drug_position})"
+            )
+
+    def _verify_reactive_group(self):
+        """Verify thiol group is present for disulfide bond formation"""
+        for atom in self.linker_payload_mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:
+                if atom.GetTotalNumHs() >= 1 and atom.GetDegree() == 1:
+                    print(" DSEB-SN38 loaded with free thiol reactive group")
+                    break
+        else:
+            print(" Warning: No free thiol group detected in DSEB-SN38")
+
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        if self.linker_payload_mol.HasSubstructMatch(ester_pattern):
+            print(" Ester bond (linker-to-drug) detected")
+        else:
+            print(" Warning: No ester bond detected")
+
+    # ------------------------------------------------------------------
+    # Conjugation site finding (Cysteine thiol)
+    # ------------------------------------------------------------------
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """Find all free cysteine thiol groups in peptide"""
+        return self.find_peptide_cysteines(peptide_smiles)
+
+    def find_peptide_cysteines(self, peptide_smiles):
+        """Find all free cysteine residues (thiol groups) in peptide"""
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        cysteines = []
+        for atom in peptide_mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:
+                total_hs = atom.GetTotalNumHs()
+                degree = atom.GetDegree()
+                if total_hs >= 1 and degree == 1:
+                    cysteines.append({
+                        'atom_idx': atom.GetIdx(),
+                        'num_hydrogens': total_hs,
+                        'neighbors': [peptide_mol.GetAtomWithIdx(n.GetIdx()).GetSymbol()
+                                      for n in atom.GetNeighbors()],
+                        'type': 'Cys_thiol'
+                    })
+        return cysteines
+
+    # ------------------------------------------------------------------
+    # Conjugation reaction (disulfide bond formation)
+    # ------------------------------------------------------------------
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """Form disulfide bond between peptide Cys thiol and linker-drug thiol."""
+        all_conjugates = []
+        seen = set()
+
+        for variant_smiles in self._variants:
+            variant_mol = Chem.MolFromSmiles(variant_smiles)
+            if variant_mol is None:
+                continue
+
+            disulfide_rxn = AllChem.ReactionFromSmarts(
+                '[C:1][SH:2].[SH:3][C:4]>>[C:1][S:2][S:3][C:4]'
+            )
+
+            products = disulfide_rxn.RunReactants((peptide_mol, variant_mol))
+            if not products:
+                continue
+
+            for product_tuple in products:
+                try:
+                    conjugate_mol = product_tuple[0]
+                    Chem.SanitizeMol(conjugate_mol)
+                    conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+                    if conjugate_smiles not in seen:
+                        seen.add(conjugate_smiles)
+                        all_conjugates.append(conjugate_smiles)
+                except Exception:
+                    continue
+
+        if not all_conjugates:
+            raise ValueError("Disulfide bond formation failed - no products generated")
+
+        return all_conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate DSEB-SN38 to peptide via disulfide bond at Cys thiol.
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Cys)
+            exclude_n_terminus: Not used for Cys conjugation (kept for interface)
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            List of all possible conjugate SMILES
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Cys thiol site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        pos_label = f"C-{self.drug_position}" if self.drug_position != 'all' else "C-10 + C-20"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [Cys disulfide] [DSEB {pos_label}]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """Batch conjugate multiple peptides."""
+        results = {
+            'conjugates': [], 'successful': [], 'failed': [], 'total_conjugates': 0
+        }
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles, exclude_n_terminus=exclude_n_terminus, return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        print(f"\n=== Summary ===")
+        print(f"Peptides: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, "
+              f"Failed: {len(results['failed'])}")
+        print(f"Total conjugates generated: {results['total_conjugates']}")
+        return results
+
+    def flatten_batch_results(self, results):
+        """Flatten batch conjugation results into a single list"""
+        all_conjugates = []
+        for conjugate_list in results['conjugates']:
+            if conjugate_list is not None:
+                if isinstance(conjugate_list, list):
+                    all_conjugates.extend(conjugate_list)
+                else:
+                    all_conjugates.append(conjugate_list)
+        return all_conjugates
+
+    # ========================================================================
+    # Validation
+    # ========================================================================
+
+    def has_disulfide_bond(self, mol):
+        """Check if molecule has disulfide bond (S-S)"""
+        disulfide_pattern = Chem.MolFromSmarts('[SX2]-[SX2]')
+        return mol.HasSubstructMatch(disulfide_pattern)
+
+    def has_free_thiol(self, mol):
+        """Check if molecule has unconjugated (free) thiol groups"""
+        for atom in mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:
+                if atom.GetTotalNumHs() >= 1 and atom.GetDegree() == 1:
+                    return True
+        return False
+
+    def has_ester_bond_to_drug(self, mol):
+        """Check for ester bond connecting linker to drug"""
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        return mol.HasSubstructMatch(ester_pattern)
+
+    def has_sn38_core(self, mol):
+        """Check for SN-38 core structure"""
+        cpt_core = Chem.MolFromSmarts('c1ccc2ncc3C(=O)OCc3c2c1')
+        return mol.HasSubstructMatch(cpt_core)
+
+    def check_attachment_site_compatibility(self, conjugate_smiles):
+        """Check attachment site compatibility for Cys-SS-DSEB-SN38 conjugation."""
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        if mol is None:
+            return False, {'parseable': False, 'error': 'Could not parse SMILES'}
+
+        details = {
+            'parseable': True,
+            'validation_method': 'cys_disulfide_ester',
+            'linker': 'LIN00168 (4-(2-Disulfanylethoxy)-4-oxobutanoic acid)',
+            'drug': 'SN-38 (7-Ethyl-10-hydroxycamptothecin)',
+        }
+        details['has_disulfide_bond'] = self.has_disulfide_bond(mol)
+        details['no_free_thiol'] = not self.has_free_thiol(mol)
+        details['has_ester_bond'] = self.has_ester_bond_to_drug(mol)
+
+        passed = all([details['has_disulfide_bond'], details['no_free_thiol']])
+        return passed, details
+
+    def _basic_validation(self, mol):
+        """Basic validation checks for DSEB-SN38 conjugates."""
+        checks = {}
+        checks['has_disulfide_bond'] = mol.HasSubstructMatch(Chem.MolFromSmarts('[#16]-[#16]'))
+        checks['no_free_thiol'] = not self.has_free_thiol(mol)
+        checks['has_ester_bond'] = self.has_ester_bond_to_drug(mol)
+        checks['has_sn38_core'] = self.has_sn38_core(mol)
+        checks['_passed'] = all([checks['has_disulfide_bond'], checks['no_free_thiol']])
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """Validate that the conjugate is chemically reasonable."""
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+        return passed, checks
+
+
+# ============================================================================
+# Main - Test Examples
+# ============================================================================
+
+if __name__ == "__main__":
+
+    print("=" * 70)
+    print("DSEB-SN38 PDC Builder - Test")
+    print("Reference: NRP1 + LIN00168 + SN-38")
+    print("=" * 70 + "\n")
+
+    test_peptides = [
+        "NCC(=O)N[C@@H](CS)C(=O)NCC(=O)O",  # Gly-Cys-Gly
+    ]
+
+    builder = DisulfanylethOxyButanoicAcidSN38_PDC_Builder()
+
+    for i, peptide in enumerate(test_peptides):
+        print(f"\n--- Test peptide {i+1} ---")
+        print(f"Input: {peptide}")
+
+        try:
+            conjugates = builder.conjugate_to_peptide(peptide)
+            print(f" Generated {len(conjugates)} conjugate(s)")
+            for j, conj in enumerate(conjugates):
+                valid, checks = builder.validate_conjugate(conj)
+                print(f"  Conjugate {j+1}: {'PASSED' if valid else 'FAILED'}")
+        except Exception as e:
+            print(f" FAILED: {e}")
+            import traceback
+            traceback.print_exc()

smiles/conjugation/DSECCpt.py

@@ -0,0 +1,522 @@
+"""
+Disulfanylethyl Carbonate - Camptothecin PDC Builder (DSECCpt)
+
+Conjugation chemistry: Lysine ε-amine + Cys-SS-Linker-Camptothecin (amide bond formation)
+
+Reference PDC:
+- PDC_00337 (SMAC-P1 CPT conjugate 7)
+
+Receptor: ERBB2 (Receptor tyrosine-protein kinase erbB-2)
+Linker: LIN00039 - 2-(Disulfanyl)ethyl hydrogen carbonate (disulfide + carbonate)
+Payload: DRG00009 - Camptothecin
+Architecture: Monomer, Lys side chain conjugation
+
+Chemistry Note:
+===============
+LIN00039 (2-(Disulfanyl)ethyl hydrogen carbonate) provides a disulfide-carbonate linkage:
+    HS─S─CH₂─CH₂─O─C(═O)─OH
+
+Linker-to-Drug Connection (Carbonate):
+    Cys─S─S─CH₂─CH₂─O─C(═O)─O─Camptothecin
+                       │
+                   carbonate bond (C-20 OH of CPT)
+
+Full PDC Architecture:
+    Peptide──Lys──NH──C(═O)──Cys──S─S──CH₂─CH₂─O──C(═O)─O──Camptothecin
+                  └── amide ──┘       └── disulfide ──┘  └── carbonate ──┘
+
+Note: This linker chemistry is equivalent to the SSCpt carbonate variant (Z8/PDC_00295).
+
+Usage:
+======
+    builder = DisulfanylethylCarbonateCpt_PDC_Builder()
+    conjugates = builder.conjugate_to_peptide(peptide_smiles)
+"""
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, Descriptors
+
+from .pdc_builder import PDC_Builder
+
+
+# ============================================================================
+# Disulfanylethyl Carbonate - Camptothecin PDC Builder
+# ============================================================================
+
+class DisulfanylethylCarbonateCpt_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for 2-(Disulfanyl)ethyl Carbonate + Camptothecin conjugates
+
+    Conjugation chemistry (amide bond formation):
+    - Peptide site: Lysine ε-amine (side chain -NH2)
+    - Linker-payload site: Carboxylic acid of Cys (Cys-SS-Linker-Drug)
+    - Reaction: Lys-NH2 + HOOC-Cys-SS-Linker-Drug → Lys-NH-CO-Cys-SS-Linker-Drug
+
+    Linker: LIN00039 - 2-(Disulfanyl)ethyl hydrogen carbonate
+        SMILES: C(CSS)OC(=O)O
+        IUPAC: 2-(disulfanyl)ethyl hydrogen carbonate
+        Connectivity: HS-S-CH2-CH2-O-C(=O)-OH
+        MW: 154.21, MF: C3H6O3S2
+
+    Drug: DRG00009 - Camptothecin
+        Canonical SMILES: CCC1(O)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1
+        Connection: Carbonate at C-20 hydroxyl
+
+    Reference: PDC_00337 (SMAC-P1 CPT conjugate 7, ERBB2)
+    """
+
+    # ========================================================================
+    # Cys-SS-Linker-Drug SMILES
+    # ========================================================================
+    #
+    # Pre-formed Cys-DisulfanylethylCarbonate-Camptothecin assembly:
+    #   HOOC─Cys(NH₂)─CH₂─S─S─CH₂─CH₂─O─C(═O)─O─CPT(C-20)
+    #
+    # Cys provides COOH for amide bond with Lys ε-amine
+    # Disulfide bond (S-S) is cleavable inside cells
+    # Carbonate bond (O-C(=O)-O) connects linker to CPT C-20 hydroxyl
+    #
+    # ========================================================================
+
+    CYS_SS_LINKER_CPT_SMILES = (
+        "N[C@H](CSSCCOC(=O)OC1(CC)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1)C(=O)O"
+    )
+
+    def __init__(self, cys_linker_drug_smiles=None):
+        """
+        Initialize Disulfanylethyl Carbonate-Camptothecin builder
+
+        Args:
+            cys_linker_drug_smiles: Optional custom Cys-SS-Linker-Drug SMILES.
+                                   If None, uses default LIN00039+CPT assembly.
+                                   Must have free carboxylic acid for conjugation.
+        """
+        if cys_linker_drug_smiles is None:
+            cys_linker_drug_smiles = self.CYS_SS_LINKER_CPT_SMILES
+
+        super().__init__(
+            cys_linker_drug_smiles,
+            component_name="Cys-SS-DisulfanylethylCarbonate-Camptothecin (LIN00039+DRG00009)"
+        )
+
+    def _verify_reactive_group(self):
+        """Verify carboxylic acid group is present for amide bond formation"""
+        # Pattern: Carboxylic acid (COOH)
+        carboxylic_acid_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+
+        if not self.linker_payload_mol.HasSubstructMatch(carboxylic_acid_pattern):
+            print(" Warning: No carboxylic acid group detected in Cys-SS-Linker-Drug")
+        else:
+            print(" Cys-SS-DSEC-Camptothecin loaded with carboxylic acid reactive group")
+
+        # Verify disulfide bond is present
+        disulfide_pattern = Chem.MolFromSmarts('[#16]-[#16]')
+        if self.linker_payload_mol.HasSubstructMatch(disulfide_pattern):
+            print(" Disulfide bond (S-S) detected")
+        else:
+            print(" Warning: No disulfide bond detected - check structure")
+
+        # Verify Cys α-amine is present
+        cys_amine_pattern = Chem.MolFromSmarts('[NH2][CH]([CH2]S)C(=O)')
+        if self.linker_payload_mol.HasSubstructMatch(cys_amine_pattern):
+            print(" Cys α-amine detected (will be free in conjugate)")
+        else:
+            print(" Warning: Cys α-amine pattern not detected")
+
+        # Verify carbonate bond to drug (O-C(=O)-O pattern)
+        carbonate_pattern = Chem.MolFromSmarts('[OX2][CX3](=[OX1])[OX2]')
+        if self.linker_payload_mol.HasSubstructMatch(carbonate_pattern):
+            print(" Carbonate bond (linker-to-drug) detected")
+        else:
+            print(" Warning: No carbonate bond detected - check linker-drug connection")
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find all lysine ε-amine groups in peptide
+
+        Returns list of dicts with site information including atom index
+        """
+        return self.find_peptide_lysines(peptide_smiles)
+
+    def find_peptide_lysines(self, peptide_smiles):
+        """
+        Find all lysine residues with free ε-amine in peptide
+
+        Lysine side chain: -CH2-CH2-CH2-CH2-NH2 (CCCCN)
+
+        Returns:
+            List of dicts with 'atom_idx' (nitrogen index) and metadata
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        lysines = []
+
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        matches = peptide_mol.GetSubstructMatches(lys_pattern)
+
+        for match in matches:
+            nitrogen_idx = match[4]
+            lysines.append({
+                'atom_idx': nitrogen_idx,
+                'side_chain_atoms': list(match),
+                'type': 'Lys_epsilon_amine'
+            })
+
+        n_terminus_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        n_term_matches = peptide_mol.GetSubstructMatches(n_terminus_pattern)
+
+        if n_term_matches:
+            print(f"  (Also found {len(n_term_matches)} N-terminus α-amine(s) - excluded by default)")
+
+        return lysines
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """
+        Perform amide bond conjugation between NH2 groups and Cys carboxylic acid
+
+        Args:
+            peptide_mol: RDKit Mol object
+            site_info: Not used (kept for compatibility)
+            exclude_n_terminus: If True, skip N-terminus conjugation
+
+        Returns: List of all possible conjugate SMILES
+        """
+        if exclude_n_terminus:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:5][CH2][CH2][CH2][CH2:4][NH2:1].[C:2](=[O:3])[OH]>>[*:5][CH2][CH2][CH2][CH2:4][NH:1][C:2](=[O:3])'
+            )
+        else:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:4][NH2:1].[C:2](=[O:3])[OH]>>[*:4][NH:1][C:2](=[O:3])'
+            )
+
+        products = amide_rxn.RunReactants((peptide_mol, self.linker_payload_mol))
+
+        if not products:
+            raise ValueError("Amide bond formation reaction failed - no products generated")
+
+        conjugates = []
+        seen_smiles = set()
+
+        for product_tuple in products:
+            try:
+                conjugate_mol = product_tuple[0]
+                Chem.SanitizeMol(conjugate_mol)
+                conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+
+                if conjugate_smiles not in seen_smiles:
+                    seen_smiles.add(conjugate_smiles)
+                    conjugates.append(conjugate_smiles)
+            except Exception:
+                continue
+
+        if not conjugates:
+            raise ValueError("All products failed sanitization")
+
+        return conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Cys-SS-DSEC-Camptothecin to peptide via amide bond at NH2 groups
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Lys with NH2)
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            List of all possible conjugate SMILES
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Lys ε-amine site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2 (including N-terminus)"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [{site_type}] [DSEC carbonate]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides
+
+        Args:
+            peptide_smiles_list: List of peptide SMILES strings
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates per peptide
+
+        Returns:
+            Dict with 'conjugates', 'successful', 'failed', 'total_conjugates'
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        print(f"\n=== Summary ===")
+        print(f"Peptides: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, Failed: {len(results['failed'])}")
+        print(f"Total conjugates generated: {results['total_conjugates']}")
+
+        return results
+
+    def flatten_batch_results(self, results):
+        """Flatten batch conjugation results into a single list"""
+        all_conjugates = []
+        for conjugate_list in results['conjugates']:
+            if conjugate_list is not None:
+                if isinstance(conjugate_list, list):
+                    all_conjugates.extend(conjugate_list)
+                else:
+                    all_conjugates.append(conjugate_list)
+        return all_conjugates
+
+    # ========================================================================
+    # Attachment Site Compatibility Checks
+    # ========================================================================
+
+    def has_lys_cys_amide_conjugation(self, mol):
+        """Check if molecule has Lys-Cys amide conjugation bond"""
+        lys_cys_amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        return mol.HasSubstructMatch(lys_cys_amide_pattern)
+
+    def has_disulfide_bond(self, mol):
+        """Check if molecule has disulfide bond (S-S)"""
+        disulfide_pattern = Chem.MolFromSmarts('[SX2]-[SX2]')
+        return mol.HasSubstructMatch(disulfide_pattern)
+
+    def has_carbonate_bond(self, mol):
+        """
+        Check if molecule has carbonate bond (O-C(=O)-O)
+
+        This specifically checks for the carbonate linkage between linker and drug.
+        """
+        carbonate_pattern = Chem.MolFromSmarts('[OX2][CX3](=[OX1])[OX2]')
+        return mol.HasSubstructMatch(carbonate_pattern)
+
+    def has_unreacted_lys_epsilon_amine(self, mol):
+        """Check for unreacted Lys ε-amine (free NH2 on 4-carbon side chain)"""
+        lys_amine_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        return mol.HasSubstructMatch(lys_amine_pattern)
+
+    def has_unreacted_cys_carboxylic_acid(self, mol):
+        """Check for unreacted Cys carboxylic acid from Cys-SS-Linker-Drug"""
+        cys_cooh_pattern = Chem.MolFromSmarts('[NH2][CH]([CH2][SX2][SX2])[CX3](=[OX1])[OX2H1]')
+        return mol.HasSubstructMatch(cys_cooh_pattern)
+
+    def find_amide_conjugation_sites(self, mol):
+        """Find all Lys-Cys amide conjugation sites in the molecule"""
+        lys_cys_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        matches = mol.GetSubstructMatches(lys_cys_pattern)
+
+        sites = []
+        for match in matches:
+            nitrogen_idx = match[0]
+            carbonyl_c_idx = match[1]
+            cys_alpha_idx = match[3]
+            sites.append((nitrogen_idx, carbonyl_c_idx, cys_alpha_idx))
+
+        return sites
+
+    def check_valency_at_conjugation_site(self, mol, site_indices):
+        """Check that valency is correct at Lys-Cys conjugation site"""
+        nitrogen_idx, carbonyl_c_idx, cys_alpha_idx = site_indices
+
+        try:
+            n_atom = mol.GetAtomWithIdx(nitrogen_idx)
+            if n_atom.GetTotalDegree() != 3:
+                return False
+
+            c_atom = mol.GetAtomWithIdx(carbonyl_c_idx)
+            if c_atom.GetTotalDegree() != 3:
+                return False
+
+            cys_c_atom = mol.GetAtomWithIdx(cys_alpha_idx)
+            if cys_c_atom.GetTotalDegree() != 4:
+                return False
+
+            return True
+        except Exception:
+            return False
+
+    def check_attachment_site_compatibility(self, conjugate_smiles):
+        """
+        Check attachment site compatibility for Lys-Cys-SS-DSEC-CPT conjugation
+
+        Validates:
+        1. Lys-Cys amide conjugation is present
+        2. Disulfide bond (S-S) is intact
+        3. Carbonate bond is present (linker-to-drug)
+        4. Valency at conjugation sites is correct
+        5. No unreacted Cys-COOH (unconjugated linker-drug)
+
+        Returns:
+            Tuple of (passed: bool, details: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'parseable': False, 'error': 'Could not parse SMILES'}
+
+        details = {
+            'parseable': True,
+            'validation_method': 'lys_cys_amide_disulfide_carbonate',
+            'linker': 'LIN00039 (2-(Disulfanyl)ethyl hydrogen carbonate)',
+            'drug': 'DRG00009 (Camptothecin)',
+            'pdc_ref': 'PDC_00337'
+        }
+
+        details['has_lys_cys_amide'] = self.has_lys_cys_amide_conjugation(mol)
+        details['has_disulfide_bond'] = self.has_disulfide_bond(mol)
+        details['has_carbonate_bond'] = self.has_carbonate_bond(mol)
+
+        conjugation_sites = self.find_amide_conjugation_sites(mol)
+        details['conjugation_sites'] = len(conjugation_sites)
+
+        valency_ok = True
+        if conjugation_sites:
+            for site in conjugation_sites:
+                if not self.check_valency_at_conjugation_site(mol, site):
+                    valency_ok = False
+                    break
+        details['valency_correct'] = valency_ok
+
+        details['no_unreacted_cys_cooh'] = not self.has_unreacted_cys_carboxylic_acid(mol)
+        details['has_unreacted_lys_amine'] = self.has_unreacted_lys_epsilon_amine(mol)
+
+        passed = all([
+            details['has_lys_cys_amide'],
+            details['has_disulfide_bond'],
+            details['has_carbonate_bond'],
+            details['valency_correct'],
+            details['no_unreacted_cys_cooh']
+        ])
+
+        return passed, details
+
+    def _basic_validation(self, mol):
+        """
+        Basic validation checks for DSEC-Camptothecin conjugates
+
+        Checks:
+        1. Amide bond present (Lys-Cys conjugation)
+        2. Disulfide bond present (S-S linker)
+        3. Specific Lys-Cys-SS conjugation pattern
+        4. Carbonate bond present (linker-to-drug)
+        5. Camptothecin core present
+        """
+        checks = {}
+
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[#6]')
+        checks['has_amide_bond'] = mol.HasSubstructMatch(amide_pattern)
+
+        disulfide_pattern = Chem.MolFromSmarts('[#16]-[#16]')
+        checks['has_disulfide_bond'] = mol.HasSubstructMatch(disulfide_pattern)
+
+        lys_cys_ss_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        checks['has_lys_cys_ss_conjugation'] = mol.HasSubstructMatch(lys_cys_ss_pattern)
+
+        carbonate_pattern = Chem.MolFromSmarts('[OX2][CX3](=[OX1])[OX2]')
+        checks['has_carbonate_bond'] = mol.HasSubstructMatch(carbonate_pattern)
+
+        camptothecin_core = Chem.MolFromSmarts('c1ccc2ncc3C(=O)OCc3c2c1')
+        checks['has_camptothecin_core'] = mol.HasSubstructMatch(camptothecin_core)
+
+        checks['_passed'] = all([
+            checks['has_amide_bond'],
+            checks['has_disulfide_bond'],
+            checks['has_lys_cys_ss_conjugation'],
+            checks['has_carbonate_bond']
+        ])
+
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """
+        Validate that the conjugate is chemically reasonable
+
+        Returns:
+            (passed: bool, checks: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+
+        return passed, checks
+
+
+# ============================================================================
+# Main - Test Examples
+# ============================================================================
+
+if __name__ == "__main__":
+
+    print("=" * 70)
+    print("Disulfanylethyl Carbonate-Camptothecin PDC Builder - Test")
+    print("Reference: PDC_00337 (SMAC-P1 CPT conjugate 7, LIN00039)")
+    print("=" * 70 + "\n")
+
+    # Test peptides
+    test_peptides = [
+        # Gly-Lys-Gly (has N-terminus NH2 + Lys ε-amine)
+        "NCC(=O)N[C@@H](CCCCN)C(=O)NCC(=O)O",
+        # Lys-Arg
+        "N[C@@H](CCCCN)C(=O)N[C@@H](CCCNC(=N)N)C(=O)O",
+    ]
+
+    builder = DisulfanylethylCarbonateCpt_PDC_Builder()
+
+    for i, peptide in enumerate(test_peptides):
+        print(f"\n--- Test peptide {i+1} ---")
+        print(f"Input: {peptide}")
+
+        try:
+            conjugates = builder.conjugate_to_peptide(peptide, exclude_n_terminus=True)
+            print(f" Generated {len(conjugates)} conjugate(s)")
+
+            for j, conj in enumerate(conjugates):
+                valid, checks = builder.validate_conjugate(conj)
+                print(f"  Conjugate {j+1}: {'PASSED' if valid else 'FAILED'} ({len(conj)} chars)")
+
+        except Exception as e:
+            print(f" FAILED: {e}")
+            import traceback
+            traceback.print_exc()

smiles/conjugation/DSEPSN38.py

@@ -0,0 +1,373 @@
+"""
+5-(2-Disulfanylethoxy)-5-oxopentanoic Acid - SN-38 PDC Builder (DSEPSN38)
+
+Conjugation chemistry: Cysteine thiol + Linker thiol → Disulfide bond (S-S)
+
+Reference PDCs:
+- PDC_02101 (SN38(10)-5-CRGDK)
+- PDC_02106 (SN38(10)-5-(Ac)CRGDK)
+
+Receptor: NRP1 (Neuropilin-1)
+Linker: LIN00169 - 5-(2-Disulfanylethoxy)-5-oxopentanoic acid
+Payload: 7-Ethyl-10-hydroxycamptothecin (SN-38)
+Architecture: Monomer, Cys thiol conjugation (disulfide bond)
+
+Chemistry Note:
+===============
+LIN00169 provides a glutaric acid-based disulfide-ester linkage:
+    HOOC─CH₂─CH₂─CH₂─C(═O)─O─CH₂─CH₂─S─SH
+
+            pentanoic acid             ethoxy-disulfanyl
+
+Linker-to-Drug Connection:
+    SN-38(10-O)─C(═O)─CH₂─CH₂─CH₂─C(═O)─O─CH₂─CH₂─SH
+
+Full PDC Architecture:
+    Peptide──Cys──S─S──CH₂─CH₂─O─C(═O)─CH₂─CH₂─CH₂─C(═O)─O──SN-38
+                  └── disulfide ─┘         └── linker ──────┘  └ ester ┘
+
+Usage:
+======
+    builder = DisulfanylethOxyPentanoicAcidSN38_PDC_Builder()
+    conjugates = builder.conjugate_to_peptide(peptide_smiles)
+"""
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, Descriptors
+
+from .pdc_builder import PDC_Builder
+
+
+# ============================================================================
+# 5-(2-Disulfanylethoxy)-5-oxopentanoic Acid - SN-38 PDC Builder
+# ============================================================================
+
+class DisulfanylethOxyPentanoicAcidSN38_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for 5-(2-Disulfanylethoxy)-5-oxopentanoic Acid + SN-38 conjugates
+
+    Conjugation chemistry (disulfide bond formation):
+    - Peptide site: Cysteine thiol (side chain -SH)
+    - Linker-payload site: Terminal thiol (-SH) on linker
+    - Reaction: Peptide-Cys-SH + HS-Linker-Drug → Peptide-Cys-S-S-Linker-Drug
+
+    Linker: LIN00169 - 5-(2-Disulfanylethoxy)-5-oxopentanoic acid
+        IUPAC: 5-(2-(disulfanyl)ethoxy)-5-oxopentanoic acid
+        Connectivity: HOOC-CH2-CH2-CH2-C(=O)-O-CH2-CH2-S-SH
+        Glutaric acid monoester with 2-(disulfanyl)ethanol
+
+    Drug: 7-Ethyl-10-hydroxycamptothecin (SN-38)
+        Connection: Ester at C-10 hydroxyl (default) or C-20 hydroxyl
+
+    Reference: PDC_02101, PDC_02106 (NRP1)
+    """
+
+    # ========================================================================
+    # Linker-Drug SMILES Variants
+    # ========================================================================
+
+    # SN-38 with DSEP ester at C-10 (phenolic OH), free thiol
+    # Drug(10-O)-C(=O)-CH2-CH2-CH2-C(=O)-O-CH2-CH2-SH
+    POS10_LINKER_SN38_SMILES = (
+        "CCC1(O)C(=O)OCc2c1cc1n(c2=O)Cc2c(CC)c3cc(OC(=O)CCCC(=O)OCCS)ccc3nc2-1"
+    )
+
+    # SN-38 with DSEP ester at C-20 (tertiary OH), free thiol
+    POS20_LINKER_SN38_SMILES = (
+        "CCC1(OC(=O)CCCC(=O)OCCS)C(=O)OCc2c1cc1n(c2=O)Cc2c(CC)c3cc(O)ccc3nc2-1"
+    )
+
+    DEFAULT_LINKER_SN38_SMILES = POS10_LINKER_SN38_SMILES
+
+    DRUG_POSITIONS = {
+        '10': POS10_LINKER_SN38_SMILES,
+        '20': POS20_LINKER_SN38_SMILES,
+    }
+
+    def __init__(self, linker_drug_smiles=None, drug_position='all'):
+        """
+        Initialize DSEP-SN38 builder
+
+        Args:
+            linker_drug_smiles: Optional custom Linker-Drug SMILES.
+            drug_position: 'all' (default), '10', or '20'.
+        """
+        self.drug_position = drug_position
+
+        if linker_drug_smiles is not None:
+            self._variants = [linker_drug_smiles]
+            self._variant_names = ['custom']
+            super().__init__(
+                linker_drug_smiles,
+                component_name="Custom-DSEP-SN38 (NRP1)"
+            )
+        elif drug_position == 'all':
+            self._variants = [self.POS10_LINKER_SN38_SMILES, self.POS20_LINKER_SN38_SMILES]
+            self._variant_names = ['pos10', 'pos20']
+            super().__init__(
+                self.POS10_LINKER_SN38_SMILES,
+                component_name="DSEP-SN38 (LIN00169, all positions)"
+            )
+        else:
+            if drug_position not in self.DRUG_POSITIONS:
+                raise ValueError(f"Unknown drug_position '{drug_position}'. "
+                                 f"Choose from: 'all', '10', '20'")
+            smiles = self.DRUG_POSITIONS[drug_position]
+            self._variants = [smiles]
+            self._variant_names = [f'pos{drug_position}']
+            super().__init__(
+                smiles,
+                component_name=f"DSEP-SN38 (LIN00169, C-{drug_position})"
+            )
+
+    def _verify_reactive_group(self):
+        """Verify thiol group is present for disulfide bond formation"""
+        for atom in self.linker_payload_mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:
+                if atom.GetTotalNumHs() >= 1 and atom.GetDegree() == 1:
+                    print(" DSEP-SN38 loaded with free thiol reactive group")
+                    break
+        else:
+            print(" Warning: No free thiol group detected in DSEP-SN38")
+
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        if self.linker_payload_mol.HasSubstructMatch(ester_pattern):
+            print(" Ester bond (linker-to-drug) detected")
+        else:
+            print(" Warning: No ester bond detected")
+
+    # ------------------------------------------------------------------
+    # Conjugation site finding (Cysteine thiol)
+    # ------------------------------------------------------------------
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """Find all free cysteine thiol groups in peptide"""
+        return self.find_peptide_cysteines(peptide_smiles)
+
+    def find_peptide_cysteines(self, peptide_smiles):
+        """Find all free cysteine residues (thiol groups) in peptide"""
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        cysteines = []
+        for atom in peptide_mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:
+                total_hs = atom.GetTotalNumHs()
+                degree = atom.GetDegree()
+                if total_hs >= 1 and degree == 1:
+                    cysteines.append({
+                        'atom_idx': atom.GetIdx(),
+                        'num_hydrogens': total_hs,
+                        'neighbors': [peptide_mol.GetAtomWithIdx(n.GetIdx()).GetSymbol()
+                                      for n in atom.GetNeighbors()],
+                        'type': 'Cys_thiol'
+                    })
+        return cysteines
+
+    # ------------------------------------------------------------------
+    # Conjugation reaction (disulfide bond formation)
+    # ------------------------------------------------------------------
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """Form disulfide bond between peptide Cys thiol and linker-drug thiol."""
+        all_conjugates = []
+        seen = set()
+
+        for variant_smiles in self._variants:
+            variant_mol = Chem.MolFromSmiles(variant_smiles)
+            if variant_mol is None:
+                continue
+
+            disulfide_rxn = AllChem.ReactionFromSmarts(
+                '[C:1][SH:2].[SH:3][C:4]>>[C:1][S:2][S:3][C:4]'
+            )
+
+            products = disulfide_rxn.RunReactants((peptide_mol, variant_mol))
+            if not products:
+                continue
+
+            for product_tuple in products:
+                try:
+                    conjugate_mol = product_tuple[0]
+                    Chem.SanitizeMol(conjugate_mol)
+                    conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+                    if conjugate_smiles not in seen:
+                        seen.add(conjugate_smiles)
+                        all_conjugates.append(conjugate_smiles)
+                except Exception:
+                    continue
+
+        if not all_conjugates:
+            raise ValueError("Disulfide bond formation failed - no products generated")
+
+        return all_conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate DSEP-SN38 to peptide via disulfide bond at Cys thiol.
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Cys)
+            exclude_n_terminus: Not used for Cys conjugation (kept for interface)
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            List of all possible conjugate SMILES
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Cys thiol site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        pos_label = f"C-{self.drug_position}" if self.drug_position != 'all' else "C-10 + C-20"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [Cys disulfide] [DSEP {pos_label}]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """Batch conjugate multiple peptides."""
+        results = {
+            'conjugates': [], 'successful': [], 'failed': [], 'total_conjugates': 0
+        }
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles, exclude_n_terminus=exclude_n_terminus, return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        print(f"\n=== Summary ===")
+        print(f"Peptides: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, "
+              f"Failed: {len(results['failed'])}")
+        print(f"Total conjugates generated: {results['total_conjugates']}")
+        return results
+
+    def flatten_batch_results(self, results):
+        """Flatten batch conjugation results into a single list"""
+        all_conjugates = []
+        for conjugate_list in results['conjugates']:
+            if conjugate_list is not None:
+                if isinstance(conjugate_list, list):
+                    all_conjugates.extend(conjugate_list)
+                else:
+                    all_conjugates.append(conjugate_list)
+        return all_conjugates
+
+    # ========================================================================
+    # Validation
+    # ========================================================================
+
+    def has_disulfide_bond(self, mol):
+        """Check if molecule has disulfide bond (S-S)"""
+        disulfide_pattern = Chem.MolFromSmarts('[SX2]-[SX2]')
+        return mol.HasSubstructMatch(disulfide_pattern)
+
+    def has_free_thiol(self, mol):
+        """Check if molecule has unconjugated (free) thiol groups"""
+        for atom in mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:
+                if atom.GetTotalNumHs() >= 1 and atom.GetDegree() == 1:
+                    return True
+        return False
+
+    def has_ester_bond_to_drug(self, mol):
+        """Check for ester bond connecting linker to drug"""
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        return mol.HasSubstructMatch(ester_pattern)
+
+    def has_sn38_core(self, mol):
+        """Check for SN-38 core structure"""
+        cpt_core = Chem.MolFromSmarts('c1ccc2ncc3C(=O)OCc3c2c1')
+        return mol.HasSubstructMatch(cpt_core)
+
+    def check_attachment_site_compatibility(self, conjugate_smiles):
+        """Check attachment site compatibility for Cys-SS-DSEP-SN38 conjugation."""
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        if mol is None:
+            return False, {'parseable': False, 'error': 'Could not parse SMILES'}
+
+        details = {
+            'parseable': True,
+            'validation_method': 'cys_disulfide_ester',
+            'linker': 'LIN00169 (5-(2-Disulfanylethoxy)-5-oxopentanoic acid)',
+            'drug': 'SN-38 (7-Ethyl-10-hydroxycamptothecin)',
+        }
+        details['has_disulfide_bond'] = self.has_disulfide_bond(mol)
+        details['no_free_thiol'] = not self.has_free_thiol(mol)
+        details['has_ester_bond'] = self.has_ester_bond_to_drug(mol)
+
+        passed = all([details['has_disulfide_bond'], details['no_free_thiol']])
+        return passed, details
+
+    def _basic_validation(self, mol):
+        """Basic validation checks for DSEP-SN38 conjugates."""
+        checks = {}
+        checks['has_disulfide_bond'] = mol.HasSubstructMatch(Chem.MolFromSmarts('[#16]-[#16]'))
+        checks['no_free_thiol'] = not self.has_free_thiol(mol)
+        checks['has_ester_bond'] = self.has_ester_bond_to_drug(mol)
+        checks['has_sn38_core'] = self.has_sn38_core(mol)
+        checks['_passed'] = all([checks['has_disulfide_bond'], checks['no_free_thiol']])
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """Validate that the conjugate is chemically reasonable."""
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+        return passed, checks
+
+
+# ============================================================================
+# Main - Test Examples
+# ============================================================================
+
+if __name__ == "__main__":
+
+    print("=" * 70)
+    print("DSEP-SN38 PDC Builder - Test")
+    print("Reference: NRP1 + LIN00169 + SN-38")
+    print("=" * 70 + "\n")
+
+    test_peptides = [
+        "NCC(=O)N[C@@H](CS)C(=O)NCC(=O)O",  # Gly-Cys-Gly
+    ]
+
+    builder = DisulfanylethOxyPentanoicAcidSN38_PDC_Builder()
+
+    for i, peptide in enumerate(test_peptides):
+        print(f"\n--- Test peptide {i+1} ---")
+        print(f"Input: {peptide}")
+
+        try:
+            conjugates = builder.conjugate_to_peptide(peptide)
+            print(f" Generated {len(conjugates)} conjugate(s)")
+            for j, conj in enumerate(conjugates):
+                valid, checks = builder.validate_conjugate(conj)
+                print(f"  Conjugate {j+1}: {'PASSED' if valid else 'FAILED'}")
+        except Exception as e:
+            print(f" FAILED: {e}")
+            import traceback
+            traceback.print_exc()

smiles/conjugation/DSPACpt.py

@@ -0,0 +1,517 @@
+"""
+Disulfanylpropanoic Acid - Camptothecin PDC Builder (DSPACpt)
+
+Conjugation chemistry: Lysine ε-amine + Cys-SS-Linker-Camptothecin (amide bond formation)
+
+Reference PDC:
+- PDC_00335 (SMAC-P1 CPT conjugate 5)
+
+Receptor: ERBB2 (Receptor tyrosine-protein kinase erbB-2)
+Linker: LIN00012 - 3-Disulfanylpropanoic Acid (disulfide + ester)
+Payload: DRG00009 - Camptothecin
+Architecture: Monomer, Lys side chain conjugation
+
+Chemistry Note:
+===============
+LIN00012 (3-Disulfanylpropanoic Acid) provides a disulfide-ester linkage:
+    HS─S─CH₂─CH₂─C(═O)─OH
+
+Linker-to-Drug Connection (Ester):
+    Cys─S─S─CH₂─CH₂─C(═O)─O─Camptothecin
+                      │
+                    ester bond (C-20 OH of CPT)
+
+Full PDC Architecture:
+    Peptide──Lys──NH──C(═O)──Cys──S─S──CH₂─CH₂─C(═O)─O──Camptothecin
+                  └── amide ──┘       └── disulfide ──┘  └── ester ──┘
+
+Note: This linker chemistry is equivalent to the SSCpt ester variant (Z7/PDC_00294).
+
+Usage:
+======
+    builder = DisulfanylpropanoicAcidCpt_PDC_Builder()
+    conjugates = builder.conjugate_to_peptide(peptide_smiles)
+"""
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, Descriptors
+
+from .pdc_builder import PDC_Builder
+
+
+# ============================================================================
+# Disulfanylpropanoic Acid - Camptothecin PDC Builder
+# ============================================================================
+
+class DisulfanylpropanoicAcidCpt_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for 3-Disulfanylpropanoic Acid + Camptothecin conjugates
+
+    Conjugation chemistry (amide bond formation):
+    - Peptide site: Lysine ε-amine (side chain -NH2)
+    - Linker-payload site: Carboxylic acid of Cys (Cys-SS-Linker-Drug)
+    - Reaction: Lys-NH2 + HOOC-Cys-SS-Linker-Drug → Lys-NH-CO-Cys-SS-Linker-Drug
+
+    Linker: LIN00012 - 3-Disulfanylpropanoic Acid
+        SMILES: C(CSS)C(=O)O
+        IUPAC: 3-(disulfanyl)propanoic acid
+        Connectivity: HS-S-CH2-CH2-C(=O)-OH
+        MW: 138.21, MF: C3H6O2S2
+
+    Drug: DRG00009 - Camptothecin
+        Canonical SMILES: CCC1(O)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1
+        Connection: Ester at C-20 hydroxyl
+
+    Reference: PDC_00335 (SMAC-P1 CPT conjugate 5, ERBB2)
+    """
+
+    # ========================================================================
+    # Cys-SS-Linker-Drug SMILES
+    # ========================================================================
+    #
+    # Pre-formed Cys-Disulfanylpropanoic acid-Camptothecin assembly:
+    #   HOOC─Cys(NH₂)─CH₂─S─S─CH₂─CH₂─C(═O)─O─CPT(C-20)
+    #
+    # Cys provides COOH for amide bond with Lys ε-amine
+    # Disulfide bond (S-S) is cleavable inside cells
+    # Ester bond connects linker to CPT C-20 hydroxyl
+    #
+    # ========================================================================
+
+    CYS_SS_LINKER_CPT_SMILES = (
+        "N[C@H](CSSCCC(=O)OC1(CC)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1)C(=O)O"
+    )
+
+    def __init__(self, cys_linker_drug_smiles=None):
+        """
+        Initialize Disulfanylpropanoic Acid-Camptothecin builder
+
+        Args:
+            cys_linker_drug_smiles: Optional custom Cys-SS-Linker-Drug SMILES.
+                                   If None, uses default LIN00012+CPT assembly.
+                                   Must have free carboxylic acid for conjugation.
+        """
+        if cys_linker_drug_smiles is None:
+            cys_linker_drug_smiles = self.CYS_SS_LINKER_CPT_SMILES
+
+        super().__init__(
+            cys_linker_drug_smiles,
+            component_name="Cys-SS-DisulfanylpropanoicAcid-Camptothecin (LIN00012+DRG00009)"
+        )
+
+    def _verify_reactive_group(self):
+        """Verify carboxylic acid group is present for amide bond formation"""
+        # Pattern: Carboxylic acid (COOH)
+        carboxylic_acid_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+
+        if not self.linker_payload_mol.HasSubstructMatch(carboxylic_acid_pattern):
+            print(" Warning: No carboxylic acid group detected in Cys-SS-Linker-Drug")
+        else:
+            print(" Cys-SS-DSPA-Camptothecin loaded with carboxylic acid reactive group")
+
+        # Verify disulfide bond is present
+        disulfide_pattern = Chem.MolFromSmarts('[#16]-[#16]')
+        if self.linker_payload_mol.HasSubstructMatch(disulfide_pattern):
+            print(" Disulfide bond (S-S) detected")
+        else:
+            print(" Warning: No disulfide bond detected - check structure")
+
+        # Verify Cys α-amine is present
+        cys_amine_pattern = Chem.MolFromSmarts('[NH2][CH]([CH2]S)C(=O)')
+        if self.linker_payload_mol.HasSubstructMatch(cys_amine_pattern):
+            print(" Cys α-amine detected (will be free in conjugate)")
+        else:
+            print(" Warning: Cys α-amine pattern not detected")
+
+        # Verify ester bond to drug
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        if self.linker_payload_mol.HasSubstructMatch(ester_pattern):
+            print(" Ester bond (linker-to-drug) detected")
+        else:
+            print(" Warning: No ester bond detected - check linker-drug connection")
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find all lysine ε-amine groups in peptide
+
+        Returns list of dicts with site information including atom index
+        """
+        return self.find_peptide_lysines(peptide_smiles)
+
+    def find_peptide_lysines(self, peptide_smiles):
+        """
+        Find all lysine residues with free ε-amine in peptide
+
+        Lysine side chain: -CH2-CH2-CH2-CH2-NH2 (CCCCN)
+
+        Returns:
+            List of dicts with 'atom_idx' (nitrogen index) and metadata
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        lysines = []
+
+        # SMARTS pattern for Lysine ε-amine: 4 CH2 groups ending in NH2
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        matches = peptide_mol.GetSubstructMatches(lys_pattern)
+
+        for match in matches:
+            nitrogen_idx = match[4]
+            lysines.append({
+                'atom_idx': nitrogen_idx,
+                'side_chain_atoms': list(match),
+                'type': 'Lys_epsilon_amine'
+            })
+
+        # Count N-terminus for informational purposes
+        n_terminus_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        n_term_matches = peptide_mol.GetSubstructMatches(n_terminus_pattern)
+
+        if n_term_matches:
+            print(f"  (Also found {len(n_term_matches)} N-terminus α-amine(s) - excluded by default)")
+
+        return lysines
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """
+        Perform amide bond conjugation between NH2 groups and Cys carboxylic acid
+
+        Args:
+            peptide_mol: RDKit Mol object
+            site_info: Not used (kept for compatibility)
+            exclude_n_terminus: If True, skip N-terminus conjugation
+
+        Returns: List of all possible conjugate SMILES
+        """
+        if exclude_n_terminus:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:5][CH2][CH2][CH2][CH2:4][NH2:1].[C:2](=[O:3])[OH]>>[*:5][CH2][CH2][CH2][CH2:4][NH:1][C:2](=[O:3])'
+            )
+        else:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:4][NH2:1].[C:2](=[O:3])[OH]>>[*:4][NH:1][C:2](=[O:3])'
+            )
+
+        products = amide_rxn.RunReactants((peptide_mol, self.linker_payload_mol))
+
+        if not products:
+            raise ValueError("Amide bond formation reaction failed - no products generated")
+
+        conjugates = []
+        seen_smiles = set()
+
+        for product_tuple in products:
+            try:
+                conjugate_mol = product_tuple[0]
+                Chem.SanitizeMol(conjugate_mol)
+                conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+
+                if conjugate_smiles not in seen_smiles:
+                    seen_smiles.add(conjugate_smiles)
+                    conjugates.append(conjugate_smiles)
+            except Exception:
+                continue
+
+        if not conjugates:
+            raise ValueError("All products failed sanitization")
+
+        return conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Cys-SS-DSPA-Camptothecin to peptide via amide bond at NH2 groups
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Lys with NH2)
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            List of all possible conjugate SMILES
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Lys ε-amine site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2 (including N-terminus)"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [{site_type}] [DSPA ester]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides
+
+        Args:
+            peptide_smiles_list: List of peptide SMILES strings
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates per peptide
+
+        Returns:
+            Dict with 'conjugates', 'successful', 'failed', 'total_conjugates'
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        print(f"\n=== Summary ===")
+        print(f"Peptides: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, Failed: {len(results['failed'])}")
+        print(f"Total conjugates generated: {results['total_conjugates']}")
+
+        return results
+
+    def flatten_batch_results(self, results):
+        """Flatten batch conjugation results into a single list"""
+        all_conjugates = []
+        for conjugate_list in results['conjugates']:
+            if conjugate_list is not None:
+                if isinstance(conjugate_list, list):
+                    all_conjugates.extend(conjugate_list)
+                else:
+                    all_conjugates.append(conjugate_list)
+        return all_conjugates
+
+    # ========================================================================
+    # Attachment Site Compatibility Checks
+    # ========================================================================
+
+    def has_lys_cys_amide_conjugation(self, mol):
+        """Check if molecule has Lys-Cys amide conjugation bond"""
+        lys_cys_amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        return mol.HasSubstructMatch(lys_cys_amide_pattern)
+
+    def has_disulfide_bond(self, mol):
+        """Check if molecule has disulfide bond (S-S)"""
+        disulfide_pattern = Chem.MolFromSmarts('[SX2]-[SX2]')
+        return mol.HasSubstructMatch(disulfide_pattern)
+
+    def has_unreacted_lys_epsilon_amine(self, mol):
+        """Check for unreacted Lys ε-amine (free NH2 on 4-carbon side chain)"""
+        lys_amine_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        return mol.HasSubstructMatch(lys_amine_pattern)
+
+    def has_unreacted_cys_carboxylic_acid(self, mol):
+        """Check for unreacted Cys carboxylic acid from Cys-SS-Linker-Drug"""
+        cys_cooh_pattern = Chem.MolFromSmarts('[NH2][CH]([CH2][SX2][SX2])[CX3](=[OX1])[OX2H1]')
+        return mol.HasSubstructMatch(cys_cooh_pattern)
+
+    def has_ester_bond_to_drug(self, mol):
+        """Check for ester bond connecting linker to drug"""
+        ester_pattern = Chem.MolFromSmarts('[SX2][SX2][CH2][CH2][CX3](=[OX1])[OX2][C]')
+        return mol.HasSubstructMatch(ester_pattern)
+
+    def find_amide_conjugation_sites(self, mol):
+        """Find all Lys-Cys amide conjugation sites in the molecule"""
+        lys_cys_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        matches = mol.GetSubstructMatches(lys_cys_pattern)
+
+        sites = []
+        for match in matches:
+            nitrogen_idx = match[0]
+            carbonyl_c_idx = match[1]
+            cys_alpha_idx = match[3]
+            sites.append((nitrogen_idx, carbonyl_c_idx, cys_alpha_idx))
+
+        return sites
+
+    def check_valency_at_conjugation_site(self, mol, site_indices):
+        """Check that valency is correct at Lys-Cys conjugation site"""
+        nitrogen_idx, carbonyl_c_idx, cys_alpha_idx = site_indices
+
+        try:
+            n_atom = mol.GetAtomWithIdx(nitrogen_idx)
+            if n_atom.GetTotalDegree() != 3:
+                return False
+
+            c_atom = mol.GetAtomWithIdx(carbonyl_c_idx)
+            if c_atom.GetTotalDegree() != 3:
+                return False
+
+            cys_c_atom = mol.GetAtomWithIdx(cys_alpha_idx)
+            if cys_c_atom.GetTotalDegree() != 4:
+                return False
+
+            return True
+        except Exception:
+            return False
+
+    def check_attachment_site_compatibility(self, conjugate_smiles):
+        """
+        Check attachment site compatibility for Lys-Cys-SS-DSPA-CPT conjugation
+
+        Validates:
+        1. Lys-Cys amide conjugation is present
+        2. Disulfide bond (S-S) is intact
+        3. Valency at conjugation sites is correct
+        4. No unreacted Cys-COOH (unconjugated linker-drug)
+        5. Ester bond to drug is present
+
+        Returns:
+            Tuple of (passed: bool, details: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'parseable': False, 'error': 'Could not parse SMILES'}
+
+        details = {
+            'parseable': True,
+            'validation_method': 'lys_cys_amide_disulfide_ester',
+            'linker': 'LIN00012 (3-Disulfanylpropanoic Acid)',
+            'drug': 'DRG00009 (Camptothecin)',
+            'pdc_ref': 'PDC_00335'
+        }
+
+        details['has_lys_cys_amide'] = self.has_lys_cys_amide_conjugation(mol)
+        details['has_disulfide_bond'] = self.has_disulfide_bond(mol)
+
+        conjugation_sites = self.find_amide_conjugation_sites(mol)
+        details['conjugation_sites'] = len(conjugation_sites)
+
+        valency_ok = True
+        if conjugation_sites:
+            for site in conjugation_sites:
+                if not self.check_valency_at_conjugation_site(mol, site):
+                    valency_ok = False
+                    break
+        details['valency_correct'] = valency_ok
+
+        details['no_unreacted_cys_cooh'] = not self.has_unreacted_cys_carboxylic_acid(mol)
+        details['has_unreacted_lys_amine'] = self.has_unreacted_lys_epsilon_amine(mol)
+
+        passed = all([
+            details['has_lys_cys_amide'],
+            details['has_disulfide_bond'],
+            details['valency_correct'],
+            details['no_unreacted_cys_cooh']
+        ])
+
+        return passed, details
+
+    def _basic_validation(self, mol):
+        """
+        Basic validation checks for DSPA-Camptothecin conjugates
+
+        Checks:
+        1. Amide bond present (Lys-Cys conjugation)
+        2. Disulfide bond present (S-S linker)
+        3. Specific Lys-Cys-SS conjugation pattern
+        4. Camptothecin core present
+        5. Ester bond to drug present
+        """
+        checks = {}
+
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[#6]')
+        checks['has_amide_bond'] = mol.HasSubstructMatch(amide_pattern)
+
+        disulfide_pattern = Chem.MolFromSmarts('[#16]-[#16]')
+        checks['has_disulfide_bond'] = mol.HasSubstructMatch(disulfide_pattern)
+
+        lys_cys_ss_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        checks['has_lys_cys_ss_conjugation'] = mol.HasSubstructMatch(lys_cys_ss_pattern)
+
+        camptothecin_core = Chem.MolFromSmarts('c1ccc2ncc3C(=O)OCc3c2c1')
+        checks['has_camptothecin_core'] = mol.HasSubstructMatch(camptothecin_core)
+
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        checks['has_ester_bond'] = mol.HasSubstructMatch(ester_pattern)
+
+        checks['_passed'] = all([
+            checks['has_amide_bond'],
+            checks['has_disulfide_bond'],
+            checks['has_lys_cys_ss_conjugation']
+        ])
+
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """
+        Validate that the conjugate is chemically reasonable
+
+        Returns:
+            (passed: bool, checks: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+
+        return passed, checks
+
+
+# ============================================================================
+# Main - Test Examples
+# ============================================================================
+
+if __name__ == "__main__":
+
+    print("=" * 70)
+    print("Disulfanylpropanoic Acid-Camptothecin PDC Builder - Test")
+    print("Reference: PDC_00335 (SMAC-P1 CPT conjugate 5, LIN00012)")
+    print("=" * 70 + "\n")
+
+    # Test peptides
+    test_peptides = [
+        # Gly-Lys-Gly (has N-terminus NH2 + Lys ε-amine)
+        "NCC(=O)N[C@@H](CCCCN)C(=O)NCC(=O)O",
+        # Lys-Arg
+        "N[C@@H](CCCCN)C(=O)N[C@@H](CCCNC(=N)N)C(=O)O",
+    ]
+
+    builder = DisulfanylpropanoicAcidCpt_PDC_Builder()
+
+    for i, peptide in enumerate(test_peptides):
+        print(f"\n--- Test peptide {i+1} ---")
+        print(f"Input: {peptide}")
+
+        try:
+            conjugates = builder.conjugate_to_peptide(peptide, exclude_n_terminus=True)
+            print(f" Generated {len(conjugates)} conjugate(s)")
+
+            for j, conj in enumerate(conjugates):
+                valid, checks = builder.validate_conjugate(conj)
+                print(f"  Conjugate {j+1}: {'PASSED' if valid else 'FAILED'} ({len(conj)} chars)")
+
+        except Exception as e:
+            print(f" FAILED: {e}")
+            import traceback
+            traceback.print_exc()

smiles/conjugation/DSPASN38.py

@@ -0,0 +1,515 @@
+"""
+Disulfanylpropanoic Acid - SN-38 PDC Builder (DSPASN38)
+
+Conjugation chemistry: Cysteine thiol + Linker thiol → Disulfide bond (S-S)
+
+Reference PDCs:
+- PDC_02099 (SN38(10)-CRGDK)
+- PDC_02104 (SN38(10)-(Ac)CRGDK)
+- PDC_02102 (SN38(10)-(Ac)CGVQIRGDK)
+- PDC_02107 (SN38(10)-CRRGDK)
+- PDC_02110 (SN38(10)-(Ac)CDRGDK)
+- PDC_02112 (SN38(20)-CRGDK)
+- PDC_02113 ((MOM)SN38(20)-CRGDK)
+- PDC_02114 ((MOM)SN38(20)-(Ac)CRGDK)
+- PDC_02115 ((MOM)SN38(20)-(Ac)CGVQIRGDK)
+
+Receptor: NRP1 (Neuropilin-1)
+Linker: LIN00149 - 3-Disulfanylpropanoic Acid (disulfide + ester)
+Payload: DRG - 7-Ethyl-10-hydroxycamptothecin (SN-38)
+Architecture: Monomer, Cys thiol conjugation (disulfide bond)
+
+Chemistry Note:
+===============
+LIN00149 (3-Disulfanylpropanoic Acid) connects via:
+    HS─CH₂─CH₂─C(═O)─OH
+
+Linker-to-Drug Connection (Ester):
+    SN-38(10-O)─C(═O)─CH₂─CH₂─SH     (position 10, default)
+    SN-38(20-O)─C(═O)─CH₂─CH₂─SH     (position 20)
+
+Full PDC Architecture:
+    Peptide──Cys──S─S──CH₂─CH₂─C(═O)─O──SN-38
+                  └── disulfide ──┘  └── ester ──┘
+
+Usage:
+======
+    builder = DisulfanylpropanoicAcidSN38_PDC_Builder()
+    conjugates = builder.conjugate_to_peptide(peptide_smiles)
+"""
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, Descriptors
+
+from .pdc_builder import PDC_Builder
+
+
+# ============================================================================
+# Disulfanylpropanoic Acid - SN-38 PDC Builder
+# ============================================================================
+
+class DisulfanylpropanoicAcidSN38_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for 3-Disulfanylpropanoic Acid + SN-38 conjugates
+
+    Conjugation chemistry (disulfide bond formation):
+    - Peptide site: Cysteine thiol (side chain -SH)
+    - Linker-payload site: Terminal thiol (-SH) on DSPA linker
+    - Reaction: Peptide-Cys-SH + HS-Linker-Drug → Peptide-Cys-S-S-Linker-Drug
+
+    Linker: LIN00149 - 3-Disulfanylpropanoic Acid
+        SMILES: OC(=O)CCSS
+        IUPAC: 3-(disulfanyl)propanoic acid
+        Connectivity: HS-S-CH2-CH2-C(=O)-OH
+
+    Drug: 7-Ethyl-10-hydroxycamptothecin (SN-38)
+        Canonical SMILES: CCC1(O)C(=O)OCc2c1cc1n(c2=O)Cc2c(CC)c3cc(O)ccc3nc2-1
+        Connection: Ester at C-10 hydroxyl (default) or C-20 hydroxyl
+
+    Supported drug conjugation positions:
+    - '10' (default): Ester at 10-OH (phenolic, most common in PDCdb)
+    - '20': Ester at 20-OH (tertiary aliphatic)
+    - 'all': Both positions, conjugates pooled
+
+    Reference: PDC_02099, PDC_02112 (NRP1)
+    """
+
+    # ========================================================================
+    # Linker-Drug SMILES Variants
+    # ========================================================================
+    #
+    # DSPA linker ester to SN-38 with free terminal thiol for Cys conjugation
+    #
+    # Position 10 (phenolic OH):
+    #   HS─CH₂─CH₂─C(═O)─O─SN38(10)
+    #
+    # Position 20 (tertiary OH):
+    #   HS─CH₂─CH₂─C(═O)─O─SN38(20)
+    #
+    # ========================================================================
+
+    # SN-38 with DSPA ester at C-10 (phenolic OH), free thiol
+    POS10_LINKER_SN38_SMILES = (
+        "CCC1(O)C(=O)OCc2c1cc1n(c2=O)Cc2c(CC)c3cc(OC(=O)CCS)ccc3nc2-1"
+    )
+
+    # SN-38 with DSPA ester at C-20 (tertiary OH), free thiol
+    POS20_LINKER_SN38_SMILES = (
+        "CCC1(OC(=O)CCS)C(=O)OCc2c1cc1n(c2=O)Cc2c(CC)c3cc(O)ccc3nc2-1"
+    )
+
+    # Default: position 10 (most common)
+    DEFAULT_LINKER_SN38_SMILES = POS10_LINKER_SN38_SMILES
+
+    # Available drug positions
+    DRUG_POSITIONS = {
+        '10': POS10_LINKER_SN38_SMILES,
+        '20': POS20_LINKER_SN38_SMILES,
+    }
+
+    def __init__(self, linker_drug_smiles=None, drug_position='all'):
+        """
+        Initialize DSPA-SN38 builder
+
+        Args:
+            linker_drug_smiles: Optional custom Linker-Drug SMILES.
+                               If provided, overrides drug_position.
+                               Must have free thiol for Cys disulfide conjugation.
+            drug_position: 'all' (default), '10', or '20'.
+                          '10': ester at C-10 phenolic OH (most common)
+                          '20': ester at C-20 tertiary OH
+                          'all': both positions, conjugates pooled
+                          Ignored if linker_drug_smiles is provided.
+        """
+        self.drug_position = drug_position
+
+        if linker_drug_smiles is not None:
+            self._variants = [linker_drug_smiles]
+            self._variant_names = ['custom']
+            super().__init__(
+                linker_drug_smiles,
+                component_name="Custom-DSPA-SN38 (NRP1)"
+            )
+        elif drug_position == 'all':
+            self._variants = [self.POS10_LINKER_SN38_SMILES, self.POS20_LINKER_SN38_SMILES]
+            self._variant_names = ['pos10', 'pos20']
+            super().__init__(
+                self.POS10_LINKER_SN38_SMILES,
+                component_name="DSPA-SN38 (LIN00149, all positions)"
+            )
+        else:
+            if drug_position not in self.DRUG_POSITIONS:
+                raise ValueError(f"Unknown drug_position '{drug_position}'. "
+                                 f"Choose from: 'all', '10', '20'")
+            smiles = self.DRUG_POSITIONS[drug_position]
+            self._variants = [smiles]
+            self._variant_names = [f'pos{drug_position}']
+            super().__init__(
+                smiles,
+                component_name=f"DSPA-SN38 (LIN00149, C-{drug_position})"
+            )
+
+    def _verify_reactive_group(self):
+        """Verify thiol group is present for disulfide bond formation"""
+        # Check for free thiol (SH)
+        for atom in self.linker_payload_mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:  # Sulfur
+                if atom.GetTotalNumHs() >= 1 and atom.GetDegree() == 1:
+                    print(" DSPA-SN38 loaded with free thiol reactive group")
+                    break
+        else:
+            print(" Warning: No free thiol group detected in DSPA-SN38")
+
+        # Verify ester bond to drug
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][c]')
+        if self.linker_payload_mol.HasSubstructMatch(ester_pattern):
+            print(" Ester bond (linker-to-drug) detected")
+        else:
+            # Try broader ester pattern
+            ester_broad = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+            if self.linker_payload_mol.HasSubstructMatch(ester_broad):
+                print(" Ester bond (linker-to-drug) detected")
+            else:
+                print(" Warning: No ester bond detected - check linker-drug connection")
+
+    # ------------------------------------------------------------------
+    # Conjugation site finding (Cysteine thiol)
+    # ------------------------------------------------------------------
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find all free cysteine thiol groups in peptide
+
+        Returns list of dicts with site information including atom index
+        """
+        return self.find_peptide_cysteines(peptide_smiles)
+
+    def find_peptide_cysteines(self, peptide_smiles):
+        """
+        Find all free cysteine residues (thiol groups) in peptide
+
+        Iterates over all atoms to find S with free H (thiol, not disulfide).
+
+        Returns:
+            List of dicts with 'atom_idx' (sulfur index) and metadata
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        cysteines = []
+
+        for atom in peptide_mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:  # Sulfur
+                total_hs = atom.GetTotalNumHs()
+                degree = atom.GetDegree()
+
+                # Thiol: S with 1 H and bonded to 1 carbon (not disulfide S-S)
+                if total_hs >= 1 and degree == 1:
+                    cysteines.append({
+                        'atom_idx': atom.GetIdx(),
+                        'num_hydrogens': total_hs,
+                        'neighbors': [peptide_mol.GetAtomWithIdx(n.GetIdx()).GetSymbol()
+                                      for n in atom.GetNeighbors()],
+                        'type': 'Cys_thiol'
+                    })
+
+        return cysteines
+
+    # ------------------------------------------------------------------
+    # Conjugation reaction (disulfide bond formation)
+    # ------------------------------------------------------------------
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """
+        Form disulfide bond between peptide Cys thiol and linker-drug thiol.
+
+        Args:
+            peptide_mol: RDKit Mol object
+            site_info: Not used (kept for compatibility)
+            exclude_n_terminus: Not used for Cys conjugation (kept for interface)
+
+        Returns: List of all possible conjugate SMILES
+        """
+        all_conjugates = []
+        seen = set()
+
+        for variant_smiles in self._variants:
+            variant_mol = Chem.MolFromSmiles(variant_smiles)
+            if variant_mol is None:
+                continue
+
+            # SMARTS: form disulfide between any two thiols
+            # Reactant 0 = peptide (Cys-CH2-SH), Reactant 1 = linker-drug (-CH2-SH)
+            disulfide_rxn = AllChem.ReactionFromSmarts(
+                '[C:1][SH:2].[SH:3][C:4]>>[C:1][S:2][S:3][C:4]'
+            )
+
+            products = disulfide_rxn.RunReactants((peptide_mol, variant_mol))
+
+            if not products:
+                continue
+
+            for product_tuple in products:
+                try:
+                    conjugate_mol = product_tuple[0]
+                    Chem.SanitizeMol(conjugate_mol)
+                    conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+
+                    if conjugate_smiles not in seen:
+                        seen.add(conjugate_smiles)
+                        all_conjugates.append(conjugate_smiles)
+                except Exception:
+                    continue
+
+        if not all_conjugates:
+            raise ValueError("Disulfide bond formation failed - no products generated")
+
+        return all_conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate DSPA-SN38 to peptide via disulfide bond at Cys thiol.
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Cys with free thiol)
+            exclude_n_terminus: Not used for Cys conjugation (kept for interface)
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            List of all possible conjugate SMILES
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Cys thiol site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        pos_label = f"C-{self.drug_position}" if self.drug_position != 'all' else "C-10 + C-20"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [Cys disulfide] [DSPA {pos_label}]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides.
+
+        Returns:
+            Dict with 'conjugates', 'successful', 'failed', 'total_conjugates'
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        print(f"\n=== Summary ===")
+        print(f"Peptides: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, "
+              f"Failed: {len(results['failed'])}")
+        print(f"Total conjugates generated: {results['total_conjugates']}")
+
+        return results
+
+    def flatten_batch_results(self, results):
+        """Flatten batch conjugation results into a single list"""
+        all_conjugates = []
+        for conjugate_list in results['conjugates']:
+            if conjugate_list is not None:
+                if isinstance(conjugate_list, list):
+                    all_conjugates.extend(conjugate_list)
+                else:
+                    all_conjugates.append(conjugate_list)
+        return all_conjugates
+
+    # ========================================================================
+    # Attachment Site Compatibility Checks (Cys disulfide)
+    # ========================================================================
+
+    def has_disulfide_bond(self, mol):
+        """Check if molecule has disulfide bond (S-S)"""
+        disulfide_pattern = Chem.MolFromSmarts('[SX2]-[SX2]')
+        return mol.HasSubstructMatch(disulfide_pattern)
+
+    def has_free_thiol(self, mol):
+        """Check if molecule has unconjugated (free) thiol groups"""
+        for atom in mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:
+                if atom.GetTotalNumHs() >= 1 and atom.GetDegree() == 1:
+                    return True
+        return False
+
+    def has_ester_bond_to_drug(self, mol):
+        """Check for ester bond connecting linker to drug"""
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        return mol.HasSubstructMatch(ester_pattern)
+
+    def has_sn38_core(self, mol):
+        """Check for SN-38 core structure (camptothecin with 7-ethyl)"""
+        # Camptothecin-like core: fused ring system with lactone
+        cpt_core = Chem.MolFromSmarts('c1ccc2ncc3C(=O)OCc3c2c1')
+        return mol.HasSubstructMatch(cpt_core)
+
+    def check_attachment_site_compatibility(self, conjugate_smiles):
+        """
+        Check attachment site compatibility for Cys-SS-DSPA-SN38 conjugation.
+
+        Validates:
+        1. Disulfide bond (S-S) is present
+        2. No free thiol remaining
+        3. Ester bond to drug is present
+
+        Returns:
+            Tuple of (passed: bool, details: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'parseable': False, 'error': 'Could not parse SMILES'}
+
+        details = {
+            'parseable': True,
+            'validation_method': 'cys_disulfide_ester',
+            'linker': 'LIN00149 (3-Disulfanylpropanoic Acid)',
+            'drug': 'SN-38 (7-Ethyl-10-hydroxycamptothecin)',
+        }
+
+        details['has_disulfide_bond'] = self.has_disulfide_bond(mol)
+        details['no_free_thiol'] = not self.has_free_thiol(mol)
+        details['has_ester_bond'] = self.has_ester_bond_to_drug(mol)
+
+        passed = all([
+            details['has_disulfide_bond'],
+            details['no_free_thiol'],
+        ])
+
+        return passed, details
+
+    def _basic_validation(self, mol):
+        """
+        Basic validation checks for DSPA-SN38 conjugates.
+
+        Checks:
+        1. Disulfide bond present (Cys-Linker conjugation)
+        2. No free thiol remaining
+        3. Ester bond to drug present
+        4. SN-38 core present
+        """
+        checks = {}
+
+        disulfide_pattern = Chem.MolFromSmarts('[#16]-[#16]')
+        checks['has_disulfide_bond'] = mol.HasSubstructMatch(disulfide_pattern)
+
+        checks['no_free_thiol'] = not self.has_free_thiol(mol)
+
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        checks['has_ester_bond'] = mol.HasSubstructMatch(ester_pattern)
+
+        checks['has_sn38_core'] = self.has_sn38_core(mol)
+
+        checks['_passed'] = all([
+            checks['has_disulfide_bond'],
+            checks['no_free_thiol'],
+        ])
+
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """
+        Validate that the conjugate is chemically reasonable.
+
+        Returns:
+            (passed: bool, checks: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+
+        return passed, checks
+
+
+# ============================================================================
+# Main - Test Examples
+# ============================================================================
+
+if __name__ == "__main__":
+
+    print("=" * 70)
+    print("DSPA-SN38 PDC Builder - Test")
+    print("Reference: NRP1 + LIN00149 + SN-38")
+    print("=" * 70 + "\n")
+
+    # Test peptides with free cysteine (NRP1-targeting sequences)
+    test_peptides = [
+        # Cys-Arg-Gly-Asp-Lys (CRGDK) - reference NRP1 peptide
+        "N[C@@H](CS)C(=O)N[C@@H](CCCNC(=N)N)C(=O)NCC(=O)N[C@@H](CC(=O)O)C(=O)N[C@@H](CCCCN)C(=O)O",
+        # Simple Gly-Cys-Gly
+        "NCC(=O)N[C@@H](CS)C(=O)NCC(=O)O",
+    ]
+
+    print("Test 1: Default mode (all positions: C-10 + C-20)")
+    print("-" * 50)
+
+    builder = DisulfanylpropanoicAcidSN38_PDC_Builder()
+
+    for i, peptide in enumerate(test_peptides):
+        print(f"\n--- Test peptide {i+1} ---")
+        print(f"Input: {peptide[:60]}...")
+
+        try:
+            conjugates = builder.conjugate_to_peptide(peptide)
+            print(f" Generated {len(conjugates)} conjugate(s)")
+
+            for j, conj in enumerate(conjugates):
+                valid, checks = builder.validate_conjugate(conj)
+                print(f"  Conjugate {j+1}: {'PASSED' if valid else 'FAILED'} "
+                      f"({len(conj)} chars)")
+
+        except Exception as e:
+            print(f" FAILED: {e}")
+            import traceback
+            traceback.print_exc()
+
+    print("\n\nTest 2: Position-10 only")
+    print("-" * 50)
+
+    builder_10 = DisulfanylpropanoicAcidSN38_PDC_Builder(drug_position='10')
+
+    try:
+        conjugates = builder_10.conjugate_to_peptide(test_peptides[1])
+        print(f" Generated {len(conjugates)} conjugate(s) [C-10 only]")
+    except Exception as e:
+        print(f" FAILED: {e}")

smiles/conjugation/DauFactory.py

@@ -0,0 +1,298 @@
+"""
+Daunorubicin PDC Factory - Drug-level builder with linker sampling
+
+Instead of specifying a particular linker-drug combo (e.g. AoaDau, ValAlaDau),
+the user specifies only the drug (Daunorubicin) and a linker is sampled
+according to GNRHR+Dau entry frequencies from PDCdb.
+
+Linker frequencies (GNRHR + Daunorubicin):
+===========================================
+    Aminooxyacetic acid (Aoa)              (AoaDau)       - 21 entries
+    Val-Ala-PABC + glutaryl spacer         (ValAlaDau)    -  2 entries
+    Val-Cit-PABC + glutaryl spacer         (ValCitDau)    -  2 entries
+    Glutaric acid                          (GlutarylDau)  -  2 entries
+                                                     Total: 27 entries
+
+Sampling probabilities:
+    AoaDau:       21/27 ≈ 0.778
+    ValAlaDau:     2/27 ≈ 0.074
+    ValCitDau:     2/27 ≈ 0.074
+    GlutarylDau:   2/27 ≈ 0.074
+
+Usage:
+======
+    # Drug-only mode - linker is sampled per peptide
+    factory = DaunorubicinPDC_Factory()
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+
+    # Fixed linker mode - use a specific linker
+    factory = DaunorubicinPDC_Factory(linker='AoaDau')
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+
+    # All linkers mode - generate conjugates with every linker
+    factory = DaunorubicinPDC_Factory(linker='all')
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+"""
+
+import numpy as np
+import contextlib
+import io
+
+from .AoaDau import AoaDaunorubicin_PDC_Builder
+from .ValAlaDau import ValAlaPABCDaunorubicin_PDC_Builder
+from .ValCitDau import ValCitPABCDaunorubicin_PDC_Builder
+from .GlutarylDau import GlutarylDaunorubicin_PDC_Builder
+
+
+# ============================================================================
+# Linker registry: name → (builder_class, PDCdb frequency, description)
+# ============================================================================
+
+GNRHR_DAU_LINKER_REGISTRY = {
+    'AoaDau': {
+        'class': AoaDaunorubicin_PDC_Builder,
+        'frequency': 21,
+        'linker_id': 'LIN00019',
+        'description': 'Aminooxyacetic acid (Aoa) - oxime bond to Dau ketone',
+        'pdc_refs': ['PDC_00011', 'PDC_00208'],
+    },
+    'ValAlaDau': {
+        'class': ValAlaPABCDaunorubicin_PDC_Builder,
+        'frequency': 2,
+        'linker_id': 'LIN00028',
+        'description': 'Val-Ala dipeptide + PABC self-immolative + glutaryl spacer',
+        'pdc_refs': ['PDC_00185', 'PDC_00193'],
+    },
+    'ValCitDau': {
+        'class': ValCitPABCDaunorubicin_PDC_Builder,
+        'frequency': 2,
+        'linker_id': 'LIN00005',
+        'description': 'Val-Cit dipeptide + PABC self-immolative + glutaryl spacer',
+        'pdc_refs': ['PDC_00189', 'PDC_00194'],
+    },
+    'GlutarylDau': {
+        'class': GlutarylDaunorubicin_PDC_Builder,
+        'frequency': 2,
+        'linker_id': 'LIN00033',
+        'description': 'Glutaric acid - simple dicarboxylic acid spacer',
+        'pdc_refs': ['PDC_00190', 'PDC_00195'],
+    },
+}
+
+
+def _compute_linker_probabilities(registry=None):
+    """Compute sampling probabilities from frequency counts."""
+    if registry is None:
+        registry = GNRHR_DAU_LINKER_REGISTRY
+    names = list(registry.keys())
+    freqs = np.array([registry[n]['frequency'] for n in names], dtype=np.float64)
+    probs = freqs / freqs.sum()
+    return names, probs
+
+
+class DaunorubicinPDC_Factory:
+    """
+    Drug-level PDC factory for Daunorubicin conjugates.
+
+    Wraps all four GNRHR+Dau linker builders and provides:
+    - 'sample' mode (default): randomly pick a linker per call according to
+      PDCdb entry frequencies
+    - 'all' mode: generate conjugates with every linker
+    - fixed linker mode: always use a specific linker (e.g. 'AoaDau')
+
+    The interface mirrors the individual builder's `conjugate_to_peptide` so it
+    can be used as a drop-in replacement in the Conjugation classifier.
+    """
+
+    def __init__(self, linker='sample', rng_seed=None):
+        """
+        Args:
+            linker: One of:
+                - 'sample' (default): sample linker per peptide according to
+                  GNRHR+Dau frequency weights
+                - 'all': use every linker and pool conjugates
+                - A specific linker name ('AoaDau', 'ValAlaDau', 'ValCitDau',
+                  'GlutarylDau')
+            rng_seed: Optional seed for reproducible linker sampling
+        """
+        self.linker_mode = linker
+        self.registry = GNRHR_DAU_LINKER_REGISTRY
+        self.rng = np.random.default_rng(rng_seed)
+
+        # Pre-compute probabilities
+        self._linker_names, self._linker_probs = _compute_linker_probabilities(self.registry)
+
+        # Validate fixed linker
+        if linker not in ('sample', 'all') and linker not in self.registry:
+            raise ValueError(
+                f"Unknown linker '{linker}'. "
+                f"Choose from: 'sample', 'all', {list(self.registry.keys())}"
+            )
+
+        # Pre-instantiate builders (suppress verbose init output)
+        self._builders = {}
+        for name, entry in self.registry.items():
+            with contextlib.redirect_stdout(io.StringIO()):
+                self._builders[name] = entry['class']()
+
+        # Track which linker was last sampled (useful for logging)
+        self.last_sampled_linker = None
+
+    # ------------------------------------------------------------------
+    # Linker selection
+    # ------------------------------------------------------------------
+
+    def sample_linker(self):
+        """Sample a linker name according to frequency weights."""
+        idx = self.rng.choice(len(self._linker_names), p=self._linker_probs)
+        return self._linker_names[idx]
+
+    def get_builder(self, linker_name):
+        """Return the pre-instantiated builder for a given linker name."""
+        return self._builders[linker_name]
+
+    @property
+    def linker_probabilities(self):
+        """Return dict of linker name → sampling probability."""
+        return dict(zip(self._linker_names, self._linker_probs.tolist()))
+
+    # ------------------------------------------------------------------
+    # Conjugation interface (mirrors individual builders)
+    # ------------------------------------------------------------------
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Daunorubicin to peptide, selecting linker by mode.
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Lys -NH2)
+            exclude_n_terminus: If True, only conjugate at Lys ε-amines
+            return_all: If True, return all positional conjugates
+
+        Returns:
+            List of conjugate SMILES.
+            In 'sample' mode: conjugates from the sampled linker.
+            In 'all' mode: conjugates from every linker (pooled).
+            In fixed mode: conjugates from the specified linker.
+        """
+        if self.linker_mode == 'all':
+            self.last_sampled_linker = 'all'
+            return self._conjugate_all_linkers(peptide_smiles, exclude_n_terminus, return_all)
+        elif self.linker_mode == 'sample':
+            chosen = self.sample_linker()
+            self.last_sampled_linker = chosen
+            return self._conjugate_with_linker(chosen, peptide_smiles, exclude_n_terminus, return_all)
+        else:
+            self.last_sampled_linker = self.linker_mode
+            return self._conjugate_with_linker(self.linker_mode, peptide_smiles, exclude_n_terminus, return_all)
+
+    def _conjugate_with_linker(self, linker_name, peptide_smiles, exclude_n_terminus, return_all):
+        """Run conjugation with a single, specific linker."""
+        builder = self._builders[linker_name]
+        with contextlib.redirect_stdout(io.StringIO()):
+            conjugates = builder.conjugate_to_peptide(
+                peptide_smiles,
+                exclude_n_terminus=exclude_n_terminus,
+                return_all=return_all,
+            )
+        # Normalize to list
+        if not isinstance(conjugates, list):
+            conjugates = [conjugates] if conjugates else []
+        return conjugates
+
+    def _conjugate_all_linkers(self, peptide_smiles, exclude_n_terminus, return_all):
+        """Generate conjugates with every linker and pool them."""
+        all_conjugates = []
+        seen = set()
+        for name in self._linker_names:
+            try:
+                conjs = self._conjugate_with_linker(name, peptide_smiles, exclude_n_terminus, return_all)
+                for c in conjs:
+                    if c not in seen:
+                        seen.add(c)
+                        all_conjugates.append(c)
+            except Exception:
+                continue
+        if not all_conjugates:
+            raise ValueError("No conjugates generated with any linker")
+        return all_conjugates
+
+    # ------------------------------------------------------------------
+    # Validation (delegates to the appropriate builder)
+    # ------------------------------------------------------------------
+
+    def validate_conjugate(self, conjugate_smiles, linker_name=None):
+        """
+        Validate a conjugate SMILES.
+
+        If linker_name is given, validate with that builder.
+        Otherwise, try all builders and return the best result.
+        """
+        if linker_name and linker_name in self._builders:
+            return self._builders[linker_name].validate_conjugate(conjugate_smiles)
+
+        # Try each builder; return first passing result
+        for name, builder in self._builders.items():
+            passed, checks = builder.validate_conjugate(conjugate_smiles)
+            if passed:
+                checks['matched_linker'] = name
+                return True, checks
+
+        # None passed - return last failure
+        return False, checks
+
+    def check_attachment_site_compatibility(self, conjugate_smiles, linker_name=None):
+        """Delegate to appropriate builder's attachment check."""
+        if linker_name and linker_name in self._builders:
+            return self._builders[linker_name].check_attachment_site_compatibility(conjugate_smiles)
+
+        for name, builder in self._builders.items():
+            passed, details = builder.check_attachment_site_compatibility(conjugate_smiles)
+            if passed:
+                details['matched_linker'] = name
+                return True, details
+        return False, details
+
+    # ------------------------------------------------------------------
+    # Batch interface
+    # ------------------------------------------------------------------
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides.
+
+        In 'sample' mode each peptide gets a freshly sampled linker.
+
+        Returns:
+            Dict with 'conjugates', 'successful', 'failed', 'total_conjugates',
+            and 'linkers_used' (list of linker names, one per peptide).
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0,
+            'linkers_used': [],
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all,
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+                results['linkers_used'].append(self.last_sampled_linker)
+            except Exception as e:
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+                results['linkers_used'].append(None)
+
+        return results

smiles/conjugation/GlutarylDau.py

@@ -0,0 +1,362 @@
+"""
+Glutaryl-Daunorubicin PDC Builder
+
+Conjugation chemistry: Lysine ε-amine + Glutaryl-Daunorubicin (amide bond formation)
+
+Reference PDCs:
+- PDC_00190 (GnRH-III-[2ΔHis-3D-Tic,8Lys(glutaryl-Val-Cit-PABC-diamine-PTX)])
+- PDC_00195 (GnRH-III-[2His-3Trp,8Lys(glutaryl-Val-Cit-PABC-diamine-PTX)])
+- Receptor: GNRHR
+- Linker: Glutaric acid (LIN00033)
+- Payload: Daunorubicin (DRG00010)
+
+Architecture:
+    Peptide──Lys──NH──C(=O)──(CH₂)₃──C(=O)──NH──Daunorubicin
+                  │              │
+              amide bond    glutaryl spacer    amide bond (pre-formed)
+
+The glutaric acid acts as a simple dicarboxylic acid spacer:
+- One COOH forms amide with Dau 3'-amine (pre-formed in linker-payload)
+- Other COOH forms amide with Lys ε-amine (conjugation reaction)
+"""
+
+from rdkit import Chem
+from rdkit.Chem import RWMol, AllChem
+
+from .pdc_builder import PDC_Builder
+
+
+class GlutarylDaunorubicin_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for Glutaric acid (Glutaryl) - Daunorubicin conjugates
+
+    Conjugation chemistry (amide bond formation):
+    - Peptide site: Lysine ε-amine (side chain -NH2)
+    - Linker-payload site: Carboxylic acid of glutaryl spacer (already conjugated
+      to Dau 3'-amine via amide)
+    - Reaction: Amide bond formation (Lys-NH2 + HOOC-glutaryl-Dau → Lys-NH-CO-glutaryl-Dau)
+
+    Architecture:
+        Peptide──Lys──NH──C(=O)──(CH₂)₃──C(=O)──NH──Daunorubicin
+                      └── amide ──┘  └── glutaryl ──┘  └── amide (pre-formed) ──┘
+
+    Reference: PDC_00190, PDC_00195 (GnRH-III conjugates, GNRHR)
+    """
+
+    # Glutaryl-Daunorubicin SMILES with free carboxylic acid
+    # Structure: HOOC-(CH2)3-CO-NH-Daunorubicin
+    # The amide bond between glutaric acid and Dau 3'-amine is pre-formed
+    DEFAULT_GLUTARYL_DAU_SMILES = (
+        "C[C@H]1[C@H]([C@H](C[C@@H](O1)O[C@H]2C[C@@](CC3=C2C(=C4C(=C3O)"
+        "C(=O)C5=C(C4=O)C(=CC=C5)OC)O)(C(=O)C)O)NC(=O)CCCC(=O)O)O"
+    )
+
+    def __init__(self, glutaryl_dau_smiles=None):
+        """
+        Initialize Glutaryl-Daunorubicin builder
+
+        Args:
+            glutaryl_dau_smiles: Optional custom Glutaryl-Dau SMILES.
+                                If None, uses default.
+                                Must have free carboxylic acid for conjugation.
+        """
+        if glutaryl_dau_smiles is None:
+            glutaryl_dau_smiles = self.DEFAULT_GLUTARYL_DAU_SMILES
+
+        super().__init__(glutaryl_dau_smiles, component_name="Glutaryl-Daunorubicin (linker-payload)")
+
+    def _verify_reactive_group(self):
+        """Verify carboxylic acid group is present for amide bond formation"""
+        carboxylic_acid_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+
+        if not self.linker_payload_mol.HasSubstructMatch(carboxylic_acid_pattern):
+            print(" Warning: No carboxylic acid group detected in Glutaryl-Dau")
+        else:
+            print(" Glutaryl-Daunorubicin loaded with carboxylic acid reactive group")
+
+        # Verify Dau 3'-amide bond is present (glutaryl-Dau linkage)
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH2]')
+        if self.linker_payload_mol.HasSubstructMatch(amide_pattern):
+            print(" Amide bond (Glutaryl-Dau linkage) detected")
+        else:
+            print(" Warning: No amide bond detected - check Glutaryl-Dau structure")
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find all lysine ε-amine groups in peptide
+
+        Returns list of dicts with site information including atom index
+        """
+        return self._find_peptide_lysines(peptide_smiles)
+
+    def _find_peptide_lysines(self, peptide_smiles):
+        """
+        Find all lysine residues with free ε-amine in peptide
+
+        Lysine side chain: -CH2-CH2-CH2-CH2-NH2 (CCCCN)
+
+        Returns:
+            List of dicts with 'atom_idx' (nitrogen index) and metadata
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        lysines = []
+
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        matches = peptide_mol.GetSubstructMatches(lys_pattern)
+
+        for match in matches:
+            nitrogen_idx = match[4]
+            lysines.append({
+                'atom_idx': nitrogen_idx,
+                'side_chain_atoms': list(match),
+                'type': 'Lys_epsilon_amine'
+            })
+
+        n_terminus_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        n_term_matches = peptide_mol.GetSubstructMatches(n_terminus_pattern)
+        if n_term_matches:
+            print(f"  (Also found {len(n_term_matches)} N-terminus α-amine(s) - excluded by default)")
+
+        return lysines
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """
+        Perform amide bond conjugation between NH2 groups and glutaryl carboxylic acid
+
+        Args:
+            peptide_mol: RDKit Mol object
+            site_info: Not used (kept for compatibility)
+            exclude_n_terminus: If True, skip N-terminus conjugation
+
+        Returns: List of all possible conjugate SMILES
+        """
+        return self._manual_conjugation(peptide_mol, site_info, exclude_n_terminus)
+
+    def _manual_conjugation(self, peptide_mol, lysine_info, exclude_n_terminus=True):
+        """
+        Create ALL possible conjugates by forming amide bond at NH2 sites
+
+        Chemistry:
+        - Any-NH2 + HOOC-glutaryl-Dau → Any-NH-C(=O)-glutaryl-Dau + H2O
+        """
+        if exclude_n_terminus:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:5][CH2][CH2][CH2][CH2:4][NH2:1].[C:2](=[O:3])[OH]>>'
+                '[*:5][CH2][CH2][CH2][CH2:4][NH:1][C:2](=[O:3])'
+            )
+        else:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:4][NH2:1].[C:2](=[O:3])[OH]>>[*:4][NH:1][C:2](=[O:3])'
+            )
+
+        products = amide_rxn.RunReactants((peptide_mol, self.linker_payload_mol))
+
+        if not products:
+            raise ValueError("Amide bond formation reaction failed - no products generated")
+
+        conjugates = []
+        seen_smiles = set()
+
+        for product_tuple in products:
+            try:
+                conjugate_mol = product_tuple[0]
+                Chem.SanitizeMol(conjugate_mol)
+                conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+
+                if conjugate_smiles not in seen_smiles:
+                    seen_smiles.add(conjugate_smiles)
+                    conjugates.append(conjugate_smiles)
+            except Exception:
+                continue
+
+        if not conjugates:
+            raise ValueError("All products failed sanitization")
+
+        return conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Glutaryl-Daunorubicin to peptide via amide bond at NH2 groups
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain NH2)
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            If return_all=True: List of all possible conjugate SMILES
+            If return_all=False: Single conjugate SMILES (first one)
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Lys ε-amine site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2 (including N-terminus)"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [{site_type}]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """Batch conjugate multiple peptides"""
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        return results
+
+    # =========================================================================
+    # Attachment Site Compatibility Checks
+    # =========================================================================
+
+    def has_glutaryl_amide_conjugation(self, mol):
+        """
+        Check if molecule has the glutaryl conjugation pattern
+
+        Pattern: N-C(=O)-(CH2)3-C(=O)-N (amide-glutaryl-amide linkage)
+        This is the signature of successful Lys-Glutaryl-Dau conjugation.
+        """
+        glutaryl_pattern = Chem.MolFromSmarts(
+            '[NX3][CX3](=[OX1])[CH2][CH2][CH2][CX3](=[OX1])[NX3]'
+        )
+        if glutaryl_pattern:
+            return mol.HasSubstructMatch(glutaryl_pattern)
+        return False
+
+    def has_dau_amide_bond(self, mol):
+        """
+        Check if molecule has the glutaryl-Dau amide bond intact
+
+        The pre-formed amide: -(CH2)3-C(=O)-NH-Dau
+        """
+        amide_pattern = Chem.MolFromSmarts('[CH2][CX3](=[OX1])[NX3]')
+        if amide_pattern:
+            return mol.HasSubstructMatch(amide_pattern)
+        return False
+
+    def has_unreacted_glutaryl_cooh(self, mol):
+        """
+        Check if molecule has unreacted glutaryl carboxylic acid
+
+        Specific pattern: Free COOH on the glutaryl chain ending
+        """
+        unreacted_pattern = Chem.MolFromSmarts(
+            '[CX3](=[OX1])[OX2H1][CH2][CH2][CH2][CX3](=[OX1])[NX3]'
+        )
+        if unreacted_pattern:
+            return mol.HasSubstructMatch(unreacted_pattern)
+        return False
+
+    def check_attachment_site_compatibility(self, mol):
+        """
+        Check if Lys-Glutaryl-Dau attachment site chemistry is valid
+
+        Validates:
+        1. Has glutaryl bis-amide conjugation pattern
+        2. Has Dau amide bond intact
+        3. No unreacted glutaryl-COOH
+        """
+        if not self.has_glutaryl_amide_conjugation(mol):
+            return False
+
+        if not self.has_dau_amide_bond(mol):
+            return False
+
+        if self.has_unreacted_glutaryl_cooh(mol):
+            return False
+
+        return True
+
+    def _basic_validation(self, mol):
+        """Basic validation checks for Glutaryl-Daunorubicin conjugates"""
+        checks = {}
+
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[#6]')
+        checks['has_amide_bond'] = mol.HasSubstructMatch(amide_pattern)
+
+        checks['has_glutaryl_conjugation'] = self.has_glutaryl_amide_conjugation(mol)
+        checks['has_dau_amide'] = self.has_dau_amide_bond(mol)
+        checks['no_unreacted_cooh'] = not self.has_unreacted_glutaryl_cooh(mol)
+        checks['attachment_site_valid'] = self.check_attachment_site_compatibility(mol)
+
+        anthraquinone = Chem.MolFromSmarts(
+            '[#6]1[#6][#6][#6]2[#6](=[#8])[#6]3[#6][#6][#6][#6][#6]3[#6](=[#8])[#6]2[#6]1'
+        )
+        checks['has_daunorubicin_core'] = mol.HasSubstructMatch(anthraquinone)
+
+        aminosugar = Chem.MolFromSmarts('[OX2]1[#6][#6][#6]([NX3])[#6][#6]1')
+        checks['has_aminosugar'] = mol.HasSubstructMatch(aminosugar)
+
+        checks['_passed'] = checks['attachment_site_valid']
+
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """Validate that the conjugate is chemically reasonable"""
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+
+        return passed, checks
+
+
+if __name__ == "__main__":
+    print("=" * 70)
+    print("Glutaryl-Daunorubicin PDC Builder - Test")
+    print("=" * 70 + "\n")
+
+    builder = GlutarylDaunorubicin_PDC_Builder()
+
+    test_peptides = [
+        "NCC(=O)N[C@@H](CCCCN)C(=O)NCC(=O)O",  # Gly-Lys-Gly
+    ]
+
+    for pep in test_peptides:
+        try:
+            conjugates = builder.conjugate_to_peptide(pep, exclude_n_terminus=True)
+            print(f" {len(conjugates)} conjugate(s)")
+            for c in conjugates:
+                valid, checks = builder.validate_conjugate(c)
+                print(f"  Valid: {valid}")
+        except Exception as e:
+            print(f" Failed: {e}")

smiles/conjugation/MalHexCpt.py

@@ -0,0 +1,560 @@
+"""
+Maleimide-Hexanoate-Camptothecin PDC Builder (MalHexCpt)
+
+Conjugation chemistry: Lysine ε-amine + Cys-S-Linker-Camptothecin (amide bond formation)
+
+Reference PDC:
+- PDC_00336 (SMAC-P1 CPT conjugate 6)
+
+Receptor: ERBB2 (Receptor tyrosine-protein kinase erbB-2)
+Linker: LIN00031 - 6-(2,5-dioxo-3-sulfanylpyrrolidin-1-yl)hexanoic acid
+Payload: DRG00009 - Camptothecin
+Architecture: Monomer, Lys side chain conjugation
+
+Chemistry Note:
+===============
+LIN00031 is a maleimide-thiol adduct (succinimide-thioether) with a hexanoic acid spacer:
+
+    Succinimide ring           Hexanoic spacer
+    ┌──────────────┐     ┌─────────────────────┐
+    O═C─CH₂─CH(SH)─C═O─N─CH₂─CH₂─CH₂─CH₂─CH₂─C(═O)─OH
+        └─────────────┘
+
+Linker-to-Drug Connection (Ester via hexanoic acid):
+    Cys─S─C(succinimide)─N─(CH₂)₅─C(═O)─O─Camptothecin
+                                       │
+                                   ester bond (C-20 OH of CPT)
+
+Full PDC Architecture:
+    Peptide──Lys──NH──C(═O)──Cys──S──[Succinimide]──N──(CH₂)₅──C(═O)─O──CPT
+                  └── amide ──┘      └── thioether ──┘     └── hexanoic spacer ──┘
+
+This is a NOVEL linker chemistry distinct from the simple disulfide-ester/carbonate
+linkers in SSCpt. The maleimide reacts with Cys thiol via Michael addition to form
+a thioether (C-S) bond. The succinimide ring provides structural rigidity and the
+hexanoic spacer provides distance between the conjugation site and the drug.
+
+Usage:
+======
+    builder = MaleimideHexanoateCpt_PDC_Builder()
+    conjugates = builder.conjugate_to_peptide(peptide_smiles)
+"""
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, Descriptors
+
+from .pdc_builder import PDC_Builder
+
+
+# ============================================================================
+# Maleimide-Hexanoate-Camptothecin PDC Builder
+# ============================================================================
+
+class MaleimideHexanoateCpt_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for Maleimide-Hexanoate + Camptothecin conjugates
+
+    Conjugation chemistry (amide bond formation):
+    - Peptide site: Lysine ε-amine (side chain -NH2)
+    - Linker-payload site: Carboxylic acid of Cys (Cys-S-Linker-Drug)
+    - Reaction: Lys-NH2 + HOOC-Cys-S-Linker-Drug → Lys-NH-CO-Cys-S-Linker-Drug
+
+    Linker: LIN00031 - 6-(2,5-dioxo-3-sulfanylpyrrolidin-1-yl)hexanoic acid
+        SMILES: C1C(C(=O)N(C1=O)CCCCCC(=O)O)S
+        IUPAC: 6-(2,5-dioxo-3-sulfanylpyrrolidin-1-yl)hexanoic acid
+        MW: 245.3, MF: C10H15NO4S
+        Features: Succinimide ring (maleimide-thiol adduct) + hexanoic acid spacer
+
+    Drug: DRG00009 - Camptothecin
+        Canonical SMILES: CCC1(O)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1
+        Connection: Ester at C-20 hydroxyl
+
+    Reference: PDC_00336 (SMAC-P1 CPT conjugate 6, ERBB2)
+    """
+
+    # ========================================================================
+    # Cys-S-Linker-Drug SMILES (thioether via Michael addition)
+    # ========================================================================
+    #
+    # Pre-formed Cys-MaleimideHexanoate-Camptothecin assembly:
+    #
+    #   HOOC─Cys(NH₂)─CH₂─S─C(succinimide)─N─(CH₂)₅─C(═O)─O─CPT(C-20)
+    #
+    # Construction from components:
+    #   Cys:     N[C@H](CS...)C(=O)O
+    #   S-C:     ...S...  (thioether, Michael addition of Cys-SH to maleimide)
+    #   LIN00031 ring: ...C4CC(=O)N(...)C4=O  (succinimide from maleimide-thiol adduct)
+    #   Spacer:  N(CCCCCC(=O)O...)  (hexanoic acid, COOH → ester to CPT)
+    #   CPT:     ...OC1(CC)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1
+    #
+    # ========================================================================
+
+    CYS_S_LINKER_CPT_SMILES = (
+        "N[C@H](CSC4CC(=O)N(CCCCCC(=O)OC1(CC)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1)C4=O)C(=O)O"
+    )
+
+    def __init__(self, cys_linker_drug_smiles=None):
+        """
+        Initialize Maleimide-Hexanoate-Camptothecin builder
+
+        Args:
+            cys_linker_drug_smiles: Optional custom Cys-S-Linker-Drug SMILES.
+                                   If None, uses default LIN00031+CPT assembly.
+                                   Must have free carboxylic acid for conjugation.
+        """
+        if cys_linker_drug_smiles is None:
+            cys_linker_drug_smiles = self.CYS_S_LINKER_CPT_SMILES
+
+        super().__init__(
+            cys_linker_drug_smiles,
+            component_name="Cys-S-MaleimideHexanoate-Camptothecin (LIN00031+DRG00009)"
+        )
+
+    def _verify_reactive_group(self):
+        """Verify carboxylic acid group is present for amide bond formation"""
+        # Pattern: Carboxylic acid (COOH)
+        carboxylic_acid_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+
+        if not self.linker_payload_mol.HasSubstructMatch(carboxylic_acid_pattern):
+            print(" Warning: No carboxylic acid group detected in Cys-S-Linker-Drug")
+        else:
+            print(" Cys-S-MalHex-Camptothecin loaded with carboxylic acid reactive group")
+
+        # Verify thioether bond is present (Cys-S-C, from Michael addition)
+        thioether_pattern = Chem.MolFromSmarts('[CH2][SX2][CH]')
+        if self.linker_payload_mol.HasSubstructMatch(thioether_pattern):
+            print(" Thioether bond (C-S-C) detected (maleimide-thiol Michael addition)")
+        else:
+            print(" Warning: No thioether bond detected - check structure")
+
+        # Verify Cys α-amine is present
+        cys_amine_pattern = Chem.MolFromSmarts('[NH2][CH]([CH2][SX2])C(=O)')
+        if self.linker_payload_mol.HasSubstructMatch(cys_amine_pattern):
+            print(" Cys α-amine detected (will be free in conjugate)")
+        else:
+            print(" Warning: Cys α-amine pattern not detected")
+
+        # Verify succinimide ring (maleimide-thiol adduct)
+        succinimide_pattern = Chem.MolFromSmarts('[CH2]C(=O)N(*)C(=O)')
+        if self.linker_payload_mol.HasSubstructMatch(succinimide_pattern):
+            print(" Succinimide ring (maleimide-thiol adduct) detected")
+        else:
+            print(" Warning: No succinimide ring detected - check linker structure")
+
+        # Verify ester bond to drug
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        if self.linker_payload_mol.HasSubstructMatch(ester_pattern):
+            print(" Ester bond (linker-to-drug) detected")
+        else:
+            print(" Warning: No ester bond detected - check linker-drug connection")
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find all lysine ε-amine groups in peptide
+
+        Returns list of dicts with site information including atom index
+        """
+        return self.find_peptide_lysines(peptide_smiles)
+
+    def find_peptide_lysines(self, peptide_smiles):
+        """
+        Find all lysine residues with free ε-amine in peptide
+
+        Lysine side chain: -CH2-CH2-CH2-CH2-NH2 (CCCCN)
+
+        Returns:
+            List of dicts with 'atom_idx' (nitrogen index) and metadata
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        lysines = []
+
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        matches = peptide_mol.GetSubstructMatches(lys_pattern)
+
+        for match in matches:
+            nitrogen_idx = match[4]
+            lysines.append({
+                'atom_idx': nitrogen_idx,
+                'side_chain_atoms': list(match),
+                'type': 'Lys_epsilon_amine'
+            })
+
+        n_terminus_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        n_term_matches = peptide_mol.GetSubstructMatches(n_terminus_pattern)
+
+        if n_term_matches:
+            print(f"  (Also found {len(n_term_matches)} N-terminus α-amine(s) - excluded by default)")
+
+        return lysines
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """
+        Perform amide bond conjugation between NH2 groups and Cys carboxylic acid
+
+        Args:
+            peptide_mol: RDKit Mol object
+            site_info: Not used (kept for compatibility)
+            exclude_n_terminus: If True, skip N-terminus conjugation
+
+        Returns: List of all possible conjugate SMILES
+        """
+        if exclude_n_terminus:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:5][CH2][CH2][CH2][CH2:4][NH2:1].[C:2](=[O:3])[OH]>>[*:5][CH2][CH2][CH2][CH2:4][NH:1][C:2](=[O:3])'
+            )
+        else:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:4][NH2:1].[C:2](=[O:3])[OH]>>[*:4][NH:1][C:2](=[O:3])'
+            )
+
+        products = amide_rxn.RunReactants((peptide_mol, self.linker_payload_mol))
+
+        if not products:
+            raise ValueError("Amide bond formation reaction failed - no products generated")
+
+        conjugates = []
+        seen_smiles = set()
+
+        for product_tuple in products:
+            try:
+                conjugate_mol = product_tuple[0]
+                Chem.SanitizeMol(conjugate_mol)
+                conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+
+                if conjugate_smiles not in seen_smiles:
+                    seen_smiles.add(conjugate_smiles)
+                    conjugates.append(conjugate_smiles)
+            except Exception:
+                continue
+
+        if not conjugates:
+            raise ValueError("All products failed sanitization")
+
+        return conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Cys-S-MalHex-Camptothecin to peptide via amide bond at NH2 groups
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Lys with NH2)
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            List of all possible conjugate SMILES
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Lys ε-amine site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2 (including N-terminus)"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [{site_type}] [maleimide-hexanoate ester]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides
+
+        Args:
+            peptide_smiles_list: List of peptide SMILES strings
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates per peptide
+
+        Returns:
+            Dict with 'conjugates', 'successful', 'failed', 'total_conjugates'
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        print(f"\n=== Summary ===")
+        print(f"Peptides: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, Failed: {len(results['failed'])}")
+        print(f"Total conjugates generated: {results['total_conjugates']}")
+
+        return results
+
+    def flatten_batch_results(self, results):
+        """Flatten batch conjugation results into a single list"""
+        all_conjugates = []
+        for conjugate_list in results['conjugates']:
+            if conjugate_list is not None:
+                if isinstance(conjugate_list, list):
+                    all_conjugates.extend(conjugate_list)
+                else:
+                    all_conjugates.append(conjugate_list)
+        return all_conjugates
+
+    # ========================================================================
+    # Attachment Site Compatibility Checks
+    # ========================================================================
+
+    def has_lys_cys_amide_conjugation(self, mol):
+        """Check if molecule has Lys-Cys amide conjugation bond"""
+        lys_cys_amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2])[NX3]')
+        return mol.HasSubstructMatch(lys_cys_amide_pattern)
+
+    def has_thioether_bond(self, mol):
+        """Check if molecule has thioether bond (C-S-C, from maleimide-thiol Michael addition)"""
+        thioether_pattern = Chem.MolFromSmarts('[CH2][SX2][CH]')
+        return mol.HasSubstructMatch(thioether_pattern)
+
+    def has_succinimide_ring(self, mol):
+        """
+        Check if molecule has succinimide ring (maleimide-thiol adduct)
+
+        The succinimide ring is the product of maleimide-thiol Michael addition.
+        Pattern: 5-membered ring with N and two C=O groups.
+        """
+        succinimide_pattern = Chem.MolFromSmarts('C1CC(=O)N(*)C1=O')
+        return mol.HasSubstructMatch(succinimide_pattern)
+
+    def has_hexanoic_spacer(self, mol):
+        """
+        Check if molecule has hexanoic acid spacer chain
+
+        Pattern: N-CH2-CH2-CH2-CH2-CH2-C(=O)-O (6-carbon chain from succinimide N)
+        """
+        hexanoic_pattern = Chem.MolFromSmarts('[NX3]CCCCCC(=O)[OX2]')
+        return mol.HasSubstructMatch(hexanoic_pattern)
+
+    def has_unreacted_lys_epsilon_amine(self, mol):
+        """Check for unreacted Lys ε-amine (free NH2 on 4-carbon side chain)"""
+        lys_amine_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        return mol.HasSubstructMatch(lys_amine_pattern)
+
+    def has_unreacted_cys_carboxylic_acid(self, mol):
+        """Check for unreacted Cys carboxylic acid from Cys-S-Linker-Drug"""
+        cys_cooh_pattern = Chem.MolFromSmarts('[NH2][CH]([CH2][SX2])[CX3](=[OX1])[OX2H1]')
+        return mol.HasSubstructMatch(cys_cooh_pattern)
+
+    def find_amide_conjugation_sites(self, mol):
+        """Find all Lys-Cys amide conjugation sites in the molecule"""
+        lys_cys_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2])[NX3]')
+        matches = mol.GetSubstructMatches(lys_cys_pattern)
+
+        sites = []
+        for match in matches:
+            nitrogen_idx = match[0]
+            carbonyl_c_idx = match[1]
+            cys_alpha_idx = match[3]
+            sites.append((nitrogen_idx, carbonyl_c_idx, cys_alpha_idx))
+
+        return sites
+
+    def check_valency_at_conjugation_site(self, mol, site_indices):
+        """Check that valency is correct at Lys-Cys conjugation site"""
+        nitrogen_idx, carbonyl_c_idx, cys_alpha_idx = site_indices
+
+        try:
+            n_atom = mol.GetAtomWithIdx(nitrogen_idx)
+            if n_atom.GetTotalDegree() != 3:
+                return False
+
+            c_atom = mol.GetAtomWithIdx(carbonyl_c_idx)
+            if c_atom.GetTotalDegree() != 3:
+                return False
+
+            cys_c_atom = mol.GetAtomWithIdx(cys_alpha_idx)
+            if cys_c_atom.GetTotalDegree() != 4:
+                return False
+
+            return True
+        except Exception:
+            return False
+
+    def check_attachment_site_compatibility(self, conjugate_smiles):
+        """
+        Check attachment site compatibility for Lys-Cys-S-MalHex-CPT conjugation
+
+        Validates:
+        1. Lys-Cys amide conjugation is present
+        2. Thioether bond (C-S-C) is intact (maleimide-thiol Michael addition)
+        3. Succinimide ring is present (maleimide-thiol adduct)
+        4. Hexanoic spacer is present
+        5. Valency at conjugation sites is correct
+        6. No unreacted Cys-COOH (unconjugated linker-drug)
+
+        Returns:
+            Tuple of (passed: bool, details: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'parseable': False, 'error': 'Could not parse SMILES'}
+
+        details = {
+            'parseable': True,
+            'validation_method': 'lys_cys_amide_thioether_maleimide_hexanoate_ester',
+            'linker': 'LIN00031 (6-(2,5-dioxo-3-sulfanylpyrrolidin-1-yl)hexanoic acid)',
+            'drug': 'DRG00009 (Camptothecin)',
+            'pdc_ref': 'PDC_00336'
+        }
+
+        details['has_lys_cys_amide'] = self.has_lys_cys_amide_conjugation(mol)
+        details['has_thioether_bond'] = self.has_thioether_bond(mol)
+        details['has_succinimide_ring'] = self.has_succinimide_ring(mol)
+        details['has_hexanoic_spacer'] = self.has_hexanoic_spacer(mol)
+
+        conjugation_sites = self.find_amide_conjugation_sites(mol)
+        details['conjugation_sites'] = len(conjugation_sites)
+
+        valency_ok = True
+        if conjugation_sites:
+            for site in conjugation_sites:
+                if not self.check_valency_at_conjugation_site(mol, site):
+                    valency_ok = False
+                    break
+        details['valency_correct'] = valency_ok
+
+        details['no_unreacted_cys_cooh'] = not self.has_unreacted_cys_carboxylic_acid(mol)
+        details['has_unreacted_lys_amine'] = self.has_unreacted_lys_epsilon_amine(mol)
+
+        passed = all([
+            details['has_lys_cys_amide'],
+            details['has_thioether_bond'],
+            details['has_succinimide_ring'],
+            details['valency_correct'],
+            details['no_unreacted_cys_cooh']
+        ])
+
+        return passed, details
+
+    def _basic_validation(self, mol):
+        """
+        Basic validation checks for MaleimideHexanoate-Camptothecin conjugates
+
+        Checks:
+        1. Amide bond present (Lys-Cys conjugation)
+        2. Thioether bond present (C-S-C, maleimide-thiol Michael addition)
+        3. Specific Lys-Cys-S conjugation pattern
+        4. Succinimide ring present (maleimide-thiol adduct)
+        5. Camptothecin core present
+        6. Ester bond to drug present
+        """
+        checks = {}
+
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[#6]')
+        checks['has_amide_bond'] = mol.HasSubstructMatch(amide_pattern)
+
+        thioether_pattern = Chem.MolFromSmarts('[CH2][SX2][CH]')
+        checks['has_thioether_bond'] = mol.HasSubstructMatch(thioether_pattern)
+
+        lys_cys_s_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2])[NX3]')
+        checks['has_lys_cys_s_conjugation'] = mol.HasSubstructMatch(lys_cys_s_pattern)
+
+        succinimide_pattern = Chem.MolFromSmarts('C1CC(=O)N(*)C1=O')
+        checks['has_succinimide_ring'] = mol.HasSubstructMatch(succinimide_pattern)
+
+        camptothecin_core = Chem.MolFromSmarts('c1ccc2ncc3C(=O)OCc3c2c1')
+        checks['has_camptothecin_core'] = mol.HasSubstructMatch(camptothecin_core)
+
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        checks['has_ester_bond'] = mol.HasSubstructMatch(ester_pattern)
+
+        checks['_passed'] = all([
+            checks['has_amide_bond'],
+            checks['has_thioether_bond'],
+            checks['has_lys_cys_s_conjugation'],
+            checks['has_succinimide_ring']
+        ])
+
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """
+        Validate that the conjugate is chemically reasonable
+
+        Returns:
+            (passed: bool, checks: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+
+        return passed, checks
+
+
+# ============================================================================
+# Main - Test Examples
+# ============================================================================
+
+if __name__ == "__main__":
+
+    print("=" * 70)
+    print("Maleimide-Hexanoate-Camptothecin PDC Builder - Test")
+    print("Reference: PDC_00336 (SMAC-P1 CPT conjugate 6, LIN00031)")
+    print("=" * 70 + "\n")
+
+    # Test peptides
+    test_peptides = [
+        # Gly-Lys-Gly (has N-terminus NH2 + Lys ε-amine)
+        "NCC(=O)N[C@@H](CCCCN)C(=O)NCC(=O)O",
+        # Lys-Arg
+        "N[C@@H](CCCCN)C(=O)N[C@@H](CCCNC(=N)N)C(=O)O",
+    ]
+
+    builder = MaleimideHexanoateCpt_PDC_Builder()
+
+    for i, peptide in enumerate(test_peptides):
+        print(f"\n--- Test peptide {i+1} ---")
+        print(f"Input: {peptide}")
+
+        try:
+            conjugates = builder.conjugate_to_peptide(peptide, exclude_n_terminus=True)
+            print(f" Generated {len(conjugates)} conjugate(s)")
+
+            for j, conj in enumerate(conjugates):
+                valid, checks = builder.validate_conjugate(conj)
+                print(f"  Conjugate {j+1}: {'PASSED' if valid else 'FAILED'} ({len(conj)} chars)")
+                if not valid:
+                    for k, v in checks.items():
+                        if v is False:
+                            print(f"     {k}")
+
+        except Exception as e:
+            print(f" FAILED: {e}")
+            import traceback
+            traceback.print_exc()

smiles/conjugation/SN38Factory.py

@@ -0,0 +1,308 @@
+"""
+SN-38 PDC Factory - Drug-level builder with linker sampling
+
+Instead of specifying a particular linker-drug combo (e.g. DSPASN38, DSEBSN38),
+the user specifies only the drug (SN-38) and a linker is sampled
+according to NRP1+SN38 entry frequencies from PDCdb.
+
+Linker frequencies (NRP1 + SN-38):
+===========================================
+    3-Disulfanylpropanoic acid               (DSPASN38) - 9 entries
+    4-(2-Disulfanylethoxy)-4-oxobutanoic acid (DSEBSN38) - 6 entries
+    5-(2-Disulfanylethoxy)-5-oxopentanoic acid (DSEPSN38) - 2 entries
+                                                    Total: 17 entries
+
+Sampling probabilities:
+    DSPASN38:  9/17 ≈ 0.529
+    DSEBSN38:  6/17 ≈ 0.353
+    DSEPSN38:  2/17 ≈ 0.118
+
+Usage:
+======
+    # Drug-only mode - linker is sampled per peptide
+    factory = SN38PDC_Factory()
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+
+    # Fixed linker mode - use a specific linker
+    factory = SN38PDC_Factory(linker='DSPASN38')
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+
+    # All linkers mode - generate conjugates with every linker
+    factory = SN38PDC_Factory(linker='all')
+    conjugates = factory.conjugate_to_peptide(peptide_smiles)
+"""
+
+import numpy as np
+import contextlib
+import io
+
+from .DSPASN38 import DisulfanylpropanoicAcidSN38_PDC_Builder
+from .DSEBSN38 import DisulfanylethOxyButanoicAcidSN38_PDC_Builder
+from .DSEPSN38 import DisulfanylethOxyPentanoicAcidSN38_PDC_Builder
+
+
+# ============================================================================
+# Linker registry: name → (builder_class, PDCdb frequency, description)
+# ============================================================================
+
+NRP1_SN38_LINKER_REGISTRY = {
+    'DSPASN38': {
+        'class': DisulfanylpropanoicAcidSN38_PDC_Builder,
+        'frequency': 9,
+        'linker_id': 'LIN00149',
+        'description': '3-Disulfanylpropanoic acid (disulfide + ester)',
+        'pdc_refs': ['PDC_02099', 'PDC_02104', 'PDC_02102', 'PDC_02107',
+                     'PDC_02110', 'PDC_02112', 'PDC_02113', 'PDC_02114', 'PDC_02115'],
+    },
+    'DSEBSN38': {
+        'class': DisulfanylethOxyButanoicAcidSN38_PDC_Builder,
+        'frequency': 6,
+        'linker_id': 'LIN00168',
+        'description': '4-(2-Disulfanylethoxy)-4-oxobutanoic acid (succinic spacer)',
+        'pdc_refs': ['PDC_02100', 'PDC_02103', 'PDC_02105',
+                     'PDC_02108', 'PDC_02109', 'PDC_02111'],
+    },
+    'DSEPSN38': {
+        'class': DisulfanylethOxyPentanoicAcidSN38_PDC_Builder,
+        'frequency': 2,
+        'linker_id': 'LIN00169',
+        'description': '5-(2-Disulfanylethoxy)-5-oxopentanoic acid (glutaric spacer)',
+        'pdc_refs': ['PDC_02101', 'PDC_02106'],
+    },
+}
+
+
+def _compute_linker_probabilities(registry=None):
+    """Compute sampling probabilities from frequency counts."""
+    if registry is None:
+        registry = NRP1_SN38_LINKER_REGISTRY
+    names = list(registry.keys())
+    freqs = np.array([registry[n]['frequency'] for n in names], dtype=np.float64)
+    probs = freqs / freqs.sum()
+    return names, probs
+
+
+class SN38PDC_Factory:
+    """
+    Drug-level PDC factory for SN-38 (7-Ethyl-10-hydroxycamptothecin) conjugates.
+
+    Wraps all three NRP1+SN38 linker builders and provides:
+    - 'sample' mode (default): randomly pick a linker per call according to
+      PDCdb entry frequencies
+    - 'all' mode: generate conjugates with every linker
+    - fixed linker mode: always use a specific linker (e.g. 'DSPASN38')
+
+    Conjugation chemistry: Cysteine thiol + Linker thiol → Disulfide bond
+    Target receptor: NRP1 (Neuropilin-1)
+
+    The interface mirrors the individual builder's `conjugate_to_peptide` so it
+    can be used as a drop-in replacement in the Conjugation classifier.
+    """
+
+    def __init__(self, linker='sample', rng_seed=None):
+        """
+        Args:
+            linker: One of:
+                - 'sample' (default): sample linker per peptide according to
+                  NRP1+SN38 frequency weights
+                - 'all': use every linker and pool conjugates
+                - A specific linker name ('DSPASN38', 'DSEBSN38', 'DSEPSN38')
+            rng_seed: Optional seed for reproducible linker sampling
+        """
+        self.linker_mode = linker
+        self.registry = NRP1_SN38_LINKER_REGISTRY
+        self.rng = np.random.default_rng(rng_seed)
+
+        # Pre-compute probabilities
+        self._linker_names, self._linker_probs = _compute_linker_probabilities(self.registry)
+
+        # Validate fixed linker
+        if linker not in ('sample', 'all') and linker not in self.registry:
+            raise ValueError(
+                f"Unknown linker '{linker}'. "
+                f"Choose from: 'sample', 'all', {list(self.registry.keys())}"
+            )
+
+        # Pre-instantiate builders (suppress verbose init output)
+        self._builders = {}
+        for name, entry in self.registry.items():
+            with contextlib.redirect_stdout(io.StringIO()):
+                self._builders[name] = entry['class']()
+
+        # Track which linker was last sampled (useful for logging)
+        self.last_sampled_linker = None
+
+    # ------------------------------------------------------------------
+    # Linker selection
+    # ------------------------------------------------------------------
+
+    def sample_linker(self):
+        """Sample a linker name according to frequency weights."""
+        idx = self.rng.choice(len(self._linker_names), p=self._linker_probs)
+        return self._linker_names[idx]
+
+    def get_builder(self, linker_name):
+        """Return the pre-instantiated builder for a given linker name."""
+        return self._builders[linker_name]
+
+    @property
+    def linker_probabilities(self):
+        """Return dict of linker name → sampling probability."""
+        return dict(zip(self._linker_names, self._linker_probs.tolist()))
+
+    # ------------------------------------------------------------------
+    # Conjugation interface (mirrors individual builders)
+    # ------------------------------------------------------------------
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate SN-38 to peptide, selecting linker by mode.
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Cys thiol)
+            exclude_n_terminus: Not used for Cys conjugation (kept for interface)
+            return_all: If True, return all positional conjugates
+
+        Returns:
+            List of conjugate SMILES.
+            In 'sample' mode: conjugates from the sampled linker.
+            In 'all' mode: conjugates from every linker (pooled).
+            In fixed mode: conjugates from the specified linker.
+        """
+        if self.linker_mode == 'all':
+            self.last_sampled_linker = 'all'
+            return self._conjugate_all_linkers(peptide_smiles, exclude_n_terminus, return_all)
+        elif self.linker_mode == 'sample':
+            chosen = self.sample_linker()
+            self.last_sampled_linker = chosen
+            return self._conjugate_with_linker(chosen, peptide_smiles, exclude_n_terminus, return_all)
+        else:
+            self.last_sampled_linker = self.linker_mode
+            return self._conjugate_with_linker(self.linker_mode, peptide_smiles, exclude_n_terminus, return_all)
+
+    def _conjugate_with_linker(self, linker_name, peptide_smiles, exclude_n_terminus, return_all):
+        """Run conjugation with a single, specific linker."""
+        builder = self._builders[linker_name]
+        with contextlib.redirect_stdout(io.StringIO()):
+            conjugates = builder.conjugate_to_peptide(
+                peptide_smiles,
+                exclude_n_terminus=exclude_n_terminus,
+                return_all=return_all,
+            )
+        # Normalize to list
+        if not isinstance(conjugates, list):
+            conjugates = [conjugates] if conjugates else []
+        return conjugates
+
+    def _conjugate_all_linkers(self, peptide_smiles, exclude_n_terminus, return_all):
+        """Generate conjugates with every linker and pool them."""
+        all_conjugates = []
+        seen = set()
+        for name in self._linker_names:
+            try:
+                conjs = self._conjugate_with_linker(name, peptide_smiles, exclude_n_terminus, return_all)
+                for c in conjs:
+                    if c not in seen:
+                        seen.add(c)
+                        all_conjugates.append(c)
+            except Exception:
+                continue
+        if not all_conjugates:
+            raise ValueError("No conjugates generated with any linker")
+        return all_conjugates
+
+    # ------------------------------------------------------------------
+    # Validation (delegates to the appropriate builder)
+    # ------------------------------------------------------------------
+
+    def validate_conjugate(self, conjugate_smiles, linker_name=None):
+        """
+        Validate a conjugate SMILES.
+
+        If linker_name is given, validate with that builder.
+        Otherwise, try all builders and return the best result.
+        """
+        if linker_name and linker_name in self._builders:
+            return self._builders[linker_name].validate_conjugate(conjugate_smiles)
+
+        # Try each builder; return first passing result
+        checks = None
+        for name, builder in self._builders.items():
+            passed, checks = builder.validate_conjugate(conjugate_smiles)
+            if passed:
+                checks['matched_linker'] = name
+                return True, checks
+
+        # None passed - return last failure
+        return False, checks
+
+    def check_attachment_site_compatibility(self, conjugate_smiles, linker_name=None):
+        """Delegate to appropriate builder's attachment check."""
+        if linker_name and linker_name in self._builders:
+            return self._builders[linker_name].check_attachment_site_compatibility(conjugate_smiles)
+
+        details = None
+        for name, builder in self._builders.items():
+            passed, details = builder.check_attachment_site_compatibility(conjugate_smiles)
+            if passed:
+                details['matched_linker'] = name
+                return True, details
+        return False, details
+
+    # ------------------------------------------------------------------
+    # Batch interface
+    # ------------------------------------------------------------------
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides.
+
+        In 'sample' mode each peptide gets a freshly sampled linker.
+
+        Returns:
+            Dict with 'conjugates', 'successful', 'failed', 'total_conjugates',
+            and 'linkers_used' (list of linker names, one per peptide).
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0,
+            'linkers_used': [],
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all,
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+                results['linkers_used'].append(
+                    self.linker_mode if self.linker_mode not in ('sample', 'all') else None
+                )
+            except Exception as e:
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+                results['linkers_used'].append(None)
+
+        return results
+
+    # ------------------------------------------------------------------
+    # Informational
+    # ------------------------------------------------------------------
+
+    def __repr__(self):
+        probs = self.linker_probabilities
+        lines = [f"SN38PDC_Factory(mode='{self.linker_mode}')"]
+        for name, prob in probs.items():
+            desc = self.registry[name]['description']
+            lines.append(f"  {name:12s}  p={prob:.3f}  ({desc})")
+        return '\n'.join(lines)

smiles/conjugation/SSCpt.py

@@ -0,0 +1,891 @@
+"""
+Disulfide-Camptothecin PDC Builder (SSCpt)
+
+Conjugation chemistry: Lysine ε-amine + Cys-SS-Linker-Camptothecin (amide bond formation)
+
+Reference PDCs:
+- PDC_00294 (Z7): Ester linker chemistry
+- PDC_00295 (Z8): Carbonate linker chemistry
+
+Receptor: ERBB2
+Linker: Disulfide bond (S-S) - cleavable inside cells
+Payload: Camptothecin
+Architecture: Monomer, Lys side chain conjugation
+
+Supported Linker-to-Drug Chemistries:
+======================================
+
+ESTER (Z7):
+    Cys─S─S─CH2─CH2─C(═O)─O─Camptothecin
+                     │
+                   ester
+
+CARBONATE (Z8):
+    Cys─S─S─CH2─CH2─O─C(═O)─O─Camptothecin
+                      │
+                   carbonate
+
+Usage:
+======
+    # DEFAULT: Both chemistries (returns combined list)
+    builder = DisulfideCamptothecin_PDC_Builder()
+    conjugates = builder.conjugate_to_peptide(peptide_smiles)
+    # Returns: [ester_conj1, ester_conj2, ..., carbonate_conj1, carbonate_conj2, ...]
+    
+    # Single chemistry only
+    builder = DisulfideCamptothecin_PDC_Builder(linker_chemistry='ester')
+    builder = DisulfideCamptothecin_PDC_Builder(linker_chemistry='carbonate')
+"""
+
+from rdkit import Chem
+from rdkit.Chem import AllChem, Descriptors
+
+from .pdc_builder import PDC_Builder
+
+
+# ============================================================================
+# Disulfide-Camptothecin PDC Builder
+# ============================================================================
+
+class DisulfideCamptothecin_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for Disulfide-Camptothecin conjugates
+    
+    Conjugation chemistry (amide bond formation):
+    - Peptide site: Lysine ε-amine (side chain -NH2)
+    - Linker-payload site: Carboxylic acid of Cys (Cys-SS-Linker-Drug)
+    - Reaction: Lys-NH2 + HOOC-Cys-SS-Linker-Drug → Lys-NH-CO-Cys-SS-Linker-Drug
+    
+    Supported linker-to-drug chemistries:
+    - 'ester' (Z7, PDC_00294): Cys-S-S-CH2-CH2-C(=O)-O-CPT
+    - 'carbonate' (Z8, PDC_00295): Cys-S-S-CH2-CH2-O-C(=O)-O-CPT
+    
+    Architecture:
+        Peptide──Lys──NH──C(=O)──Cys──S─S──Linker──Camptothecin
+                      └── amide ──┘      └── disulfide ──┘
+    
+    Usage:
+        # Default: BOTH chemistries (returns combined list)
+        builder = DisulfideCamptothecin_PDC_Builder()
+        conjugates = builder.conjugate_to_peptide(peptide_smiles)
+        # For 1 Lys: returns 2 conjugates (1 ester + 1 carbonate)
+        
+        # Single chemistry only
+        builder = DisulfideCamptothecin_PDC_Builder(linker_chemistry='ester')
+    """
+    
+    # ========================================================================
+    # Linker Chemistry Variants
+    # ========================================================================
+    # 
+    # Two linker-to-drug chemistries are supported:
+    #
+    # ESTER (Z7, PDC_00294):
+    #   Cys─S─S─CH2─CH2─C(═O)─O─Camptothecin
+    #                    │
+    #                  ester
+    #
+    # CARBONATE (Z8, PDC_00295):
+    #   Cys─S─S─CH2─CH2─O─C(═O)─O─Camptothecin
+    #                     │
+    #                  carbonate
+    #
+    # ========================================================================
+    
+    # Ester variant (Z7, PDC_00294)
+    # Structure: HOOC-Cys(α-amine)-CH2-S-S-CH2-CH2-C(=O)-O-Camptothecin
+    ESTER_CYS_SS_LINKER_CPT_SMILES = (
+        "N[C@H](CSSCCC(=O)OC1(CC)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1)C(=O)O"
+    )
+    
+    # Carbonate variant (Z8, PDC_00295)
+    # Structure: HOOC-Cys(α-amine)-CH2-S-S-CH2-CH2-O-C(=O)-O-Camptothecin
+    CARBONATE_CYS_SS_LINKER_CPT_SMILES = (
+        "N[C@H](CSSCCOC(=O)OC1(CC)C(=O)OCc2c1cc1n(c2=O)Cc2cc3ccccc3nc2-1)C(=O)O"
+    )
+    
+    # Default (backward compatibility)
+    DEFAULT_CYS_SS_LINKER_CPT_SMILES = ESTER_CYS_SS_LINKER_CPT_SMILES
+    
+    # Available linker chemistries
+    LINKER_CHEMISTRIES = {
+        'ester': ESTER_CYS_SS_LINKER_CPT_SMILES,
+        'carbonate': CARBONATE_CYS_SS_LINKER_CPT_SMILES,
+    }
+    
+    def __init__(self, cys_linker_drug_smiles=None, linker_chemistry='all'):
+        """
+        Initialize Disulfide-Camptothecin builder
+        
+        Args:
+            cys_linker_drug_smiles: Optional custom Cys-SS-Linker-Drug SMILES. 
+                                   If provided, overrides linker_chemistry.
+                                   Must have free carboxylic acid for conjugation.
+            linker_chemistry: 'all' (default), 'ester' (Z7), or 'carbonate' (Z8).
+                             'all' generates conjugates with both chemistries.
+                             Ignored if cys_linker_drug_smiles is provided.
+        """
+        self.linker_chemistry = linker_chemistry
+        
+        if cys_linker_drug_smiles is not None:
+            # Custom SMILES provided - use single chemistry mode
+            self.linker_chemistry = 'custom'
+            super().__init__(cys_linker_drug_smiles, component_name="Cys-SS-Linker-Camptothecin (custom)")
+        elif linker_chemistry == 'all':
+            # All chemistries mode - initialize with ester as base, but will use both
+            super().__init__(self.ESTER_CYS_SS_LINKER_CPT_SMILES, component_name="Cys-SS-Linker-Camptothecin (all)")
+        else:
+            # Single chemistry mode
+            if linker_chemistry not in self.LINKER_CHEMISTRIES:
+                raise ValueError(f"Unknown linker_chemistry: {linker_chemistry}. "
+                               f"Choose from: 'all', {list(self.LINKER_CHEMISTRIES.keys())}")
+            cys_linker_drug_smiles = self.LINKER_CHEMISTRIES[linker_chemistry]
+            super().__init__(cys_linker_drug_smiles, component_name=f"Cys-SS-Linker-Camptothecin ({linker_chemistry})")
+    
+    @classmethod
+    def generate_all_chemistry_variants(cls, peptide_smiles, exclude_n_terminus=True):
+        """
+        Generate conjugates with ALL available linker chemistries (Ester + Carbonate)
+        
+        This is useful for expanding candidate diversity.
+        
+        Args:
+            peptide_smiles: SMILES string of peptide
+            exclude_n_terminus: If True, only conjugate at Lys ε-amines
+        
+        Returns:
+            Dict with:
+                'conjugates': List of all conjugate SMILES
+                'by_chemistry': Dict mapping chemistry name to list of conjugates
+                'by_site': Dict mapping site index to list of (chemistry, smiles) tuples
+                'total': Total number of conjugates
+        """
+        results = {
+            'conjugates': [],
+            'by_chemistry': {},
+            'by_site': {},
+            'total': 0
+        }
+        
+        for chemistry_name in cls.LINKER_CHEMISTRIES.keys():
+            builder = cls(linker_chemistry=chemistry_name)
+            
+            try:
+                conjugates = builder.conjugate_to_peptide(
+                    peptide_smiles, 
+                    exclude_n_terminus=exclude_n_terminus
+                )
+                
+                results['by_chemistry'][chemistry_name] = conjugates
+                
+                for i, conj in enumerate(conjugates):
+                    results['conjugates'].append(conj)
+                    
+                    # Track by site
+                    if i not in results['by_site']:
+                        results['by_site'][i] = []
+                    results['by_site'][i].append((chemistry_name, conj))
+                
+            except Exception as e:
+                print(f" {chemistry_name} variant failed: {e}")
+                results['by_chemistry'][chemistry_name] = []
+        
+        results['total'] = len(results['conjugates'])
+        
+        print(f"\n=== All Chemistry Variants ===")
+        print(f"Total conjugates: {results['total']}")
+        for chem, conjs in results['by_chemistry'].items():
+            print(f"  {chem}: {len(conjs)} conjugate(s)")
+        
+        return results
+    
+    @classmethod
+    def batch_generate_all_variants(cls, peptide_smiles_list, exclude_n_terminus=True):
+        """
+        Batch generate all chemistry variants for multiple peptides
+        
+        Args:
+            peptide_smiles_list: List of peptide SMILES strings
+            exclude_n_terminus: If True, only conjugate at Lys ε-amines
+        
+        Returns:
+            Dict with:
+                'all_conjugates': Flat list of all conjugate SMILES
+                'by_peptide': List of results dicts (one per peptide)
+                'summary': Overall statistics
+        """
+        results = {
+            'all_conjugates': [],
+            'by_peptide': [],
+            'summary': {
+                'peptides': len(peptide_smiles_list),
+                'successful': 0,
+                'failed': 0,
+                'total_conjugates': 0
+            }
+        }
+        
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            print(f"\n--- Peptide {i+1}/{len(peptide_smiles_list)} ---")
+            try:
+                pep_results = cls.generate_all_chemistry_variants(
+                    peptide_smiles, 
+                    exclude_n_terminus=exclude_n_terminus
+                )
+                results['by_peptide'].append(pep_results)
+                results['all_conjugates'].extend(pep_results['conjugates'])
+                results['summary']['successful'] += 1
+                results['summary']['total_conjugates'] += pep_results['total']
+            except Exception as e:
+                print(f" Peptide {i} failed completely: {e}")
+                results['by_peptide'].append(None)
+                results['summary']['failed'] += 1
+        
+        print(f"\n=== Batch Summary ===")
+        print(f"Peptides: {results['summary']['peptides']}")
+        print(f"Successful: {results['summary']['successful']}")
+        print(f"Failed: {results['summary']['failed']}")
+        print(f"Total conjugates: {results['summary']['total_conjugates']}")
+        
+        return results
+    
+    def _verify_reactive_group(self):
+        """Verify carboxylic acid group is present for amide bond formation"""
+        # Pattern: Carboxylic acid (COOH)
+        carboxylic_acid_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+        
+        if not self.linker_payload_mol.HasSubstructMatch(carboxylic_acid_pattern):
+            print(" Warning: No carboxylic acid group detected in Cys-SS-Linker-Drug")
+        else:
+            print(" Cys-SS-Linker-Camptothecin loaded with carboxylic acid reactive group")
+        
+        # Verify disulfide bond is present
+        disulfide_pattern = Chem.MolFromSmarts('[#16]-[#16]')
+        if self.linker_payload_mol.HasSubstructMatch(disulfide_pattern):
+            print(" Disulfide bond (S-S) detected")
+        else:
+            print(" Warning: No disulfide bond detected - check structure")
+        
+        # Verify Cys α-amine is present
+        cys_amine_pattern = Chem.MolFromSmarts('[NH2][CH]([CH2]S)C(=O)')
+        if self.linker_payload_mol.HasSubstructMatch(cys_amine_pattern):
+            print(" Cys α-amine detected (will be free in conjugate)")
+        else:
+            print(" Warning: Cys α-amine pattern not detected")
+    
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find all lysine ε-amine groups in peptide
+        
+        Returns list of dicts with site information including atom index
+        """
+        return self.find_peptide_lysines(peptide_smiles)
+    
+    def find_peptide_lysines(self, peptide_smiles):
+        """
+        Find all lysine residues with free ε-amine in peptide
+        
+        Lysine side chain: -CH2-CH2-CH2-CH2-NH2 (CCCCN)
+        
+        Returns:
+            List of dicts with 'atom_idx' (nitrogen index) and metadata
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+        
+        lysines = []
+        
+        # SMARTS pattern for Lysine ε-amine: 4 CH2 groups ending in NH2
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        matches = peptide_mol.GetSubstructMatches(lys_pattern)
+        
+        for match in matches:
+            nitrogen_idx = match[4]
+            
+            lysines.append({
+                'atom_idx': nitrogen_idx,
+                'side_chain_atoms': list(match),
+                'type': 'Lys_epsilon_amine'
+            })
+        
+        # Also count N-terminus for informational purposes
+        n_terminus_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        n_term_matches = peptide_mol.GetSubstructMatches(n_terminus_pattern)
+        
+        if n_term_matches:
+            print(f"  (Also found {len(n_term_matches)} N-terminus α-amine(s) - excluded by default)")
+        
+        return lysines
+    
+    def count_all_nh2_sites(self, peptide_smiles):
+        """
+        Count all NH2 sites in peptide (both Lys ε-amines and N-terminus)
+        
+        Returns:
+            Dict with counts: {'lys_epsilon': n, 'n_terminus': n, 'other': n, 'total': n}
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+        
+        counts = {'lys_epsilon': 0, 'n_terminus': 0, 'other': 0, 'total': 0}
+        
+        # Lys ε-amine: CH2-CH2-CH2-CH2-NH2 (4-carbon side chain)
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        counts['lys_epsilon'] = len(peptide_mol.GetSubstructMatches(lys_pattern))
+        
+        # N-terminus: NH2 attached to backbone carbon
+        n_term_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        counts['n_terminus'] = len(peptide_mol.GetSubstructMatches(n_term_pattern))
+        
+        # Total NH2 (any primary amine)
+        all_nh2_pattern = Chem.MolFromSmarts('[NH2]')
+        total_nh2 = len(peptide_mol.GetSubstructMatches(all_nh2_pattern))
+        
+        counts['other'] = total_nh2 - counts['lys_epsilon'] - counts['n_terminus']
+        counts['total'] = total_nh2
+        
+        return counts
+    
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """
+        Perform amide bond conjugation between NH2 groups and Cys carboxylic acid
+        
+        Args:
+            peptide_mol: RDKit Mol object
+            site_info: Not used (kept for compatibility)
+            exclude_n_terminus: If True, skip N-terminus conjugation
+        
+        Returns: List of all possible conjugate SMILES
+        """
+        return self._manual_conjugation(peptide_mol, site_info, exclude_n_terminus)
+    
+    def _manual_conjugation(self, peptide_mol, lysine_info, exclude_n_terminus=True):
+        """
+        Create ALL possible conjugates by forming amide bond at NH2 sites
+        
+        Chemistry:
+        - Lys-NH2 + HOOC-Cys-SS-Linker-Drug → Lys-NH-C(=O)-Cys-SS-Linker-Drug + H2O
+        
+        Args:
+            peptide_mol: RDKit Mol object of peptide
+            lysine_info: Not used (kept for interface compatibility)
+            exclude_n_terminus: If True, only conjugate at Lys ε-amines (4-carbon chain),
+                               not at N-terminus α-amine. Default: True
+        
+        Returns:
+            List of all possible conjugate SMILES
+        """
+        if exclude_n_terminus:
+            # SMARTS that matches only Lys ε-amine: requires 4 CH2 chain before NH2
+            # Include wildcard (*) to preserve connection to peptide backbone
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:5][CH2][CH2][CH2][CH2:4][NH2:1].[C:2](=[O:3])[OH]>>[*:5][CH2][CH2][CH2][CH2:4][NH:1][C:2](=[O:3])'
+            )
+        else:
+            # Match ANY primary amine (includes N-terminus)
+            # Include wildcard to preserve the rest of the molecule
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:4][NH2:1].[C:2](=[O:3])[OH]>>[*:4][NH:1][C:2](=[O:3])'
+            )
+        
+        # Run the reaction - returns all possible products
+        products = amide_rxn.RunReactants((peptide_mol, self.linker_payload_mol))
+        
+        if not products:
+            raise ValueError("Amide bond formation reaction failed - no products generated")
+        
+        # Collect ALL valid conjugates
+        conjugates = []
+        seen_smiles = set()  # Deduplicate
+        
+        for product_tuple in products:
+            try:
+                conjugate_mol = product_tuple[0]
+                Chem.SanitizeMol(conjugate_mol)
+                conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+                
+                # Deduplicate (same SMILES = same molecule)
+                if conjugate_smiles not in seen_smiles:
+                    seen_smiles.add(conjugate_smiles)
+                    conjugates.append(conjugate_smiles)
+            except Exception as e:
+                # Skip invalid products
+                continue
+        
+        if not conjugates:
+            raise ValueError("All products failed sanitization")
+        
+        return conjugates
+    
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Cys-SS-Linker-Camptothecin to peptide via amide bond at NH2 groups
+        
+        If linker_chemistry='all' (default), generates conjugates with BOTH ester and 
+        carbonate chemistries, returning a combined list.
+        
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain Lys with NH2)
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines,
+                               not at N-terminus. Set to False to include N-terminus.
+            return_all: If True, return all possible conjugates. If False, return first only.
+        
+        Returns:
+            List of all possible conjugate SMILES (from all chemistries if mode='all')
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+        
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Lys ε-amine site(s)")
+        
+        # If 'all' mode, generate conjugates with both chemistries
+        if self.linker_chemistry == 'all':
+            all_conjugates = []
+            seen_smiles = set()
+            
+            for chem_name, chem_smiles in self.LINKER_CHEMISTRIES.items():
+                # Temporarily set linker_payload_mol to this chemistry
+                temp_mol = Chem.MolFromSmiles(chem_smiles)
+                original_mol = self.linker_payload_mol
+                self.linker_payload_mol = temp_mol
+                
+                try:
+                    conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+                    
+                    for conj in conjugates:
+                        if conj not in seen_smiles:
+                            seen_smiles.add(conj)
+                            all_conjugates.append(conj)
+                    
+                    print(f"  {chem_name}: {len(conjugates)} conjugate(s)")
+                    
+                except Exception as e:
+                    print(f"  {chem_name}: failed - {e}")
+                finally:
+                    # Restore original
+                    self.linker_payload_mol = original_mol
+            
+            site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2"
+            print(f"Generated {len(all_conjugates)} total unique conjugate(s) [{site_type}] [all chemistries]")
+            
+            if not all_conjugates:
+                raise ValueError("No conjugates generated - no valid conjugation sites found")
+            
+            if return_all:
+                return all_conjugates
+            else:
+                return all_conjugates[0] if all_conjugates else None
+        
+        else:
+            # Single chemistry mode
+            conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+            
+            site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2 (including N-terminus)"
+            print(f"Generated {len(conjugates)} unique conjugate(s) [{site_type}] [{self.linker_chemistry}]")
+            
+            if return_all:
+                return conjugates
+            else:
+                return conjugates[0] if conjugates else None
+    
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """
+        Batch conjugate multiple peptides
+        
+        Args:
+            peptide_smiles_list: List of peptide SMILES strings
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates per peptide
+        
+        Returns:
+            Dict with:
+                'conjugates': List of lists (if return_all=True) or list of SMILES
+                'successful': List of successful peptide indices
+                'failed': List of (index, error) tuples
+                'total_conjugates': Total number of conjugates generated
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+        
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles, 
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+        
+        print(f"\n=== Summary ===")
+        print(f"Peptides: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, Failed: {len(results['failed'])}")
+        print(f"Total conjugates generated: {results['total_conjugates']}")
+        
+        return results
+    
+    def flatten_batch_results(self, results):
+        """
+        Flatten batch conjugation results into a single list of all conjugates
+        
+        Args:
+            results: Dict returned by batch_conjugate()
+        
+        Returns:
+            List of all conjugate SMILES (flattened)
+        """
+        all_conjugates = []
+        for conjugate_list in results['conjugates']:
+            if conjugate_list is not None:
+                if isinstance(conjugate_list, list):
+                    all_conjugates.extend(conjugate_list)
+                else:
+                    all_conjugates.append(conjugate_list)
+        return all_conjugates
+    
+    # ========================================================================
+    # Attachment Site Compatibility Checks (Lys ε-amine + Cys-COOH amide bond)
+    # ========================================================================
+    
+    def has_lys_cys_amide_conjugation(self, mol):
+        """
+        Check if molecule has Lys-Cys amide conjugation bond
+        
+        Pattern: Lys-ε-N-C(=O)-Cα(Cys)-CH2-S-S
+        The amide bond connects Lys side chain to Cys α-carbon
+        
+        Returns:
+            bool: True if Lys-Cys amide conjugation is present
+        """
+        # Pattern: N-C(=O)-C(N)(C-S-S) - amide to Cys with disulfide
+        lys_cys_amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        return mol.HasSubstructMatch(lys_cys_amide_pattern)
+    
+    def has_disulfide_bond(self, mol):
+        """
+        Check if molecule has disulfide bond (S-S)
+        
+        Returns:
+            bool: True if disulfide bond is present
+        """
+        disulfide_pattern = Chem.MolFromSmarts('[SX2]-[SX2]')
+        return mol.HasSubstructMatch(disulfide_pattern)
+    
+    def has_unreacted_lys_epsilon_amine(self, mol):
+        """
+        Check for unreacted Lys ε-amine (free NH2 on 4-carbon side chain)
+        
+        In a fully conjugated PDC, all Lys ε-amines should be consumed.
+        This checks for remaining free NH2 groups on Lys side chains.
+        
+        Returns:
+            bool: True if unreacted Lys ε-amine is present (indicates incomplete conjugation)
+        """
+        # Lys ε-amine: 4 CH2 chain ending in NH2
+        lys_amine_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        return mol.HasSubstructMatch(lys_amine_pattern)
+    
+    def has_unreacted_cys_carboxylic_acid(self, mol):
+        """
+        Check for unreacted Cys carboxylic acid (from Cys-SS-Linker-Drug)
+        
+        Pattern: Free COOH on Cys α-carbon adjacent to CH2-S-S
+        
+        Returns:
+            bool: True if unreacted Cys-COOH is present (indicates unconjugated linker-drug)
+        """
+        # Cys with free COOH: H2N-CH(CH2-S-S)-COOH
+        cys_cooh_pattern = Chem.MolFromSmarts('[NH2][CH]([CH2][SX2][SX2])[CX3](=[OX1])[OX2H1]')
+        return mol.HasSubstructMatch(cys_cooh_pattern)
+    
+    def find_amide_conjugation_sites(self, mol):
+        """
+        Find all Lys-Cys amide conjugation sites in the molecule
+        
+        Returns:
+            List of tuples: (nitrogen_idx, carbonyl_carbon_idx, cys_alpha_carbon_idx)
+        """
+        # Pattern: N-C(=O)-C(N)(C-S-S) 
+        lys_cys_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        matches = mol.GetSubstructMatches(lys_cys_pattern)
+        
+        sites = []
+        for match in matches:
+            nitrogen_idx = match[0]       # Lys nitrogen
+            carbonyl_c_idx = match[1]     # Carbonyl carbon
+            cys_alpha_idx = match[3]      # Cys α-carbon
+            sites.append((nitrogen_idx, carbonyl_c_idx, cys_alpha_idx))
+        
+        return sites
+    
+    def check_valency_at_conjugation_site(self, mol, site_indices):
+        """
+        Check that valency is correct at Lys-Cys conjugation site
+        
+        Args:
+            mol: RDKit Mol object
+            site_indices: Tuple of (nitrogen_idx, carbonyl_c_idx, cys_alpha_idx)
+        
+        Returns:
+            bool: True if valency is correct
+        """
+        nitrogen_idx, carbonyl_c_idx, cys_alpha_idx = site_indices
+        
+        try:
+            # Lys nitrogen should have 3 bonds (NH connected to amide)
+            n_atom = mol.GetAtomWithIdx(nitrogen_idx)
+            if n_atom.GetTotalDegree() != 3:
+                return False
+            
+            # Carbonyl carbon should have 3 bonds (C=O and two single bonds)
+            c_atom = mol.GetAtomWithIdx(carbonyl_c_idx)
+            if c_atom.GetTotalDegree() != 3:
+                return False
+            
+            # Cys α-carbon should have 4 bonds
+            cys_c_atom = mol.GetAtomWithIdx(cys_alpha_idx)
+            if cys_c_atom.GetTotalDegree() != 4:
+                return False
+            
+            return True
+            
+        except Exception:
+            return False
+    
+    def check_attachment_site_compatibility(self, conjugate_smiles):
+        """
+        Check attachment site compatibility for Lys-Cys-SS-Linker-Drug conjugation
+        
+        Validates:
+        1. Lys-Cys amide conjugation is present
+        2. Disulfide bond (S-S) is intact
+        3. Valency at conjugation sites is correct
+        4. No unreacted Cys-COOH (unconjugated linker-drug)
+        
+        Args:
+            conjugate_smiles: SMILES string of the conjugate
+        
+        Returns:
+            Tuple of (passed: bool, details: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        
+        if mol is None:
+            return False, {
+                'parseable': False,
+                'error': 'Could not parse SMILES'
+            }
+        
+        details = {
+            'parseable': True,
+            'validation_method': 'lys_cys_amide_disulfide'
+        }
+        
+        # Check 1: Lys-Cys amide conjugation present
+        details['has_lys_cys_amide'] = self.has_lys_cys_amide_conjugation(mol)
+        
+        # Check 2: Disulfide bond intact
+        details['has_disulfide_bond'] = self.has_disulfide_bond(mol)
+        
+        # Check 3: Find conjugation sites and validate valency
+        conjugation_sites = self.find_amide_conjugation_sites(mol)
+        details['conjugation_sites'] = len(conjugation_sites)
+        
+        valency_ok = True
+        if conjugation_sites:
+            for site in conjugation_sites:
+                if not self.check_valency_at_conjugation_site(mol, site):
+                    valency_ok = False
+                    break
+        details['valency_correct'] = valency_ok
+        
+        # Check 4: No unreacted Cys-COOH (would indicate unconjugated linker-drug)
+        details['no_unreacted_cys_cooh'] = not self.has_unreacted_cys_carboxylic_acid(mol)
+        
+        # Optional info: Check for unreacted Lys ε-amines (for multi-Lys peptides)
+        details['has_unreacted_lys_amine'] = self.has_unreacted_lys_epsilon_amine(mol)
+        
+        # Overall pass/fail
+        passed = all([
+            details['has_lys_cys_amide'],
+            details['has_disulfide_bond'],
+            details['valency_correct'],
+            details['no_unreacted_cys_cooh']
+        ])
+        
+        return passed, details
+    
+    def _basic_validation(self, mol):
+        """
+        Basic validation checks for Disulfide-Camptothecin conjugates
+        
+        Checks:
+        1. Amide bond present (Lys-Cys conjugation)
+        2. Disulfide bond present (S-S linker)
+        3. Specific Lys-Cys-SS conjugation pattern
+        4. Camptothecin core present
+        5. Ester bond to drug present
+        """
+        checks = {}
+        
+        # Check 1: Has amide bond (general check)
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[#6]')
+        checks['has_amide_bond'] = mol.HasSubstructMatch(amide_pattern)
+        
+        # Check 2: Disulfide bond present (S-S)
+        disulfide_pattern = Chem.MolFromSmarts('[#16]-[#16]')
+        checks['has_disulfide_bond'] = mol.HasSubstructMatch(disulfide_pattern)
+        
+        # Check 3: Specific Lys-Cys-SS conjugation pattern
+        # Looking for: N-C(=O)-C-C-S-S (amide to Cys to disulfide)
+        lys_cys_ss_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH]([CH2][SX2][SX2])[NX3]')
+        checks['has_lys_cys_ss_conjugation'] = mol.HasSubstructMatch(lys_cys_ss_pattern)
+        
+        # Check 4: Camptothecin core (quinoline fused ring system)
+        # Simplified pattern for the fused ring system
+        camptothecin_core = Chem.MolFromSmarts('c1ccc2ncc3C(=O)OCc3c2c1')
+        checks['has_camptothecin_core'] = mol.HasSubstructMatch(camptothecin_core)
+        
+        # Check 5: Ester bond (linker to drug)
+        ester_pattern = Chem.MolFromSmarts('[#6][CX3](=[OX1])[OX2][#6]')
+        checks['has_ester_bond'] = mol.HasSubstructMatch(ester_pattern)
+        
+        # Overall pass/fail - based on essential conjugation checks
+        checks['_passed'] = all([
+            checks['has_amide_bond'],
+            checks['has_disulfide_bond'],
+            checks['has_lys_cys_ss_conjugation']
+        ])
+        
+        return checks
+    
+    def validate_conjugate(self, conjugate_smiles):
+        """
+        Validate that the conjugate is chemically reasonable
+        
+        Returns:
+            (passed: bool, checks: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+        
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+        
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+        
+        return passed, checks
+
+
+# ============================================================================
+# Main - Test Examples  
+# ============================================================================
+
+if __name__ == "__main__":
+    
+    print("=" * 70)
+    print("Disulfide-Camptothecin PDC Builder - Test")
+    print("Reference: PDC_00294 (Z7, Ester) & PDC_00295 (Z8, Carbonate)")
+    print("=" * 70 + "\n")
+    
+    # Test peptides
+    test_peptides = [
+        # Simple test: Gly-Lys-Gly (has N-terminus NH2 + Lys ε-amine)
+        "NCC(=O)N[C@@H](CCCCN)C(=O)NCC(=O)O",
+        # Lys-Arg
+        "N[C@@H](CCCCN)C(=O)N[C@@H](CCCNC(=N)N)C(=O)O",
+    ]
+    
+    print("\n" + "=" * 70)
+    print("Test 1: DEFAULT MODE (all chemistries: Ester + Carbonate)")
+    print("=" * 70 + "\n")
+    
+    # Default: linker_chemistry='all'
+    builder = DisulfideCamptothecin_PDC_Builder()
+    
+    peptide = test_peptides[0]
+    print(f"Input peptide (Gly-Lys-Gly): {peptide}")
+    
+    try:
+        conjugates = builder.conjugate_to_peptide(peptide, exclude_n_terminus=True)
+        print(f"\n Total conjugates: {len(conjugates)} (should be 2: 1 ester + 1 carbonate)")
+        
+        for i, conj in enumerate(conjugates):
+            valid, checks = builder.validate_conjugate(conj)
+            print(f"  Conjugate {i+1}: {'PASSED' if valid else 'FAILED'} ({len(conj)} chars)")
+            
+    except Exception as e:
+        print(f"\n FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    print("\n" + "=" * 70)
+    print("Test 2: Single Chemistry Mode (Ester only)")
+    print("=" * 70 + "\n")
+    
+    builder_ester = DisulfideCamptothecin_PDC_Builder(linker_chemistry='ester')
+    
+    print(f"Input peptide (Gly-Lys-Gly): {peptide}")
+    
+    try:
+        conjugates = builder_ester.conjugate_to_peptide(peptide, exclude_n_terminus=True)
+        print(f"\n Ester only: {len(conjugates)} conjugate(s)")
+            
+    except Exception as e:
+        print(f"\n FAILED: {e}")
+    
+    print("\n" + "=" * 70)
+    print("Test 3: Single Chemistry Mode (Carbonate only)")
+    print("=" * 70 + "\n")
+    
+    builder_carbonate = DisulfideCamptothecin_PDC_Builder(linker_chemistry='carbonate')
+    
+    print(f"Input peptide (Gly-Lys-Gly): {peptide}")
+    
+    try:
+        conjugates = builder_carbonate.conjugate_to_peptide(peptide, exclude_n_terminus=True)
+        print(f"\n Carbonate only: {len(conjugates)} conjugate(s)")
+            
+    except Exception as e:
+        print(f"\n FAILED: {e}")
+    
+    print("\n" + "=" * 70)
+    print("Test 4: Batch Conjugation (default: all chemistries)")
+    print("=" * 70 + "\n")
+    
+    builder = DisulfideCamptothecin_PDC_Builder()  # default: all
+    
+    try:
+        results = builder.batch_conjugate(test_peptides, exclude_n_terminus=True)
+        
+        print(f"\nResults breakdown:")
+        for i in results['successful']:
+            conjugates = results['conjugates'][i]
+            num_valid = sum(1 for c in conjugates if builder.validate_conjugate(c)[0])
+            name = ["Gly-Lys-Gly", "Lys-Arg"][i]
+            print(f"  {name}: {len(conjugates)} conjugate(s), {num_valid} valid")
+            
+    except Exception as e:
+        print(f"\n FAILED: {e}")
+        import traceback
+        traceback.print_exc()

smiles/conjugation/ValAlaDau.py

@@ -0,0 +1,348 @@
+"""
+Val-Ala-PABC-Daunorubicin PDC Builder
+
+Conjugation chemistry: Lysine ε-amine + Glutaryl-Val-Ala-PABC-Daunorubicin (amide bond formation)
+
+Reference PDCs:
+- PDC_00185 (GnRH-III-[2ΔHis-3D-Tic, 8Lys(glutaryl-Val-Ala-PABC-Dau)])
+- PDC_00193 (GnRH-III-[2His-3Trp,8Lys(glutaryl-Val-Ala-PABC-Dau)])
+- Receptor: GNRHR
+- Linker: Val-Ala (LIN00028)
+- Payload: Daunorubicin (DRG00010)
+
+Architecture:
+    Peptide──Lys──NH──C(=O)──(CH₂)₃──C(=O)──Val──Ala──PABC──O──C(=O)──NH──Dau
+                  │              │              │         │
+              amide bond    glutaryl     dipeptide     carbamate
+              (reaction)    spacer       cleavable      self-immolative
+
+The cleavable dipeptide Val-Ala is recognized by cathepsin B (lysosomal protease).
+After cleavage, PABC self-immolation releases free Daunorubicin.
+"""
+
+from rdkit import Chem
+from rdkit.Chem import RWMol, AllChem
+
+from .pdc_builder import PDC_Builder
+
+
+class ValAlaPABCDaunorubicin_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for Glutaryl-Val-Ala-PABC-Daunorubicin conjugates
+
+    Conjugation chemistry (amide bond formation):
+    - Peptide site: Lysine ε-amine (side chain -NH2)
+    - Linker-payload site: Free carboxylic acid on glutaryl spacer
+    - Reaction: Amide bond (Lys-NH2 + HOOC-glutaryl-Val-Ala-PABC-Dau)
+
+    Architecture:
+        Peptide──Lys──NH──CO──(CH₂)₃──CO──Val──Ala──PABC──OCO──NH──Dau
+                      └── amide ──┘         └── cleavable ──┘  └─ carbamate ─┘
+
+    Reference: PDC_00185, PDC_00193 (GnRH-III conjugates, GNRHR)
+    """
+
+    # Val-Ala-PABC-Daunorubicin with free glutaryl COOH
+    # Structure: HOOC-(CH2)3-CO-Val-Ala-PABC-OCO-NH-Dau
+    # - Glutaryl spacer connects to Val N-terminus via amide
+    # - Val-Ala dipeptide (cathepsin B cleavable)
+    # - PABC (para-aminobenzylcarbamate) self-immolative spacer
+    # - Carbamate bond from PABC to Dau 3'-amine
+    DEFAULT_VALA_PABC_DAU_SMILES = (
+        "C[C@H]1[C@H]([C@H](C[C@@H](O1)O[C@H]2C[C@@](CC3=C2C(=C4C(=C3O)"
+        "C(=O)C5=C(C4=O)C(=CC=C5)OC)O)(C(=O)C)O)NC(=O)OCc6ccc(NC(=O)"
+        "[C@@H](C)NC(=O)[C@@H](C(C)C)NC(=O)CCCC(=O)O)cc6)O"
+    )
+
+    def __init__(self, vala_pabc_dau_smiles=None):
+        """
+        Initialize Val-Ala-PABC-Daunorubicin builder
+
+        Args:
+            vala_pabc_dau_smiles: Optional custom SMILES. If None, uses default.
+                                 Must have free carboxylic acid for conjugation.
+        """
+        if vala_pabc_dau_smiles is None:
+            vala_pabc_dau_smiles = self.DEFAULT_VALA_PABC_DAU_SMILES
+
+        super().__init__(vala_pabc_dau_smiles,
+                         component_name="Val-Ala-PABC-Daunorubicin (linker-payload)")
+
+    def _verify_reactive_group(self):
+        """Verify carboxylic acid group is present for amide bond formation"""
+        carboxylic_acid_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+
+        if not self.linker_payload_mol.HasSubstructMatch(carboxylic_acid_pattern):
+            print(" Warning: No carboxylic acid group detected in Val-Ala-PABC-Dau")
+        else:
+            print(" Val-Ala-PABC-Daunorubicin loaded with carboxylic acid reactive group")
+
+        # Verify PABC benzyl carbamate
+        pabc_pattern = Chem.MolFromSmarts('c1ccc(COC(=O)N)cc1')
+        if self.linker_payload_mol.HasSubstructMatch(pabc_pattern):
+            print(" PABC self-immolative spacer detected")
+        else:
+            print(" Warning: No PABC pattern detected")
+
+        # Verify dipeptide amide bonds
+        dipeptide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH][NX3][CX3](=[OX1])')
+        if self.linker_payload_mol.HasSubstructMatch(dipeptide_pattern):
+            print(" Dipeptide (Val-Ala) linkage detected")
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """Find all lysine ε-amine groups in peptide"""
+        return self._find_peptide_lysines(peptide_smiles)
+
+    def _find_peptide_lysines(self, peptide_smiles):
+        """Find all lysine residues with free ε-amine in peptide"""
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        lysines = []
+
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        matches = peptide_mol.GetSubstructMatches(lys_pattern)
+
+        for match in matches:
+            nitrogen_idx = match[4]
+            lysines.append({
+                'atom_idx': nitrogen_idx,
+                'side_chain_atoms': list(match),
+                'type': 'Lys_epsilon_amine'
+            })
+
+        n_terminus_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        n_term_matches = peptide_mol.GetSubstructMatches(n_terminus_pattern)
+        if n_term_matches:
+            print(f"  (Also found {len(n_term_matches)} N-terminus α-amine(s) - excluded by default)")
+
+        return lysines
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """Perform amide bond conjugation"""
+        return self._manual_conjugation(peptide_mol, site_info, exclude_n_terminus)
+
+    def _manual_conjugation(self, peptide_mol, lysine_info, exclude_n_terminus=True):
+        """
+        Create ALL possible conjugates by forming amide bond at NH2 sites
+
+        Chemistry:
+        - Any-NH2 + HOOC-glutaryl-Val-Ala-PABC-Dau → amide bond + H2O
+        """
+        if exclude_n_terminus:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:5][CH2][CH2][CH2][CH2:4][NH2:1].[C:2](=[O:3])[OH]>>'
+                '[*:5][CH2][CH2][CH2][CH2:4][NH:1][C:2](=[O:3])'
+            )
+        else:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:4][NH2:1].[C:2](=[O:3])[OH]>>[*:4][NH:1][C:2](=[O:3])'
+            )
+
+        products = amide_rxn.RunReactants((peptide_mol, self.linker_payload_mol))
+
+        if not products:
+            raise ValueError("Amide bond formation reaction failed - no products generated")
+
+        conjugates = []
+        seen_smiles = set()
+
+        for product_tuple in products:
+            try:
+                conjugate_mol = product_tuple[0]
+                Chem.SanitizeMol(conjugate_mol)
+                conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+
+                if conjugate_smiles not in seen_smiles:
+                    seen_smiles.add(conjugate_smiles)
+                    conjugates.append(conjugate_smiles)
+            except Exception:
+                continue
+
+        if not conjugates:
+            raise ValueError("All products failed sanitization")
+
+        return conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Val-Ala-PABC-Daunorubicin to peptide via amide bond
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain NH2)
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            If return_all=True: List of all possible conjugate SMILES
+            If return_all=False: Single conjugate SMILES
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Lys ε-amine site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2 (including N-terminus)"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [{site_type}]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """Batch conjugate multiple peptides"""
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        return results
+
+    # =========================================================================
+    # Attachment Site Compatibility Checks
+    # =========================================================================
+
+    def has_pabc_carbamate(self, mol):
+        """Check for PABC benzyl carbamate: Ar-CH2-O-C(=O)-N"""
+        pabc_pattern = Chem.MolFromSmarts('c1ccc(COC(=O)[NX3])cc1')
+        if pabc_pattern:
+            return mol.HasSubstructMatch(pabc_pattern)
+        return False
+
+    def has_val_ala_dipeptide(self, mol):
+        """
+        Check for Val-Ala dipeptide fragment
+
+        Pattern: isopropyl-CH-C(=O)-NH-CH(CH3)-C(=O)
+        """
+        # General dipeptide amide pattern
+        dipeptide = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[CH][NX3][CX3](=[OX1])')
+        if dipeptide:
+            return mol.HasSubstructMatch(dipeptide)
+        return False
+
+    def has_glutaryl_spacer(self, mol):
+        """Check for glutaryl spacer: -C(=O)-(CH2)3-C(=O)-"""
+        glutaryl = Chem.MolFromSmarts('[CX3](=[OX1])[CH2][CH2][CH2][CX3](=[OX1])')
+        if glutaryl:
+            return mol.HasSubstructMatch(glutaryl)
+        return False
+
+    def has_unreacted_glutaryl_cooh(self, mol):
+        """Check if free COOH on glutaryl end remains (should be consumed)"""
+        unreacted = Chem.MolFromSmarts(
+            '[OX2H1][CX3](=[OX1])[CH2][CH2][CH2][CX3](=[OX1])[NX3]'
+        )
+        if unreacted:
+            return mol.HasSubstructMatch(unreacted)
+        return False
+
+    def check_attachment_site_compatibility(self, mol):
+        """
+        Check if attachment site chemistry is valid
+
+        Validates:
+        1. Has PABC carbamate (intact self-immolative spacer)
+        2. Has glutaryl spacer
+        3. No unreacted glutaryl-COOH
+        """
+        if not self.has_pabc_carbamate(mol):
+            return False
+
+        if not self.has_glutaryl_spacer(mol):
+            return False
+
+        if self.has_unreacted_glutaryl_cooh(mol):
+            return False
+
+        return True
+
+    def _basic_validation(self, mol):
+        """Basic validation checks for Val-Ala-PABC-Daunorubicin conjugates"""
+        checks = {}
+
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[#6]')
+        checks['has_amide_bond'] = mol.HasSubstructMatch(amide_pattern)
+
+        checks['has_pabc_carbamate'] = self.has_pabc_carbamate(mol)
+        checks['has_val_ala_dipeptide'] = self.has_val_ala_dipeptide(mol)
+        checks['has_glutaryl_spacer'] = self.has_glutaryl_spacer(mol)
+        checks['no_unreacted_cooh'] = not self.has_unreacted_glutaryl_cooh(mol)
+        checks['attachment_site_valid'] = self.check_attachment_site_compatibility(mol)
+
+        anthraquinone = Chem.MolFromSmarts(
+            '[#6]1[#6][#6][#6]2[#6](=[#8])[#6]3[#6][#6][#6][#6][#6]3[#6](=[#8])[#6]2[#6]1'
+        )
+        checks['has_daunorubicin_core'] = mol.HasSubstructMatch(anthraquinone)
+
+        aminosugar = Chem.MolFromSmarts('[OX2]1[#6][#6][#6]([NX3])[#6][#6]1')
+        checks['has_aminosugar'] = mol.HasSubstructMatch(aminosugar)
+
+        checks['_passed'] = checks['attachment_site_valid']
+
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """Validate that the conjugate is chemically reasonable"""
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+
+        return passed, checks
+
+
+if __name__ == "__main__":
+    print("=" * 70)
+    print("Val-Ala-PABC-Daunorubicin PDC Builder - Test")
+    print("=" * 70 + "\n")
+
+    builder = ValAlaPABCDaunorubicin_PDC_Builder()
+
+    test_peptides = [
+        "NCC(=O)N[C@@H](CCCCN)C(=O)NCC(=O)O",  # Gly-Lys-Gly
+    ]
+
+    for pep in test_peptides:
+        try:
+            conjugates = builder.conjugate_to_peptide(pep, exclude_n_terminus=True)
+            print(f" {len(conjugates)} conjugate(s)")
+            for c in conjugates:
+                valid, checks = builder.validate_conjugate(c)
+                print(f"  Valid: {valid}")
+        except Exception as e:
+            print(f" Failed: {e}")

smiles/conjugation/ValCitDau.py

@@ -0,0 +1,358 @@
+"""
+Val-Cit-PABC-Daunorubicin PDC Builder
+
+Conjugation chemistry: Lysine ε-amine + Glutaryl-Val-Cit-PABC-Daunorubicin (amide bond formation)
+
+Reference PDCs:
+- PDC_00189 (GnRH-III-[2ΔHis-3D-Tic,8Lys(glutaryl-Val-Cit-PABC-Dau)])
+- PDC_00194 (GnRH-III-[2His-3Trp,8Lys(glutaryl-Val-Cit-PABC-Dau)])
+- Receptor: GNRHR
+- Linker: Val-Cit (LIN00005)
+- Payload: Daunorubicin (DRG00010)
+
+Architecture:
+    Peptide──Lys──NH──C(=O)──(CH₂)₃──C(=O)──Val──Cit──PABC──O──C(=O)──NH──Dau
+                  │              │              │          │
+              amide bond    glutaryl     dipeptide      carbamate
+              (reaction)    spacer       cleavable       self-immolative
+
+The cleavable dipeptide Val-Cit (valine-citrulline) is recognized by
+cathepsin B (lysosomal protease). After cleavage, PABC self-immolation
+releases free Daunorubicin.
+
+Citrulline (Cit) is a non-coded amino acid with a ureido side chain:
+  -CH₂-CH₂-CH₂-NH-C(=O)-NH₂
+"""
+
+from rdkit import Chem
+from rdkit.Chem import RWMol, AllChem
+
+from .pdc_builder import PDC_Builder
+
+
+class ValCitPABCDaunorubicin_PDC_Builder(PDC_Builder):
+    """
+    PDC builder for Glutaryl-Val-Cit-PABC-Daunorubicin conjugates
+
+    Conjugation chemistry (amide bond formation):
+    - Peptide site: Lysine ε-amine (side chain -NH2)
+    - Linker-payload site: Free carboxylic acid on glutaryl spacer
+    - Reaction: Amide bond (Lys-NH2 + HOOC-glutaryl-Val-Cit-PABC-Dau)
+
+    Architecture:
+        Peptide──Lys──NH──CO──(CH₂)₃──CO──Val──Cit──PABC──OCO──NH──Dau
+                      └── amide ──┘          └── cleavable ──┘  └─ carbamate ─┘
+
+    Reference: PDC_00189, PDC_00194 (GnRH-III conjugates, GNRHR)
+    """
+
+    # Val-Cit-PABC-Daunorubicin with free glutaryl COOH
+    # Structure: HOOC-(CH2)3-CO-Val-Cit-PABC-OCO-NH-Dau
+    # - Glutaryl spacer connects to Val N-terminus via amide
+    # - Val-Cit dipeptide (cathepsin B cleavable, Cit has ureido side chain)
+    # - PABC (para-aminobenzylcarbamate) self-immolative spacer
+    # - Carbamate bond from PABC to Dau 3'-amine
+    DEFAULT_VALCIT_PABC_DAU_SMILES = (
+        "C[C@H]1[C@H]([C@H](C[C@@H](O1)O[C@H]2C[C@@](CC3=C2C(=C4C(=C3O)"
+        "C(=O)C5=C(C4=O)C(=CC=C5)OC)O)(C(=O)C)O)NC(=O)OCc6ccc(NC(=O)"
+        "[C@@H](CCCNC(=O)N)NC(=O)[C@@H](C(C)C)NC(=O)CCCC(=O)O)cc6)O"
+    )
+
+    def __init__(self, valcit_pabc_dau_smiles=None):
+        """
+        Initialize Val-Cit-PABC-Daunorubicin builder
+
+        Args:
+            valcit_pabc_dau_smiles: Optional custom SMILES. If None, uses default.
+                                   Must have free carboxylic acid for conjugation.
+        """
+        if valcit_pabc_dau_smiles is None:
+            valcit_pabc_dau_smiles = self.DEFAULT_VALCIT_PABC_DAU_SMILES
+
+        super().__init__(valcit_pabc_dau_smiles,
+                         component_name="Val-Cit-PABC-Daunorubicin (linker-payload)")
+
+    def _verify_reactive_group(self):
+        """Verify carboxylic acid group is present for amide bond formation"""
+        carboxylic_acid_pattern = Chem.MolFromSmarts('[CX3](=O)[OX2H1]')
+
+        if not self.linker_payload_mol.HasSubstructMatch(carboxylic_acid_pattern):
+            print(" Warning: No carboxylic acid group detected in Val-Cit-PABC-Dau")
+        else:
+            print(" Val-Cit-PABC-Daunorubicin loaded with carboxylic acid reactive group")
+
+        # Verify PABC benzyl carbamate
+        pabc_pattern = Chem.MolFromSmarts('c1ccc(COC(=O)N)cc1')
+        if self.linker_payload_mol.HasSubstructMatch(pabc_pattern):
+            print(" PABC self-immolative spacer detected")
+        else:
+            print(" Warning: No PABC pattern detected")
+
+        # Verify citrulline ureido group
+        urea_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[NX3]')
+        if self.linker_payload_mol.HasSubstructMatch(urea_pattern):
+            print(" Citrulline ureido group detected")
+        else:
+            print(" Warning: No citrulline ureido group detected")
+
+    def find_conjugation_sites(self, peptide_smiles):
+        """Find all lysine ε-amine groups in peptide"""
+        return self._find_peptide_lysines(peptide_smiles)
+
+    def _find_peptide_lysines(self, peptide_smiles):
+        """Find all lysine residues with free ε-amine in peptide"""
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+
+        lysines = []
+
+        lys_pattern = Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][NH2]')
+        matches = peptide_mol.GetSubstructMatches(lys_pattern)
+
+        for match in matches:
+            nitrogen_idx = match[4]
+            lysines.append({
+                'atom_idx': nitrogen_idx,
+                'side_chain_atoms': list(match),
+                'type': 'Lys_epsilon_amine'
+            })
+
+        n_terminus_pattern = Chem.MolFromSmarts('[NH2][CH,CH2][C](=O)')
+        n_term_matches = peptide_mol.GetSubstructMatches(n_terminus_pattern)
+        if n_term_matches:
+            print(f"  (Also found {len(n_term_matches)} N-terminus α-amine(s) - excluded by default)")
+
+        return lysines
+
+    def _perform_conjugation(self, peptide_mol, site_info, exclude_n_terminus=True):
+        """Perform amide bond conjugation"""
+        return self._manual_conjugation(peptide_mol, site_info, exclude_n_terminus)
+
+    def _manual_conjugation(self, peptide_mol, lysine_info, exclude_n_terminus=True):
+        """
+        Create ALL possible conjugates by forming amide bond at NH2 sites
+
+        Chemistry:
+        - Any-NH2 + HOOC-glutaryl-Val-Cit-PABC-Dau → amide bond + H2O
+        """
+        if exclude_n_terminus:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:5][CH2][CH2][CH2][CH2:4][NH2:1].[C:2](=[O:3])[OH]>>'
+                '[*:5][CH2][CH2][CH2][CH2:4][NH:1][C:2](=[O:3])'
+            )
+        else:
+            amide_rxn = AllChem.ReactionFromSmarts(
+                '[*:4][NH2:1].[C:2](=[O:3])[OH]>>[*:4][NH:1][C:2](=[O:3])'
+            )
+
+        products = amide_rxn.RunReactants((peptide_mol, self.linker_payload_mol))
+
+        if not products:
+            raise ValueError("Amide bond formation reaction failed - no products generated")
+
+        conjugates = []
+        seen_smiles = set()
+
+        for product_tuple in products:
+            try:
+                conjugate_mol = product_tuple[0]
+                Chem.SanitizeMol(conjugate_mol)
+                conjugate_smiles = Chem.MolToSmiles(conjugate_mol)
+
+                if conjugate_smiles not in seen_smiles:
+                    seen_smiles.add(conjugate_smiles)
+                    conjugates.append(conjugate_smiles)
+            except Exception:
+                continue
+
+        if not conjugates:
+            raise ValueError("All products failed sanitization")
+
+        return conjugates
+
+    def conjugate_to_peptide(self, peptide_smiles, exclude_n_terminus=True, return_all=True):
+        """
+        Conjugate Val-Cit-PABC-Daunorubicin to peptide via amide bond
+
+        Args:
+            peptide_smiles: SMILES string of peptide (must contain NH2)
+            exclude_n_terminus: If True (default), only conjugate at Lys ε-amines
+            return_all: If True, return all possible conjugates
+
+        Returns:
+            If return_all=True: List of all possible conjugate SMILES
+            If return_all=False: Single conjugate SMILES
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+
+        sites = self.find_conjugation_sites(peptide_smiles)
+        print(f"Found {len(sites)} Lys ε-amine site(s)")
+
+        conjugates = self._perform_conjugation(peptide_mol, None, exclude_n_terminus)
+
+        site_type = "Lys ε-amine only" if exclude_n_terminus else "all NH2 (including N-terminus)"
+        print(f"Generated {len(conjugates)} unique conjugate(s) [{site_type}]")
+
+        if return_all:
+            return conjugates
+        else:
+            return conjugates[0] if conjugates else None
+
+    def batch_conjugate(self, peptide_smiles_list, exclude_n_terminus=True, return_all=True):
+        """Batch conjugate multiple peptides"""
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+            'total_conjugates': 0
+        }
+
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugates = self.conjugate_to_peptide(
+                    peptide_smiles,
+                    exclude_n_terminus=exclude_n_terminus,
+                    return_all=return_all
+                )
+                if return_all:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += len(conjugates)
+                else:
+                    results['conjugates'].append(conjugates)
+                    results['total_conjugates'] += 1
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+
+        return results
+
+    # =========================================================================
+    # Attachment Site Compatibility Checks
+    # =========================================================================
+
+    def has_pabc_carbamate(self, mol):
+        """Check for PABC benzyl carbamate: Ar-CH2-O-C(=O)-N"""
+        pabc_pattern = Chem.MolFromSmarts('c1ccc(COC(=O)[NX3])cc1')
+        if pabc_pattern:
+            return mol.HasSubstructMatch(pabc_pattern)
+        return False
+
+    def has_val_cit_dipeptide(self, mol):
+        """
+        Check for Val-Cit dipeptide fragment
+
+        Citrulline has a ureido side chain: -CH₂CH₂CH₂-NH-C(=O)-NH₂
+        """
+        # Check for ureido (citrulline marker)
+        urea_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[NX3]')
+        if urea_pattern:
+            return mol.HasSubstructMatch(urea_pattern)
+        return False
+
+    def has_glutaryl_spacer(self, mol):
+        """Check for glutaryl spacer: -C(=O)-(CH2)3-C(=O)-"""
+        glutaryl = Chem.MolFromSmarts('[CX3](=[OX1])[CH2][CH2][CH2][CX3](=[OX1])')
+        if glutaryl:
+            return mol.HasSubstructMatch(glutaryl)
+        return False
+
+    def has_unreacted_glutaryl_cooh(self, mol):
+        """Check if free COOH on glutaryl end remains (should be consumed)"""
+        unreacted = Chem.MolFromSmarts(
+            '[OX2H1][CX3](=[OX1])[CH2][CH2][CH2][CX3](=[OX1])[NX3]'
+        )
+        if unreacted:
+            return mol.HasSubstructMatch(unreacted)
+        return False
+
+    def check_attachment_site_compatibility(self, mol):
+        """
+        Check if attachment site chemistry is valid
+
+        Validates:
+        1. Has PABC carbamate
+        2. Has glutaryl spacer
+        3. Has citrulline ureido group
+        4. No unreacted glutaryl-COOH
+        """
+        if not self.has_pabc_carbamate(mol):
+            return False
+
+        if not self.has_glutaryl_spacer(mol):
+            return False
+
+        if not self.has_val_cit_dipeptide(mol):
+            return False
+
+        if self.has_unreacted_glutaryl_cooh(mol):
+            return False
+
+        return True
+
+    def _basic_validation(self, mol):
+        """Basic validation checks for Val-Cit-PABC-Daunorubicin conjugates"""
+        checks = {}
+
+        amide_pattern = Chem.MolFromSmarts('[NX3][CX3](=[OX1])[#6]')
+        checks['has_amide_bond'] = mol.HasSubstructMatch(amide_pattern)
+
+        checks['has_pabc_carbamate'] = self.has_pabc_carbamate(mol)
+        checks['has_val_cit_dipeptide'] = self.has_val_cit_dipeptide(mol)
+        checks['has_glutaryl_spacer'] = self.has_glutaryl_spacer(mol)
+        checks['no_unreacted_cooh'] = not self.has_unreacted_glutaryl_cooh(mol)
+        checks['attachment_site_valid'] = self.check_attachment_site_compatibility(mol)
+
+        anthraquinone = Chem.MolFromSmarts(
+            '[#6]1[#6][#6][#6]2[#6](=[#8])[#6]3[#6][#6][#6][#6][#6]3[#6](=[#8])[#6]2[#6]1'
+        )
+        checks['has_daunorubicin_core'] = mol.HasSubstructMatch(anthraquinone)
+
+        aminosugar = Chem.MolFromSmarts('[OX2]1[#6][#6][#6]([NX3])[#6][#6]1')
+        checks['has_aminosugar'] = mol.HasSubstructMatch(aminosugar)
+
+        checks['_passed'] = checks['attachment_site_valid']
+
+        return checks
+
+    def validate_conjugate(self, conjugate_smiles):
+        """Validate that the conjugate is chemically reasonable"""
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+
+        checks = self._basic_validation(mol)
+        checks['parseable'] = True
+
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']
+
+        return passed, checks
+
+
+if __name__ == "__main__":
+    print("=" * 70)
+    print("Val-Cit-PABC-Daunorubicin PDC Builder - Test")
+    print("=" * 70 + "\n")
+
+    builder = ValCitPABCDaunorubicin_PDC_Builder()
+
+    test_peptides = [
+        "NCC(=O)N[C@@H](CCCCN)C(=O)NCC(=O)O",  # Gly-Lys-Gly
+    ]
+
+    for pep in test_peptides:
+        try:
+            conjugates = builder.conjugate_to_peptide(pep, exclude_n_terminus=True)
+            print(f" {len(conjugates)} conjugate(s)")
+            for c in conjugates:
+                valid, checks = builder.validate_conjugate(c)
+                print(f"  Valid: {valid}")
+        except Exception as e:
+            print(f" Failed: {e}")

smiles/conjugation/VcMMAE.py

@@ -0,0 +1,457 @@
+from rdkit import Chem
+from rdkit.Chem import RWMol
+
+from .pdc_builder import PDC_Builder
+
+
+class VcMMAE_PDC_Builder(PDC_Builder):
+    """
+    Specialized builder for vcMMAE (valine-citrulline-MMAE) conjugates
+    
+    vcMMAE consists of:
+    - Linker: Val-Cit-PABC (valine-citrulline-para-aminobenzylcarbamate) - cleavable dipeptide
+    - Payload: MMAE (monomethyl auristatin E) - cytotoxic drug
+    - Reactive group: Maleimide for cysteine conjugation
+    
+    Conjugation chemistry (thiol-maleimide Michael addition):
+    - Peptide site: Free cysteine thiol (-SH)
+    - vcMMAE site: Maleimide group (C=C in 5-membered ring)
+    - Reaction: Thiol adds to maleimide C=C, forming stable thioether bond
+    - Product: Peptide-S-vcMMAE conjugate
+    """
+    
+    def __init__(self, vcmmae_smiles=None):
+        """
+        Initialize vcMMAE builder
+        
+        Args:
+            vcmmae_smiles: Optional custom vcMMAE (linker-payload) SMILES. If None, uses default.
+        """
+        # Default vcMMAE SMILES: maleimide-(6-carbon spacer)-Val-Cit-PABC-MMAE
+        # Structure: Maleimide-C6-Val-Cit-PABC-MMAE
+        if vcmmae_smiles is None:
+            vcmmae_smiles = "O[C@@H](C1=CC=CC=C1)[C@@H](C)NC([C@H](C)[C@H]([C@@]2([H])N(CCC2)C(C[C@@H](OC)[C@H]([C@@H](C)CC)N(C)C([C@H](C(C)C)NC([C@H](C(C)C)N(C)C(OCC3=CC=C(C=C3)NC([C@H](CCCNC(N)=O)NC([C@H](C(C)C)NC(CCCCCN4C(C=CC4=O)=O)=O)=O)=O)=O)=O)=O)=O)OC)=O"
+        
+        # Initialize base class
+        super().__init__(vcmmae_smiles, component_name="vcMMAE (linker-payload)")
+    
+    def _verify_reactive_group(self):
+        """Verify maleimide group is present"""
+        maleimide_pattern = Chem.MolFromSmarts('[#7]1[#6](=[#8])[#6]=[#6][#6]1=[#8]')
+        if not self.linker_payload_mol.HasSubstructMatch(maleimide_pattern):
+            print(" Warning: No maleimide group detected")
+        else:
+            print(" vcMMAE (linker-payload) loaded with maleimide reactive group")
+        
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find all free cysteine residues (thiol groups) in peptide
+        
+        Returns list of dicts with site information including atom index
+        """
+        return self.find_peptide_cysteines(peptide_smiles)
+    
+    def find_peptide_cysteines(self, peptide_smiles):
+        """
+        Find all free cysteine residues in peptide
+        
+        FIX: Iterate over all atoms instead of using restrictive SMARTS
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError("Invalid peptide SMILES")
+        
+        cysteines = []
+        
+        # FIX: Directly check each sulfur atom
+        for atom in peptide_mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:  # Sulfur
+                # Check if it's a thiol (has hydrogen attached)
+                total_hs = atom.GetTotalNumHs()  # Includes implicit H
+                degree = atom.GetDegree()  # Number of bonds to other heavy atoms
+                
+                # Thiol: S with 1 H and bonded to 1 carbon
+                # (vs disulfide: S-S with no H)
+                if total_hs >= 1 and degree == 1:
+                    cysteines.append({
+                        'atom_idx': atom.GetIdx(),
+                        'num_hydrogens': total_hs,
+                        'neighbors': [peptide_mol.GetAtomWithIdx(n.GetIdx()).GetSymbol() 
+                                     for n in atom.GetNeighbors()]
+                    })
+        
+        return cysteines
+    
+    def _find_maleimide_atoms(self, mol):
+        """
+        Find the maleimide C=C double bond atoms
+        
+        Returns: (carbon1_idx, carbon2_idx) of the C=C in maleimide
+        """
+        # Pattern: N in ring, with two C=O, and C=C
+        # O=C-C=C-C=O with N closing the ring
+        maleimide_pattern = Chem.MolFromSmarts('[#7]1[#6](=[#8])[#6]=[#6][#6]1=[#8]')
+        matches = mol.GetSubstructMatches(maleimide_pattern)
+        
+        if not matches:
+            return None, None
+        
+        # Pattern indices: 0=N, 1=C(=O), 2=C(of C=C), 3=C(of C=C), 4=C(=O)
+        match = matches[0]
+        return match[2], match[3]
+    
+    def _perform_conjugation(self, peptide_mol, site_info):
+        """
+        Perform thiol-maleimide conjugation
+        
+        Uses manual bond formation for reliability with complex molecules
+        """
+        return self._manual_conjugation(peptide_mol, site_info)
+    
+    def conjugate_to_peptide(self, peptide_smiles, cys_position=0):
+        """
+        Conjugate vcMMAE to peptide via thiol-maleimide Michael addition
+        
+        Compatibility wrapper - calls parent class method with renamed parameter
+        """
+        return super().conjugate_to_peptide(peptide_smiles, site_position=cys_position)
+    
+    def batch_conjugate(self, peptide_smiles_list, cys_position=0):
+        """
+        Batch conjugate multiple peptides
+        
+        Compatibility wrapper - calls parent class method with renamed parameter
+        """
+        return super().batch_conjugate(peptide_smiles_list, site_position=cys_position)
+    
+    # =========================================================================
+    # Attachment Site Compatibility Checks (thiol-maleimide specific)
+    # =========================================================================
+    
+    def has_free_thiol(self, mol):
+        """
+        Check if molecule has unconjugated (free) thiol groups
+        
+        A valid PDC should NOT have free thiols - they should be conjugated
+        """
+        # SMARTS pattern check
+        free_thiol_pattern = Chem.MolFromSmarts('[CH2][SH1]')
+        free_thiol_alt = Chem.MolFromSmarts('[CH2][S;H1;v2]')
+        
+        if free_thiol_pattern and mol.HasSubstructMatch(free_thiol_pattern):
+            return True
+        if free_thiol_alt and mol.HasSubstructMatch(free_thiol_alt):
+            return True
+        
+        # Also check by iterating atoms (more reliable)
+        for atom in mol.GetAtoms():
+            if atom.GetAtomicNum() == 16:  # Sulfur
+                if atom.GetTotalNumHs() >= 1 and atom.GetDegree() == 1:
+                    # Sulfur with H and only one bond = free thiol
+                    return True
+        
+        return False
+    
+    def has_unreacted_maleimide(self, mol):
+        """
+        Check if molecule has unreacted maleimide (C=C in ring)
+        
+        After thiol-maleimide conjugation, the C=C should become C-C
+        """
+        maleimide_unreacted = Chem.MolFromSmarts('[#7]1[#6](=[#8])[#6]=[#6][#6]1=[#8]')
+        if maleimide_unreacted:
+            return mol.HasSubstructMatch(maleimide_unreacted)
+        return False
+    
+    def find_thioether_conjugation_sites(self, mol):
+        """
+        Identify thioether (C-S-C) conjugation sites in the molecule
+        
+        For thiol-maleimide conjugation, the sulfur should be bonded to:
+        - One carbon from the peptide cysteine
+        - One carbon from the (former) maleimide ring
+        
+        Returns list of sulfur atom indices that are conjugation sites
+        """
+        thioether_pattern = Chem.MolFromSmarts('[#6][#16][#6]')
+        conjugation_sites = []
+        
+        if thioether_pattern:
+            matches = mol.GetSubstructMatches(thioether_pattern)
+            for match in matches:
+                sulfur_idx = match[1]  # S is the middle atom in C-S-C
+                sulfur_atom = mol.GetAtomWithIdx(sulfur_idx)
+                
+                # Check if one neighbor is in a 5-membered ring (succinimide)
+                neighbors = [mol.GetAtomWithIdx(n.GetIdx()) for n in sulfur_atom.GetNeighbors()]
+                for neighbor in neighbors:
+                    if neighbor.IsInRingSize(5):
+                        conjugation_sites.append(sulfur_idx)
+                        break
+                else:
+                    # Even if not clearly in succinimide, thioether is likely conjugation
+                    conjugation_sites.append(sulfur_idx)
+        
+        return list(set(conjugation_sites))  # Remove duplicates
+    
+    def check_valency_at_conjugation_site(self, mol, conjugation_atom_idx):
+        """
+        Check valency specifically at a sulfur conjugation site atom
+        
+        Valid thioether states for thiol-maleimide conjugation:
+        - Thioether: S with 2 single bonds to C, no H (C-S-C)
+        - NOT: free thiol, sulfonium, hypervalent sulfur
+        """
+        atom = mol.GetAtomWithIdx(conjugation_atom_idx)
+        
+        if atom.GetSymbol() != 'S':
+            return True  # Not sulfur, skip this check
+        
+        degree = atom.GetDegree()  # Number of bonds to other heavy atoms
+        total_hs = atom.GetTotalNumHs()
+        formal_charge = atom.GetFormalCharge()
+        
+        # Get neighbor information
+        neighbors = [mol.GetAtomWithIdx(n.GetIdx()) for n in atom.GetNeighbors()]
+        neighbor_symbols = [n.GetSymbol() for n in neighbors]
+        
+        # Count double bonds to oxygen (for sulfoxide/sulfone)
+        double_bonded_oxygens = 0
+        for neighbor in neighbors:
+            if neighbor.GetSymbol() == 'O':
+                bond = mol.GetBondBetweenAtoms(conjugation_atom_idx, neighbor.GetIdx())
+                if bond and bond.GetBondType() == Chem.BondType.DOUBLE:
+                    double_bonded_oxygens += 1
+        
+        # ========== Valid Sulfur States ==========
+        
+        # 1. Thioether: S with 2 single bonds to C, no H
+        #    C-S-C (degree=2, no O, no charge)
+        if degree == 2 and 'O' not in neighbor_symbols and formal_charge == 0:
+            if total_hs == 0:
+                return True  # Valid thioether
+            else:
+                return False  # Has H = still partially a thiol
+        
+        # 2. Sulfoxide: S with 1 double-bonded O
+        #    C-S(=O)-C (degree=3, one =O)
+        if degree == 3 and double_bonded_oxygens == 1 and formal_charge == 0:
+            return True
+        
+        # 3. Sulfone: S with 2 double-bonded O
+        #    C-S(=O)(=O)-C (degree=4, two =O)
+        if degree == 4 and double_bonded_oxygens == 2 and formal_charge == 0:
+            return True
+        
+        # ========== Invalid Sulfur States ==========
+        
+        # Sulfonium (positively charged S with 3+ bonds) - INVALID
+        if degree >= 3 and formal_charge > 0:
+            return False
+        
+        # Hypervalent sulfur without proper oxygen - INVALID
+        if degree > 2 and double_bonded_oxygens == 0:
+            return False
+        
+        # If none of the valid patterns matched, it's invalid
+        return False
+    
+    def check_attachment_site_compatibility(self, mol):
+        """
+        Check if thiol-maleimide attachment site chemistry is valid
+        
+        For vcMMAE conjugation, validates:
+        1. No free thiols (conjugation must have happened)
+        2. Has thioether bond (C-S-C)
+        3. No unreacted maleimide (C=C should be reduced)
+        4. Valid valency at conjugation sites
+        
+        Args:
+            mol: RDKit Mol object of the conjugate
+        
+        Returns:
+            bool: True if attachment site is valid, False otherwise
+        """
+        # ===== CRITICAL CHECK: Reject if free thiol present =====
+        # A valid PDC should have the thiol CONSUMED by conjugation
+        if self.has_free_thiol(mol):
+            return False
+        
+        # ===== Check for thioether (C-S-C) =====
+        thioether_pattern = Chem.MolFromSmarts('[#6][#16][#6]')
+        if not thioether_pattern or not mol.HasSubstructMatch(thioether_pattern):
+            return False  # No thioether = no conjugation
+        
+        # ===== Check valency at conjugation sites =====
+        conjugation_sites = self.find_thioether_conjugation_sites(mol)
+        
+        for site_idx in conjugation_sites:
+            if not self.check_valency_at_conjugation_site(mol, site_idx):
+                return False
+        
+        return True
+    
+    # =========================================================================
+    # Manual Conjugation
+    # =========================================================================
+    
+    def _manual_conjugation(self, peptide_mol, cysteine_info):
+        """
+        Create conjugate by manually forming bonds
+        
+        Chemistry:
+        - Maleimide C=C becomes C-C (single bond)
+        - Sulfur bonds to one of these carbons
+        - Proton transfers (S loses H, C gains H)
+        
+        Result: Peptide-Cys-S-(linker-payload)
+        """
+        
+        # Find maleimide C=C in vcMMAE linker-payload
+        maleimide_c1, maleimide_c2 = self._find_maleimide_atoms(self.linker_payload_mol)
+        
+        if maleimide_c1 is None:
+            raise ValueError("Could not locate maleimide in vcMMAE linker-payload")
+        
+        sulfur_idx = cysteine_info['atom_idx']
+        
+        # Step 1: Combine molecules
+        combined = Chem.CombineMols(self.linker_payload_mol, peptide_mol)
+        combined_rw = RWMol(combined)
+        
+        # Step 2: Adjust indices (peptide atoms are shifted)
+        linker_payload_atom_count = self.linker_payload_mol.GetNumAtoms()
+        new_sulfur_idx = sulfur_idx + linker_payload_atom_count
+        
+        # Step 3: Change maleimide C=C to C-C
+        bond = combined_rw.GetBondBetweenAtoms(maleimide_c1, maleimide_c2)
+        if bond is not None:
+            bond.SetBondType(Chem.BondType.SINGLE)
+        
+        # Step 4: Add S-C bond (sulfur to one maleimide carbon)
+        combined_rw.AddBond(new_sulfur_idx, maleimide_c2, Chem.BondType.SINGLE)
+        
+        # Step 5: Sanitize and clean up
+        try:
+            mol = combined_rw.GetMol()
+            Chem.SanitizeMol(mol)
+            conjugate_smiles = Chem.MolToSmiles(mol)
+            return conjugate_smiles
+            
+        except Exception as e:
+            # Try alternative sanitization
+            try:
+                mol = combined_rw.GetMol()
+                mol = Chem.RemoveAllHs(mol)
+                Chem.SanitizeMol(mol)
+                conjugate_smiles = Chem.MolToSmiles(mol)
+                return conjugate_smiles
+            except Exception as e2:
+                raise ValueError(f"Sanitization failed: {e2}")
+    
+    def _basic_validation(self, mol):
+        """
+        Basic validation checks for thiol-maleimide conjugates
+        
+        Includes attachment site compatibility check with valency validation.
+        
+        Returns dict of check_name: bool pairs
+        """
+        checks = {}
+        
+        # Check 1: Has thioether (the new C-S-C bond)
+        thioether = Chem.MolFromSmarts('[#6]-[#16]-[#6]')
+        checks['has_thioether'] = mol.HasSubstructMatch(thioether)
+        
+        # Check 2: Maleimide C=C is gone (should be succinimide now)
+        checks['maleimide_consumed'] = not self.has_unreacted_maleimide(mol)
+        
+        # Check 3: Has succinimide (reduced maleimide)
+        succinimide = Chem.MolFromSmarts('[#7]1[#6](=[#8])[#6][#6][#6]1=[#8]')
+        checks['has_succinimide'] = mol.HasSubstructMatch(succinimide)
+        
+        # Check 4: No free thiol remaining
+        checks['no_free_thiol'] = not self.has_free_thiol(mol)
+        
+        # Check 5: Full attachment site compatibility (includes valency)
+        checks['attachment_site_valid'] = self.check_attachment_site_compatibility(mol)
+        
+        # Find and report conjugation sites for diagnostics
+        conjugation_sites = self.find_thioether_conjugation_sites(mol)
+        checks['conjugation_sites'] = conjugation_sites
+        
+        # Overall pass/fail - attachment_site_valid includes thioether and no_free_thiol checks
+        checks['_passed'] = checks['attachment_site_valid']
+        
+        return checks
+
+
+
+# ============================================================================
+# Main - Test Examples  
+# ============================================================================
+
+if __name__ == "__main__":
+    
+    print("=" * 70)
+    print("PDC Builder Framework - vcMMAE Example")
+    print("=" * 70 + "\n")
+    
+    builder = VcMMAE_PDC_Builder()
+    
+    # Test peptides with free cysteine
+    test_peptides = [
+        "NCC(=O)N[C@@H](CS)C(=O)NCC(=O)O",  # Gly-Cys-Gly
+        "N[C@@H](CS)C(=O)N[C@@H](CCCNC(=N)N)C(=O)O",  # Cys-Arg
+        "N[C@@H](CCCNC(=N)N)C(=O)N[C@@H](CS)C(=O)N[C@@H](CC(C)C)C(=O)O",  # Arg-Cys-Leu
+    ]
+    
+    print("\n" + "=" * 70)
+    print("Testing Single Conjugation")
+    print("=" * 70 + "\n")
+    
+    peptide = test_peptides[0]
+    print(f"Input peptide: {peptide}")
+    
+    try:
+        conjugate = builder.conjugate_to_peptide(peptide)
+        print(f"\n SUCCESS!")
+        print(f"Conjugate SMILES: {conjugate[:80]}...")
+        
+        valid, checks = builder.validate_conjugate(conjugate, use_analyzer=False)
+        print(f"\n--- Basic Validation ---")
+        print(f"Result: {'PASSED' if valid else 'FAILED'}")
+        for k, v in checks.items():
+            print(f"  {k}: {v}")
+
+    except Exception as e:
+        print(f"\n FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+    
+    print("\n" + "=" * 70)
+    print("Batch Conjugation & Validation")
+    print("=" * 70 + "\n")
+    
+    # Batch conjugate
+    results = builder.batch_conjugate(test_peptides)
+    
+    # Get successful conjugates
+    successful_conjugates = [c for c in results['conjugates'] if c is not None]
+    
+    if len(successful_conjugates) > 0:
+        print(f"\nValidating {len(successful_conjugates)} conjugates...")
+        
+        # Basic validation
+        print("\n--- Basic Validation ---")
+        basic_results = builder.batch_validate(successful_conjugates, use_analyzer=False)
+        
+        for i, (passed, checks) in enumerate(basic_results):
+            status = " VALID" if passed else " INVALID"
+            print(f"  Conjugate {i+1}: {status}")
+            
+            if not passed:
+                # Show which checks failed
+                for check_name, check_val in checks.items():
+                    if check_name not in ['validation_method', 'parseable', 'conjugation_sites'] and not check_val:
+                        print(f"    Failed: {check_name}")

smiles/conjugation/__init__.py

@@ -0,0 +1 @@
+"""PDC conjugation builders."""

smiles/conjugation/pdc_builder.py

@@ -0,0 +1,223 @@
+"""
+PDC Builders - General and Specialized
+"""
+
+from rdkit import Chem
+from rdkit.Chem import Descriptors
+from abc import ABC, abstractmethod
+
+
+class PDC_Builder(ABC):
+    """
+    Base class for Peptide-Drug Conjugate (PDC) builders
+    
+    This abstract class defines the interface for building PDCs with different
+    conjugation chemistries. Subclasses must implement specific conjugation rules.
+    
+    The conjugation component can be:
+    - Just a drug payload with reactive group
+    - A linker-payload combination (e.g., vcMMAE = Val-Cit linker + MMAE payload)
+    - A linker alone (for later payload attachment)
+    """
+    
+    def __init__(self, linker_payload_smiles, component_name="Drug-Linker"):
+        """
+        Initialize PDC builder with a linker-payload component
+        
+        Args:
+            linker_payload_smiles: SMILES string of the linker-payload component
+            component_name: Name of the component (for display)
+        """
+        self.linker_payload_smiles = linker_payload_smiles
+        self.component_name = component_name
+        self.linker_payload_mol = Chem.MolFromSmiles(linker_payload_smiles)
+        
+        if self.linker_payload_mol is None:
+            raise ValueError(f"Could not parse {component_name} SMILES: {linker_payload_smiles}")
+        
+        # Verify reactive group is present
+        self._verify_reactive_group()
+    
+    @abstractmethod
+    def _verify_reactive_group(self):
+        """Verify that the linker-payload component has the expected reactive group"""
+        pass
+    
+    @abstractmethod
+    def find_conjugation_sites(self, peptide_smiles):
+        """
+        Find conjugation sites in the peptide
+        
+        Args:
+            peptide_smiles: SMILES string of the peptide
+        
+        Returns:
+            List of dicts with conjugation site information
+        """
+        pass
+    
+    @abstractmethod
+    def _perform_conjugation(self, peptide_mol, site_info):
+        """
+        Perform the actual conjugation chemistry
+        
+        Args:
+            peptide_mol: RDKit Mol object of the peptide
+            site_info: Dict with conjugation site information
+        
+        Returns:
+            SMILES string of the conjugate
+        """
+        pass
+    
+    @abstractmethod
+    def check_attachment_site_compatibility(self, mol):
+        """
+        Check if attachment site chemistry is valid for this conjugation type
+        
+        This is conjugation-method specific and must be implemented by subclasses.
+        Checks should include:
+        - Presence of expected conjugation bond (e.g., thioether for thiol-maleimide)
+        - Absence of unreacted precursors (e.g., free thiol, unreacted maleimide)
+        - Valency at conjugation sites
+        
+        Args:
+            mol: RDKit Mol object of the conjugate
+        
+        Returns:
+            bool: True if attachment site is valid, False otherwise
+        """
+        pass
+    
+    def conjugate_to_peptide(self, peptide_smiles, site_position=0):
+        """
+        Conjugate linker-payload to peptide at specified conjugation site
+        
+        Args:
+            peptide_smiles: SMILES string of the peptide
+            site_position: Index of the conjugation site to use (default: 0)
+        
+        Returns:
+            SMILES string of the PDC conjugate
+        """
+        peptide_mol = Chem.MolFromSmiles(peptide_smiles)
+        if peptide_mol is None:
+            raise ValueError(f"Invalid peptide SMILES: {peptide_smiles}")
+        
+        sites = self.find_conjugation_sites(peptide_smiles)
+        
+        if len(sites) == 0:
+            raise ValueError(f"No conjugation sites found in peptide")
+        
+        if site_position >= len(sites):
+            raise ValueError(f"Site position {site_position} out of range (found {len(sites)} sites)")
+        
+        print(f"Found {len(sites)} conjugation site(s) in peptide")
+        print(f"Conjugating at site position {site_position}")
+        
+        target_site = sites[site_position]
+        return self._perform_conjugation(peptide_mol, target_site)
+    
+    def batch_conjugate(self, peptide_smiles_list, site_position=0):
+        """
+        Conjugate multiple peptides in batch
+        
+        Args:
+            peptide_smiles_list: List of peptide SMILES strings
+            site_position: Index of conjugation site to use
+        
+        Returns:
+            Dict with 'conjugates', 'successful', and 'failed' keys
+        """
+        results = {
+            'conjugates': [],
+            'successful': [],
+            'failed': [],
+        }
+        
+        for i, peptide_smiles in enumerate(peptide_smiles_list):
+            try:
+                conjugate = self.conjugate_to_peptide(peptide_smiles, site_position)
+                results['conjugates'].append(conjugate)
+                results['successful'].append(i)
+            except Exception as e:
+                print(f" Peptide {i} failed: {e}")
+                results['conjugates'].append(None)
+                results['failed'].append((i, str(e)))
+        
+        print(f"\n=== Summary ===")
+        print(f"Total: {len(peptide_smiles_list)}, Success: {len(results['successful'])}, Failed: {len(results['failed'])}")
+        
+        return results
+    
+    def get_conjugate_properties(self, conjugate_smiles):
+        """Calculate molecular properties of the conjugate"""
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        if mol is None:
+            return None
+        
+        return {
+            'molecular_weight': Descriptors.MolWt(mol),
+            'num_heavy_atoms': mol.GetNumHeavyAtoms(),
+            'logp': Descriptors.MolLogP(mol),
+            'tpsa': Descriptors.TPSA(mol),  # Topological Polar Surface Area
+        }
+    
+    def validate_conjugate(self, conjugate_smiles):
+        """
+        Validate that the conjugate is chemically reasonable
+        
+        Args:
+            conjugate_smiles: SMILES string of the conjugate
+        
+        Returns:
+            (passed: bool, checks: dict)
+        """
+        mol = Chem.MolFromSmiles(conjugate_smiles)
+        
+        if mol is None:
+            return False, {'error': 'Could not parse SMILES'}
+        
+        # Basic validation (delegated to subclass)
+        checks = self._basic_validation(mol)
+        checks['validation_method'] = 'basic'
+        checks['parseable'] = True
+        
+        # Get pass/fail from subclass implementation
+        passed = checks.get('_passed', True)
+        if '_passed' in checks:
+            del checks['_passed']  # Remove internal flag from output
+        
+        return passed, checks
+    
+    @abstractmethod
+    def _basic_validation(self, mol):
+        """
+        Perform basic validation checks specific to conjugation type
+        
+        Should include check_attachment_site_compatibility() and any other
+        conjugation-specific checks.
+        
+        Args:
+            mol: RDKit Mol object of the conjugate
+        
+        Returns:
+            Dict of check_name: bool pairs. Include '_passed': bool for overall result.
+        """
+        pass
+    
+    def batch_validate(self, conjugate_smiles_list):
+        """Validate multiple conjugates"""
+        results = []
+        for smiles in conjugate_smiles_list:
+            if smiles is None:
+                results.append((False, {'error': 'None SMILES'}))
+            else:
+                passed, checks = self.validate_conjugate(smiles)
+                results.append((passed, checks))
+        
+        num_passed = sum(1 for passed, _ in results if passed)
+        print(f"\nValidation Summary: {num_passed}/{len(results)} passed")
+        
+        return results
+

smiles/evaluation/__init__.py

@@ -0,0 +1 @@
+"""PDC evaluation and structure-analysis utilities."""

smiles/evaluation/metrics.py

@@ -0,0 +1,570 @@
+import torch
+import warnings
+import numpy as np
+import pandas as pd
+from rdkit import Chem, rdBase, DataStructs
+from rdkit.Chem import AllChem
+import os
+import sys
+
+# Add current directory to path for imports
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+# Import custom analyzers and conjugation builders
+from smiles.evaluation.peptide_analyzer import PeptideAnalyzer
+from smiles.evaluation.pdc_analyzer import PDCAnalyzer
+from smiles.conjugation.AoaDau import AoaDaunorubicin_PDC_Builder
+from smiles.conjugation.SSCpt import DisulfideCamptothecin_PDC_Builder
+from smiles.conjugation.DSPASN38 import DisulfanylpropanoicAcidSN38_PDC_Builder
+from smiles.scoring.classifiers import BindingAffinity
+
+rdBase.DisableLog('rdApp.error')
+warnings.filterwarnings("ignore", category=DeprecationWarning)
+warnings.filterwarnings("ignore", category=UserWarning)
+warnings.filterwarnings("ignore", category=FutureWarning)
+
+def fingerprints_from_mol(molecule, radius=3, size=2048, hashed=False):
+    """ Create ECFP fingerprint of a molecule """
+    if hashed:
+        fp_bits = AllChem.GetHashedMorganFingerprint(molecule, radius, nBits=size)
+    else:
+        fp_bits = AllChem.GetMorganFingerprintAsBitVect(molecule, radius, nBits=size)
+    fp_np = np.zeros((1,))
+    DataStructs.ConvertToNumpyArray(fp_bits, fp_np)
+    return fp_np.reshape(1, -1)
+
+def average_agg_tanimoto(stock_vecs, gen_vecs, batch_size=5000, agg='max', device='cuda', p=1):
+    """
+    For each molecule in gen_vecs finds closest molecule in stock_vecs.
+    Returns average tanimoto score for between these molecules
+
+    Parameters:
+        stock_vecs: numpy array <n_vectors x dim>
+        gen_vecs: numpy array <n_vectors' x dim>
+        agg: max or mean
+        p: power for averaging: (mean x^p)^(1/p)
+    """
+    assert agg in ['max', 'mean'], "Can aggregate only max or mean"
+    agg_tanimoto = np.zeros(len(gen_vecs))
+    total = np.zeros(len(gen_vecs))
+    for j in range(0, stock_vecs.shape[0], batch_size):
+        x_stock = torch.tensor(stock_vecs[j:j + batch_size]).to(device).float()
+        for i in range(0, gen_vecs.shape[0], batch_size):
+            y_gen = torch.tensor(gen_vecs[i:i + batch_size]).to(device).float()
+            y_gen = y_gen.transpose(0, 1)
+            tp = torch.mm(x_stock, y_gen)
+            jac = (tp / (x_stock.sum(1, keepdim=True) + y_gen.sum(0, keepdim=True) - tp)).cpu().numpy()
+            jac[np.isnan(jac)] = 1
+            if p != 1:
+                jac = jac**p
+            if agg == 'max':
+                agg_tanimoto[i:i + y_gen.shape[1]] = np.maximum(agg_tanimoto[i:i + y_gen.shape[1]], jac.max(0))
+            elif agg == 'mean':
+                agg_tanimoto[i:i + y_gen.shape[1]] += jac.sum(0)
+                total[i:i + y_gen.shape[1]] += jac.shape[0]
+    if agg == 'mean':
+        agg_tanimoto /= total
+    if p != 1:
+        agg_tanimoto = (agg_tanimoto)**(1/p)
+    return np.mean(agg_tanimoto)
+
+def get_mol(smiles_or_mol):
+    '''
+    Loads SMILES/molecule into RDKit's object
+    '''
+    if isinstance(smiles_or_mol, str):
+        if len(smiles_or_mol) == 0:
+            return None
+        mol = Chem.MolFromSmiles(smiles_or_mol)
+        if mol is None:
+            return None
+        try:
+            Chem.SanitizeMol(mol)
+        except ValueError:
+            return None
+        return mol
+    return smiles_or_mol
+
+def canonic_smiles(smiles_or_mol):
+    mol = get_mol(smiles_or_mol)
+    if mol is None:
+        return None
+    return Chem.MolToSmiles(mol)
+
+
+class Metrics:
+    
+    def __init__(self, prior_path='./data/smiles/30K_all.csv', n_jobs=100, input_type='smiles', device='cuda'):
+        self.device = device
+        self.input_type = input_type  # 'smiles', 'peptide', or 'conjugate'
+        self.n_jobs = n_jobs
+        
+        # Initialize analyzers
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.pdc_analyzer = PDCAnalyzer(device)
+        
+        # Initialize conjugation builders
+        self.aoa_dau_builder = AoaDaunorubicin_PDC_Builder()
+        self.ss_cpt_builder = DisulfideCamptothecin_PDC_Builder()
+        self.dspa_sn38_builder = DisulfanylpropanoicAcidSN38_PDC_Builder()
+        
+        # Load reference set for similarity calculations
+        train_set_cano_smi = pd.read_csv(prior_path)['SMILES'].astype(str).tolist()
+        
+        for smi in train_set_cano_smi:
+            mol = Chem.MolFromSmiles(smi)
+            
+            if mol is None:
+                print(f"Invalid SMILES: {smi}")
+            
+        self.train_set_cano_smi = train_set_cano_smi
+        
+        self.ref_fps = np.vstack([
+            fingerprints_from_mol(Chem.MolFromSmiles(smi)) 
+            for smi in train_set_cano_smi 
+            if Chem.MolFromSmiles(smi) is not None
+        ])
+
+    def _determine_target_from_filename(self, generated_path):
+        """Determine target protein from filename"""
+        filename = os.path.basename(generated_path).upper()
+        if 'GNRHR' in filename:
+            return 'GNRHR'
+        elif 'ERBB2' in filename:
+            return 'ERBB2'
+        elif 'NRP1' in filename:
+            return 'NRP1'
+        else:
+            raise ValueError(f"Cannot determine target from filename: {filename}. Expected GNRHR / ERBB2 / NRP1 in filename.")
+
+    def _get_target_sequence(self, target):
+        """Get target protein sequence for binding affinity calculation"""
+        # Define target sequences (you may need to update these)
+        target_sequences = {
+            'GNRHR': 'MANSASPEQNQNHCSAINNSIPLMQGNLPTLTLSGKIRVTVTFFLFLLSATFNASFLLKLQKWTQKKEKGKKLSRMKLLLKHLTLANLLETLIVMPLDGMWNITVQWYAGELLCKVLSYLKLFSMYAPAFMMVVISLDRSLAITRPLALKSNSKVGQSMVGLAWILSSVFAGPQLYIFRMIHLADSSGQTKVFSQCVTHCSFSQWWHQAFYNFFTFSCLFIIPLFIMLICNAKIIFTLTRVLHQDPHELQLNQSKNNIPRARLKTLKMTVAFATSFTVCWTPYYVLGIWYWFDPEMLNRLSDPVNHFFFLFAFLNPCFDPLIYGYFSL',
+            'ERBB2': 'MELAALCRWGLLLALLPPGAASTQVCTGTDMKLRLPASPETHLDMLRHLYQGCQVVQGNLELTYLPTNASLSFLQDIQEVQGYVLIAHNQVRQVPLQRLRIVRGTQLFEDNYALAVLDNGDPLNNTTPVTGASPGGLRELQLRSLTEILKGGVLIQRNPQLCYQDTILWKDIFHKNNQLALTLIDTNRSRACHPCSPMCKGSRCWGESSEDCQSLTRTVCAGGCARCKGPLPTDCCHEQCAAGCTGPKHSDCLACLHFNHSGICELHCPALVTYNTDTFESMPNPEGRYTFGASCVTACPYNYLSTDVGSCTLVCPLHNQEVTAEDGTQRCEKCSKPCARVCYGLGMEHLREVRAVTSANIQEFAGCKKIFGSLAFLPESFDGDPASNTAPLQPEQLQVFETLEEITGYLYISAWPDSLPDLSVFQNLQVIRGRILHNGAYSLTLQGLGISWLGLRSLRELGSGLALIHHNTHLCFVHTVPWDQLFRNPHQALLHTANRPEDECVGEGLACHQLCARGHCWGPGPTQCVNCSQFLRGQECVEECRVLQGLPREYVNARHCLPCHPECQPQNGSVTCFGPEADQCVACAHYKDPPFCVARCPSGVKPDLSYMPIWKFPDEEGACQPCPINCTHSCVDLDDKGCPAEQRASPLTSIISAVVGILLVVVLGVVFGILIKRRQQKIRKYTMRRLLQETELVEPLTPSGAMPNQAQMRILKETELRKVKVLGSGAFGTVYKGIWIPDGENVKIPVAIKVLRENTSPKANKEILDEAYVMAGVGSPYVSRLLGICLTSTVQLVTQLMPYGCLLDHVRENRGRLGSQDLLNWCMQIAKGMSYLEDVRLVHRDLAARNVLVKSPNHVKITDFGLARLLDIDETEYHADGGKVPIKWMALESILRRRFTHQSDVWSYGVTVWELMTFGAKPYDGIPAREIPDLLEKGERLPQPPICTIDVYMIMVKCWMIDSECRPRFRELVSEFSRMARDPQRFVVIQNEDLGPASPLDSTFYRSLLEDDDMGDLVDAEEYLVPQQGFFCPDPAPGAGGMVHHRHRSSSTRSGGGDLTLGLEPSEEEAPRSPLAPSEGAGSDVFDGDLGMGAAKGLQSLPTHDPSPLQRYSEDPTVPLPSETDGYVAPLTCSPQPEYVNQPDVRPQPPSPREGPLPAARPAGATLERPKTLSPGKNGVVKDVFAFGGAVENPEYLTPQGGAAPQPHPPPAFSPAFDNLYYWDQDPPERGAPPSTFKGTPTAENPEYLGLDVPV',
+            'NRP1': 'MERGLPLLCAVLALVLAPAGAFRNDKCGDTIKIESPGYLTSPGYPHSYHPSEKCEWLIQAPDPYQRIMINFNPHFDLEDRDCKYDYVEVFDGENENGHFRGKFCGKIAPPPVVSSGPFLFIKFVSDYETHGAGFSIRYEIFKRGPECSQNYTTPSGVIKSPGFPEKYPNSLECTYIVFVPKMSEIILEFESFDLEPDSNPPGGMFCRYDRLEIWDGFPDVGPHIGRYCGQKTPGRIRSSSGILSMVFYTDSAIAKEGFSANYSVLQSSVSEDFKCMEALGMESGEIHSDQITASSQYSTNWSAERSRLNYPENGWTPGEDSYREWIQVDLGLLRFVTAVGTQGAISKETKKKYYVKTYKIDVSSNGEDWITIKEGNKPVLFQGNTNPTDVVVAVFPKPLITRFVRIKPATWETGISMRFEVYGCKITDYPCSGMLGMVSGLISDSQITSSNQGDRNWMPENIRLVTSRSGWALPPAPHSYINEWLQIDLGEEKIVRGIIIQGGKHRENKVFMRKFKIGYSNNGSDWKMIMDDSKRKAKSFEGNNNYDTPELRTFPALSTRFIRIYPERATHGGLGLRMELLGCEVEAPTAGPTTPNGNLVDECDDDQANCHSGTGDDFQLTGGTTVLATEKPTVIDSTIQSEFPTYGFNCEFGWGSHKTFCHWEHDNHVQLKWSVLTSKTGPIQDHTGDGNFIYSQADENQKGKVARLVSPVVYSQNSAHCMTFWYHMSGSHVGTLRVKLRYQKPEEYDQLVWMAIGHQGDHWKEGRVLLHKSLKLYQVIFEGEIGKGNLGGIAVDDISINNHISQEDCAKPADLDKKNPEIKIDETGSTPGYEGEGEGDKNISRKPGNVLKTLDPILITIIAMSALGVLLGAVCGVVLYCACWHNGMSERNLSALENYNFELVDGVKLKKDKLNTQSTYSEA'
+        }
+        return target_sequences.get(target, '')
+
+    def get_metrics_for_peptides(self, generated_path, save_csv=True):
+        """
+        Analyze peptide sequences with conjugation, validity, binding affinity, and drug-likeness
+        """
+        print(f"Analyzing peptides from: {generated_path}")
+        
+        # Determine target from filename
+        target = self._determine_target_from_filename(generated_path)
+        print(f"Detected target: {target}")
+        
+        # Load data
+        df = pd.read_csv(generated_path)
+        if 'sequence' in df.columns:
+            peptide_sequences = df['sequence'].astype(str).tolist()
+        elif 'SMILES' in df.columns:
+            peptide_sequences = df['SMILES'].astype(str).tolist()
+        else:
+            raise ValueError("No 'sequence' or 'SMILES' column found in the input file")
+        
+        total_sequences = len(peptide_sequences)
+        print(f"Total sequences: {total_sequences}")
+        
+        # Step 1: Filter valid peptides
+        print("Step 1: Filtering valid peptides...")
+        valid_peptides = []
+        valid_indices = []
+        
+        for i, seq in enumerate(peptide_sequences):
+            if self.peptide_analyzer.is_peptide(seq):
+                valid_peptides.append(seq)
+                valid_indices.append(i)
+        
+        print(f"Valid peptides: {len(valid_peptides)}/{total_sequences}")
+        
+        if len(valid_peptides) == 0:
+            print("No valid peptides found!")
+            return {
+                'total_sequences': total_sequences,
+                'valid_peptides': 0,
+                'validity_rate': 0.0,
+                'conjugation_success_rate': 0.0,
+                'binding_affinity_mean': 0.0,
+                'binding_affinity_std': 0.0,
+            }
+        
+        # Step 2: Conjugate peptides
+        print("Step 2: Conjugating peptides...")
+        if target == 'GNRHR':
+            print("Using Aoa-Dau conjugation builder for GNRHR")
+            builder = self.aoa_dau_builder 
+        elif target == 'ERBB2':
+            builder = self.ss_cpt_builder
+        elif target == 'NRP1':
+            builder = self.dspa_sn38_builder
+
+        conjugate_results = []
+        
+        for i, peptide in enumerate(valid_peptides):
+            try:
+                print(f"  Conjugating peptide {i+1}/{len(valid_peptides)}: {peptide[:50]}...")
+                # Try different methods to get conjugates
+                try:
+                    conjugates = builder.conjugate_to_peptide(peptide, exclude_n_terminus=True)
+                    if conjugates and len(conjugates) > 0:
+                        print(f"    Success: Generated {len(conjugates)} conjugate(s)")
+                        first_conjugate = conjugates[0]
+                        print(f"    DEBUG: First conjugate length: {len(first_conjugate)} chars")
+                        print(f"    DEBUG: First conjugate preview: {first_conjugate[:100]}...")
+                        print(f"    DEBUG: Expected peptide in conjugate: {peptide[:30] in first_conjugate}")
+                        conjugate_results.append(first_conjugate)  # Take first conjugate
+                    else:
+                        print(f"    No conjugates generated")
+                        conjugate_results.append(None)
+                except AttributeError:
+                    # Fallback to batch_conjugate without return_all parameter
+                    result = builder.batch_conjugate([peptide])
+                    if result and 'conjugates' in result and len(result['conjugates']) > 0:
+                        conjugates = result['conjugates'][0]  # First peptide's conjugates
+                        if conjugates:
+                            print(f"    Success: Generated {len(conjugates)} conjugate(s)")
+                            first_conjugate = conjugates[0]
+                            print(f"    DEBUG: Fallback conjugate length: {len(first_conjugate)} chars")
+                            print(f"    DEBUG: Fallback conjugate preview: {first_conjugate[:100]}...")
+                            conjugate_results.append(first_conjugate)  # Take first conjugate
+                        else:
+                            print(f"    No conjugates generated")
+                            conjugate_results.append(None)
+                    else:
+                        print(f"    Batch conjugation failed")
+                        conjugate_results.append(None)
+                        
+            except Exception as e:
+                print(f"    Conjugation failed for peptide {i+1}: {e}")
+                conjugate_results.append(None)
+        
+        print(f"Total conjugation attempts: {len(conjugate_results)}")
+        print(f"Successful conjugations: {sum(1 for c in conjugate_results if c is not None)}")
+        
+        # Step 3: Analyze conjugates with PDC analyzer
+        print("Step 3: Analyzing conjugates...")
+        valid_conjugates = []
+        conjugate_properties = []
+        
+        for i, conjugate in enumerate(conjugate_results):
+            if conjugate is not None:
+                print(f"  Analyzing conjugate {i+1}: {conjugate[:60]}...")
+                # Check if SMILES is valid
+                mol = Chem.MolFromSmiles(conjugate)
+                if mol is not None:
+                    print(f"    Valid SMILES structure")
+                    
+                    # Try PDC analyzer validation
+                    try:
+                        linker_feasible = self.pdc_analyzer.check_linker_feasibility(mol)
+                        repr_valid = self.pdc_analyzer.check_representational_validity(mol)
+                        print(f"    PDC checks - Linker feasible: {linker_feasible}, Repr valid: {repr_valid}")
+
+                        is_valid = linker_feasible and repr_valid
+                        if not is_valid:
+                            print(f"    PDC validation failed, dropped this molecule")
+                            is_valid = False 
+
+                    except Exception as e:
+                        print(f"    PDC validation error: {e}, dropped this molecule")
+                        is_valid = False
+                    
+                    if is_valid:
+                        valid_conjugates.append(conjugate)
+                        try:
+                            # DEBUG: Check the molecule being analyzed  
+                            print(f"    DEBUG: SMILES to analyze: {conjugate[:100]}...")
+                            
+                            # CRITICAL FIX: Create molecule from conjugate SMILES
+                            conjugate_mol = Chem.MolFromSmiles(conjugate)
+                            conjugate_formula = Chem.rdMolDescriptors.CalcMolFormula(conjugate_mol)
+                            conjugate_mw = Chem.Descriptors.MolWt(conjugate_mol)
+                            print(f"    DEBUG: Analyzing molecule - Formula: {conjugate_formula}, MW: {conjugate_mw:.2f}")
+                            print(f"    DEBUG: SMILES length: {len(conjugate)} chars")
+                            
+                            # DEBUG: Double-check we're analyzing the right molecule
+                            if conjugate_formula == 'C33H40N3O11':
+                                print(f"     WARNING: This is the linker-drug formula, not full conjugate!")
+                            else:
+                                print(f"     This appears to be the full conjugate formula")
+                            
+                            # Use the correctly created molecule for property calculation
+                            props = self.pdc_analyzer.calculate_molecular_properties(conjugate_mol)
+                            # Ensure molecular weight is included
+                            if 'molecular_weight' not in props:
+                                from rdkit.Chem import Descriptors
+                                props['molecular_weight'] = Descriptors.MolWt(conjugate_mol)
+                            print(f"    Calculated properties successfully: {props.get('mf', 'No MF')}, MW: {props.get('molecular_weight', 'N/A'):.2f}")
+                        except Exception as e:
+                            print(f"    Property calculation failed: {e}, using basic properties")
+                            # Fallback to basic RDKit properties - use conjugate molecule
+                            conjugate_mol = Chem.MolFromSmiles(conjugate)
+                            if conjugate_mol:
+                                from rdkit.Chem import Descriptors
+                                props = {
+                                    'molecular_weight': Descriptors.MolWt(conjugate_mol),
+                                    'logp': Descriptors.MolLogP(conjugate_mol),
+                                    'num_atoms': conjugate_mol.GetNumHeavyAtoms(),
+                                    'num_bonds': conjugate_mol.GetNumBonds()
+                                }
+                            else:
+                                props = {
+                                    'molecular_weight': 0,
+                                    'logp': 0,
+                                    'num_atoms': 0,
+                                    'num_bonds': 0
+                                }
+                        
+                        # Get amino acid sequence from peptide analyzer
+                        try:
+                            aa_seq, seq_length = self.peptide_analyzer.analyze_structure(valid_peptides[i])
+                        except Exception as e:
+                            print(f"    Failed to analyze structure: {e}")
+                            aa_seq, seq_length = None, None
+                        
+                        conjugate_properties.append({
+                            'peptide_index': valid_indices[i],
+                            'peptide_smiles': valid_peptides[i],
+                            'aa_seq': aa_seq,
+                            'seq_length': seq_length,
+                            'conjugate_smiles': conjugate,
+                            **props
+                        })
+                        print(f"     Added to valid conjugates")
+                    else:
+                        print(f"     Failed validation")
+                else:
+                    print(f"     Invalid SMILES structure")
+            else:
+                print(f"  Conjugate {i+1}: None (conjugation failed)")
+        
+        print(f"Valid conjugates: {len(valid_conjugates)}")
+        if len(valid_conjugates) > 0:
+            print(f"First few valid conjugates:")
+            for i, conj in enumerate(valid_conjugates[:3]):
+                print(f"  {i+1}: {conj[:80]}...")
+        else:
+            print("No valid conjugates found!")
+        
+        # Step 4: Calculate binding affinity
+        print("Step 4: Calculating binding affinity...")
+        target_seq = self._get_target_sequence(target)
+        if not target_seq:
+            print(f"Warning: No target sequence found for {target}")
+            affinity_scores_all = []
+            affinity_scores_conjugatable = []
+            # Add None for binding affinity when no target sequence
+            for prop in conjugate_properties:
+                prop['binding_affinity'] = None
+        else:
+            print(f"Using target sequence (length: {len(target_seq)})")
+            try:
+                # Initialize binding affinity model
+                from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+                tokenizer = SMILES_SPE_Tokenizer('./smiles/smiles_tokenizer/new_vocab.txt', './smiles/smiles_tokenizer/new_splits.txt')
+                binding_model = BindingAffinity(target_seq, self.device)
+                print(f"Binding model initialized successfully")
+                
+                # Tokenize valid peptides
+                valid_peptide_tokens = []
+                for peptide in valid_peptides:
+                    tokens = tokenizer(peptide)
+                    # Handle BatchEncoding object from tokenizer
+                    if hasattr(tokens, 'input_ids'):
+                        tokens = tokens.input_ids  # Extract token IDs from BatchEncoding
+                    if isinstance(tokens, list):
+                        tokens = torch.tensor(tokens)
+                    valid_peptide_tokens.append(tokens)
+                
+                print(f"Tokenized {len(valid_peptide_tokens)} valid peptides")
+                
+                # Calculate affinity for all valid peptides
+                if valid_peptide_tokens:
+                    # Pad sequences
+                    max_len = max(len(tokens) for tokens in valid_peptide_tokens)
+                    padded_tokens = torch.zeros((len(valid_peptide_tokens), max_len), dtype=torch.long)
+                    for i, tokens in enumerate(valid_peptide_tokens):
+                        padded_tokens[i, :len(tokens)] = tokens
+                    
+                    print(f"Calculating affinity for {len(valid_peptide_tokens)} peptides...")
+                    affinity_scores_all = binding_model(padded_tokens.to(self.device)).cpu().numpy()
+                    print(f"Calculated {len(affinity_scores_all)} affinity scores")
+                else:
+                    affinity_scores_all = []
+                
+                # ===== NEW CODE STARTS HERE =====
+                # Create mapping from valid_indices to affinity scores
+                affinity_map = {}
+                for i, valid_idx in enumerate(valid_indices):
+                    if i < len(affinity_scores_all):
+                        affinity_map[valid_idx] = float(affinity_scores_all[i])
+                
+                # Add binding affinity to conjugate_properties
+                print(f"Adding binding affinity to conjugate properties...")
+                for prop in conjugate_properties:
+                    peptide_idx = prop['peptide_index']
+                    if peptide_idx in affinity_map:
+                        prop['binding_affinity'] = affinity_map[peptide_idx]
+                        print(f"  Added affinity {affinity_map[peptide_idx]:.4f} for peptide index {peptide_idx}")
+                    else:
+                        prop['binding_affinity'] = None
+                        print(f"  No affinity score for peptide index {peptide_idx}")
+
+                # Calculate affinity for peptides with valid conjugates
+                # NOTE: pre-initialize so the variable always exists even when
+                # conjugate_properties is empty (zero valid conjugates).
+                # Also: this block runs ONCE here, not once per prop above.
+                affinity_scores_conjugatable = []
+                conjugatable_indices = [prop['peptide_index'] for prop in conjugate_properties]
+                valid_conjugatable_indices = [
+                    valid_indices.index(idx)
+                    for idx in conjugatable_indices
+                    if idx in valid_indices
+                ]
+
+                print(f"Found {len(valid_conjugatable_indices)} conjugatable peptides")
+
+                if valid_conjugatable_indices:
+                    conjugatable_tokens = [valid_peptide_tokens[i] for i in valid_conjugatable_indices]
+                    max_len = max(len(tokens) for tokens in conjugatable_tokens)
+                    padded_tokens = torch.zeros((len(conjugatable_tokens), max_len), dtype=torch.long)
+                    for i, tokens in enumerate(conjugatable_tokens):
+                        padded_tokens[i, :len(tokens)] = tokens
+
+                    affinity_scores_conjugatable = binding_model(padded_tokens.to(self.device)).cpu().numpy()
+                    print(f"Calculated {len(affinity_scores_conjugatable)} conjugatable affinity scores")
+            
+            except Exception as e:
+                print(f"Binding affinity calculation failed: {e}")
+                import traceback
+                traceback.print_exc()
+                affinity_scores_all = []
+                affinity_scores_conjugatable = []
+                # ===== NEW CODE: Handle errors =====
+                # Still add None values for binding affinity
+                for prop in conjugate_properties:
+                    prop['binding_affinity'] = None
+                # ===== NEW CODE ENDS =====
+    
+        
+        # Save detailed results to CSV
+        csv_path = None
+        if save_csv and conjugate_properties:
+            output_dir = os.path.dirname(generated_path)
+            base_name = os.path.splitext(os.path.basename(generated_path))[0]
+            csv_path = os.path.join(output_dir, f"{base_name}_conjugate_analysis.csv")
+            
+            df_results = pd.DataFrame(conjugate_properties)
+            df_results.to_csv(csv_path, index=False)
+            print(f"Detailed results saved to: {csv_path}")
+        
+        # Calculate summary statistics
+        validity_rate = len(valid_peptides) / total_sequences
+        conjugation_success_rate = len(valid_conjugates) / total_sequences
+        
+        results = {
+            'total_sequences': total_sequences,
+            'valid_peptides': len(valid_peptides),
+            'validity_rate': validity_rate,
+            'valid_conjugates': len(valid_conjugates),
+            'conjugation_success_rate': conjugation_success_rate,
+            'target': target,
+            'csv_path': csv_path,  # Add CSV path to results
+        }
+        
+        # Binding affinity statistics
+        if len(affinity_scores_all) > 0:
+            results.update({
+                'binding_affinity_all_mean': float(np.mean(affinity_scores_all)),
+                'binding_affinity_all_std': float(np.std(affinity_scores_all)),
+                'binding_affinity_all_min': float(np.min(affinity_scores_all)),
+                'binding_affinity_all_max': float(np.max(affinity_scores_all)),
+            })
+        
+        if len(affinity_scores_conjugatable) > 0:
+            results.update({
+                'binding_affinity_conjugatable_mean': float(np.mean(affinity_scores_conjugatable)),
+                'binding_affinity_conjugatable_std': float(np.std(affinity_scores_conjugatable)),
+                'binding_affinity_conjugatable_min': float(np.min(affinity_scores_conjugatable)),
+                'binding_affinity_conjugatable_max': float(np.max(affinity_scores_conjugatable)),
+            })
+    
+        
+        # Print summary
+        print("\n" + "="*60)
+        print("PEPTIDE ANALYSIS SUMMARY")
+        print("="*60)
+        print(f"Target: {target}")
+        print(f"Total sequences: {total_sequences}")
+        print(f"Valid peptides: {len(valid_peptides)} ({validity_rate:.1%})")
+        print(f"Valid conjugates: {len(valid_conjugates)} ({conjugation_success_rate:.1%})")
+        
+        print(f"\nBINDING AFFINITY ANALYSIS:")
+        if len(affinity_scores_all) > 0:
+            print(f"  All valid peptides ({len(affinity_scores_all)} peptides):")
+            print(f"    Mean: {np.mean(affinity_scores_all):.4f} ± {np.std(affinity_scores_all):.4f}")
+            print(f"    Range: {np.min(affinity_scores_all):.4f} to {np.max(affinity_scores_all):.4f}")
+        else:
+            print(f"  All valid peptides: No scores calculated (binding affinity model failed)")
+            
+        if len(affinity_scores_conjugatable) > 0:
+            print(f"  Conjugatable peptides ({len(affinity_scores_conjugatable)} peptides):")
+            print(f"    Mean: {np.mean(affinity_scores_conjugatable):.4f} ± {np.std(affinity_scores_conjugatable):.4f}")
+            print(f"    Range: {np.min(affinity_scores_conjugatable):.4f} to {np.max(affinity_scores_conjugatable):.4f}")
+        else:
+            print(f"  Conjugatable peptides: No scores calculated")
+        
+        return results
+    
+    def get_metrics(self, generated_path):
+        """Main method to get metrics based on input type"""
+        if self.input_type == 'peptide':
+            return self.get_metrics_for_peptides(generated_path)
+        elif self.input_type == 'conjugate':
+            # Placeholder for future conjugate analysis
+            print("Conjugate analysis not yet implemented")
+            return self._get_basic_metrics(generated_path)
+        else:
+            # Default SMILES analysis
+            return self._get_basic_metrics(generated_path)
+
+    def _get_basic_metrics(self, generated_path):
+        """Original metrics calculation for general SMILES"""
+        generated_smi = pd.read_csv(generated_path)['SMILES'].astype(str).tolist()
+        #generated_smi = pd.read_csv(generated_path, usecols=['Generated SMILES'])['Generated SMILES'].tolist()
+        #mols = [Chem.MolFromSmiles(smi) if smi else None for smi in generated_smi]
+        mols = [Chem.MolFromSmiles(smi) if smi else None for smi in generated_smi]
+        is_valid = [1 if mol else 0 for mol in mols]
+        validity = sum(is_valid) / len(is_valid)
+
+        valid_canon_smiles = [Chem.MolToSmiles(mol) for mol in mols if mol]
+        uniqueness = len(set(valid_canon_smiles)) / len(valid_canon_smiles)
+
+        uniq_smis = list(set(valid_canon_smiles))
+        uniq_mols = [Chem.MolFromSmiles(smi) for smi in uniq_smis]
+
+        fps = np.vstack([fingerprints_from_mol(mol) for mol in uniq_mols])
+        diversity = 1 - (average_agg_tanimoto(fps, fps, agg='mean', p=1)).mean()
+
+        snn = average_agg_tanimoto(self.ref_fps, fps, agg='max', p=1)
+        
+        # gen_smiles = mapper(self.n_jobs)(canonic_smiles, valid_canon_smiles)
+        # gen_smiles_set = set(gen_smiles) - {None}
+        # train_set = set(self.train_set_cano_smi)
+        # novelty = len(gen_smiles_set - train_set) / len(gen_smiles_set)
+
+        # print(f"validity\tuniqueness\tdiversity\tsnn\tnovelty")
+        # print(f"{validity:.3f}\t{uniqueness:.3f}\t{diversity:.3f}\t{snn:.3f}\t{novelty:.3f}")
+        print(f"validity\tuniqueness\tdiversity\tsnn")
+        print(f"{validity:.3f}\t{uniqueness:.3f}\t{diversity:.3f}\t{snn:.3f}")
+        return {
+            "validity": validity,
+            "uniqueness": uniqueness,
+            "diversity": diversity,
+            "snn": snn,  # "structural novelty"
+            # "novelty": novelty,
+        }

smiles/evaluation/pdc_analyzer.py

@@ -0,0 +1,431 @@
+"""
+PDC Analyzer
+Analyzes Peptide-Drug Conjugates (PDCs) for structural validity
+
+Hard constraints (validity = 1 if all pass, 0 if any fails):
+1. Linker feasibility: ring-closure constraints, long carbon chains
+2. Representational validity: SMILES parsing, sanitization, chemistry problems
+
+Note: Attachment/conjugation site checks are conjugation-method specific 
+and should be handled in the conjugation PDC builder (e.g., pdc_builder.py)
+"""
+
+import torch
+from rdkit import Chem
+from rdkit.Chem import Descriptors, AllChem, rdMolDescriptors
+
+
+class PDCAnalyzer:
+    """
+    Analyzes Peptide-Drug Conjugates (PDCs) for structural validity
+    
+    Hard constraint checks (validity = 1 if all pass, 0 if any fails):
+    1. Linker feasibility: ring-closure constraints, long carbon chains
+    2. Representational validity: SMILES parsing, sanitization, chemistry problems
+    
+    Optional molecular properties (toggle on/off):
+    - mf: molecular formula
+    - polararea: topological polar surface area
+    - complexity: molecular complexity (BertzCT)
+    - xlogp: calculated LogP (Wildman-Crippen)
+    - heavycnt: heavy atom count
+    - hbonddonor: hydrogen bond donors
+    - hbondacc: hydrogen bond acceptors
+    - rotbonds: rotatable bonds
+    
+    Note: Attachment/conjugation site checks are conjugation-method specific
+    and should be handled in the conjugation PDC builder.
+    """
+
+    def __init__(self, device, 
+                 max_linker_chain_length=16,
+                 calculate_properties=True):
+        self.device = device
+        self.max_linker_chain_length = max_linker_chain_length
+        self.calculate_properties = calculate_properties
+        
+        # Property calculation flags (all on by default)
+        self.property_flags = {
+            'mf': True,           # molecular formula
+            'polararea': True,    # topological polar surface area
+            'complexity': True,   # BertzCT complexity
+            'xlogp': True,        # Wildman-Crippen LogP
+            'heavycnt': True,     # heavy atom count
+            'hbonddonor': True,   # H-bond donors
+            'hbondacc': True,     # H-bond acceptors
+            'rotbonds': True,     # rotatable bonds
+        }
+        
+        # Pre-compile SMARTS patterns for efficiency
+        self._compile_patterns()
+    
+    def set_property_flags(self, **kwargs):
+        """
+        Enable/disable specific property calculations
+        
+        Args:
+            **kwargs: property_name=True/False pairs
+            
+        Example:
+            analyzer.set_property_flags(mf=False, xlogp=True, rotbonds=False)
+        """
+        for key, value in kwargs.items():
+            if key in self.property_flags:
+                self.property_flags[key] = bool(value)
+            else:
+                raise ValueError(f"Unknown property: {key}. Valid properties: {list(self.property_flags.keys())}")
+    
+    def enable_all_properties(self):
+        """Enable calculation of all optional properties"""
+        for key in self.property_flags:
+            self.property_flags[key] = True
+    
+    def disable_all_properties(self):
+        """Disable calculation of all optional properties"""
+        for key in self.property_flags:
+            self.property_flags[key] = False
+    
+    def _compile_patterns(self):
+        """Pre-compile SMARTS patterns used frequently"""
+        self.patterns = {
+            # Long carbon chain (indicates potentially problematic linker)
+            # 16 consecutive CH2 groups
+            'long_chain': Chem.MolFromSmarts('[CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2][CH2]'),
+        }
+    
+    def find_linker_rings(self, mol):
+        """
+        Identify rings that are part of the linker (not peptide backbone or drug core)
+        
+        Linker rings typically include:
+        - Succinimide (from maleimide after thiol addition)
+        - Small saturated rings used as spacers
+        - NOT aromatic rings (usually part of drug)
+        """
+        ring_info = mol.GetRingInfo()
+        atom_rings = ring_info.AtomRings()
+        
+        linker_rings = []
+        
+        for ring in atom_rings:
+            ring_size = len(ring)
+            ring_atoms = [mol.GetAtomWithIdx(idx) for idx in ring]
+            ring_symbols = [atom.GetSymbol() for atom in ring_atoms]
+            
+            # Skip aromatic rings (likely drug component)
+            if all(atom.GetIsAromatic() for atom in ring_atoms):
+                continue
+            
+            # Check for succinimide (5-membered, N + 2 C=O)
+            if ring_size == 5 and 'N' in ring_symbols:
+                carbonyl_count = 0
+                for idx in ring:
+                    atom = mol.GetAtomWithIdx(idx)
+                    if atom.GetSymbol() == 'C':
+                        for neighbor in atom.GetNeighbors():
+                            if neighbor.GetSymbol() == 'O' and neighbor.GetIdx() not in ring:
+                                bond = mol.GetBondBetweenAtoms(idx, neighbor.GetIdx())
+                                if bond and bond.GetBondType() == Chem.BondType.DOUBLE:
+                                    carbonyl_count += 1
+                
+                if carbonyl_count >= 2:
+                    # This is a succinimide - valid linker ring
+                    linker_rings.append(ring)
+                    continue
+            
+            # Check for small saturated rings (potential linker spacers)
+            if ring_size <= 6:
+                is_saturated = not any(atom.GetIsAromatic() for atom in ring_atoms)
+                if is_saturated:
+                    linker_rings.append(ring)
+            
+            # Flag very large non-aromatic rings (unusual in linkers)
+            if ring_size > 12:
+                linker_rings.append(ring)  # Will be checked for validity
+        
+        return linker_rings
+    
+    def check_ring_closure_at_linker(self, mol):
+        """
+        Check ring-closure constraints in the linker region
+        """
+        linker_rings = self.find_linker_rings(mol)
+        
+        if len(linker_rings) == 0:
+            return True  # No linker rings is fine
+        
+        for ring in linker_rings:
+            ring_size = len(ring)
+            
+            # Impossible rings
+            if ring_size < 3:
+                return False
+            
+            # 3-membered rings: check for severe strain
+            if ring_size == 3:
+                for atom_idx in ring:
+                    atom = mol.GetAtomWithIdx(atom_idx)
+                    # sp2 carbon in 3-ring is very strained (except epoxides/aziridines)
+                    if atom.GetSymbol() == 'C':
+                        if atom.GetHybridization() == Chem.HybridizationType.SP2:
+                            # Check if it's part of a C=C (bad) vs just geometry
+                            for neighbor in atom.GetNeighbors():
+                                bond = mol.GetBondBetweenAtoms(atom_idx, neighbor.GetIdx())
+                                if bond and bond.GetBondType() == Chem.BondType.DOUBLE:
+                                    if neighbor.GetSymbol() == 'C':
+                                        return False  # C=C in 3-ring is invalid
+            
+            # 4-membered rings: check for excessive sp2 character
+            if ring_size == 4:
+                sp2_count = sum(1 for idx in ring 
+                              if mol.GetAtomWithIdx(idx).GetHybridization() == Chem.HybridizationType.SP2)
+                if sp2_count > 2:
+                    return False
+            
+            # Very large rings in linker are suspicious
+            if ring_size > 12:
+                return False
+        
+        return True
+    
+    def check_linker_feasibility(self, mol):
+        """
+        Check if linker properties are reasonable
+        
+        Hard constraints only:
+        - Ring-closure constraints
+        - Long carbon chain check
+        """
+        # Check ring-closure constraints
+        if not self.check_ring_closure_at_linker(mol):
+            return False
+        
+        # Check for excessively long carbon chains
+        if self.patterns['long_chain'] and mol.HasSubstructMatch(self.patterns['long_chain']):
+            return False
+        
+        return True
+    
+    def check_representational_validity(self, mol):
+        """
+        Check if molecule can be properly represented
+        """
+        # Sanitization
+        try:
+            Chem.SanitizeMol(mol)
+        except:
+            return False
+        
+        # Round-trip SMILES
+        try:
+            smiles = Chem.MolToSmiles(mol)
+            mol_reconstructed = Chem.MolFromSmiles(smiles)
+            if mol_reconstructed is None:
+                return False
+        except:
+            return False
+        
+        # Check for chemistry problems
+        try:
+            problems = Chem.DetectChemistryProblems(mol)
+            for problem in problems:
+                problem_type = problem.GetType()
+                if problem_type in ['AtomValenceException', 'KekulizeException']:
+                    return False
+        except:
+            return False
+        
+        # 2D coordinates (basic structural validity)
+        try:
+            AllChem.Compute2DCoords(mol)
+        except:
+            return False
+        
+        return True
+    
+    def calculate_molecular_properties(self, mol):
+        """
+        Calculate optional molecular properties based on property_flags
+        
+        Returns dict of property_name: value for enabled properties
+        """
+        properties = {}
+        
+        if not self.calculate_properties:
+            return properties
+        
+        if self.property_flags.get('mf'):
+            properties['mf'] = rdMolDescriptors.CalcMolFormula(mol)
+        
+        if self.property_flags.get('polararea'):
+            properties['polararea'] = round(Descriptors.TPSA(mol), 2)
+        
+        if self.property_flags.get('complexity'):
+            properties['complexity'] = round(Descriptors.BertzCT(mol), 2)
+        
+        if self.property_flags.get('xlogp'):
+            properties['xlogp'] = round(Descriptors.MolLogP(mol), 2)
+        
+        if self.property_flags.get('heavycnt'):
+            properties['heavycnt'] = mol.GetNumHeavyAtoms()
+        
+        if self.property_flags.get('hbonddonor'):
+            properties['hbonddonor'] = Descriptors.NumHDonors(mol)
+        
+        if self.property_flags.get('hbondacc'):
+            properties['hbondacc'] = Descriptors.NumHAcceptors(mol)
+        
+        if self.property_flags.get('rotbonds'):
+            properties['rotbonds'] = Descriptors.NumRotatableBonds(mol)
+        
+        return properties
+    
+    def get_scores(self, input_seqs_list):
+        """
+        Validate PDC structures based on hard constraints
+        
+        Hard constraints:
+        1. Linker feasibility (ring-closure, long chains)
+        2. Representational validity (SMILES, sanitization, chemistry)
+        
+        Returns list of 0/1 scores
+        """
+        results = []
+
+        for candidates in input_seqs_list:
+            for smiles in candidates:
+                mol = Chem.MolFromSmiles(smiles)
+                if mol is None:
+                    continue
+                
+                # Run hard constraint checks only
+                linker_ok = self.check_linker_feasibility(mol)
+                representational_ok = self.check_representational_validity(mol)
+                
+                is_valid = linker_ok and representational_ok
+                if is_valid:
+                    results.append(1)
+                    break
+            else:
+                results.append(0)
+
+        return results
+    
+    def get_detailed_scores(self, input_seqs_list):
+        """
+        Return detailed validation results with optional properties
+        """
+        detailed_results = []
+
+        for candidates in input_seqs_list:
+            candidate_results = []
+            
+            for smiles in candidates:
+                mol = Chem.MolFromSmiles(smiles)
+                if mol is None:
+                    candidate_results.append({
+                        'valid': 0,
+                        'smiles': smiles,
+                        'error': 'Cannot parse SMILES'
+                    })
+                    continue
+            
+                # Run hard constraint checks
+                checks = {
+                    'linker_feasible': self.check_linker_feasibility(mol),
+                    'representational': self.check_representational_validity(mol)
+                }
+                
+                is_valid = all(checks.values())
+                
+                # Calculate optional properties if enabled
+                properties = self.calculate_molecular_properties(mol)
+                
+                # Add linker ring info for diagnostics
+                linker_rings = self.find_linker_rings(mol)
+                checks['linker_ring_sizes'] = [len(ring) for ring in linker_rings]
+                
+                candidate_results.append({
+                    'valid': 1 if is_valid else 0,
+                    'smiles': smiles,
+                    'checks': checks,
+                    'properties': properties
+                })
+            
+            detailed_results.append(candidate_results)
+
+        return detailed_results
+    
+    def __call__(self, input_seqs):
+        scores = self.get_scores(input_seqs)
+        return torch.tensor(scores)
+
+
+# ============================================================================
+# Testing
+# ============================================================================
+
+if __name__ == "__main__":
+    device = 'cpu'
+    pdc_analyzer = PDCAnalyzer(device)
+    
+    print("=" * 70)
+    print("PDC Analyzer - Structural Validity Checker")
+    print("=" * 70)
+    print("\nHard constraints checked:")
+    print("  1. Linker feasibility (ring-closure, long carbon chains)")
+    print("  2. Representational validity (SMILES, sanitization)")
+    print("\nNote: Conjugation site checks are method-specific")
+    print("      and should be in the conjugation PDC builder.\n")
+    
+    # Test property flag toggling
+    print("Testing property flags:")
+    print(f"  Default flags: {pdc_analyzer.property_flags}")
+    
+    pdc_analyzer.set_property_flags(mf=False, complexity=False)
+    print(f"  After disabling mf, complexity: {pdc_analyzer.property_flags}")
+    
+    pdc_analyzer.enable_all_properties()
+    print(f"  After enable_all: {pdc_analyzer.property_flags}")
+    
+    print("\n" + "=" * 70)
+    
+    test_smiles = [
+        # Test 1: Valid structure
+        ("NCC(=O)N[C@@H](CSC1CC(=O)N(CCOCCOCCOCCOCCN)C1=O)C(=O)NCC(=O)O",
+         "Valid structure with succinimide ring"),
+        
+        # Test 2: Invalid - strained 3-membered ring with C=C
+        ("NCC(=O)N[C@@H](CSC1=CC1CCCCN)C(=O)NCC(=O)O",
+         "Invalid: Strained 3-membered ring with sp2 carbon"),
+        
+        # Test 3: Simple valid peptide
+        ("NCC(=O)N[C@@H](C)C(=O)NCC(=O)O",
+         "Simple valid peptide structure"),
+    ]
+    
+    print("\nTest Results")
+    print("=" * 70 + "\n")
+    
+    for i, (smiles, description) in enumerate(test_smiles):
+        print(f"Test {i+1}: {description}")
+        print(f"SMILES: {smiles[:60]}{'...' if len(smiles) > 60 else ''}")
+        
+        detailed = pdc_analyzer.get_detailed_scores([smiles])[0]
+        
+        valid_str = " VALID" if detailed['valid'] == 1 else " INVALID"
+        print(f"Result: {valid_str}")
+        
+        if 'checks' in detailed:
+            print("\nHard constraint checks:")
+            for check_name in ['linker_feasible', 'representational']:
+                if check_name in detailed['checks']:
+                    status = "" if detailed['checks'][check_name] else ""
+                    print(f"  {status} {check_name}")
+        
+        if 'properties' in detailed and detailed['properties']:
+            print("\nMolecular properties:")
+            for prop, value in detailed['properties'].items():
+                print(f"  {prop}: {value}")
+        
+        print("\n" + "-" * 70 + "\n")

smiles/evaluation/peptide_analyzer.py

@@ -0,0 +1,1259 @@
+import os
+import re
+import pandas as pd
+from io import StringIO
+import rdkit
+from rdkit import Chem
+from rdkit.Chem import AllChem, Draw
+import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+from io import BytesIO
+import tempfile
+from rdkit import Chem, RDLogger
+RDLogger.DisableLog('rdApp.*')
+
+import pdb
+
+class PeptideAnalyzer:
+    def __init__(self):
+        self.bond_patterns = [
+            (r'OC\(=O\)', 'ester'),  # Ester bond
+            (r'N\(C\)C\(=O\)', 'n_methyl'),  # N-methylated peptide bond
+            (r'N[0-9]C\(=O\)', 'proline'),  # Proline peptide bond
+            (r'NC\(=O\)', 'peptide'),  # Standard peptide bond
+            (r'C\(=O\)N\(C\)', 'n_methyl_reverse'),  # Reverse N-methylated
+            (r'C\(=O\)N[12]?', 'peptide_reverse')  # Reverse peptide bond
+        ]
+        # Three to one letter code mapping
+        self.three_to_one = {
+            'Ala': 'A', 'Cys': 'C', 'Asp': 'D', 'Glu': 'E',
+            'Phe': 'F', 'Gly': 'G', 'His': 'H', 'Ile': 'I',
+            'Lys': 'K', 'Leu': 'L', 'Met': 'M', 'Asn': 'N',
+            'Pro': 'P', 'Gln': 'Q', 'Arg': 'R', 'Ser': 'S',
+            'Thr': 'T', 'Val': 'V', 'Trp': 'W', 'Tyr': 'Y'
+        }
+    
+    def is_peptide(self, smiles):
+        """Check if the SMILES represents a peptide structure"""
+        # pdb.set_trace()
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return False
+            
+        # Look for peptide bonds: NC(=O) pattern
+        peptide_bond_pattern = Chem.MolFromSmarts('[NH][C](=O)')
+        if mol.HasSubstructMatch(peptide_bond_pattern):
+            return True
+            
+        # Look for N-methylated peptide bonds: N(C)C(=O) pattern
+        n_methyl_pattern = Chem.MolFromSmarts('[N;H0;$(NC)](C)[C](=O)')
+        if mol.HasSubstructMatch(n_methyl_pattern):
+            return True
+        
+        return False
+
+    def is_cyclic(self, smiles):
+        """Improved cyclic peptide detection"""
+        # Check for C-terminal carboxyl
+        if smiles.endswith('C(=O)O'):
+            return False, [], []
+            
+        # Find all numbers used in ring closures
+        ring_numbers = re.findall(r'(?:^|[^c])[0-9](?=[A-Z@\(\)])', smiles)
+        
+        # Find aromatic ring numbers
+        aromatic_matches = re.findall(r'c[0-9](?:ccccc|c\[nH\]c)[0-9]', smiles)
+        aromatic_cycles = []
+        for match in aromatic_matches:
+            numbers = re.findall(r'[0-9]', match)
+            aromatic_cycles.extend(numbers)
+        
+        # Numbers that aren't part of aromatic rings are peptide cycles
+        peptide_cycles = [n for n in ring_numbers if n not in aromatic_cycles]
+        
+        is_cyclic = len(peptide_cycles) > 0 and not smiles.endswith('C(=O)O')
+        return is_cyclic, peptide_cycles, aromatic_cycles
+    
+    def split_on_bonds(self, smiles):
+        """Split SMILES into segments with simplified Pro handling"""
+        positions = []
+        used = set()
+        
+        # Find Gly pattern first
+        gly_pattern = r'NCC\(=O\)'
+        for match in re.finditer(gly_pattern, smiles):
+            if not any(p in range(match.start(), match.end()) for p in used):
+                positions.append({
+                    'start': match.start(),
+                    'end': match.end(),
+                    'type': 'gly',
+                    'pattern': match.group()
+                })
+                used.update(range(match.start(), match.end()))
+        
+        for pattern, bond_type in self.bond_patterns:
+            for match in re.finditer(pattern, smiles):
+                if not any(p in range(match.start(), match.end()) for p in used):
+                    positions.append({
+                        'start': match.start(),
+                        'end': match.end(),
+                        'type': bond_type,
+                        'pattern': match.group()
+                    })
+                    used.update(range(match.start(), match.end()))
+
+        # Sort by position
+        positions.sort(key=lambda x: x['start'])
+        
+        # Create segments
+        segments = []
+        
+        if positions:
+            # First segment
+            if positions[0]['start'] > 0:
+                segments.append({
+                    'content': smiles[0:positions[0]['start']],
+                    'bond_after': positions[0]['pattern']
+                })
+            
+            # Process segments
+            for i in range(len(positions)-1):
+                current = positions[i]
+                next_pos = positions[i+1]
+                
+                if current['type'] == 'gly':
+                    segments.append({
+                        'content': 'NCC(=O)',
+                        'bond_before': positions[i-1]['pattern'] if i > 0 else None,
+                        'bond_after': next_pos['pattern']
+                    })
+                else:
+                    content = smiles[current['end']:next_pos['start']]
+                    if content:
+                        segments.append({
+                            'content': content,
+                            'bond_before': current['pattern'],
+                            'bond_after': next_pos['pattern']
+                        })
+            
+            # Last segment
+            if positions[-1]['end'] < len(smiles):
+                segments.append({
+                    'content': smiles[positions[-1]['end']:],
+                    'bond_before': positions[-1]['pattern']
+                })
+        
+        return segments
+
+    def clean_terminal_carboxyl(self, segment):
+        """Remove C-terminal carboxyl only if it's the true terminus"""
+        content = segment['content']
+        
+        # Only clean if:
+        # 1. Contains C(=O)O
+        # 2. No bond_after exists (meaning it's the last segment)
+        # 3. C(=O)O is at the end of the content
+        if 'C(=O)O' in content and not segment.get('bond_after'):
+            print('recognized?')
+            # Remove C(=O)O pattern regardless of position
+            cleaned = re.sub(r'\(C\(=O\)O\)', '', content)
+            # Remove any leftover empty parentheses
+            cleaned = re.sub(r'\(\)', '', cleaned)
+            print(cleaned)
+            return cleaned
+        return content
+    
+    def identify_residue(self, segment):
+        """Identify residue with Pro reconstruction"""
+        # Only clean terminal carboxyl if this is the last segment
+        content = self.clean_terminal_carboxyl(segment)
+        mods = self.get_modifications(segment)
+        
+        # UAA pattern matching section - before regular residues
+        # Phenylglycine and derivatives
+        if 'c1ccccc1' in content:
+            if '[C@@H](c1ccccc1)' in content or '[C@H](c1ccccc1)' in content:
+                return '4', mods  # Base phenylglycine
+                
+        # 4-substituted phenylalanines
+        if 'Cc1ccc' in content:
+            if 'OMe' in content or 'OCc1ccc' in content:
+                return '0A1', mods  # 4-methoxy-Phenylalanine
+            elif 'Clc1ccc' in content:
+                return '200', mods  # 4-chloro-Phenylalanine
+            elif 'Brc1ccc' in content:
+                return '4BF', mods  # 4-Bromo-phenylalanine
+            elif 'C#Nc1ccc' in content:
+                return '4CF', mods  # 4-cyano-phenylalanine
+            elif 'Ic1ccc' in content:
+                return 'PHI', mods  # 4-Iodo-phenylalanine
+            elif 'Fc1ccc' in content:
+                return 'PFF', mods  # 4-Fluoro-phenylalanine
+                
+        # Modified tryptophans
+        if 'c[nH]c2' in content:
+            if 'Oc2cccc2' in content:
+                return '0AF', mods  # 7-hydroxy-tryptophan
+            elif 'Fc2cccc2' in content:
+                return '4FW', mods  # 4-fluoro-tryptophan
+            elif 'Clc2cccc2' in content:
+                return '6CW', mods  # 6-chloro-tryptophan
+            elif 'Brc2cccc2' in content:
+                return 'BTR', mods  # 6-bromo-tryptophan
+            elif 'COc2cccc2' in content:
+                return 'MOT5', mods  # 5-Methoxy-tryptophan
+            elif 'Cc2cccc2' in content:
+                return 'MTR5', mods  # 5-Methyl-tryptophan
+                
+        # Special amino acids
+        if 'CC(C)(C)[C@@H]' in content or 'CC(C)(C)[C@H]' in content:
+            return 'BUG', mods  # Tertleucine
+            
+        if 'CCCNC(=N)N' in content:
+            return 'CIR', mods  # Citrulline
+            
+        if '[SeH]' in content:
+            return 'CSE', mods  # Selenocysteine
+            
+        if '[NH3]CC[C@@H]' in content or '[NH3]CC[C@H]' in content:
+            return 'DAB', mods  # Diaminobutyric acid
+            
+        if 'C1CCCCC1' in content:
+            if 'C1CCCCC1[C@@H]' in content or 'C1CCCCC1[C@H]' in content:
+                return 'CHG', mods  # Cyclohexylglycine
+            elif 'C1CCCCC1C[C@@H]' in content or 'C1CCCCC1C[C@H]' in content:
+                return 'ALC', mods  # 3-cyclohexyl-alanine
+                
+        # Naphthalene derivatives
+        if 'c1cccc2c1cccc2' in content:
+            if 'c1cccc2c1cccc2[C@@H]' in content or 'c1cccc2c1cccc2[C@H]' in content:
+                return 'NAL', mods  # 2-Naphthyl-alanine
+                
+        # Heteroaromatic derivatives
+        if 'c1cncc' in content:
+            return 'PYR4', mods  # 3-(4-Pyridyl)-alanine
+        if 'c1cscc' in content:
+            return 'THA3', mods  # 3-(3-thienyl)-alanine
+        if 'c1nnc' in content:
+            return 'TRZ4', mods  # 3-(1,2,4-Triazol-1-yl)-alanine
+            
+        # Modified serines and threonines
+        if 'OP(O)(O)O' in content:
+            if '[C@@H](COP' in content or '[C@H](COP' in content:
+                return 'SEP', mods  # phosphoserine
+            elif '[C@@H](OP' in content or '[C@H](OP' in content:
+                return 'TPO', mods  # phosphothreonine
+            
+        # Specialized ring systems
+        if 'c1c2ccccc2cc2c1cccc2' in content:
+            return 'ANTH', mods  # 3-(9-anthryl)-alanine
+        if 'c1csc2c1cccc2' in content:
+            return 'BTH3', mods  # 3-(3-benzothienyl)-alanine
+        if '[C@]12C[C@H]3C[C@@H](C2)C[C@@H](C1)C3' in content:
+            return 'ADAM', mods  # Adamanthane
+
+        # Fluorinated derivatives
+        if 'FC(F)(F)' in content:
+            if 'CC(F)(F)F' in content:
+                return 'FLA', mods  # Trifluoro-alanine
+            if 'C(F)(F)F)c1' in content:
+                if 'c1ccccc1C(F)(F)F' in content:
+                    return 'TFG2', mods  # 2-(Trifluoromethyl)-phenylglycine
+                if 'c1cccc(c1)C(F)(F)F' in content:
+                    return 'TFG3', mods  # 3-(Trifluoromethyl)-phenylglycine
+                if 'c1ccc(cc1)C(F)(F)F' in content:
+                    return 'TFG4', mods  # 4-(Trifluoromethyl)-phenylglycine
+        
+        # Multiple halogen patterns
+        if 'F' in content and 'c1' in content:
+            if 'c1ccc(c(c1)F)F' in content:
+                return 'F2F', mods  # 3,4-Difluoro-phenylalanine
+            if 'cc(F)cc(c1)F' in content:
+                return 'WFP', mods  # 3,5-Difluoro-phenylalanine
+        if 'Cl' in content and 'c1' in content:
+            if 'c1ccc(cc1Cl)Cl' in content:
+                return 'CP24', mods  # 2,4-dichloro-phenylalanine
+            if 'c1ccc(c(c1)Cl)Cl' in content:
+                return 'CP34', mods  # 3,4-dichloro-phenylalanine
+
+        # Hydroxy and amino derivatives
+        if 'O' in content and 'c1' in content:
+            if 'c1cc(O)cc(c1)O' in content:
+                return '3FG', mods  # (2s)-amino(3,5-dihydroxyphenyl)-ethanoic acid
+            if 'c1ccc(c(c1)O)O' in content:
+                return 'DAH', mods  # 3,4-Dihydroxy-phenylalanine
+        
+        # Cyclic amino acids
+        if 'C1CCCC1' in content:
+            return 'CPA3', mods  # 3-Cyclopentyl-alanine
+        if 'C1CCCCC1' in content:
+            if 'CC1CCCCC1' in content:
+                return 'ALC', mods  # 3-cyclohexyl-alanine
+            else:
+                return 'CHG', mods  # Cyclohexylglycine
+
+        # Chain-length variants
+        if 'CCC[C@@H]' in content or 'CCC[C@H]' in content:
+            return 'NLE', mods  # Norleucine
+        if 'CC[C@@H]' in content or 'CC[C@H]' in content:
+            if not any(x in content for x in ['CC(C)', 'COC', 'CN(']):
+                return 'ABA', mods  # 2-Aminobutyric acid
+        
+        # Modified histidines
+        if 'c1cnc' in content:
+            if '[C@@H]1CN[C@@H](N1)F' in content:
+                return '2HF', mods  # 2-fluoro-l-histidine
+            if 'c1cnc([nH]1)F' in content:
+                return '2HF1', mods  # 2-fluoro-l-histidine variant
+            if 'c1c[nH]c(n1)F' in content:
+                return '2HF2', mods  # 2-fluoro-l-histidine variant
+
+        # Sulfur and selenium containing
+        if '[SeH]' in content:
+            return 'CSE', mods  # Selenocysteine
+        if 'S' in content:
+            if 'CSCc1ccccc1' in content:
+                return 'BCS', mods  # benzylcysteine
+            if 'CCSC' in content:
+                return 'ESC', mods  # Ethionine
+            if 'CCS' in content:
+                return 'HCS', mods  # homocysteine
+
+        # Additional modifications
+        if 'CN=[N]=N' in content:
+            return 'AZDA', mods  # azido-alanine
+        if '[NH]=[C](=[NH2])=[NH2]' in content:
+            if 'CCC[NH]=' in content:
+                return 'AGM', mods  # 5-methyl-arginine
+            if 'CC[NH]=' in content:
+                return 'GDPR', mods  # 2-Amino-3-guanidinopropionic acid
+
+        if 'CCON' in content:
+            return 'CAN', mods  # canaline
+        if '[C@@H]1C=C[C@@H](C=C1)' in content:
+            return 'ACZ', mods  # cis-amiclenomycin
+        if 'CCC(=O)[NH3]' in content:
+            return 'ONL', mods  # 5-oxo-l-norleucine
+        if 'c1ccncc1' in content:
+            return 'PYR4', mods  # 3-(4-Pyridyl)-alanine
+        if 'c1ccco1' in content:
+            return 'FUA2', mods  # (2-furyl)-alanine
+            
+        if 'c1ccc' in content:
+            if 'c1ccc(cc1)c1ccccc1' in content:
+                return 'BIF', mods  # 4,4-biphenylalanine
+            if 'c1ccc(cc1)C(=O)c1ccccc1' in content:
+                return 'PBF', mods  # 4-benzoyl-phenylalanine
+            if 'c1ccc(cc1)C(C)(C)C' in content:
+                return 'TBP4', mods  # 4-tert-butyl-phenylalanine
+            if 'c1ccc(cc1)[C](=[NH2])=[NH2]' in content:
+                return '0BN', mods  # 4-carbamimidoyl-l-phenylalanine
+            if 'c1cccc(c1)[C](=[NH2])=[NH2]' in content:
+                return 'APM', mods  # m-amidinophenyl-3-alanine
+
+        # Multiple hydroxy patterns
+        if 'O' in content:
+            if '[C@H]([C@H](C)O)O' in content:
+                return 'ILX', mods  # 4,5-dihydroxy-isoleucine
+            if '[C@H]([C@@H](C)O)O' in content:
+                return 'ALO', mods  # Allo-threonine
+            if '[C@H](COP(O)(O)O)' in content:
+                return 'SEP', mods  # phosphoserine
+            if '[C@H]([C@@H](C)OP(O)(O)O)' in content:
+                return 'TPO', mods  # phosphothreonine
+            if '[C@H](c1ccc(O)cc1)O' in content:
+                return 'OMX', mods  # (betar)-beta-hydroxy-l-tyrosine
+            if '[C@H](c1ccc(c(Cl)c1)O)O' in content:
+                return 'OMY', mods  # (betar)-3-chloro-beta-hydroxy-l-tyrosine
+
+        # Heterocyclic patterns
+        if 'n1' in content:
+            if 'n1cccn1' in content:
+                return 'PYZ1', mods  # 3-(1-Pyrazolyl)-alanine
+            if 'n1nncn1' in content:
+                return 'TEZA', mods  # 3-(2-Tetrazolyl)-alanine
+            if 'c2c(n1)cccc2' in content:
+                return 'QU32', mods  # 3-(2-Quinolyl)-alanine
+            if 'c1cnc2c(c1)cccc2' in content:
+                return 'QU33', mods  # 3-(3-quinolyl)-alanine
+            if 'c1ccnc2c1cccc2' in content:
+                return 'QU34', mods  # 3-(4-quinolyl)-alanine
+            if 'c1ccc2c(c1)nccc2' in content:
+                return 'QU35', mods  # 3-(5-Quinolyl)-alanine
+            if 'c1ccc2c(c1)cncc2' in content:
+                return 'QU36', mods  # 3-(6-Quinolyl)-alanine
+            if 'c1cnc2c(n1)cccc2' in content:
+                return 'QX32', mods  # 3-(2-quinoxalyl)-alanine
+
+        # Multiple nitrogen patterns
+        if 'N' in content:
+            if '[NH3]CC[C@@H]' in content:
+                return 'DAB', mods  # Diaminobutyric acid
+            if '[NH3]C[C@@H]' in content:
+                return 'DPP', mods  # 2,3-Diaminopropanoic acid
+            if '[NH3]CCCCCC[C@@H]' in content:
+                return 'HHK', mods  # (2s)-2,8-diaminooctanoic acid
+            if 'CCC[NH]=[C](=[NH2])=[NH2]' in content:
+                return 'GBUT', mods  # 2-Amino-4-guanidinobutryric acid
+            if '[NH]=[C](=S)=[NH2]' in content:
+                return 'THIC', mods  # Thio-citrulline
+
+        # Chain modified amino acids
+        if 'CC' in content:
+            if 'CCCC[C@@H]' in content:
+                return 'AHP', mods  # 2-Aminoheptanoic acid
+            if 'CCC([C@@H])(C)C' in content:
+                return 'I2M', mods  # 3-methyl-l-alloisoleucine
+            if 'CC[C@H]([C@@H])C' in content:
+                return 'IIL', mods  # Allo-Isoleucine
+            if '[C@H](CCC(C)C)' in content:
+                return 'HLEU', mods  # Homoleucine
+            if '[C@@H]([C@@H](C)O)C' in content:
+                return 'HLU', mods  # beta-hydroxyleucine
+
+        # Modified glutamate/aspartate patterns
+        if '[C@@H]' in content:
+            if '[C@@H](C[C@@H](F))' in content:
+                return 'FGA4', mods  # 4-Fluoro-glutamic acid
+            if '[C@@H](C[C@@H](O))' in content:
+                return '3GL', mods  # 4-hydroxy-glutamic-acid
+            if '[C@@H](C[C@H](C))' in content:
+                return 'LME', mods  # (3r)-3-methyl-l-glutamic acid
+            if '[C@@H](CC[C@H](C))' in content:
+                return 'MEG', mods  # (3s)-3-methyl-l-glutamic acid
+
+        # Sulfur and selenium modifications
+        if 'S' in content:
+            if 'SCC[C@@H]' in content:
+                return 'HSER', mods  # homoserine
+            if 'SCCN' in content:
+                return 'SLZ', mods  # thialysine
+            if 'SC(=O)' in content:
+                return 'CSA', mods  # s-acetonylcysteine
+            if '[S@@](=O)' in content:
+                return 'SME', mods  # Methionine sulfoxide
+            if 'S(=O)(=O)' in content:
+                return 'OMT', mods  # Methionine sulfone
+
+        # Double bond containing
+        if 'C=' in content:
+            if 'C=C[C@@H]' in content:
+                return '2AG', mods  # 2-Allyl-glycine
+            if 'C=C[C@@H]' in content:
+                return 'LVG', mods  # vinylglycine
+            if 'C=Cc1ccccc1' in content:
+                return 'STYA', mods  # Styrylalanine
+
+        # Special cases
+        if '[C@@H]1Cc2c(C1)cccc2' in content:
+            return 'IGL', mods  # alpha-amino-2-indanacetic acid
+        if '[C](=[C](=O)=O)=O' in content:
+            return '26P', mods  # 2-amino-6-oxopimelic acid
+        if '[C](=[C](=O)=O)=C' in content:
+            return '2NP', mods  # l-2-amino-6-methylene-pimelic acid
+        if 'c2cnc[nH]2' in content:
+            return 'HIS', mods  # histidine core
+        if 'c1cccc2c1cc(O)cc2' in content:
+            return 'NAO1', mods  # 5-hydroxy-1-naphthalene
+        if 'c1ccc2c(c1)cc(O)cc2' in content:
+            return 'NAO2', mods  # 6-hydroxy-2-naphthalene
+            
+        # Proline (P) - flexible ring numbers
+        if any([
+            # Check for any ring number in bond patterns
+            (segment.get('bond_after', '').startswith(f'N{n}C(=O)') and 'CCC' in content and 
+            any(f'[C@@H]{n}' in content or f'[C@H]{n}' in content for n in '123456789'))
+            for n in '123456789'
+        ]) or any([
+            # Check ending patterns with any ring number
+            (f'CCCN{n}' in content and content.endswith('=O') and 
+            any(f'[C@@H]{n}' in content or f'[C@H]{n}' in content for n in '123456789'))
+            for n in '123456789'
+        ]) or any([
+            # Handle CCC[C@H]n patterns
+            (content == f'CCC[C@H]{n}' and segment.get('bond_before', '').startswith(f'C(=O)N{n}')) or
+            (content == f'CCC[C@@H]{n}' and segment.get('bond_before', '').startswith(f'C(=O)N{n}')) or
+            # N-terminal Pro with any ring number
+            (f'N{n}CCC[C@H]{n}' in content) or
+            (f'N{n}CCC[C@@H]{n}' in content)
+            for n in '123456789'
+        ]):
+            return 'Pro', mods
+        
+        # Tryptophan (W) - more specific indole pattern
+        if re.search(r'c[0-9]c\[nH\]c[0-9]ccccc[0-9][0-9]', content) and \
+        'c[nH]c' in content.replace(' ', ''):
+            return 'Trp', mods
+        
+        # Lysine (K) - both patterns
+        if '[C@@H](CCCCN)' in content or '[C@H](CCCCN)' in content:
+            return 'Lys', mods
+            
+        # Arginine (R) - both patterns
+        if '[C@@H](CCCNC(=N)N)' in content or '[C@H](CCCNC(=N)N)' in content:
+            return 'Arg', mods
+        
+        if ('C[C@H](CCCC)' in content or 'C[C@@H](CCCC)' in content) and 'CC(C)' not in content:
+            return 'Nle', mods
+            
+        # Ornithine (Orn) - 3-carbon chain with NH2
+        if ('C[C@H](CCCN)' in content or 'C[C@@H](CCCN)' in content) and 'CC(C)' not in content:
+            return 'Orn', mods
+            
+        # 2-Naphthylalanine (2Nal) - distinct from Phe pattern
+        if ('Cc3cc2ccccc2c3' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return '2Nal', mods
+            
+        # Cyclohexylalanine (Cha) - already in your code but moved here for clarity
+        if 'N2CCCCC2' in content or 'CCCCC2' in content:
+            return 'Cha', mods
+            
+        # Aminobutyric acid (Abu) - 2-carbon chain
+        if ('C[C@H](CC)' in content or 'C[C@@H](CC)' in content) and not any(p in content for p in ['CC(C)', 'CCCC', 'CCC(C)']):
+            return 'Abu', mods
+            
+        # Pipecolic acid (Pip) - 6-membered ring like Pro
+        if ('N3CCCCC3' in content or 'CCCCC3' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return 'Pip', mods
+
+        # Cyclohexylglycine (Chg) - direct cyclohexyl without CH2
+        if ('C[C@H](C1CCCCC1)' in content or 'C[C@@H](C1CCCCC1)' in content):
+            return 'Chg', mods
+            
+        # 4-Fluorophenylalanine (4F-Phe)
+        if ('Cc2ccc(F)cc2' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return '4F-Phe', mods
+        
+        # Regular residue identification
+        if ('NCC(=O)' in content) or (content == 'C'):
+            # Middle case - between bonds
+            if segment.get('bond_before') and segment.get('bond_after'):
+                if ('C(=O)N' in segment['bond_before'] or 'C(=O)N(C)' in segment['bond_before']):
+                    return 'Gly', mods
+            # Terminal case - at the end
+            elif segment.get('bond_before') and segment.get('bond_before').startswith('C(=O)N'):
+                return 'Gly', mods
+            
+        if 'CC(C)C[C@H]' in content or 'CC(C)C[C@@H]' in content:
+            return 'Leu', mods
+        if '[C@@H](CC(C)C)' in content or '[C@H](CC(C)C)' in content:
+            return 'Leu', mods
+
+        if '[C@@H]([C@@H](C)O)' in content or '[C@H]([C@H](C)O)' in content:
+            return 'Thr', mods
+        
+        if '[C@H](Cc2ccccc2)' in content or '[C@@H](Cc2ccccc2)' in content:
+            return 'Phe', mods
+            
+        if ('[C@H](C(C)C)' in content or       # With outer parentheses
+            '[C@@H](C(C)C)' in content or      # With outer parentheses
+            '[C@H]C(C)C' in content or         # Without outer parentheses 
+            '[C@@H]C(C)C' in content):         # Without outer parentheses
+            if not any(p in content for p in ['CC(C)C[C@H]', 'CC(C)C[C@@H]']):  # Still check not Leu
+                return 'Val', mods
+        
+        if '[C@H](COC(C)(C)C)' in content or '[C@@H](COC(C)(C)C)' in content:
+            return 'O-tBu', mods
+        
+        if any([
+            'CC[C@H](C)' in content,
+            'CC[C@@H](C)' in content,
+            'C(C)C[C@H]' in content and 'CC(C)C' not in content,
+            'C(C)C[C@@H]' in content and 'CC(C)C' not in content
+        ]):
+            return 'Ile', mods
+        
+        if ('[C@H](C)' in content or '[C@@H](C)' in content):
+            if not any(p in content for p in ['C(C)C', 'COC', 'CN(', 'C(C)O', 'CC[C@H]', 'CC[C@@H]']):
+                return 'Ala', mods
+        
+        # Tyrosine (Tyr) - 4-hydroxybenzyl side chain
+        if re.search(r'Cc[0-9]ccc\(O\)cc[0-9]', content):
+            return 'Tyr', mods
+
+        
+        # Serine (Ser) - Hydroxymethyl side chain
+        if '[C@H](CO)' in content or '[C@@H](CO)' in content:
+            if not ('C(C)O' in content or 'COC' in content):
+                return 'Ser', mods
+        
+        # Threonine (Thr) - 1-hydroxyethyl side chain
+        if '[C@@H]([C@@H](C)O)' in content or '[C@H]([C@H](C)O)' in content or '[C@@H](C)O' in content or '[C@H](C)O' in content:
+            return 'Thr', mods
+        
+        # Cysteine (Cys) - Thiol side chain
+        if '[C@H](CS)' in content or '[C@@H](CS)' in content:
+            return 'Cys', mods
+        
+        # Methionine (Met) - Methylthioethyl side chain
+        if ('C[C@H](CCSC)' in content or 'C[C@@H](CCSC)' in content):
+            return 'Met', mods
+        
+        # Asparagine (Asn) - Carbamoylmethyl side chain
+        if ('CC(=O)N' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return 'Asn', mods
+        
+        # Glutamine (Gln) - Carbamoylethyl side chain
+        if ('CCC(=O)N' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return 'Gln', mods
+        
+        # Aspartic acid (Asp) - Carboxymethyl side chain
+        if ('CC(=O)O' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return 'Asp', mods
+        
+        # Glutamic acid (Glu) - Carboxyethyl side chainCC(C(=O)O)N
+        if ('CCC(=O)O' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return 'Glu', mods
+        
+        # Arginine (Arg) - 3-guanidinopropyl side chain
+        if ('CCCNC(=N)N' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return 'Arg', mods
+        
+        # Histidine (His) - Imidazole side chain
+        if ('Cc2cnc[nH]2' in content) and ('C[C@H]' in content or 'C[C@@H]' in content):
+            return 'His', mods
+            
+        return None, mods
+
+    def get_modifications(self, segment):
+        """Get modifications based on bond types"""
+        mods = []
+        if segment.get('bond_after'):
+            if 'N(C)' in segment['bond_after'] or segment['bond_after'].startswith('C(=O)N(C)'):
+                mods.append('N-Me')
+            if 'OC(=O)' in segment['bond_after']:
+                mods.append('O-linked')
+        return mods
+
+    def analyze_structure(self, smiles):
+        """Main analysis function with debug output"""
+        print("\nAnalyzing structure:", smiles)
+        
+        # Split into segments
+        segments = self.split_on_bonds(smiles)
+        
+        print("\nSegment Analysis:")
+        sequence = []
+        for i, segment in enumerate(segments):
+            print(f"\nSegment {i}:")
+            print(f"Content: {segment['content']}")
+            print(f"Bond before: {segment.get('bond_before', 'None')}")
+            print(f"Bond after: {segment.get('bond_after', 'None')}")
+            
+            residue, mods = self.identify_residue(segment)
+            if residue:
+                if mods:
+                    sequence.append(f"{residue}({','.join(mods)})")
+                else:
+                    sequence.append(residue)
+                print(f"Identified as: {residue}")
+                print(f"Modifications: {mods}")
+            else:
+                print(f"Warning: Could not identify residue in segment: {segment['content']}")
+        
+        # Check if cyclic
+        is_cyclic, peptide_cycles, aromatic_cycles = self.is_cyclic(smiles)
+        three_letter = '-'.join(sequence)
+        one_letter = ''.join(self.three_to_one.get(aa.split('(')[0], 'X') for aa in sequence)
+        
+        if is_cyclic:
+            three_letter = f"cyclo({three_letter})"
+            one_letter = f"cyclo({one_letter})"
+        
+        print(f"\nFinal sequence: {three_letter}")
+        print(f"One-letter code: {one_letter}")
+        print(f"Is cyclic: {is_cyclic}")
+        #print(f"Peptide cycles: {peptide_cycles}")
+        #print(f"Aromatic cycles: {aromatic_cycles}")
+        
+        return three_letter, len(segments)
+        """return {
+            'three_letter': three_letter,
+            #'one_letter': one_letter,
+            'is_cyclic': is_cyclic
+        }"""
+        
+    def return_sequence(self, smiles):
+        """Main analysis function with debug output"""
+        print("\nAnalyzing structure:", smiles)
+        
+        # Split into segments
+        segments = self.split_on_bonds(smiles)
+        
+        print("\nSegment Analysis:")
+        sequence = []
+        for i, segment in enumerate(segments):
+            print(f"\nSegment {i}:")
+            print(f"Content: {segment['content']}")
+            print(f"Bond before: {segment.get('bond_before', 'None')}")
+            print(f"Bond after: {segment.get('bond_after', 'None')}")
+            
+            residue, mods = self.identify_residue(segment)
+            if residue:
+                if mods:
+                    sequence.append(f"{residue}({','.join(mods)})")
+                else:
+                    sequence.append(residue)
+                print(f"Identified as: {residue}")
+                print(f"Modifications: {mods}")
+            else:
+                print(f"Warning: Could not identify residue in segment: {segment['content']}")
+        
+        return sequence
+
+"""
+def annotate_cyclic_structure(mol, sequence):
+    '''Create annotated 2D structure with clear, non-overlapping residue labels'''
+    # Generate 2D coordinates
+    # Generate 2D coordinates
+    AllChem.Compute2DCoords(mol)
+    
+    # Create drawer with larger size for annotations
+    drawer = Draw.rdMolDraw2D.MolDraw2DCairo(2000, 2000)  # Even larger size
+    
+    # Get residue list and reverse it to match structural representation
+    if sequence.startswith('cyclo('):
+        residues = sequence[6:-1].split('-')
+    else:
+        residues = sequence.split('-')
+    residues = list(reversed(residues))  # Reverse the sequence
+    
+    # Draw molecule first to get its bounds
+    drawer.drawOptions().addAtomIndices = False
+    drawer.DrawMolecule(mol)
+    drawer.FinishDrawing()
+    
+    # Convert to PIL Image
+    img = Image.open(BytesIO(drawer.GetDrawingText()))
+    draw = ImageDraw.Draw(img)
+    
+    try:
+        # Try to use DejaVuSans as it's commonly available on Linux systems
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 60)
+        small_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 60)
+    except OSError:
+        try:
+            # Fallback to Arial if available (common on Windows)
+            font = ImageFont.truetype("arial.ttf", 60)
+            small_font = ImageFont.truetype("arial.ttf", 60)
+        except OSError:
+            # If no TrueType fonts are available, fall back to default
+            print("Warning: TrueType fonts not available, using default font")
+            font = ImageFont.load_default()
+            small_font = ImageFont.load_default()
+    # Get molecule bounds
+    conf = mol.GetConformer()
+    positions = []
+    for i in range(mol.GetNumAtoms()):
+        pos = conf.GetAtomPosition(i)
+        positions.append((pos.x, pos.y))
+    
+    x_coords = [p[0] for p in positions]
+    y_coords = [p[1] for p in positions]
+    min_x, max_x = min(x_coords), max(x_coords)
+    min_y, max_y = min(y_coords), max(y_coords)
+    
+    # Calculate scaling factors
+    scale = 150  # Increased scale factor
+    center_x = 1000  # Image center
+    center_y = 1000
+    
+    # Add residue labels in a circular arrangement around the structure
+    n_residues = len(residues)
+    radius = 700  # Distance of labels from center
+    
+    # Start from the rightmost point (3 o'clock position) and go counterclockwise
+    # Offset by -3 positions to align with structure
+    offset = 0  # Adjust this value to match the structure alignment
+    for i, residue in enumerate(residues):
+        # Calculate position in a circle around the structure
+        # Start from 0 (3 o'clock) and go counterclockwise
+        angle = -(2 * np.pi * ((i + offset) % n_residues) / n_residues)
+        
+        # Calculate label position
+        label_x = center_x + radius * np.cos(angle)
+        label_y = center_y + radius * np.sin(angle)
+        
+        # Draw residue label
+        text = f"{i+1}. {residue}"
+        bbox = draw.textbbox((label_x, label_y), text, font=font)
+        padding = 10
+        draw.rectangle([bbox[0]-padding, bbox[1]-padding, 
+                       bbox[2]+padding, bbox[3]+padding], 
+                      fill='white', outline='white')
+        draw.text((label_x, label_y), text, 
+                 font=font, fill='black', anchor="mm")
+    
+    # Add sequence at the top with white background
+    seq_text = f"Sequence: {sequence}"
+    bbox = draw.textbbox((center_x, 100), seq_text, font=small_font)
+    padding = 10
+    draw.rectangle([bbox[0]-padding, bbox[1]-padding, 
+                   bbox[2]+padding, bbox[3]+padding], 
+                  fill='white', outline='white')
+    draw.text((center_x, 100), seq_text, 
+             font=small_font, fill='black', anchor="mm")
+    
+    return img
+"""
+def annotate_cyclic_structure(mol, sequence):
+    """Create structure visualization with just the sequence header"""
+    # Generate 2D coordinates
+    AllChem.Compute2DCoords(mol)
+    
+    # Create drawer with larger size for annotations
+    drawer = Draw.rdMolDraw2D.MolDraw2DCairo(2000, 2000)
+    
+    # Draw molecule first
+    drawer.drawOptions().addAtomIndices = False
+    drawer.DrawMolecule(mol)
+    drawer.FinishDrawing()
+    
+    # Convert to PIL Image
+    img = Image.open(BytesIO(drawer.GetDrawingText()))
+    draw = ImageDraw.Draw(img)
+    try:
+        small_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 60)
+    except OSError:
+        try:
+            small_font = ImageFont.truetype("arial.ttf", 60)
+        except OSError:
+            print("Warning: TrueType fonts not available, using default font")
+            small_font = ImageFont.load_default()
+    
+    # Add just the sequence header at the top
+    seq_text = f"Sequence: {sequence}"
+    bbox = draw.textbbox((1000, 100), seq_text, font=small_font)
+    padding = 10
+    draw.rectangle([bbox[0]-padding, bbox[1]-padding, 
+                   bbox[2]+padding, bbox[3]+padding], 
+                  fill='white', outline='white')
+    draw.text((1000, 100), seq_text, 
+             font=small_font, fill='black', anchor="mm")
+    
+    return img
+
+def create_enhanced_linear_viz(sequence, smiles):
+    """Create an enhanced linear representation using PeptideAnalyzer"""
+    analyzer = PeptideAnalyzer()  # Create analyzer instance
+    
+    # Create figure with two subplots
+    fig = plt.figure(figsize=(15, 10))
+    gs = fig.add_gridspec(2, 1, height_ratios=[1, 2])
+    ax_struct = fig.add_subplot(gs[0])
+    ax_detail = fig.add_subplot(gs[1])
+    
+    # Parse sequence and get residues
+    if sequence.startswith('cyclo('):
+        residues = sequence[6:-1].split('-')
+    else:
+        residues = sequence.split('-')
+    
+    # Get segments using analyzer
+    segments = analyzer.split_on_bonds(smiles)
+    
+    # Debug print
+    print(f"Number of residues: {len(residues)}")
+    print(f"Number of segments: {len(segments)}")
+    
+    # Top subplot - Basic structure
+    ax_struct.set_xlim(0, 10)
+    ax_struct.set_ylim(0, 2)
+    
+    num_residues = len(residues)
+    spacing = 9.0 / (num_residues - 1) if num_residues > 1 else 9.0
+    
+    # Draw basic structure
+    y_pos = 1.5
+    for i in range(num_residues):
+        x_pos = 0.5 + i * spacing
+        
+        # Draw amino acid box
+        rect = patches.Rectangle((x_pos-0.3, y_pos-0.2), 0.6, 0.4, 
+                               facecolor='lightblue', edgecolor='black')
+        ax_struct.add_patch(rect)
+        
+        # Draw connecting bonds if not the last residue
+        if i < num_residues - 1:
+            segment = segments[i] if i < len(segments) else None
+            if segment:
+                # Determine bond type from segment info
+                bond_type = 'ester' if 'O-linked' in segment.get('bond_after', '') else 'peptide'
+                is_n_methylated = 'N-Me' in segment.get('bond_after', '')
+                
+                bond_color = 'red' if bond_type == 'ester' else 'black'
+                linestyle = '--' if bond_type == 'ester' else '-'
+                
+                # Draw bond line
+                ax_struct.plot([x_pos+0.3, x_pos+spacing-0.3], [y_pos, y_pos], 
+                             color=bond_color, linestyle=linestyle, linewidth=2)
+                
+                # Add bond type label
+                mid_x = x_pos + spacing/2
+                bond_label = f"{bond_type}"
+                if is_n_methylated:
+                    bond_label += "\n(N-Me)"
+                ax_struct.text(mid_x, y_pos+0.1, bond_label, 
+                             ha='center', va='bottom', fontsize=10, 
+                             color=bond_color)
+        
+        # Add residue label
+        ax_struct.text(x_pos, y_pos-0.5, residues[i], 
+                      ha='center', va='top', fontsize=14)
+    
+    # Bottom subplot - Detailed breakdown
+    ax_detail.set_ylim(0, len(segments)+1)
+    ax_detail.set_xlim(0, 1)
+    
+    # Create detailed breakdown
+    segment_y = len(segments)  # Start from top
+    for i, segment in enumerate(segments):
+        y = segment_y - i
+        
+        # Check if this is a bond or residue
+        residue, mods = analyzer.identify_residue(segment)
+        if residue:
+            text = f"Residue {i+1}: {residue}"
+            if mods:
+                text += f" ({', '.join(mods)})"
+            color = 'blue'
+        else:
+            # Must be a bond
+            text = f"Bond {i}: "
+            if 'O-linked' in segment.get('bond_after', ''):
+                text += "ester"
+            elif 'N-Me' in segment.get('bond_after', ''):
+                text += "peptide (N-methylated)"
+            else:
+                text += "peptide"
+            color = 'red'
+        
+        # Add segment analysis
+        ax_detail.text(0.05, y, text, fontsize=12, color=color)
+        ax_detail.text(0.5, y, f"SMILES: {segment.get('content', '')}", fontsize=10, color='gray')
+    
+    # If cyclic, add connection indicator
+    if sequence.startswith('cyclo('):
+        ax_struct.annotate('', xy=(9.5, y_pos), xytext=(0.5, y_pos),
+                          arrowprops=dict(arrowstyle='<->', color='red', lw=2))
+        ax_struct.text(5, y_pos+0.3, 'Cyclic Connection', 
+                      ha='center', color='red', fontsize=14)
+    
+    # Add titles and adjust layout
+    ax_struct.set_title("Peptide Structure Overview", pad=20)
+    ax_detail.set_title("Segment Analysis Breakdown", pad=20)
+    
+    # Remove axes
+    for ax in [ax_struct, ax_detail]:
+        ax.set_xticks([])
+        ax.set_yticks([])
+        ax.axis('off')
+    
+    plt.tight_layout()
+    return fig
+
+class PeptideStructureGenerator:
+    """A class to generate 3D structures of peptides using different embedding methods"""
+    
+    @staticmethod
+    def prepare_molecule(smiles):
+        """Prepare molecule with proper hydrogen handling"""
+        mol = Chem.MolFromSmiles(smiles, sanitize=False)
+        if mol is None:
+            raise ValueError("Failed to create molecule from SMILES")
+        
+        # Calculate valence for each atom
+        for atom in mol.GetAtoms():
+            atom.UpdatePropertyCache(strict=False)
+        
+        # Sanitize with reduced requirements
+        Chem.SanitizeMol(mol, 
+                        sanitizeOps=Chem.SANITIZE_FINDRADICALS|
+                                  Chem.SANITIZE_KEKULIZE|
+                                  Chem.SANITIZE_SETAROMATICITY|
+                                  Chem.SANITIZE_SETCONJUGATION|
+                                  Chem.SANITIZE_SETHYBRIDIZATION|
+                                  Chem.SANITIZE_CLEANUPCHIRALITY)
+        
+        mol = Chem.AddHs(mol)
+        return mol
+
+    @staticmethod
+    def get_etkdg_params(attempt=0):
+        """Get ETKDG parameters with optional modifications based on attempt number"""
+        params = AllChem.ETKDGv3()
+        params.randomSeed = -1
+        params.maxIterations = 200
+        params.numThreads = 4  # Reduced for web interface
+        params.useBasicKnowledge = True
+        params.enforceChirality = True
+        params.useExpTorsionAnglePrefs = True
+        params.useSmallRingTorsions = True
+        params.useMacrocycleTorsions = True
+        params.ETversion = 2
+        params.pruneRmsThresh = -1
+        params.embedRmsThresh = 0.5
+        
+        if attempt > 10:
+            params.bondLength = 1.5 + (attempt - 10) * 0.02
+            params.useExpTorsionAnglePrefs = False
+        
+        return params
+
+    def generate_structure_etkdg(self, smiles, max_attempts=20):
+        """Generate 3D structure using ETKDG without UFF optimization"""
+        success = False
+        mol = None
+        
+        for attempt in range(max_attempts):
+            try:
+                mol = self.prepare_molecule(smiles)
+                params = self.get_etkdg_params(attempt)
+                
+                if AllChem.EmbedMolecule(mol, params) == 0:
+                    success = True
+                    break
+            except Exception as e:
+                continue
+        
+        if not success:
+            raise ValueError("Failed to generate structure with ETKDG")
+        
+        return mol
+
+    def generate_structure_uff(self, smiles, max_attempts=20):
+        """Generate 3D structure using ETKDG followed by UFF optimization"""
+        best_mol = None
+        lowest_energy = float('inf')
+        
+        for attempt in range(max_attempts):
+            try:
+                test_mol = self.prepare_molecule(smiles)
+                params = self.get_etkdg_params(attempt)
+                
+                if AllChem.EmbedMolecule(test_mol, params) == 0:
+                    res = AllChem.UFFOptimizeMolecule(test_mol, maxIters=2000, 
+                                                     vdwThresh=10.0, confId=0,
+                                                     ignoreInterfragInteractions=True)
+                    
+                    if res == 0:
+                        ff = AllChem.UFFGetMoleculeForceField(test_mol)
+                        if ff:
+                            current_energy = ff.CalcEnergy()
+                            if current_energy < lowest_energy:
+                                lowest_energy = current_energy
+                                best_mol = Chem.Mol(test_mol)
+            except Exception:
+                continue
+        
+        if best_mol is None:
+            raise ValueError("Failed to generate optimized structure")
+        
+        return best_mol
+
+    @staticmethod
+    def mol_to_sdf_bytes(mol):
+        """Convert RDKit molecule to SDF file bytes"""
+        # First write to StringIO in text mode
+        sio = StringIO()
+        writer = Chem.SDWriter(sio)
+        writer.write(mol)
+        writer.close()
+        
+        # Convert the string to bytes
+        return sio.getvalue().encode('utf-8')
+
+def process_input(smiles_input=None, file_obj=None, show_linear=False, 
+                 show_segment_details=False, generate_3d=False, use_uff=False):
+    """Process input and create visualizations using PeptideAnalyzer"""
+    analyzer = PeptideAnalyzer()
+    temp_dir = tempfile.mkdtemp() if generate_3d else None
+    structure_files = []
+    
+    # Handle direct SMILES input
+    if smiles_input:
+        smiles = smiles_input.strip()
+        
+        # First check if it's a peptide using analyzer's method
+        if not analyzer.is_peptide(smiles):
+            return "Error: Input SMILES does not appear to be a peptide structure.", None, None
+                
+        try:
+            # Create molecule
+            mol = Chem.MolFromSmiles(smiles)
+            if mol is None:
+                return "Error: Invalid SMILES notation.", None, None
+            
+            # Generate 3D structures if requested
+            if generate_3d:
+                generator = PeptideStructureGenerator()
+                
+                try:
+                    # Generate ETKDG structure
+                    mol_etkdg = generator.generate_structure_etkdg(smiles)
+                    etkdg_path = os.path.join(temp_dir, "structure_etkdg.sdf")
+                    writer = Chem.SDWriter(etkdg_path)
+                    writer.write(mol_etkdg)
+                    writer.close()
+                    structure_files.append(etkdg_path)
+                    
+                    # Generate UFF structure if requested
+                    if use_uff:
+                        mol_uff = generator.generate_structure_uff(smiles)
+                        uff_path = os.path.join(temp_dir, "structure_uff.sdf")
+                        writer = Chem.SDWriter(uff_path)
+                        writer.write(mol_uff)
+                        writer.close()
+                        structure_files.append(uff_path)
+                
+                except Exception as e:
+                    return f"Error generating 3D structures: {str(e)}", None, None, None
+                
+            # Use analyzer to get sequence
+            segments = analyzer.split_on_bonds(smiles)
+            
+            # Process segments and build sequence
+            sequence_parts = []
+            output_text = ""
+            
+            # Only include segment analysis in output if requested
+            if show_segment_details:
+                output_text += "Segment Analysis:\n"
+                for i, segment in enumerate(segments):
+                    output_text += f"\nSegment {i}:\n"
+                    output_text += f"Content: {segment['content']}\n"
+                    output_text += f"Bond before: {segment.get('bond_before', 'None')}\n"
+                    output_text += f"Bond after: {segment.get('bond_after', 'None')}\n"
+                    
+                    residue, mods = analyzer.identify_residue(segment)
+                    if residue:
+                        if mods:
+                            sequence_parts.append(f"{residue}({','.join(mods)})")
+                        else:
+                            sequence_parts.append(residue)
+                        output_text += f"Identified as: {residue}\n"
+                        output_text += f"Modifications: {mods}\n"
+                    else:
+                        output_text += f"Warning: Could not identify residue in segment: {segment['content']}\n"
+                output_text += "\n"
+            else:
+                # Just build sequence without detailed analysis in output
+                for segment in segments:
+                    residue, mods = analyzer.identify_residue(segment)
+                    if residue:
+                        if mods:
+                            sequence_parts.append(f"{residue}({','.join(mods)})")
+                        else:
+                            sequence_parts.append(residue)
+            
+            # Check if cyclic using analyzer's method
+            is_cyclic, peptide_cycles, aromatic_cycles = analyzer.is_cyclic(smiles)
+            three_letter = '-'.join(sequence_parts)
+            one_letter = ''.join(analyzer.three_to_one.get(aa.split('(')[0], 'X') for aa in sequence_parts)
+        
+            if is_cyclic:
+                three_letter = f"cyclo({three_letter})"
+                one_letter = f"cyclo({one_letter})"
+            
+            # Create cyclic structure visualization
+            img_cyclic = annotate_cyclic_structure(mol, three_letter)
+            
+            # Create linear representation if requested
+            img_linear = None
+            if show_linear:
+                fig_linear = create_enhanced_linear_viz(three_letter, smiles)
+                buf = BytesIO()
+                fig_linear.savefig(buf, format='png', bbox_inches='tight', dpi=300)
+                buf.seek(0)
+                img_linear = Image.open(buf)
+                plt.close(fig_linear)
+
+            # Add summary to output
+            summary = "Summary:\n"
+            summary += f"Sequence: {three_letter}\n"
+            summary += f"One-letter code: {one_letter}\n"
+            summary += f"Is Cyclic: {'Yes' if is_cyclic else 'No'}\n"
+            #if is_cyclic:
+                #summary += f"Peptide Cycles: {', '.join(peptide_cycles)}\n"
+                #summary += f"Aromatic Cycles: {', '.join(aromatic_cycles)}\n"
+            
+            if structure_files:
+                summary += "\n3D Structures Generated:\n"
+                for filepath in structure_files:
+                    summary += f"- {os.path.basename(filepath)}\n"
+            
+            return summary + output_text, img_cyclic, img_linear, structure_files if structure_files else None
+            
+        except Exception as e:
+            return f"Error processing SMILES: {str(e)}", None, None, None
+    
+    # Handle file input
+    if file_obj is not None:
+        try:
+            # Handle file content
+            if hasattr(file_obj, 'name'):
+                with open(file_obj.name, 'r') as f:
+                    content = f.read()
+            else:
+                content = file_obj.decode('utf-8') if isinstance(file_obj, bytes) else str(file_obj)
+            
+            output_text = ""
+            for line in content.splitlines():
+                smiles = line.strip()
+                if smiles:
+                    # Check if it's a peptide
+                    if not analyzer.is_peptide(smiles):
+                        output_text += f"Skipping non-peptide SMILES: {smiles}\n"
+                        continue
+                    
+                    # Process this SMILES
+                    segments = analyzer.split_on_bonds(smiles)
+                    sequence_parts = []
+                    
+                    # Add segment details if requested
+                    if show_segment_details:
+                        output_text += f"\nSegment Analysis for SMILES: {smiles}\n"
+                        for i, segment in enumerate(segments):
+                            output_text += f"\nSegment {i}:\n"
+                            output_text += f"Content: {segment['content']}\n"
+                            output_text += f"Bond before: {segment.get('bond_before', 'None')}\n"
+                            output_text += f"Bond after: {segment.get('bond_after', 'None')}\n"
+                            residue, mods = analyzer.identify_residue(segment)
+                            if residue:
+                                if mods:
+                                    sequence_parts.append(f"{residue}({','.join(mods)})")
+                                else:
+                                    sequence_parts.append(residue)
+                                output_text += f"Identified as: {residue}\n"
+                                output_text += f"Modifications: {mods}\n"
+                    else:
+                        for segment in segments:
+                            residue, mods = analyzer.identify_residue(segment)
+                            if residue:
+                                if mods:
+                                    sequence_parts.append(f"{residue}({','.join(mods)})")
+                                else:
+                                    sequence_parts.append(residue)
+                    
+                    # Get cyclicity and create sequence
+                    is_cyclic, peptide_cycles, aromatic_cycles = analyzer.is_cyclic(smiles)
+                    sequence = f"cyclo({'-'.join(sequence_parts)})" if is_cyclic else '-'.join(sequence_parts)
+                    
+                    output_text += f"\nSummary for SMILES: {smiles}\n"
+                    output_text += f"Sequence: {sequence}\n"
+                    output_text += f"Is Cyclic: {'Yes' if is_cyclic else 'No'}\n"
+                    if is_cyclic:
+                        output_text += f"Peptide Cycles: {', '.join(peptide_cycles)}\n"
+                        #output_text += f"Aromatic Cycles: {', '.join(aromatic_cycles)}\n"
+                    output_text += "-" * 50 + "\n"
+            
+            return output_text, None, None
+            
+        except Exception as e:
+            return f"Error processing file: {str(e)}", None, None
+    
+    return "No input provided.", None, None
+
+

smiles/evaluation/run_metrics.py

@@ -0,0 +1,282 @@
+#!/usr/bin/env python3
+"""
+Script to run peptide metrics analysis on SMILES text files
+
+Usage:
+    python run_metrics.py input.txt [--target TARGET] [--output OUTPUT_DIR] [--device DEVICE]
+
+Arguments:
+    input.txt: Text file containing SMILES sequences (one per line)
+    --target: Target protein (GNRHR / ERBB2 / NRP1). If not specified, will try to detect from filename
+    --output: Output directory for results (default: same as input file directory)
+    --device: Device for computation (cuda or cpu, default: cuda if available)
+"""
+
+import argparse
+import os
+import sys
+import pandas as pd
+import tempfile
+import torch
+import shutil
+from pathlib import Path
+
+# Add current directory to path for imports
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+from smiles.evaluation.metrics import Metrics
+
+def read_smiles_file(input_path):
+    """
+    Read SMILES sequences from a text file
+    
+    Args:
+        input_path: Path to input text file
+        
+    Returns:
+        List of SMILES strings
+    """
+    smiles_list = []
+    
+    with open(input_path, 'r') as f:
+        for line_num, line in enumerate(f, 1):
+            line = line.strip()
+            if line and not line.startswith('#'):  # Skip empty lines and comments
+                smiles_list.append(line)
+    
+    print(f"Read {len(smiles_list)} SMILES sequences from {input_path}")
+    return smiles_list
+
+def create_temp_csv(smiles_list, target=None):
+    """
+    Create a temporary CSV file with SMILES data
+    
+    Args:
+        smiles_list: List of SMILES strings
+        target: Target protein name (for filename)
+        
+    Returns:
+        Path to temporary CSV file
+    """
+    # Create DataFrame
+    df = pd.DataFrame({'SMILES': smiles_list})
+    
+    # Create temporary file
+    suffix = f"_{target}" if target else ""
+    temp_file = tempfile.NamedTemporaryFile(
+        mode='w', 
+        suffix=f"{suffix}_smiles.csv", 
+        delete=False
+    )
+    
+    # Save to CSV
+    df.to_csv(temp_file.name, index=False)
+    temp_file.close()
+    
+    print(f"Created temporary CSV: {temp_file.name}")
+    return temp_file.name
+
+def determine_target(input_path, target_arg):
+    """
+    Determine target protein from filename or argument
+    
+    Args:
+        input_path: Path to input file
+        target_arg: Target specified via command line argument
+        
+    Returns:
+        Target protein name or None
+    """
+    if target_arg:
+        target = target_arg.upper()
+        if target not in ['GNRHR', 'ERBB2', 'NRP1']:
+            print(f"Warning: Unknown target '{target}'. Supported targets: GNRHR, ERBB2, NRP1")
+            return None
+        return target
+    
+    # Try to detect from filename
+    filename = os.path.basename(input_path).upper()
+    if 'GNRHR' in filename:
+        return 'GNRHR'
+    elif 'ERBB2' in filename:
+        return 'ERBB2'
+    elif 'NRP1' in filename:
+        return 'NRP1'
+    else:
+        print("Warning: Could not determine target from filename. Specify with --target")
+        print("Supported targets: GNRHR, ERBB2, NRP1")
+        return None
+
+def save_results(results, output_path):
+    """
+    Save analysis results to a text file
+    
+    Args:
+        results: Dictionary of analysis results
+        output_path: Path to output file
+    """
+    with open(output_path, 'w') as f:
+        f.write("PEPTIDE METRICS ANALYSIS RESULTS\n")
+        f.write("=" * 50 + "\n\n")
+        
+        # Basic statistics
+        f.write("BASIC STATISTICS:\n")
+        f.write(f"  Total sequences: {results.get('total_sequences', 0)}\n")
+        f.write(f"  Valid peptides: {results.get('valid_peptides', 0)}\n")
+        f.write(f"  Validity rate: {results.get('validity_rate', 0):.1%}\n")
+        f.write(f"  Valid conjugates: {results.get('valid_conjugates', 0)}\n")
+        f.write(f"  Conjugation success rate: {results.get('conjugation_success_rate', 0):.1%}\n")
+        f.write(f"  Target protein: {results.get('target', 'Unknown')}\n\n")
+        
+        # Binding affinity
+        if 'binding_affinity_all_mean' in results:
+            f.write("BINDING AFFINITY (All valid peptides):\n")
+            f.write(f"  Mean: {results['binding_affinity_all_mean']:.4f}\n")
+            f.write(f"  Std: {results['binding_affinity_all_std']:.4f}\n")
+            f.write(f"  Min: {results['binding_affinity_all_min']:.4f}\n")
+            f.write(f"  Max: {results['binding_affinity_all_max']:.4f}\n\n")
+        
+        if 'binding_affinity_conjugatable_mean' in results:
+            f.write("BINDING AFFINITY (Conjugatable peptides):\n")
+            f.write(f"  Mean: {results['binding_affinity_conjugatable_mean']:.4f}\n")
+            f.write(f"  Std: {results['binding_affinity_conjugatable_std']:.4f}\n")
+            f.write(f"  Min: {results['binding_affinity_conjugatable_min']:.4f}\n")
+            f.write(f"  Max: {results['binding_affinity_conjugatable_max']:.4f}\n\n")
+        
+        # Raw results (for programmatic access)
+        f.write("RAW RESULTS (JSON format):\n")
+        f.write(str(results))
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Run peptide metrics analysis on SMILES text files",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+    python run_metrics.py peptides.txt
+    python run_metrics.py gnrhr_peptides.txt --target GNRHR
+    python run_metrics.py peptides.txt --target ERBB2 --output ./results
+    python run_metrics.py peptides.txt --target NRP1 --device cpu
+        """
+    )
+    
+    parser.add_argument('input_file', type=str,
+                        help='Text file containing SMILES sequences (one per line)')
+    parser.add_argument('--target', type=str, choices=['GNRHR', 'ERBB2', 'NRP1'], default=None,
+                        help='Target protein (GNRHR, ERBB2, NRP1). If not specified, will try to detect from filename')
+    parser.add_argument('--output', type=str, default=None,
+                        help='Output directory for results (default: same as input file directory)')
+    parser.add_argument('--device', type=str, default='auto',
+                        help='Device for computation (cuda, cpu, or auto)')
+    parser.add_argument('--no-csv', action='store_true',
+                        help='Skip saving detailed CSV results')
+    
+    args = parser.parse_args()
+    
+    # Validate input file
+    if not os.path.exists(args.input_file):
+        print(f"Error: Input file '{args.input_file}' does not exist")
+        return 1
+    
+    # Determine device
+    if args.device == 'auto':
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    else:
+        device = args.device
+    
+    print(f"Using device: {device}")
+    
+    # Set up output directory
+    if args.output:
+        output_dir = args.output
+        os.makedirs(output_dir, exist_ok=True)
+    else:
+        output_dir = os.path.dirname(os.path.abspath(args.input_file))
+    
+    # Generate output filenames
+    base_name = Path(args.input_file).stem
+    results_file = os.path.join(output_dir, f"{base_name}_metrics.txt")
+    
+    try:
+        # Read SMILES sequences
+        print("Reading SMILES sequences...")
+        smiles_list = read_smiles_file(args.input_file)
+        
+        if not smiles_list:
+            print("Error: No SMILES sequences found in input file")
+            return 1
+        
+        # Determine target
+        target = determine_target(args.input_file, args.target)
+        if not target:
+            print("Error: Could not determine target protein. Use --target to specify.")
+            return 1
+        
+        # Create temporary CSV file
+        temp_csv_path = create_temp_csv(smiles_list, target)
+        
+        try:
+            # Initialize metrics calculator
+            print("Initializing metrics calculator...")
+            metrics_calculator = Metrics(
+                input_type='peptide',
+                device=device
+            )
+            
+            # Run analysis
+            print("Running peptide analysis...")
+            results = metrics_calculator.get_metrics_for_peptides(
+                temp_csv_path, 
+                save_csv=(not args.no_csv)
+            )
+            
+            # Save results
+            save_results(results, results_file)
+            print(f"\nResults saved to: {results_file}")
+            
+            # Handle detailed CSV if it was created
+            if not args.no_csv and results.get('csv_path'):
+                source_csv = results['csv_path']
+                
+                if os.path.exists(source_csv):
+                    final_csv = os.path.join(output_dir, f"{base_name}_conjugate_analysis.csv")
+                    
+                    # If output_dir is different from source directory, move the file
+                    if os.path.abspath(os.path.dirname(source_csv)) != os.path.abspath(output_dir):
+                        try:
+                            shutil.move(source_csv, final_csv)
+                            print(f"Detailed CSV moved to: {final_csv}")
+                        except Exception as e:
+                            # Fallback: copy then delete
+                            try:
+                                shutil.copy2(source_csv, final_csv)
+                                os.remove(source_csv)
+                                print(f"Detailed CSV copied to: {final_csv}")
+                            except Exception as e2:
+                                print(f"Warning: Failed to move CSV file: {e2}")
+                                print(f"CSV remains at: {source_csv}")
+                    else:
+                        print(f"Detailed CSV already saved at: {source_csv}")
+                else:
+                    print(f"Warning: Expected detailed CSV not found at: {source_csv}")
+            elif not args.no_csv:
+                print("Warning: No detailed CSV was created (possibly no valid conjugates found)")
+            
+        finally:
+            # Clean up temporary CSV
+            if os.path.exists(temp_csv_path):
+                os.unlink(temp_csv_path)
+                print(f"Cleaned up temporary file: {temp_csv_path}")
+        
+        print("\nAnalysis complete!")
+        return 0
+        
+    except Exception as e:
+        print(f"Error during analysis: {e}")
+        import traceback
+        traceback.print_exc()
+        return 1
+
+if __name__ == "__main__":
+    exit_code = main()
+    sys.exit(exit_code)

smiles/evaluation/run_metrics_uaa.py

@@ -0,0 +1,590 @@
+#!/usr/bin/env python3
+"""
+run_metrics_uaa.py
+==================
+
+Re-evaluates already-generated UAA peptide-drug conjugates after stripping
+Fmoc-strategy side-chain protecting groups, and reports protected-vs-deprotected
+correlation for the binding-affinity oracle.
+
+This script does NOT re-run guidance / generation. It works entirely on the
+existing `*_conjugate_analysis.csv` files produced by run_metrics.py.
+
+Pipeline per input CSV
+----------------------
+  1. Load CSV (must contain `peptide_smiles`, `conjugate_smiles`, `binding_affinity`).
+  2. Strip protecting groups from `peptide_smiles` and `conjugate_smiles`
+     using a small library of RDKit reactions (Mmt, Trt, Boc, OtBu ester,
+     OtBu/OBn ether, OBn ester, OBn-phosphate, Pbf).
+  3. Recompute molecular properties (MW, LogP, TPSA, HBD/HBA, rotbonds,
+     heavy atoms) on the deprotected conjugate.
+  4. Re-tokenize the deprotected peptide SMILES and re-score with the same
+     `BindingAffinity` model used originally. (Skip with --no-rescore.)
+  5. Optionally pair each row with the matching NCAA sequence string from the
+     original `.txt` file (--ncaa-txt or auto-detect).
+  6. Save `<input>_deprot.csv` with both protected and deprotected columns,
+     plus a `uaa_deprotection_summary.csv` aggregating across all inputs:
+        - n
+        - mean / std of binding (×10 pKd) protected and deprotected
+        - Pearson and Spearman correlation across the two
+        - per-rule strip counts on conjugates
+
+Usage
+-----
+    python run_metrics_uaa.py INPUT_CSV [INPUT_CSV ...] \
+        [--output OUTPUT_DIR] [--target {ERBB2,GNRHR,NRP1}] \
+        [--device {cuda,cpu,auto}] [--no-rescore] \
+        [--vocab PATH] [--splits PATH] [--ncaa-txt-dir PATH]
+
+Examples
+--------
+    # Re-evaluate one CSV
+    python -m smiles.evaluation.run_metrics_uaa \
+        smiles/results/uaa_results/cupofmoca_guided_uaa_ERBB2_SSCpt_k8_len10_smiles_conjugate_analysis.csv \
+        --output smiles/results/uaa_results_deprot
+
+    # Re-evaluate everything (see scripts/analyze_smiles_uaa_deprotect.sh)
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from rdkit import Chem, RDLogger
+from rdkit.Chem import AllChem, Descriptors
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+RDLogger.DisableLog("rdApp.*")
+
+
+# ---------------------------------------------------------------------------
+# Target protein sequences (must match metrics.py - copied here so this script
+# is self-contained and doesn't trigger the full Metrics(...) construction
+# which loads a 30K reference SMILES set we don't need.)
+# ---------------------------------------------------------------------------
+TARGET_SEQUENCES = {
+    'GNRHR': (
+        'MANSASPEQNQNHCSAINNSIPLMQGNLPTLTLSGKIRVTVTFFLFLLSATFNASFLLKLQKWTQKKEKGKK'
+        'LSRMKLLLKHLTLANLLETLIVMPLDGMWNITVQWYAGELLCKVLSYLKLFSMYAPAFMMVVISLDRSLAIT'
+        'RPLALKSNSKVGQSMVGLAWILSSVFAGPQLYIFRMIHLADSSGQTKVFSQCVTHCSFSQWWHQAFYNFFTF'
+        'SCLFIIPLFIMLICNAKIIFTLTRVLHQDPHELQLNQSKNNIPRARLKTLKMTVAFATSFTVCWTPYYVLGI'
+        'WYWFDPEMLNRLSDPVNHFFFLFAFLNPCFDPLIYGYFSL'
+    ),
+    'ERBB2': (
+        'MELAALCRWGLLLALLPPGAASTQVCTGTDMKLRLPASPETHLDMLRHLYQGCQVVQGNLELTYLPTNASLS'
+        'FLQDIQEVQGYVLIAHNQVRQVPLQRLRIVRGTQLFEDNYALAVLDNGDPLNNTTPVTGASPGGLRELQLRS'
+        'LTEILKGGVLIQRNPQLCYQDTILWKDIFHKNNQLALTLIDTNRSRACHPCSPMCKGSRCWGESSEDCQSLT'
+        'RTVCAGGCARCKGPLPTDCCHEQCAAGCTGPKHSDCLACLHFNHSGICELHCPALVTYNTDTFESMPNPEGR'
+        'YTFGASCVTACPYNYLSTDVGSCTLVCPLHNQEVTAEDGTQRCEKCSKPCARVCYGLGMEHLREVRAVTSAN'
+        'IQEFAGCKKIFGSLAFLPESFDGDPASNTAPLQPEQLQVFETLEEITGYLYISAWPDSLPDLSVFQNLQVIR'
+        'GRILHNGAYSLTLQGLGISWLGLRSLRELGSGLALIHHNTHLCFVHTVPWDQLFRNPHQALLHTANRPEDEC'
+        'VGEGLACHQLCARGHCWGPGPTQCVNCSQFLRGQECVEECRVLQGLPREYVNARHCLPCHPECQPQNGSVTC'
+        'FGPEADQCVACAHYKDPPFCVARCPSGVKPDLSYMPIWKFPDEEGACQPCPINCTHSCVDLDDKGCPAEQRA'
+        'SPLTSIISAVVGILLVVVLGVVFGILIKRRQQKIRKYTMRRLLQETELVEPLTPSGAMPNQAQMRILKETEL'
+        'RKVKVLGSGAFGTVYKGIWIPDGENVKIPVAIKVLRENTSPKANKEILDEAYVMAGVGSPYVSRLLGICLTS'
+        'TVQLVTQLMPYGCLLDHVRENRGRLGSQDLLNWCMQIAKGMSYLEDVRLVHRDLAARNVLVKSPNHVKITDF'
+        'GLARLLDIDETEYHADGGKVPIKWMALESILRRRFTHQSDVWSYGVTVWELMTFGAKPYDGIPAREIPDLLE'
+        'KGERLPQPPICTIDVYMIMVKCWMIDSECRPRFRELVSEFSRMARDPQRFVVIQNEDLGPASPLDSTFYRSL'
+        'LEDDDMGDLVDAEEYLVPQQGFFCPDPAPGAGGMVHHRHRSSSTRSGGGDLTLGLEPSEEEAPRSPLAPSEG'
+        'AGSDVFDGDLGMGAAKGLQSLPTHDPSPLQRYSEDPTVPLPSETDGYVAPLTCSPQPEYVNQPDVRPQPPSP'
+        'REGPLPAARPAGATLERPKTLSPGKNGVVKDVFAFGGAVENPEYLTPQGGAAPQPHPPPAFSPAFDNLYYWD'
+        'QDPPERGAPPSTFKGTPTAENPEYLGLDVPV'
+    ),
+    'NRP1': (
+        'MERGLPLLCAVLALVLAPAGAFRNDKCGDTIKIESPGYLTSPGYPHSYHPSEKCEWLIQAPDPYQRIMINFN'
+        'PHFDLEDRDCKYDYVEVFDGENENGHFRGKFCGKIAPPPVVSSGPFLFIKFVSDYETHGAGFSIRYEIFKRG'
+        'PECSQNYTTPSGVIKSPGFPEKYPNSLECTYIVFVPKMSEIILEFESFDLEPDSNPPGGMFCRYDRLEIWDG'
+        'FPDVGPHIGRYCGQKTPGRIRSSSGILSMVFYTDSAIAKEGFSANYSVLQSSVSEDFKCMEALGMESGEIHS'
+        'DQITASSQYSTNWSAERSRLNYPENGWTPGEDSYREWIQVDLGLLRFVTAVGTQGAISKETKKKYYVKTYKI'
+        'DVSSNGEDWITIKEGNKPVLFQGNTNPTDVVVAVFPKPLITRFVRIKPATWETGISMRFEVYGCKITDYPCS'
+        'GMLGMVSGLISDSQITSSNQGDRNWMPENIRLVTSRSGWALPPAPHSYINEWLQIDLGEEKIVRGIIIQGGK'
+        'HRENKVFMRKFKIGYSNNGSDWKMIMDDSKRKAKSFEGNNNYDTPELRTFPALSTRFIRIYPERATHGGLGL'
+        'RMELLGCEVEAPTAGPTTPNGNLVDECDDDQANCHSGTGDDFQLTGGTTVLATEKPTVIDSTIQSEFPTYGF'
+        'NCEFGWGSHKTFCHWEHDNHVQLKWSVLTSKTGPIQDHTGDGNFIYSQADENQKGKVARLVSPVVYSQNSAH'
+        'CMTFWYHMSGSHVGTLRVKLRYQKPEEYDQLVWMAIGHQGDHWKEGRVLLHKSLKLYQVIFEGEIGKGNLGG'
+        'IAVDDISINNHISQEDCAKPADLDKKNPEIKIDETGSTPGYEGEGEGDKNISRKPGNVLKTLDPILITIIAM'
+        'SALGVLLGAVCGVVLYCACWHNGMSERNLSALENYNFELVDGVKLKKDKLNTQSTYSEA'
+    ),
+}
+
+
+# ---------------------------------------------------------------------------
+# Deprotection rules (Fmoc SPPS side-chain protecting groups).
+# Each rule is (label, reaction-SMARTS).  Reactions are applied iteratively
+# until no rule matches; on each match we re-sanitize the molecule.
+# Patterns are intentionally specific so they don't accidentally rewrite
+# atoms in the linker/payload (Aoa, SS, DSPA, Cpt, Dau, SN-38 do not contain
+# tBu, trityl, or benzyl-phosphate groups).
+# ---------------------------------------------------------------------------
+DEPROTECTION_RULES = [
+    # 4-methoxytrityl on Cys: -S-C(Ph)(Ph)(p-anisyl) -> -SH
+    ("Mmt-S",
+     "[#16:1][C](-c1ccccc1)(-c2ccccc2)-c3ccc(O[CH3])cc3>>[#16:1]"),
+
+    # Trityl on Cys / His-N: -S-C(Ph)3 -> -SH ; -N-C(Ph)3 -> -NH
+    ("Trt-S/N",
+     "[#16,#7:1][C](-c1ccccc1)(-c2ccccc2)-c3ccccc3>>[*:1]"),
+
+    # Boc on amine: -N-C(=O)-O-C(CH3)3 -> -NH
+    ("Boc-N",
+     "[N:1][C](=O)O[C]([CH3])([CH3])[CH3]>>[N:1]"),
+
+    # OtBu ester on COOH: -C(=O)-O-C(CH3)3 -> -C(=O)OH
+    ("OtBu-ester",
+     "[C:1](=O)O[C]([CH3])([CH3])[CH3]>>[C:1](=O)O"),
+
+    # OtBu ether on Ser/Thr/Tyr: -O-C(CH3)3 -> -OH
+    ("OtBu-ether",
+     "[O:1][C]([CH3])([CH3])[CH3]>>[O:1]"),
+
+    # Bn on phosphate / phosphonate: -O-P(=O)(O)-O-CH2-Ph -> -O-P(=O)(O)-OH
+    ("Bn-phosphate",
+     "[#15:1](=O)([O])O[CH2]c1ccccc1>>[#15:1](=O)([O])O"),
+
+    # Bn ester: -C(=O)-O-CH2-Ph -> -C(=O)OH
+    ("Bn-ester",
+     "[C:1](=O)O[CH2]c1ccccc1>>[C:1](=O)O"),
+
+    # Bn ether (catch-all, after phosphate/ester): -O-CH2-Ph -> -OH
+    ("Bn-ether",
+     "[O:1][CH2]c1ccccc1>>[O:1]"),
+
+    # Pbf on Arg: -N-S(=O)(=O)-pentamethyldihydrobenzofuran -> -N
+    ("Pbf-N",
+     "[N:1][S](=O)(=O)c1c([CH3])c2c(c1)O[C]([CH3])([CH3])[C]2([CH3])[CH3]>>[N:1]"),
+
+    # 2,4-dimethoxybenzyl (Dmb) amide / amine cap: -N-CH2-Ar(2,4-OMe) -> -NH
+    ("Dmb-N",
+     "[N:1][CH2]c1cc(O[CH3])cc(O[CH3])c1>>[N:1]"),
+    ("Dmb-N-alt",
+     "[N:1][CH2]c1c(O[CH3])cc(O[CH3])cc1>>[N:1]"),
+]
+
+
+def deprotect_mol(mol, rules, max_iters=30):
+    """Apply deprotection reactions until convergence.
+
+    Returns
+    -------
+    final_mol : Chem.Mol or None
+    counts    : dict {label: int}  -- how many times each rule fired.
+    """
+    counts = {label: 0 for label, _ in rules}
+    if mol is None:
+        return None, counts
+
+    # Pre-compile reactions
+    compiled = []
+    for label, rxn_sm in rules:
+        try:
+            rxn = AllChem.ReactionFromSmarts(rxn_sm)
+            if rxn.GetNumReactantTemplates() == 1:
+                compiled.append((label, rxn))
+        except Exception:
+            pass
+
+    for _ in range(max_iters):
+        changed = False
+        for label, rxn in compiled:
+            try:
+                products = rxn.RunReactants((mol,))
+            except Exception:
+                continue
+            if not products:
+                continue
+            new_mol = products[0][0]
+            try:
+                Chem.SanitizeMol(new_mol)
+            except Exception:
+                continue
+            try:
+                if Chem.MolToSmiles(new_mol) == Chem.MolToSmiles(mol):
+                    continue
+            except Exception:
+                continue
+            mol = new_mol
+            counts[label] += 1
+            changed = True
+            break  # restart rule list from the top
+        if not changed:
+            break
+    return mol, counts
+
+
+def deprotect_smiles(smi, rules):
+    """Convenience wrapper: SMILES -> (deprot_smi, counts)."""
+    empty = {label: 0 for label, _ in rules}
+    if not isinstance(smi, str) or not smi:
+        return None, empty
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return None, empty
+    final_mol, counts = deprotect_mol(mol, rules)
+    if final_mol is None:
+        return None, counts
+    try:
+        return Chem.MolToSmiles(final_mol), counts
+    except Exception:
+        return None, counts
+
+
+def compute_props(smi):
+    if not isinstance(smi, str) or not smi:
+        return None
+    mol = Chem.MolFromSmiles(smi)
+    if mol is None:
+        return None
+    return {
+        'mw_dep':      float(Descriptors.MolWt(mol)),
+        'logp_dep':    float(Descriptors.MolLogP(mol)),
+        'tpsa_dep':    float(Descriptors.TPSA(mol)),
+        'hbd_dep':     int(Descriptors.NumHDonors(mol)),
+        'hba_dep':     int(Descriptors.NumHAcceptors(mol)),
+        'rotbonds_dep': int(Descriptors.NumRotatableBonds(mol)),
+        'heavy_dep':   int(mol.GetNumHeavyAtoms()),
+    }
+
+
+def determine_target(input_path, target_arg):
+    if target_arg:
+        t = target_arg.upper()
+        if t in TARGET_SEQUENCES:
+            return t
+    fn = os.path.basename(input_path).upper()
+    for t in TARGET_SEQUENCES:
+        if t in fn:
+            return t
+    return None
+
+
+def score_binding_smiles(model, tokenizer, smi_list, device, batch_size=32):
+    """Score a list of peptide SMILES with the BindingAffinity model.
+
+    Returns a numpy array of length len(smi_list); failed entries are NaN.
+    """
+    import torch
+    out = np.full(len(smi_list), np.nan, dtype=float)
+    if not smi_list:
+        return out
+
+    valid_idx = []
+    valid_tok = []
+    for i, s in enumerate(smi_list):
+        if not isinstance(s, str) or not s:
+            continue
+        try:
+            t = tokenizer(s)
+        except Exception:
+            continue
+        if hasattr(t, 'input_ids'):
+            t = t.input_ids
+        if isinstance(t, list):
+            t = torch.tensor(t)
+        if not torch.is_tensor(t):
+            continue
+        if t.dim() == 2:
+            t = t.squeeze(0)
+        if t.dim() != 1 or t.numel() == 0:
+            continue
+        valid_idx.append(i)
+        valid_tok.append(t.long())
+
+    if not valid_tok:
+        return out
+
+    for j in range(0, len(valid_tok), batch_size):
+        chunk = valid_tok[j:j + batch_size]
+        max_len = max(int(t.numel()) for t in chunk)
+        pad = torch.zeros((len(chunk), max_len), dtype=torch.long)
+        for k, t in enumerate(chunk):
+            pad[k, :t.numel()] = t
+        try:
+            with torch.no_grad():
+                scores = model(pad.to(device)).cpu().numpy().reshape(-1)
+        except Exception as e:
+            print(f"  [score_binding] batch {j}: {e}")
+            continue
+        for k, idx in enumerate(valid_idx[j:j + batch_size]):
+            out[idx] = float(scores[k])
+    return out
+
+
+def maybe_load_ncaa_sequences(csv_path, ncaa_dir):
+    """Pair rows of the CSV with NCAA sequence strings from the matching
+    `<base>.txt` two-column file (NCAA<TAB>SMILES). Returns a list of
+    NCAA strings (None for unmatched rows) aligned to the CSV rows.
+
+    Lookup strategy:
+      1. Read all (NCAA, SMILES) pairs from the .txt.
+      2. For each row in the CSV, match by `peptide_smiles`.
+    """
+    base = Path(csv_path).stem
+    if base.endswith('_conjugate_analysis'):
+        base = base[: -len('_conjugate_analysis')]
+    candidates = []
+    if ncaa_dir:
+        candidates.append(os.path.join(ncaa_dir, base + '.txt'))
+    csv_dir = os.path.dirname(os.path.abspath(csv_path))
+    candidates.extend([
+        os.path.join(csv_dir, base + '.txt'),
+        os.path.join(csv_dir, '..', base + '.txt'),
+        os.path.join(os.path.dirname(csv_dir), base + '.txt'),
+    ])
+    # also try stripping _smiles suffix used by some files
+    if base.endswith('_smiles'):
+        b2 = base[: -len('_smiles')]
+        for d in [ncaa_dir, csv_dir, os.path.join(csv_dir, '..'), os.path.dirname(csv_dir)]:
+            if d:
+                candidates.append(os.path.join(d, b2 + '.txt'))
+
+    for path in candidates:
+        if path and os.path.exists(path):
+            try:
+                pairs = {}
+                with open(path, 'r', encoding='utf-8', errors='replace') as f:
+                    for line in f:
+                        line = line.rstrip('\n').rstrip('\r')
+                        if not line:
+                            continue
+                        if '\t' in line:
+                            ncaa, smi = line.split('\t', 1)
+                        else:
+                            # one-column SMILES file, no NCAA available
+                            return None
+                        pairs[smi] = ncaa
+                return path, pairs
+            except Exception:
+                continue
+    return None
+
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="UAA peptide-drug conjugate deprotection + re-scoring",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    p.add_argument('inputs', nargs='+',
+                   help='UAA *_conjugate_analysis.csv file(s) produced by run_metrics.py')
+    p.add_argument('--output', '-o', default=None,
+                   help='Output directory for *_deprot.csv files and summary '
+                        '(default: same directory as each input)')
+    p.add_argument('--target', default=None, choices=list(TARGET_SEQUENCES),
+                   help='Target protein (default: auto-detect from filename)')
+    p.add_argument('--device', default='auto',
+                   help='cuda / cpu / auto (default: auto)')
+    p.add_argument('--no-rescore', action='store_true',
+                   help='Skip binding rescoring (deprotect + recompute properties only)')
+    p.add_argument('--vocab', default='./smiles/smiles_tokenizer/new_vocab.txt')
+    p.add_argument('--splits', default='./smiles/smiles_tokenizer/new_splits.txt')
+    p.add_argument('--ncaa-txt-dir', default=None,
+                   help='Directory holding the original NCAA .txt files '
+                        '(NCAA<TAB>SMILES). Default: same dir as the CSV.')
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # Lazy heavy imports (only if needed)
+    if not args.no_rescore:
+        import torch
+        if args.device == 'auto':
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        else:
+            device = args.device
+        print(f"Device: {device}")
+
+        # Load tokenizer + binding model lazily
+        try:
+            from smiles.scoring.classifiers import BindingAffinity
+            from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+        except Exception as e:
+            print(f"Cannot import binding-affinity oracle modules: {e}")
+            print("Falling back to --no-rescore.")
+            args.no_rescore = True
+
+        if not args.no_rescore:
+            if not (os.path.exists(args.vocab) and os.path.exists(args.splits)):
+                print(f"Tokenizer vocab/splits not found "
+                      f"(--vocab={args.vocab}, --splits={args.splits}).")
+                print("Run from MoCA root, or pass --vocab/--splits. Skipping rescoring.")
+                args.no_rescore = True
+            else:
+                tokenizer = SMILES_SPE_Tokenizer(args.vocab, args.splits)
+    else:
+        device = 'cpu'
+
+    # Print rule list
+    print("Deprotection rules (in order, applied iteratively):")
+    for label, _ in DEPROTECTION_RULES:
+        print(f"  - {label}")
+
+    binding_models = {}
+    summary_rows = []
+    overall_strip_counts = {label: 0 for label, _ in DEPROTECTION_RULES}
+
+    for in_csv in args.inputs:
+        print(f"\n>>> {in_csv}")
+        if not os.path.exists(in_csv):
+            print(f"  missing, skipping")
+            continue
+        df = pd.read_csv(in_csv)
+        if df.empty:
+            print(f"  empty CSV, skipping")
+            continue
+        target = determine_target(in_csv, args.target)
+        if target is None:
+            print(f"  cannot determine target, skipping")
+            continue
+        print(f"  target = {target}, n = {len(df)}")
+
+        # Try to attach NCAA sequence strings from the matching .txt
+        ncaa_seq_col = [None] * len(df)
+        ncaa_lookup = maybe_load_ncaa_sequences(in_csv, args.ncaa_txt_dir)
+        if ncaa_lookup is not None:
+            ncaa_path, pairs = ncaa_lookup
+            n_matched = 0
+            for i, row in df.iterrows():
+                smi = row.get('peptide_smiles')
+                if isinstance(smi, str) and smi in pairs:
+                    ncaa_seq_col[i] = pairs[smi]
+                    n_matched += 1
+            print(f"  paired {n_matched}/{len(df)} rows with NCAA sequences from {ncaa_path}")
+        df['ncaa_seq'] = ncaa_seq_col
+
+        # ---- Deprotection ----
+        prot_pept = df['peptide_smiles'].fillna('').astype(str).tolist()
+        prot_conj = df['conjugate_smiles'].fillna('').astype(str).tolist() \
+            if 'conjugate_smiles' in df.columns else [''] * len(df)
+
+        deprot_pept = []
+        deprot_conj = []
+        rule_counts_pept = {label: 0 for label, _ in DEPROTECTION_RULES}
+        rule_counts_conj = {label: 0 for label, _ in DEPROTECTION_RULES}
+        n_strips_pept = []
+        n_strips_conj = []
+        for sp, sc in zip(prot_pept, prot_conj):
+            dp, cp = deprotect_smiles(sp, DEPROTECTION_RULES)
+            dc, cc = deprotect_smiles(sc, DEPROTECTION_RULES)
+            deprot_pept.append(dp)
+            deprot_conj.append(dc)
+            for k, v in cp.items():
+                rule_counts_pept[k] += v
+            for k, v in cc.items():
+                rule_counts_conj[k] += v
+                overall_strip_counts[k] += v
+            n_strips_pept.append(sum(cp.values()))
+            n_strips_conj.append(sum(cc.values()))
+
+        df['peptide_smiles_deprot']   = deprot_pept
+        df['conjugate_smiles_deprot'] = deprot_conj
+        df['n_strip_pept']            = n_strips_pept
+        df['n_strip_conj']            = n_strips_conj
+
+        # ---- Properties on deprotected conjugate ----
+        props_list = [compute_props(s) for s in deprot_conj]
+        for k in ('mw_dep', 'logp_dep', 'tpsa_dep', 'hbd_dep',
+                  'hba_dep', 'rotbonds_dep', 'heavy_dep'):
+            df[k] = [p[k] if p else np.nan for p in props_list]
+
+        # ---- Binding rescoring ----
+        df['binding_protected'] = df.get('binding_affinity', np.nan)
+        if args.no_rescore:
+            df['binding_deprotected'] = np.nan
+        else:
+            if target not in binding_models:
+                print(f"  loading binding-affinity model for {target}...")
+                seq = TARGET_SEQUENCES[target]
+                binding_models[target] = BindingAffinity(seq, device=device)
+            model = binding_models[target]
+            print(f"  scoring {len(deprot_pept)} deprotected peptide SMILES...")
+            df['binding_deprotected'] = score_binding_smiles(
+                model, tokenizer, deprot_pept, device
+            )
+
+        # Convenience pKd columns (×10)
+        df['binding_protected_pKd']   = df['binding_protected']   * 10
+        df['binding_deprotected_pKd'] = df['binding_deprotected'] * 10
+
+        # ---- Save per-CSV output ----
+        out_dir = args.output or os.path.dirname(os.path.abspath(in_csv))
+        os.makedirs(out_dir, exist_ok=True)
+        out_csv = os.path.join(out_dir, Path(in_csv).stem + '_deprot.csv')
+        df.to_csv(out_csv, index=False)
+        print(f"  -> {out_csv}")
+
+        # ---- Per-CSV summary ----
+        bp = pd.to_numeric(df['binding_protected'],   errors='coerce').values
+        bd = pd.to_numeric(df['binding_deprotected'], errors='coerce').values
+        mask = np.isfinite(bp) & np.isfinite(bd)
+        n_pair = int(mask.sum())
+        if n_pair > 1:
+            try:
+                from scipy.stats import pearsonr, spearmanr
+                pr = float(pearsonr(bp[mask], bd[mask])[0])
+                sr = float(spearmanr(bp[mask], bd[mask]).correlation)
+            except Exception:
+                bp_z = bp[mask] - bp[mask].mean()
+                bd_z = bd[mask] - bd[mask].mean()
+                denom = (np.linalg.norm(bp_z) * np.linalg.norm(bd_z))
+                pr = float((bp_z * bd_z).sum() / denom) if denom else float('nan')
+                sr = float('nan')
+        else:
+            pr = float('nan')
+            sr = float('nan')
+
+        bp_v = bp[np.isfinite(bp)]
+        bd_v = bd[np.isfinite(bd)]
+        row = {
+            'file':              os.path.basename(in_csv),
+            'target':            target,
+            'n':                 len(df),
+            'n_with_both':       n_pair,
+            'protect_avg_strips_per_conj': float(np.mean(n_strips_conj)) if n_strips_conj else 0.0,
+            'binding_prot_pKd_mean':   float(bp_v.mean() * 10) if bp_v.size else float('nan'),
+            'binding_prot_pKd_std':    float(bp_v.std()  * 10) if bp_v.size else float('nan'),
+            'binding_dep_pKd_mean':    float(bd_v.mean() * 10) if bd_v.size else float('nan'),
+            'binding_dep_pKd_std':     float(bd_v.std()  * 10) if bd_v.size else float('nan'),
+            'pearson_r':         pr,
+            'spearman_r':        sr,
+        }
+        for label, _ in DEPROTECTION_RULES:
+            row[f'strips_conj__{label}'] = rule_counts_conj[label]
+        summary_rows.append(row)
+
+        # Print
+        print(f"\n  --- {os.path.basename(in_csv)} ---")
+        print(f"  binding (pKd) protected   : {row['binding_prot_pKd_mean']:.3f} ± {row['binding_prot_pKd_std']:.3f}  (n={int(np.isfinite(bp).sum())})")
+        print(f"  binding (pKd) deprotected : {row['binding_dep_pKd_mean']:.3f} ± {row['binding_dep_pKd_std']:.3f}  (n={int(np.isfinite(bd).sum())})")
+        print(f"  Pearson r  (prot vs deprot, paired n={n_pair}): {pr:.3f}")
+        print(f"  Spearman r (prot vs deprot, paired n={n_pair}): {sr:.3f}")
+        avg_per = float(np.mean(n_strips_conj)) if n_strips_conj else 0.0
+        print(f"  Total strip events on conjugates: {sum(n_strips_conj)}  ({avg_per:.2f} per molecule)")
+        for label, _ in DEPROTECTION_RULES:
+            v = rule_counts_conj[label]
+            if v:
+                print(f"    {label:<14s} {v}")
+
+    # ---- Aggregate summary ----
+    if summary_rows:
+        out_dir = args.output or os.path.dirname(os.path.abspath(args.inputs[0]))
+        os.makedirs(out_dir, exist_ok=True)
+        summary_csv = os.path.join(out_dir, 'uaa_deprotection_summary.csv')
+        pd.DataFrame(summary_rows).to_csv(summary_csv, index=False)
+        print(f"\nAggregate summary -> {summary_csv}")
+        print("\nOverall strip events on conjugates across all inputs:")
+        for label, _ in DEPROTECTION_RULES:
+            v = overall_strip_counts[label]
+            if v:
+                print(f"  {label:<14s} {v}")
+
+    return 0
+
+
+if __name__ == '__main__':
+    sys.exit(main() or 0)

smiles/generation/__init__.py

@@ -0,0 +1 @@
+"""PDC generation entry points."""

smiles/generation/cupofmoca_guided.py

@@ -0,0 +1,996 @@
+"""
+CupOfMoCA-style Guided Generation with Alternative Guidance Methods
+
+This implements Algorithm 2 (Sampling with CupOfMoCA) with different guidance strategies:
+1. Exponential Tilting (ExpTilt) - Direct implementation of Algorithm 2
+2. Classifier Guidance (CG) - Using external scores to adjust logits
+3. Best-of-K Selection (BoK) - Sample K candidates and select best
+
+Key differences from the Barker-based method in guided.py:
+- ExpTilt: Uses p'(z') ∝ p(z') · exp(β · S(z')) directly without Barker balancing
+- CG: Adjusts base logits using external classifier/score signals: logits = base + w * score_logits
+- BoK: Samples multiple candidates per step and selects the one with highest score
+
+Note on Classifier Guidance vs Classifier-Free Guidance:
+- Classifier Guidance (what we implement here): External scorers modify the generation distribution
+- Classifier-Free Guidance (NOT implemented): Would require model trained with conditioning dropout,
+  and extrapolates between conditional and unconditional model outputs without external classifiers
+"""
+
+import argparse
+import re
+import random
+from collections import Counter
+import csv
+import json
+import os
+from datetime import datetime
+from enum import Enum
+from typing import List, Optional, Tuple, Dict, Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+import numpy as np
+
+# Import from the original codebase
+from smiles.models.train import MDLMLightningModule
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+from smiles.evaluation.peptide_analyzer import PeptideAnalyzer
+from smiles.evaluation.pdc_analyzer import PDCAnalyzer
+from smiles.scoring.classifiers import Conjugation, BindingAffinity
+
+
+class GuidanceMethod(Enum):
+    """Available guidance methods"""
+    EXPONENTIAL_TILT = "exp_tilt"      # Direct exponential tilting (Algorithm 2)
+    CLASSIFIER_GUIDANCE = "cg"          # Classifier/score-based logit adjustment
+    BEST_OF_K = "bok"                   # Best-of-K selection
+    HYBRID = "hybrid"                   # Combination of methods
+
+
+class GenerationTracker:
+    """
+    Tracks detailed validity information during generation
+    """
+    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):
+        """Log information for a generation step"""
+        step_info = {
+            'step': step,
+            'scores': scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else scores,
+            'num_valid': int(sum(1 for s in scores if s > 0)) if hasattr(scores, '__iter__') else 0,
+            'improved_indices': improved_indices or [],
+            'guidance_info': guidance_info or {}
+        }
+        
+        if detailed_results:
+            step_info['detailed'] = []
+            for i, result in enumerate(detailed_results):
+                step_info['detailed'].append({
+                    'sample_idx': i,
+                    'is_valid_peptide': result.get('is_valid_peptide', False),
+                    'has_conjugation_site': result.get('has_conjugation_site', False),
+                    'conjugation_success': result.get('conjugation_success', False),
+                    'pdc_valid': result.get('pdc_valid', False),
+                    'num_conjugates': len(result.get('conjugates', [])),
+                    'error': result.get('error')
+                })
+        
+        self.step_logs.append(step_info)
+        
+    def log_final(self, sequences, detailed_results_with_props):
+        """Log final results with properties"""
+        for i, (seq, result) in enumerate(zip(sequences, detailed_results_with_props)):
+            final_info = {
+                'sample_idx': i,
+                'peptide_smiles': seq,
+                'is_valid_peptide': result.get('is_valid_peptide', False),
+                'has_conjugation_site': result.get('has_conjugation_site', False),
+                'num_conjugates': len(result.get('conjugates', [])),
+                'num_valid_pdcs': len(result.get('valid_pdcs', [])),
+                'best_pdc': None,
+                'linker_used': result.get('linker_used', ''),
+                'all_conjugates': result.get('conjugates', []),
+                'properties': result.get('properties', {})
+            }
+            
+            if result.get('best_pdc'):
+                final_info['best_pdc'] = {
+                    'smiles': result['best_pdc'].get('smiles', ''),
+                    'checks': result['best_pdc'].get('checks', {})
+                }
+            
+            self.final_results.append(final_info)
+    
+    def save(self, filename_prefix='generation'):
+        """Save all logs to files"""
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_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', 'peptide_smiles', 'is_valid_peptide', 'has_conjugation_site',
+                'linker_used', 'num_conjugates', 'num_valid_pdcs', 'best_pdc_smiles',
+                'mf', 'polararea', 'xlogp', 'heavycnt', 'rotbonds'
+            ])
+            for result in self.final_results:
+                props = result.get('properties', {})
+                writer.writerow([
+                    result['sample_idx'],
+                    result['peptide_smiles'],
+                    result['is_valid_peptide'],
+                    result['has_conjugation_site'],
+                    result.get('linker_used', ''),
+                    result['num_conjugates'],
+                    result['num_valid_pdcs'],
+                    result['best_pdc']['smiles'] if result.get('best_pdc') else '',
+                    props.get('mf', ''),
+                    props.get('polararea', ''),
+                    props.get('xlogp', ''),
+                    props.get('heavycnt', ''),
+                    props.get('rotbonds', '')
+                ])
+        
+        # Save dedicated PDC output: peptide → linker → all conjugate SMILES
+        pdc_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_pdcs.csv")
+        with open(pdc_path, 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow(['sample_idx', 'peptide_smiles', 'linker_used', 'conjugate_idx', 'pdc_smiles'])
+            for result in self.final_results:
+                for j, conj in enumerate(result.get('all_conjugates', [])):
+                    writer.writerow([
+                        result['sample_idx'],
+                        result['peptide_smiles'],
+                        result.get('linker_used', ''),
+                        j,
+                        conj
+                    ])
+        
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        print(f"  PDC conjugates: {pdc_path}")
+        
+        return step_log_path, final_path, csv_path, pdc_path
+    
+    def print_step_summary(self, step, scores, detailed_results=None, guidance_method=None):
+        """Print a summary for the current step"""
+        if isinstance(scores, torch.Tensor):
+            scores_list = scores.cpu().numpy().tolist()
+        else:
+            scores_list = list(scores)
+        
+        num_valid = sum(1 for s in scores_list if s > 0)
+        
+        print(f"\n--- Step {step} Summary ({guidance_method or 'unknown'}) ---")
+        print(f"  Valid: {num_valid}/{len(scores_list)}")
+        
+        if detailed_results:
+            valid_peptide = sum(1 for r in detailed_results if r.get('is_valid_peptide'))
+            has_site = sum(1 for r in detailed_results if r.get('has_conjugation_site'))
+            conj_success = sum(1 for r in detailed_results if r.get('conjugation_success'))
+            pdc_valid = sum(1 for r in detailed_results if r.get('pdc_valid'))
+            
+            print(f"  Breakdown:")
+            print(f"    Valid peptides: {valid_peptide}")
+            print(f"    Has conjugation site: {has_site}")
+            print(f"    Conjugation success: {conj_success}")
+            print(f"    PDC valid: {pdc_valid}")
+    
+    def print_final_summary(self):
+        """Print final summary"""
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        
+        total = len(self.final_results)
+        valid_pdcs = sum(1 for r in self.final_results if r.get('num_valid_pdcs', 0) > 0)
+        
+        print(f"\nTotal sequences: {total}")
+        print(f"With valid PDCs: {valid_pdcs}")
+        
+        print("\nDetailed results:")
+        for result in self.final_results:
+            status = "" if result.get('num_valid_pdcs', 0) > 0 else ""
+            print(f"\n  [{status}] Sample {result['sample_idx']}:")
+            print(f"      Peptide: {result['peptide_smiles'][:60]}...")
+            print(f"      Valid peptide: {result['is_valid_peptide']}")
+            print(f"      Has conjugation site: {result['has_conjugation_site']}")
+            print(f"      Conjugates: {result['num_conjugates']}")
+            print(f"      Valid PDCs: {result['num_valid_pdcs']}")
+
+
+class CupOfMoCAGenerator:
+    """
+    CupOfMoCA-style generator implementing Algorithm 2 with multiple guidance strategies.
+    
+    Key differences from Barker-based approach:
+    1. ExpTilt: Direct exponential tilting without balancing function
+    2. CFG: Classifier-free guidance by mixing logits
+    3. BoK: Best-of-K selection strategy
+    """
+    
+    def __init__(self, model, device, objectives, args, tracker=None):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.pdc_analyzer = PDCAnalyzer(device)
+        self.invalid_tokens = [0, 1, 2, 3, 4, 585, 586]  # Special tokens to avoid
+        self.tracker = tracker
+        
+        # Parse guidance method
+        self.guidance_method = GuidanceMethod(args.guidance_method)
+        
+        # Find the Conjugation objective for detailed tracking
+        self.conjugation_obj = None
+        for obj in objectives:
+            if isinstance(obj, Conjugation):
+                self.conjugation_obj = obj
+                break
+    
+    # =========================================================================
+    # Initial Generation (x0 from prior μ₀)
+    # =========================================================================
+    
+    def _generate_single_x0(self, n_steps: int = 16, temperature: float = 1.0) -> torch.Tensor:
+        """
+        Generate initial sequences using flow matching (Algorithm 2, Line 6: Sample Z_0 ~ μ_0)
+        This follows the flow from noise to data without guidance.
+        """
+        self.model.eval()
+        
+        # Start with random tokens (noise at t=0)
+        x = torch.randint(
+            0, self.model.tokenizer.vocab_size,
+            (self.args.num_samples, self.args.gen_len),
+            device=self.device
+        )
+        
+        # Time schedule: t_0 = 0 to t_K = T (Line 5 of Algorithm 2)
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+        
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Initial Flow Matching", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i + 1]
+                
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+                
+                # Get predicted distribution (Line 8: p_{k+1} ← Φ_{θ*}^{t_k→t_{k+1}}(Z_k))
+                logits = self.model(x, t_tensor)
+                logits = logits / temperature
+                logits[:, :, 586] = -1000  # Mask invalid token
+                
+                pred_x1 = torch.argmax(logits, dim=-1)
+                
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+                
+                # Interpolate between noise and prediction
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                noise = torch.randint(0, self.model.tokenizer.vocab_size, x.shape, device=self.device)
+                x = torch.where(mask, noise, pred_x1)
+        
+        return x
+    
+    def _check_conjugation_validity(self, x: torch.Tensor) -> torch.Tensor:
+        """Check if sequences can be conjugated (validity indicator 1_valid in Algorithm 2)"""
+        if self.conjugation_obj is None:
+            sequences = self.model.tokenizer.batch_decode(x)
+            return torch.tensor([self.peptide_analyzer.is_peptide(seq) for seq in sequences], 
+                              device=self.device)
+        scores = self.conjugation_obj(x)
+        return scores > 0
+    
+    def generate_x0(self, n_steps: int = 16, temperature: float = 1.0,
+                    require_valid: bool = False, max_retries: int = 10) -> torch.Tensor:
+        """
+        Generate initial x0 sequences with optional validity enforcement.
+        Implements the constraint that x0 should satisfy validity before guidance.
+        """
+        print("Generating initial sequences (x0 from prior)...")
+        
+        x = self._generate_single_x0(n_steps, temperature)
+        
+        if require_valid:
+            print("Enforcing initial validity constraint...")
+            
+            for retry in range(max_retries):
+                validity = self._check_conjugation_validity(x)
+                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 are valid!")
+                    break
+                
+                print(f"  Retry {retry + 1}/{max_retries}: {num_valid}/{self.args.num_samples} valid")
+                
+                invalid_indices = torch.where(~validity)[0]
+                original_num_samples = self.args.num_samples
+                self.args.num_samples = num_invalid
+                new_x = self._generate_single_x0(n_steps, temperature)
+                self.args.num_samples = original_num_samples
+                x[invalid_indices] = new_x
+            
+            final_validity = self._check_conjugation_validity(x)
+            if final_validity.sum().item() < self.args.num_samples:
+                print(f"⚠ After {max_retries} retries, {final_validity.sum().item()}/{self.args.num_samples} valid")
+        
+        generated_sequences = self.model.tokenizer.batch_decode(x)
+        for seq in generated_sequences:
+            print(f"  {seq[:80]}...")
+        
+        return x
+    
+    # =========================================================================
+    # Score Functions
+    # =========================================================================
+    
+    def _get_scores(self, x_batch: torch.Tensor) -> torch.Tensor:
+        """Calculate scores for all objectives: S(z) in Algorithm 2"""
+        scores = []
+        for obj_func in self.objectives:
+            scores.append(obj_func(x_batch.to(self.device)))
+        return torch.stack(scores, dim=0).to(self.device)
+    
+    def _compute_aggregated_score(self, scores: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+        """Compute weighted aggregated score S(z) for guidance"""
+        # Weighted sum or min depending on aggregation strategy
+        if self.args.score_aggregation == "weighted_sum":
+            return (weights * scores).sum(dim=0)
+        elif self.args.score_aggregation == "weighted_min":
+            return torch.min(weights * scores, dim=0).values
+        else:  # default: weighted sum
+            return (weights * scores).sum(dim=0)
+    
+    # =========================================================================
+    # Guidance Methods
+    # =========================================================================
+    
+    def _exponential_tilt_guidance(
+        self,
+        x: torch.Tensor,
+        pos_probs: torch.Tensor,
+        candidates: torch.Tensor,
+        mut_idx: int,
+        sample_idx: int,
+        beta: float,
+        weights: torch.Tensor,
+        current_score: float = None
+    ) -> Tuple[int, Dict]:
+        """
+        Exponential Tilting Guidance (Algorithm 2, Lines 10-12) with acceptance criterion
+        
+        p'(z') ∝ p_{k+1}(z') · exp(β · S(z')) · 1_valid(z')
+        
+        Added: Only accept if new_score >= current_score (prevents regression)
+        """
+        num_candidates = len(candidates)
+        current_token = x[sample_idx, mut_idx].item()
+        
+        if num_candidates == 0:
+            return current_token, {'method': 'exp_tilt', 'fallback': True, 'accepted': False}
+        
+        # Limit candidates for efficiency
+        if num_candidates > self.args.top_k:
+            candidates = candidates[:self.args.top_k]
+            num_candidates = self.args.top_k
+        
+        # Create batch of proposals
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(num_candidates, 1)
+        x_proposals[:, mut_idx] = candidates
+        
+        # Compute scores S(z') for all candidates
+        proposal_scores = self._get_scores(x_proposals)  # [num_objectives, num_candidates]
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)  # [num_candidates]
+        
+        # Validity mask: 1_valid(z')
+        validity_mask = (proposal_scores[0] > 0).float()  # Using first objective (conjugation) as validity
+        
+        # Get base probabilities from the model
+        base_probs = pos_probs[sample_idx, candidates]
+        
+        # Apply exponential tilting with numerical stability:
+        # p'(z') ∝ p(z') · exp(β · S(z')) · 1_valid(z')
+        # 
+        # To avoid overflow, we use log-space computation:
+        # log(p') = log(p) + β·S + log(1_valid)
+        # Then normalize in log-space before exp
+        
+        log_base_probs = torch.log(base_probs + 1e-10)
+        log_exp_scores = beta * aggregated_scores
+        log_validity = torch.log(validity_mask + 1e-10)  # -inf for invalid
+        
+        log_tilted = log_base_probs + log_exp_scores + log_validity
+        
+        # Subtract max for numerical stability (log-sum-exp trick)
+        log_tilted_max = log_tilted.max()
+        if torch.isinf(log_tilted_max) or torch.isnan(log_tilted_max):
+            # All candidates are invalid or have zero probability
+            # Fallback to base distribution
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+        else:
+            log_tilted_stable = log_tilted - log_tilted_max
+            tilted_probs = torch.exp(log_tilted_stable)
+            
+            # Normalize
+            prob_sum = tilted_probs.sum()
+            if prob_sum < 1e-10 or torch.isnan(prob_sum) or torch.isinf(prob_sum):
+                # Fallback to base distribution
+                fallback = True
+                final_probs = base_probs / (base_probs.sum() + 1e-9)
+            else:
+                fallback = False
+                final_probs = tilted_probs / prob_sum
+        
+        # Final safety check
+        if torch.any(torch.isnan(final_probs)) or torch.any(torch.isinf(final_probs)):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+        
+        # Ensure no negative or zero probabilities for multinomial
+        final_probs = torch.clamp(final_probs, min=1e-10)
+        final_probs = final_probs / final_probs.sum()
+        
+        selected_idx = torch.multinomial(final_probs, 1).item()
+        selected_token = candidates[selected_idx].item()
+        selected_score = aggregated_scores[selected_idx].item()
+        
+        # ACCEPTANCE CRITERION: Only accept if score improves or stays same
+        accepted = True
+        if self.args.use_acceptance and current_score is not None:
+            if selected_score < current_score:
+                # Reject: keep current token
+                selected_token = current_token
+                accepted = False
+        
+        info = {
+            'method': 'exp_tilt',
+            'beta': beta,
+            'num_candidates': num_candidates,
+            'proposal_score': aggregated_scores[selected_idx].item(),
+            'selected_score': selected_score if accepted else current_score,
+            'fallback': fallback,
+            'validity_rate': validity_mask.mean().item(),
+            'accepted': accepted
+        }
+        
+        return selected_token, info
+    
+    def _classifier_guidance(
+        self,
+        x: torch.Tensor,
+        logits: torch.Tensor,
+        mut_idx: int,
+        sample_idx: int,
+        guidance_weight: float,
+        weights: torch.Tensor
+    ) -> Tuple[int, Dict]:
+        """
+        Classifier Guidance (Score-based Logit Adjustment)
+        
+        Uses external score functions to adjust the base model's logits:
+        logits_guided = logits_base + guidance_weight * score_logits
+        
+        This differs from Classifier-Free Guidance (CFG) which requires:
+        - A model trained with conditioning dropout
+        - Extrapolation between conditional and unconditional model outputs
+        
+        Here we use external classifiers/scorers to provide guidance signal.
+        """
+        base_logits = logits[sample_idx, mut_idx, :]
+        
+        # Get top-k candidates from base distribution
+        sorted_logits, sorted_indices = torch.sort(base_logits, descending=True)
+        candidates = sorted_indices[:self.args.top_k]
+        
+        # Filter invalid tokens
+        valid_mask = ~torch.isin(candidates, torch.tensor(self.invalid_tokens, device=self.device))
+        candidates = candidates[valid_mask]
+        
+        if len(candidates) == 0:
+            return x[sample_idx, mut_idx].item(), {'method': 'cg', 'fallback': True}
+        
+        # Create proposals and score them
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(len(candidates), 1)
+        x_proposals[:, mut_idx] = candidates
+        
+        proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+        
+        # Create score-based logits adjustment
+        # Higher score -> higher probability
+        score_logits = aggregated_scores * self.args.score_temperature
+        
+        # Base logits for candidates
+        base_candidate_logits = base_logits[candidates]
+        
+        # Classifier guidance: guided = base + w * score_adjustment
+        guided_logits = base_candidate_logits + guidance_weight * score_logits
+        
+        # Convert to probabilities
+        guided_probs = F.softmax(guided_logits, dim=-1)
+        
+        # Sample
+        selected_idx = torch.multinomial(guided_probs, 1).item()
+        selected_token = candidates[selected_idx].item()
+        
+        info = {
+            'method': 'cg',
+            'guidance_weight': guidance_weight,
+            'num_candidates': len(candidates),
+            'selected_score': aggregated_scores[selected_idx].item(),
+            'fallback': False
+        }
+        
+        return selected_token, info
+    
+    def _best_of_k_guidance(
+        self,
+        x: torch.Tensor,
+        pos_probs: torch.Tensor,
+        candidates: torch.Tensor,
+        mut_idx: int,
+        sample_idx: int,
+        k: int,
+        weights: torch.Tensor
+    ) -> Tuple[int, Dict]:
+        """
+        Best-of-K Selection
+        
+        1. Sample K candidates from the base distribution
+        2. Score all K candidates
+        3. Select the one with the highest score
+        
+        This is a simple but effective strategy that doesn't require 
+        probability manipulation.
+        """
+        num_candidates = len(candidates)
+        if num_candidates == 0:
+            return x[sample_idx, mut_idx].item(), {'method': 'bok', 'fallback': True}
+        
+        # Limit k to available candidates
+        k = min(k, num_candidates)
+        
+        # Get probabilities for candidates
+        candidate_probs = pos_probs[sample_idx, candidates]
+        candidate_probs = candidate_probs / (candidate_probs.sum() + 1e-9)
+        
+        # Sample K candidates
+        if k >= num_candidates:
+            sampled_indices = torch.arange(num_candidates, device=self.device)
+        else:
+            sampled_indices = torch.multinomial(candidate_probs, k, replacement=False)
+        
+        sampled_tokens = candidates[sampled_indices]
+        
+        # Create proposals for all K samples
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(len(sampled_tokens), 1)
+        x_proposals[:, mut_idx] = sampled_tokens
+        
+        # Score all proposals
+        proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+        
+        # Select best
+        best_idx = torch.argmax(aggregated_scores).item()
+        selected_token = sampled_tokens[best_idx].item()
+        
+        info = {
+            'method': 'bok',
+            'k': k,
+            'num_candidates': num_candidates,
+            'selected_score': aggregated_scores[best_idx].item(),
+            'mean_score': aggregated_scores.mean().item(),
+            'max_score': aggregated_scores.max().item(),
+            'fallback': False
+        }
+        
+        return selected_token, info
+    
+    def _hybrid_guidance(
+        self,
+        x: torch.Tensor,
+        pos_probs: torch.Tensor,
+        logits: torch.Tensor,
+        candidates: torch.Tensor,
+        mut_idx: int,
+        sample_idx: int,
+        step: int,
+        beta: float,
+        weights: torch.Tensor
+    ) -> Tuple[int, Dict]:
+        """
+        Hybrid Guidance: Combines multiple methods adaptively
+        
+        Strategy:
+        - Early steps: Use Best-of-K for exploration
+        - Middle steps: Use Exponential Tilting for balanced guidance
+        - Late steps: Use CFG for refinement
+        """
+        total_steps = self.args.optimization_steps
+        progress = step / max(total_steps - 1, 1)
+        
+        if progress < 0.3:
+            # Early: Best-of-K for exploration
+            return self._best_of_k_guidance(
+                x, pos_probs, candidates, mut_idx, sample_idx,
+                k=self.args.bok_k, weights=weights
+            )
+        elif progress < 0.7:
+            # Middle: Exponential tilting
+            return self._exponential_tilt_guidance(
+                x, pos_probs, candidates, mut_idx, sample_idx,
+                beta=beta, weights=weights
+            )
+        else:
+            # Late: Classifier guidance for refinement
+            return self._classifier_guidance(
+                x, logits, mut_idx, sample_idx,
+                guidance_weight=self.args.cg_weight, weights=weights
+            )
+    
+    # =========================================================================
+    # Main Generation Loop (Algorithm 2)
+    # =========================================================================
+    
+    def generate(self) -> torch.Tensor:
+        """
+        Main generation loop implementing Algorithm 2 with the selected guidance method.
+        
+        Lines 7-17 of Algorithm 2:
+        for k = 0 to K-1 do
+            p_{k+1} ← Φ_{θ*}^{t_k→t_{k+1}}(Z_k)
+            if guidance is active (β > 0) then
+                for each candidate z' ∈ TopK(p_{k+1}) do
+                    p'(z') ∝ p_{k+1}(z') · exp(β · S(z')) · 1_valid(z')
+                Sample Z_{k+1} ~ p'
+            else
+                Sample Z_{k+1} ~ p_{k+1}
+        return Z_K
+        """
+        # Generate initial x0 with validity constraint
+        x = self.generate_x0(
+            require_valid=self.args.require_valid_x0,
+            max_retries=self.args.max_x0_retries
+        )
+        
+        # Setup weights
+        if self.args.weights is None:
+            weights = torch.tensor(
+                [1.0] + [1.0 / (self.num_objectives - 1)] * (self.num_objectives - 1),
+                device=self.device
+            ).view(-1, 1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+        
+        print(f"\nGuidance Method: {self.guidance_method.value}")
+        print(f"Weights: {weights.squeeze().tolist()}")
+        print(f"β range: [{self.args.beta_min}, {self.args.beta_max}]")
+        print(f"Acceptance criterion: {'ENABLED' if self.args.use_acceptance else 'DISABLED'}")
+        
+        # Main optimization loop (Line 7: for k = 0 to K-1)
+        with torch.no_grad():
+            for step in tqdm(range(self.args.optimization_steps), desc="Guided Generation"):
+                improved_indices = []
+                
+                # Anneal guidance strength β (from beta_min to beta_max)
+                progress = step / max(self.args.optimization_steps - 1, 1)
+                beta = self.args.beta_min + (self.args.beta_max - self.args.beta_min) * progress
+                
+                # Choose position to modify (can be random or sequential)
+                if self.args.position_strategy == "random":
+                    mut_idx = random.randint(1, self.args.gen_len - 2)
+                else:  # sequential
+                    mut_idx = (step % (self.args.gen_len - 2)) + 1
+                
+                # Get model prediction (Line 8: p_{k+1} ← Φ_{θ*}^{t_k→t_{k+1}}(Z_k))
+                time_t = torch.full(
+                    (self.args.num_samples,),
+                    step / self.args.optimization_steps,
+                    device=self.device
+                )
+                
+                logits = self.model(x, time_t)
+                probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+                
+                # Zero out current tokens and invalid tokens
+                pos_probs[:, x[:, mut_idx]] = 0
+                pos_probs[:, self.invalid_tokens] = 0
+                
+                # Get TopK candidates (Line 10: for each candidate z' ∈ TopK(p_{k+1}))
+                sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+                
+                # Get current scores for comparison
+                current_scores = self._get_scores(x)
+                current_aggregated = self._compute_aggregated_score(current_scores, weights)
+                
+                if step == 0:
+                    print(f"\nInitial Scores: {current_scores.cpu().numpy()}")
+                    if self.tracker and self.conjugation_obj:
+                        detailed = self.conjugation_obj.get_last_detailed_results()
+                        self.tracker.log_step(step, current_scores[0], detailed)
+                        self.tracker.print_step_summary(step, current_scores[0], detailed, 
+                                                       self.guidance_method.value)
+                
+                # Process each sample
+                final_tokens = []
+                guidance_infos = []
+                
+                for i in range(self.args.num_samples):
+                    # Get candidates for this sample
+                    sample_probs = sorted_probs[i]
+                    sample_indices = sorted_indices[i]
+                    
+                    # Top-p filtering
+                    cumsum = torch.cumsum(sample_probs, dim=0)
+                    mask = cumsum <= self.args.top_p
+                    mask[0] = True  # Always include at least one
+                    candidates = sample_indices[mask]
+                    
+                    # Current score for this sample (for acceptance criterion)
+                    current_sample_score = current_aggregated[i].item()
+                    
+                    # Apply selected guidance method (Lines 9-15)
+                    if beta > 0:  # Guidance active (Line 9)
+                        if self.guidance_method == GuidanceMethod.EXPONENTIAL_TILT:
+                            token, info = self._exponential_tilt_guidance(
+                                x, pos_probs, candidates, mut_idx, i, beta, weights,
+                                current_score=current_sample_score
+                            )
+                        elif self.guidance_method == GuidanceMethod.CLASSIFIER_GUIDANCE:
+                            token, info = self._classifier_guidance(
+                                x, logits, mut_idx, i, self.args.cg_weight, weights
+                            )
+                        elif self.guidance_method == GuidanceMethod.BEST_OF_K:
+                            token, info = self._best_of_k_guidance(
+                                x, pos_probs, candidates, mut_idx, i, self.args.bok_k, weights
+                            )
+                        elif self.guidance_method == GuidanceMethod.HYBRID:
+                            token, info = self._hybrid_guidance(
+                                x, pos_probs, logits, candidates, mut_idx, i, step, beta, weights
+                            )
+                        else:
+                            raise ValueError(f"Unknown guidance method: {self.guidance_method}")
+                    else:
+                        # No guidance: sample from p_{k+1} directly (Line 15)
+                        candidate_probs = pos_probs[i, candidates]
+                        candidate_probs = candidate_probs / (candidate_probs.sum() + 1e-9)
+                        idx = torch.multinomial(candidate_probs, 1).item()
+                        token = candidates[idx].item()
+                        info = {'method': 'unguided'}
+                    
+                    final_tokens.append(token)
+                    guidance_infos.append(info)
+                
+                # Update sequences with new tokens
+                x[:, mut_idx] = torch.tensor(final_tokens, device=self.device)
+                
+                # Check for improvements
+                new_scores = self._get_scores(x)
+                new_aggregated = self._compute_aggregated_score(new_scores, weights)
+                
+                for i in range(self.args.num_samples):
+                    if new_aggregated[i] > current_aggregated[i]:
+                        improved_indices.append(i)
+                
+                if improved_indices:
+                    print(f"\nStep {step}: Improved samples: {improved_indices}")
+                    for i in improved_indices:
+                        print(f"  Sample {i}:")
+                        # print(f"    Previous score: {current_aggregated[i].item():.4f} | Per-objective: {current_scores[:, i].cpu().numpy()}")
+                        print(f"    New score:      {new_aggregated[i].item():.4f} | Per-objective: {new_scores[:, i].cpu().numpy()}")
+                    print(f"  Sequences: {self.model.tokenizer.batch_decode(x[improved_indices])}")
+                
+                # Logging
+                if self.tracker and self.conjugation_obj and (step % 4 == 0 or step == self.args.optimization_steps - 1):
+                    step_scores = self._get_scores(x)
+                    detailed = self.conjugation_obj.get_last_detailed_results()
+                    self.tracker.log_step(
+                        step, step_scores[0], detailed, improved_indices,
+                        guidance_info={'method': self.guidance_method.value, 'beta': beta}
+                    )
+                    
+                    if step == self.args.optimization_steps - 1:
+                        self.tracker.print_step_summary(step, step_scores[0], detailed,
+                                                       self.guidance_method.value)
+        
+        return x
+
+
+# =============================================================================
+# Main Execution
+# =============================================================================
+
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    print(f"Guidance method: {args.guidance_method}")
+    
+    tokenizer = SMILES_SPE_Tokenizer(
+        './smiles/smiles_tokenizer/new_vocab.txt',
+        './smiles/smiles_tokenizer/new_splits.txt'
+    )
+    
+    # Initialize objectives
+    conjugation_model = Conjugation(tokenizer, device=device, builder_type=args.builder_type)
+    affinity_model = BindingAffinity(args.target, device)
+    OBJECTIVE_FUNCTIONS = [conjugation_model, affinity_model]
+    
+    # Initialize tracker
+    tracker = GenerationTracker(output_dir=args.output_dir)
+    
+    # Set random seeds
+    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)
+        print(f"Random seed: {args.seed}")
+    
+    # Load model
+    print(f"Loading model from: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+    model_hparams = checkpoint["hyper_parameters"]["args"]
+    
+    model = MDLMLightningModule.load_from_checkpoint(
+        args.checkpoint,
+        args=model_hparams,
+        tokenizer=tokenizer,
+        map_location=device,
+        strict=False
+    )
+    model.to(device)
+    print("Model loaded successfully.")
+    
+    # Create generator
+    generator = CupOfMoCAGenerator(model, device, OBJECTIVE_FUNCTIONS, args, tracker=tracker)
+    
+    all_sequences = []
+    all_final_results = []
+    
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'=' * 60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'=' * 60}")
+        
+        generated_tokens = generator.generate()
+        final_scores = generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+        
+        print("\n--- Batch Results ---")
+        print(f"Sequences: {sequence_str}")
+        print(f"Scores: {final_scores}")
+        
+        all_sequences.extend(sequence_str)
+        
+        # Get detailed results
+        detailed_with_props = conjugation_model.get_detailed_scores_with_properties(generated_tokens)
+        tracker.log_final(sequence_str, detailed_with_props)
+        all_final_results.extend(detailed_with_props)
+    
+    # Save results
+    tracker.print_final_summary()
+    tracker.save(filename_prefix=f'cupofmoca_{args.guidance_method}')
+    
+    with open(args.output_file, 'w') as f:
+        for seq in all_sequences:
+            f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+    print("\nGeneration complete.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="CupOfMoCA-style Guided Generation with Alternative Guidance Methods"
+    )
+    
+    # Model and data
+    parser.add_argument("--checkpoint", type=str, required=True,
+                        help="Path to trained model checkpoint")
+    parser.add_argument("--target", type=str, required=True,
+                        help="Target protein sequence for binding affinity")
+    
+    # Generation settings
+    parser.add_argument("--num_samples", type=int, default=10,
+                        help="Number of samples per batch")
+    parser.add_argument("--num_batches", type=int, default=10,
+                        help="Number of batches to generate")
+    parser.add_argument("--gen_len", type=int, default=50,
+                        help="Length of sequences to generate")
+    parser.add_argument("--optimization_steps", type=int, default=16,
+                        help="Number of guided optimization steps (K in Algorithm 2)")
+    
+    # Guidance method selection
+    parser.add_argument("--guidance_method", type=str, default="exp_tilt",
+                        choices=["exp_tilt", "cg", "bok", "hybrid"],
+                        help="Guidance method: exp_tilt (exponential tilting), "
+                             "cg (classifier/score guidance), bok (best-of-k), hybrid")
+    
+    # Guidance parameters
+    parser.add_argument("--beta_min", type=float, default=1.0,
+                        help="Minimum guidance strength β")
+    parser.add_argument("--beta_max", type=float, default=20.0,
+                        help="Maximum guidance strength β")
+    parser.add_argument("--cg_weight", type=float, default=2.0,
+                        help="Classifier guidance weight (for cg method)")
+    parser.add_argument("--use_acceptance", action="store_true",
+                        help="Use acceptance criterion: only accept if new_score >= current_score (prevents regression)")
+    parser.add_argument("--bok_k", type=int, default=8,
+                        help="Number of candidates for best-of-k selection")
+    parser.add_argument("--top_k", type=int, default=200,
+                        help="Top-K candidates to consider for guidance")
+    parser.add_argument("--top_p", type=float, default=0.9,
+                        help="Top-p threshold for candidate filtering")
+    parser.add_argument("--score_temperature", type=float, default=1.0,
+                        help="Temperature for score-based logit adjustment")
+    parser.add_argument("--score_aggregation", type=str, default="weighted_sum",
+                        choices=["weighted_sum", "weighted_min"],
+                        help="How to aggregate multi-objective scores")
+    
+    # Weights
+    parser.add_argument("--weights", type=float, nargs='+', default=None,
+                        help="Weights for objectives (e.g., 0.5 0.5)")
+    
+    # Position selection
+    parser.add_argument("--position_strategy", type=str, default="random",
+                        choices=["random", "sequential"],
+                        help="Strategy for selecting mutation positions")
+    
+    # Initial generation
+    parser.add_argument("--require_valid_x0", action="store_true",
+                        help="Regenerate x0 until valid conjugatable peptides")
+    parser.add_argument("--max_x0_retries", type=int, default=10,
+                        help="Max retries for valid x0 generation")
+    
+    # PDC builder
+    parser.add_argument("--builder_type", type=str, default="AoaDau",
+                        choices=["AoaDau", "ValAlaDau", "ValCitDau", "GlutarylDau", "Dau",
+                                 "SSCpt", "VcMMAE", "DSPACpt", "MalHexCpt", "DSECCpt", "Cpt",
+                                 "DSPASN38", "DSEBSN38", "DSEPSN38", "SN38"],
+                        help="PDC builder type for conjugation. "
+                             "'Dau' samples linker by GNRHR+Dau frequency. "
+                             "'Cpt' samples linker by ERBB2+Cpt frequency. "
+                             "'SN38' samples linker by NRP1+SN38 frequency.")
+    
+    # Output
+    parser.add_argument("--output_file", type=str, default="./smiles.txt",
+                        help="File to save generated sequences")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs",
+                        help="Directory for detailed logs")
+    
+    # Reproducibility
+    parser.add_argument("--seed", type=int, default=None,
+                        help="Random seed")
+    
+    args = parser.parse_args()
+    main(args)

smiles/generation/cupofmoca_guided_aa.py

@@ -0,0 +1,1126 @@
+"""CupOfMoCA-style guided generation with the natural-AA foundation model."""
+
+import argparse
+import csv
+import json
+import os
+import random
+import sys
+import time
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from datetime import datetime
+from enum import Enum
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from easydict import EasyDict as edict
+from tqdm import tqdm
+
+from smiles.models.train import MDLMLightningModule
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+from smiles.evaluation.peptide_analyzer import PeptideAnalyzer
+from smiles.evaluation.pdc_analyzer import PDCAnalyzer
+from smiles.scoring.classifiers import Conjugation, BindingAffinity
+
+# AA-mode shared utilities
+from smiles.models.aa_helpers import (
+    ESM2AminoTokenizer as _AAEsmTokenizer,
+    MDLM_PepReDi as _AAMDLM,
+    make_affinity,
+    DEFAULT_AA_AFFINITY_CKPT,
+    INVALID_AA_TOKEN_IDS as _AA_INVALID_IDS,
+)
+
+from third_party.SMILESToPeptide.peptide_to_smiles import aa_to_smiles
+from third_party.SMILESToPeptide.aminoacid_selective import one_letter_to_three_letter
+
+
+# =============================================================================
+# Amino Acid Foundation Model (DiT-based MDLM)
+# =============================================================================
+
+# --- AA-mode Model + Tokenizer (imported from aa_helpers) ---
+MDLM = _AAMDLM
+AminoTokenizer = _AAEsmTokenizer
+
+
+# =============================================================================
+# Profiler
+# =============================================================================
+
+class Profiler:
+    """Aggregates wall-clock time per named section. Optionally syncs CUDA.
+
+    `reset()` clears all state (use between batches for self-contained summaries).
+    `print_summary()` shows cumulative totals since the last reset, plus a
+    delta column showing time since the previous `print_summary()` call.
+    """
+
+    def __init__(self, enabled=False, cuda_sync=True):
+        self.enabled = enabled
+        self.cuda_sync = cuda_sync and torch.cuda.is_available()
+        self.totals = defaultdict(float)
+        self.counts = defaultdict(int)
+        self._last_totals = defaultdict(float)
+        self._last_counts = defaultdict(int)
+
+    @contextmanager
+    def section(self, name):
+        if not self.enabled:
+            yield
+            return
+        if self.cuda_sync:
+            torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        try:
+            yield
+        finally:
+            if self.cuda_sync:
+                torch.cuda.synchronize()
+            self.totals[name] += time.perf_counter() - t0
+            self.counts[name] += 1
+
+    def reset(self):
+        self.totals.clear()
+        self.counts.clear()
+        self._last_totals.clear()
+        self._last_counts.clear()
+
+    def print_summary(self, header=""):
+        if not self.enabled or not self.totals:
+            return
+        cumulative_total = sum(self.totals.values())
+        delta_totals = {k: self.totals[k] - self._last_totals.get(k, 0.0)
+                        for k in self.totals}
+        delta_counts = {k: self.counts[k] - self._last_counts.get(k, 0)
+                        for k in self.counts}
+        interval_total = sum(delta_totals.values())
+
+        print(f"\n=== Profiler summary {header} ===")
+        print(f"  (cumulative since reset | interval since previous summary)")
+        print(f"  {'section':<34} "
+              f"{'calls':>6} {'cum(s)':>8} {'cum_avg':>8}  "
+              f"{'Δcalls':>6} {'Δ(s)':>7} {'Δavg':>7} {'%Δ':>5}")
+        for name, tot in sorted(self.totals.items(), key=lambda kv: -kv[1]):
+            n = self.counts[name]
+            cum_avg = (tot / n) * 1000.0 if n else 0.0
+            dt = delta_totals[name]
+            dn = delta_counts[name]
+            d_avg = (dt / dn) * 1000.0 if dn else 0.0
+            d_pct = (dt / interval_total) * 100.0 if interval_total > 0 else 0.0
+            print(f"  {name:<34} "
+                  f"{n:>6d} {tot:>8.2f} {cum_avg:>7.1f}ms  "
+                  f"{dn:>6d} {dt:>7.2f} {d_avg:>6.1f}ms {d_pct:>4.1f}%")
+        print(f"  {'-' * 34} {'-' * 6} {'-' * 8} {'-' * 8}  "
+              f"{'-' * 6} {'-' * 7} {'-' * 7} {'-' * 5}")
+        print(f"  {'TOTAL':<34} {'':>6} {cumulative_total:>8.2f} {'':>8}  "
+              f"{'':>6} {interval_total:>7.2f}")
+
+        self._last_totals = dict(self.totals)
+        self._last_counts = dict(self.counts)
+
+
+# =============================================================================
+# Guidance Method Enum
+# =============================================================================
+
+class GuidanceMethod(Enum):
+    EXPONENTIAL_TILT = "exp_tilt"
+    CLASSIFIER_GUIDANCE = "cg"
+    BEST_OF_K = "bok"
+    HYBRID = "hybrid"
+
+
+# =============================================================================
+# Generation Tracker
+# =============================================================================
+
+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, 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(sum(1 for s in scores if s > 0)) if hasattr(scores, '__iter__') else 0,
+            'improved_indices': improved_indices or [],
+            'guidance_info': guidance_info or {}
+        }
+        if detailed_results:
+            step_info['detailed'] = []
+            for i, result in enumerate(detailed_results):
+                step_info['detailed'].append({
+                    'sample_idx': i,
+                    'is_valid_peptide': result.get('is_valid_peptide', False),
+                    'has_conjugation_site': result.get('has_conjugation_site', False),
+                    'conjugation_success': result.get('conjugation_success', False),
+                    'pdc_valid': result.get('pdc_valid', False),
+                    'num_conjugates': len(result.get('conjugates', [])),
+                    'error': result.get('error')
+                })
+        self.step_logs.append(step_info)
+
+    def log_final(self, sequences, detailed_results_with_props, smiles_list=None):
+        for i, (seq, result) in enumerate(zip(sequences, detailed_results_with_props)):
+            final_info = {
+                'sample_idx': i,
+                'aa_sequence': seq,
+                'peptide_smiles': smiles_list[i] if smiles_list else seq,
+                'is_valid_peptide': result.get('is_valid_peptide', False),
+                'has_conjugation_site': result.get('has_conjugation_site', False),
+                'num_conjugates': len(result.get('conjugates', [])),
+                'num_valid_pdcs': len(result.get('valid_pdcs', [])),
+                'best_pdc': None,
+                'linker_used': result.get('linker_used', ''),
+                'all_conjugates': result.get('conjugates', []),
+                'properties': result.get('properties', {})
+            }
+            if result.get('best_pdc'):
+                final_info['best_pdc'] = {
+                    'smiles': result['best_pdc'].get('smiles', ''),
+                    'checks': result['best_pdc'].get('checks', {})
+                }
+            self.final_results.append(final_info)
+
+    def save(self, filename_prefix='generation'):
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_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', 'aa_sequence', 'peptide_smiles', 'is_valid_peptide',
+                'has_conjugation_site', 'linker_used', 'num_conjugates', 'num_valid_pdcs',
+                'best_pdc_smiles', 'mf', 'polararea', 'xlogp', 'heavycnt', 'rotbonds'
+            ])
+            for result in self.final_results:
+                props = result.get('properties', {})
+                writer.writerow([
+                    result['sample_idx'],
+                    result.get('aa_sequence', ''),
+                    result['peptide_smiles'],
+                    result['is_valid_peptide'],
+                    result['has_conjugation_site'],
+                    result.get('linker_used', ''),
+                    result['num_conjugates'],
+                    result['num_valid_pdcs'],
+                    result['best_pdc']['smiles'] if result.get('best_pdc') else '',
+                    props.get('mf', ''),
+                    props.get('polararea', ''),
+                    props.get('xlogp', ''),
+                    props.get('heavycnt', ''),
+                    props.get('rotbonds', '')
+                ])
+
+        pdc_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_pdcs.csv")
+        with open(pdc_path, 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow(['sample_idx', 'aa_sequence', 'peptide_smiles', 'linker_used',
+                             'conjugate_idx', 'pdc_smiles'])
+            for result in self.final_results:
+                for j, conj in enumerate(result.get('all_conjugates', [])):
+                    writer.writerow([
+                        result['sample_idx'],
+                        result.get('aa_sequence', ''),
+                        result['peptide_smiles'],
+                        result.get('linker_used', ''),
+                        j,
+                        conj
+                    ])
+
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        print(f"  PDC conjugates: {pdc_path}")
+        return step_log_path, final_path, csv_path, pdc_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)
+        num_valid = sum(1 for s in scores_list if s > 0)
+        print(f"\n--- Step {step} Summary ({guidance_method or 'unknown'}) ---")
+        print(f"  Valid: {num_valid}/{len(scores_list)}")
+        if detailed_results:
+            print(f"  Breakdown:")
+            print(f"    Valid peptides: {sum(1 for r in detailed_results if r.get('is_valid_peptide'))}")
+            print(f"    Has conjugation site: {sum(1 for r in detailed_results if r.get('has_conjugation_site'))}")
+            print(f"    Conjugation success: {sum(1 for r in detailed_results if r.get('conjugation_success'))}")
+            print(f"    PDC valid: {sum(1 for r in detailed_results if r.get('pdc_valid'))}")
+
+    def print_final_summary(self):
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        total = len(self.final_results)
+        valid_pdcs = sum(1 for r in self.final_results if r.get('num_valid_pdcs', 0) > 0)
+        print(f"\nTotal sequences: {total}")
+        print(f"With valid PDCs: {valid_pdcs}")
+        for result in self.final_results:
+            status = "" if result.get('num_valid_pdcs', 0) > 0 else ""
+            print(f"\n  [{status}] Sample {result['sample_idx']}:")
+            print(f"      AA seq: {result.get('aa_sequence', '')}")
+            print(f"      Peptide SMILES: {result['peptide_smiles'][:60]}...")
+            print(f"      Valid peptide: {result['is_valid_peptide']}")
+            print(f"      Has conjugation site: {result['has_conjugation_site']}")
+            print(f"      Conjugates: {result['num_conjugates']}")
+            print(f"      Valid PDCs: {result['num_valid_pdcs']}")
+
+
+# =============================================================================
+# CupOfMoCA Generator (UAA-aware)
+# =============================================================================
+
+class CupOfMoCAGeneratorUAA:
+    """CupOfMoCA Algorithm 2 generator that operates in AA token space.
+
+    AA tokens are converted to SMILES on the fly for scoring with the
+    Conjugation and BindingAffinity objectives.
+    """
+
+    def __init__(self, model, device, objectives, args, tokenizer, tokenizer_type='amino',
+                 tracker=None):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.pdc_analyzer = PDCAnalyzer(device)
+        self.tracker = tracker
+        self.tokenizer = tokenizer
+        self.tokenizer_type = tokenizer_type
+        self.guidance_method = GuidanceMethod(args.guidance_method)
+        self.profiler = Profiler(enabled=getattr(args, 'profile', False))
+
+        cascade_workers = max(1, getattr(args, 'cascade_workers', 1))
+        self._cascade_pool = (
+            ThreadPoolExecutor(max_workers=cascade_workers)
+            if cascade_workers > 1 else None
+        )
+
+        if tokenizer_type == 'amino':
+            # PepReDi natural-AA vocab is 24 tokens; 0-3 = cls/pad/eos/unk,
+            # 4-23 = canonical AAs. Block specials only. (UAA used id 38
+            # for D-AA 'g'; that's out of PepReDi vocab range.)
+            self.invalid_tokens = [0, 1, 2, 3]
+            self.smiles_tokenizer = SMILES_SPE_Tokenizer(
+                args.smiles_vocab_path, args.smiles_splits_path
+            )
+        else:
+            self.invalid_tokens = [0, 1, 2, 3, 4, 585, 586]
+            self.smiles_tokenizer = None
+
+        self.conjugation_obj = None
+        for obj in objectives:
+            if isinstance(obj, Conjugation):
+                self.conjugation_obj = obj
+                break
+
+    # ------------------------------------------------------------------
+    # x0 generation (Algorithm 2, Line 6)
+    # ------------------------------------------------------------------
+    def _generate_single_x0(self, n_steps=16, temperature=1.0):
+        self.model.eval()
+
+        if self.tokenizer_type == 'amino':
+            vocab_size = self.model.vocab_size
+            valid_tokens = [i for i in range(vocab_size) if i not in self.invalid_tokens]
+            valid_tokens = torch.tensor(valid_tokens, device=self.device)
+            rand_indices = torch.randint(
+                0, len(valid_tokens),
+                (self.args.num_samples, self.args.gen_len),
+                device=self.device
+            )
+            body = valid_tokens[rand_indices]
+            cls_col = torch.full((self.args.num_samples, 1), self.tokenizer.cls_id, device=self.device)
+            eos_col = torch.full((self.args.num_samples, 1), self.tokenizer.eos_id, device=self.device)
+            x = torch.cat([cls_col, body, eos_col], dim=1)
+        else:
+            vocab_size = self.model.tokenizer.vocab_size
+            x = torch.randint(
+                0, vocab_size,
+                (self.args.num_samples, self.args.gen_len),
+                device=self.device
+            )
+
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Initial Flow Matching", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i + 1]
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+
+                logits = self.model(x, t_tensor) / temperature
+                for inv_tok in self.invalid_tokens:
+                    if inv_tok < logits.shape[-1]:
+                        logits[:, :, inv_tok] = -1000
+
+                pred_x1 = torch.argmax(logits, dim=-1)
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                noise = torch.randint(0, vocab_size, x.shape, device=self.device)
+                x = torch.where(mask, noise, pred_x1)
+
+        return x
+
+    def _aa_to_smiles_batch(self, x_batch):
+        """Decode AA tokens, convert to SMILES via three-letter codes."""
+        with self.profiler.section("aa2smiles.decode"):
+            seqs = self.tokenizer.batch_decode(x_batch)
+        with self.profiler.section("aa2smiles.one_to_three"):
+            three_letter_seqs = [one_letter_to_three_letter(s) for s in seqs]
+        with self.profiler.section("aa2smiles.aa_to_smiles"):
+            with ThreadPoolExecutor() as pool:
+                smiles_seqs = list(pool.map(
+                    lambda s: aa_to_smiles(s, cyclic=self.args.cyclic),
+                    three_letter_seqs
+                ))
+        return seqs, smiles_seqs
+
+    def _check_conjugation_validity(self, x):
+        if self.tokenizer_type == 'amino':
+            _, smiles_seqs = self._aa_to_smiles_batch(x)
+            validities = [
+                self.peptide_analyzer.is_peptide(s) if isinstance(s, str) else False
+                for s in smiles_seqs
+            ]
+            return torch.tensor(validities, device=self.device)
+
+        if self.conjugation_obj is None:
+            sequences = self.tokenizer.batch_decode(x)
+            return torch.tensor(
+                [self.peptide_analyzer.is_peptide(seq) for seq in sequences],
+                device=self.device
+            )
+        return self.conjugation_obj(x) > 0
+
+    def generate_x0(self, n_steps=16, temperature=1.0,
+                    require_valid=False, max_retries=10):
+        print("Generating initial sequences (x0 from prior)...")
+        x = self._generate_single_x0(n_steps, temperature)
+
+        if require_valid:
+            print("Enforcing initial validity constraint...")
+            for retry in range(max_retries):
+                validity = self._check_conjugation_validity(x)
+                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 are valid!")
+                    break
+                print(f"  Retry {retry + 1}/{max_retries}: {num_valid}/{self.args.num_samples} valid")
+                invalid_indices = torch.where(~validity)[0]
+                original_num = self.args.num_samples
+                self.args.num_samples = num_invalid
+                new_x = self._generate_single_x0(n_steps, temperature)
+                self.args.num_samples = original_num
+                x[invalid_indices] = new_x
+
+            final_validity = self._check_conjugation_validity(x)
+            if final_validity.sum().item() < self.args.num_samples:
+                print(f"⚠ After {max_retries} retries, "
+                      f"{final_validity.sum().item()}/{self.args.num_samples} valid")
+
+        generated_sequences = self.tokenizer.batch_decode(x)
+        for seq in generated_sequences:
+            print(f"  {seq}")
+        return x
+
+    # ------------------------------------------------------------------
+    # Scoring
+    # ------------------------------------------------------------------
+    def _get_scores(self, x_batch):
+        with self.profiler.section("get_scores.total"):
+            scores = []
+
+            if self.tokenizer_type == 'amino':
+                _, smiles_seqs = self._aa_to_smiles_batch(x_batch)
+                with self.profiler.section("get_scores.valid_mask"):
+                    valid_mask = torch.tensor(
+                        [1 if isinstance(s, str) else 0 for s in smiles_seqs],
+                        device=self.device
+                    )
+                    for i in torch.where((1 - valid_mask) == 1)[0].tolist():
+                        smiles_seqs[i] = 'C'
+
+                with self.profiler.section("get_scores.smiles_tokenize"):
+                    smiles_batch = self.smiles_tokenizer(
+                        smiles_seqs, return_tensors='pt', padding=True
+                    ).to(self.device)
+                attention_mask = smiles_batch.get("attention_mask", None)
+
+                for obj_func in self.objectives:
+                    obj_name = type(obj_func).__name__
+                    with self.profiler.section(f"get_scores.obj.{obj_name}"):
+                        if getattr(obj_func, "input_type", None) == "aa":
+                            # AA-input objective: score raw AA tokens (no
+                            # SMILES conversion, no valid_mask multiply).
+                            score = obj_func(x_batch.to(self.device)).to(self.device)
+                            scores.append(score)
+                        else:
+                            try:
+                                score = obj_func(
+                                    smiles_batch["input_ids"], attention_mask=attention_mask
+                                ).to(self.device)
+                            except TypeError:
+                                score = obj_func(smiles_batch["input_ids"]).to(self.device)
+                            scores.append(score * valid_mask)
+            else:
+                for obj_func in self.objectives:
+                    obj_name = type(obj_func).__name__
+                    with self.profiler.section(f"get_scores.obj.{obj_name}"):
+                        scores.append(obj_func(x_batch.to(self.device)))
+
+            return torch.stack(scores, dim=0).to(self.device)
+
+    def _compute_aggregated_score(self, scores, weights):
+        if self.args.score_aggregation == "weighted_min":
+            return torch.min(weights * scores, dim=0).values
+        return (weights * scores).sum(dim=0)
+
+    # ------------------------------------------------------------------
+    # Cascade scoring (cheap filter + affinity on all → full PDC on top-K)
+    # ------------------------------------------------------------------
+    def _get_scores_cascaded(self, x_batch):
+        """Run cheap conjugation pre-filter + BindingAffinity on all candidates,
+        then run the expensive PDC validation only on the top `cascade_k` by
+        affinity (among those that pass the cheap filter).
+
+        Returns [num_objectives, batch] in the same layout as `_get_scores`.
+        Falls back to full scoring if Conjugation isn't an objective or
+        tokenizer_type != 'amino'.
+        """
+        if self.tokenizer_type != 'amino' or self.conjugation_obj is None:
+            return self._get_scores(x_batch)
+
+        with self.profiler.section("cascade.total"):
+            n = x_batch.shape[0]
+            cascade_k = self.args.cascade_k
+            conj_obj = self.conjugation_obj
+
+            _, smiles_seqs = self._aa_to_smiles_batch(x_batch)
+            valid_mask = torch.tensor(
+                [1.0 if isinstance(s, str) else 0.0 for s in smiles_seqs],
+                device=self.device
+            )
+            safe_smiles = [s if isinstance(s, str) else 'C' for s in smiles_seqs]
+
+            with self.profiler.section("cascade.cheap_filter"):
+                cheap_pass = []
+                for s, vm in zip(safe_smiles, valid_mask.tolist()):
+                    if vm == 0.0:
+                        cheap_pass.append(False)
+                        continue
+                    cheap_pass.append(
+                        conj_obj.is_valid_peptide(s)
+                        and conj_obj.has_conjugation_site(s)
+                    )
+                cheap_pass_t = torch.tensor(
+                    cheap_pass, device=self.device, dtype=torch.float
+                )
+
+            with self.profiler.section("cascade.smiles_tokenize"):
+                smiles_batch = self.smiles_tokenizer(
+                    safe_smiles, return_tensors='pt', padding=True
+                ).to(self.device)
+                attention_mask = smiles_batch.get("attention_mask", None)
+
+            aff_obj = next(
+                (o for o in self.objectives if not isinstance(o, Conjugation)),
+                None
+            )
+            if aff_obj is None:
+                aff_scores = torch.zeros(n, device=self.device)
+            else:
+                with self.profiler.section(f"cascade.obj.{type(aff_obj).__name__}"):
+                    if getattr(aff_obj, "input_type", None) == "aa":
+                        # AA-input affinity: score raw AA tokens directly.
+                        # Don't multiply by valid_mask (the AA->SMILES outcome
+                        # is irrelevant for the AA-input scorer).
+                        aff_scores = aff_obj(x_batch.to(self.device)).to(self.device)
+                    else:
+                        try:
+                            aff_scores = aff_obj(
+                                smiles_batch["input_ids"], attention_mask=attention_mask
+                            ).to(self.device)
+                        except TypeError:
+                            aff_scores = aff_obj(smiles_batch["input_ids"]).to(self.device)
+                        aff_scores = aff_scores * valid_mask
+
+            eligible_idx = torch.where(cheap_pass_t > 0)[0]
+            if cascade_k > 0 and len(eligible_idx) > cascade_k:
+                topk_local = torch.topk(aff_scores[eligible_idx], k=cascade_k).indices
+                validate_idx = eligible_idx[topk_local]
+            else:
+                validate_idx = eligible_idx
+
+            conj_scores = torch.zeros(n, device=self.device)
+            with self.profiler.section("cascade.full_validate"):
+                for j in validate_idx.tolist():
+                    ok, conjs, _, _ = conj_obj.conjugate_peptide(
+                        safe_smiles[j], exclude_n_terminus=True
+                    )
+                    if not ok or not conjs:
+                        continue
+                    has_valid, _ = conj_obj.validate_conjugates(conjs)
+                    if has_valid:
+                        conj_scores[j] = 1.0
+
+            out = []
+            for obj in self.objectives:
+                if isinstance(obj, Conjugation):
+                    out.append(conj_scores)
+                else:
+                    out.append(aff_scores)
+            return torch.stack(out, dim=0).to(self.device)
+
+    # ------------------------------------------------------------------
+    # Guidance methods
+    # ------------------------------------------------------------------
+    def _exponential_tilt_guidance(self, x, pos_probs, candidates, mut_idx, sample_idx,
+                                   beta, weights, current_score=None):
+        num_candidates = len(candidates)
+        current_token = x[sample_idx, mut_idx].item()
+
+        if num_candidates == 0:
+            return current_token, {'method': 'exp_tilt', 'fallback': True, 'accepted': False}
+
+        if num_candidates > self.args.top_k:
+            candidates = candidates[:self.args.top_k]
+            num_candidates = self.args.top_k
+
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(num_candidates, 1)
+        x_proposals[:, mut_idx] = candidates
+
+        with self.profiler.section("exp_tilt.proposal_scores"):
+            if self.args.cascade_k > 0:
+                proposal_scores = self._get_scores_cascaded(x_proposals)
+            else:
+                proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+        validity_mask = (proposal_scores[0] > 0).float()
+        base_probs = pos_probs[sample_idx, candidates]
+
+        log_base_probs = torch.log(base_probs + 1e-10)
+        log_exp_scores = beta * aggregated_scores
+        log_validity = torch.log(validity_mask + 1e-10)
+        log_tilted = log_base_probs + log_exp_scores + log_validity
+
+        log_tilted_max = log_tilted.max()
+        fallback = False
+        if torch.isinf(log_tilted_max) or torch.isnan(log_tilted_max):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+        else:
+            log_tilted_stable = log_tilted - log_tilted_max
+            tilted_probs = torch.exp(log_tilted_stable)
+            prob_sum = tilted_probs.sum()
+            if prob_sum < 1e-10 or torch.isnan(prob_sum) or torch.isinf(prob_sum):
+                fallback = True
+                final_probs = base_probs / (base_probs.sum() + 1e-9)
+            else:
+                final_probs = tilted_probs / prob_sum
+
+        if torch.any(torch.isnan(final_probs)) or torch.any(torch.isinf(final_probs)):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+
+        final_probs = torch.clamp(final_probs, min=1e-10)
+        final_probs = final_probs / final_probs.sum()
+
+        selected_idx = torch.multinomial(final_probs, 1).item()
+        selected_token = candidates[selected_idx].item()
+        selected_score = aggregated_scores[selected_idx].item()
+
+        accepted = True
+        if self.args.use_acceptance and current_score is not None:
+            if selected_score < current_score:
+                selected_token = current_token
+                accepted = False
+
+        return selected_token, {
+            'method': 'exp_tilt', 'beta': beta, 'num_candidates': num_candidates,
+            'proposal_score': aggregated_scores[selected_idx].item(),
+            'selected_score': selected_score if accepted else current_score,
+            'fallback': fallback, 'validity_rate': validity_mask.mean().item(),
+            'accepted': accepted
+        }
+
+    def _classifier_guidance(self, x, logits, mut_idx, sample_idx,
+                             guidance_weight, weights):
+        base_logits = logits[sample_idx, mut_idx, :]
+        sorted_logits, sorted_indices = torch.sort(base_logits, descending=True)
+        candidates = sorted_indices[:self.args.top_k]
+
+        valid_mask = ~torch.isin(
+            candidates, torch.tensor(self.invalid_tokens, device=self.device)
+        )
+        candidates = candidates[valid_mask]
+
+        if len(candidates) == 0:
+            return x[sample_idx, mut_idx].item(), {'method': 'cg', 'fallback': True}
+
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(len(candidates), 1)
+        x_proposals[:, mut_idx] = candidates
+
+        with self.profiler.section("cg.proposal_scores"):
+            if self.args.cascade_k > 0:
+                proposal_scores = self._get_scores_cascaded(x_proposals)
+            else:
+                proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+
+        score_logits = aggregated_scores * self.args.score_temperature
+        base_candidate_logits = base_logits[candidates]
+        guided_logits = base_candidate_logits + guidance_weight * score_logits
+        guided_probs = F.softmax(guided_logits, dim=-1)
+
+        selected_idx = torch.multinomial(guided_probs, 1).item()
+        return candidates[selected_idx].item(), {
+            'method': 'cg', 'guidance_weight': guidance_weight,
+            'num_candidates': len(candidates),
+            'selected_score': aggregated_scores[selected_idx].item(), 'fallback': False
+        }
+
+    def _best_of_k_guidance(self, x, pos_probs, candidates, mut_idx, sample_idx,
+                            k, weights):
+        num_candidates = len(candidates)
+        if num_candidates == 0:
+            return x[sample_idx, mut_idx].item(), {'method': 'bok', 'fallback': True}
+
+        k = min(k, num_candidates)
+        candidate_probs = pos_probs[sample_idx, candidates]
+        candidate_probs = candidate_probs / (candidate_probs.sum() + 1e-9)
+
+        if k >= num_candidates:
+            sampled_indices = torch.arange(num_candidates, device=self.device)
+        else:
+            sampled_indices = torch.multinomial(candidate_probs, k, replacement=False)
+
+        sampled_tokens = candidates[sampled_indices]
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(len(sampled_tokens), 1)
+        x_proposals[:, mut_idx] = sampled_tokens
+
+        with self.profiler.section("bok.proposal_scores"):
+            if self.args.cascade_k > 0:
+                proposal_scores = self._get_scores_cascaded(x_proposals)
+            else:
+                proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+
+        best_idx = torch.argmax(aggregated_scores).item()
+        return sampled_tokens[best_idx].item(), {
+            'method': 'bok', 'k': k, 'num_candidates': num_candidates,
+            'selected_score': aggregated_scores[best_idx].item(),
+            'mean_score': aggregated_scores.mean().item(),
+            'max_score': aggregated_scores.max().item(), 'fallback': False
+        }
+
+    def _hybrid_guidance(self, x, pos_probs, logits, candidates, mut_idx, sample_idx,
+                         step, beta, weights):
+        progress = step / max(self.args.optimization_steps - 1, 1)
+        if progress < 0.3:
+            return self._best_of_k_guidance(
+                x, pos_probs, candidates, mut_idx, sample_idx,
+                k=self.args.bok_k, weights=weights
+            )
+        elif progress < 0.7:
+            return self._exponential_tilt_guidance(
+                x, pos_probs, candidates, mut_idx, sample_idx,
+                beta=beta, weights=weights
+            )
+        else:
+            return self._classifier_guidance(
+                x, logits, mut_idx, sample_idx,
+                guidance_weight=self.args.cg_weight, weights=weights
+            )
+
+    # ------------------------------------------------------------------
+    # Main loop
+    # ------------------------------------------------------------------
+    def generate(self):
+        self.profiler.reset()
+        x = self.generate_x0(
+            require_valid=self.args.require_valid_x0,
+            max_retries=self.args.max_x0_retries
+        )
+
+        if self.args.weights is None:
+            weights = torch.tensor(
+                [1.0] + [1.0 / (self.num_objectives - 1)] * (self.num_objectives - 1),
+                device=self.device
+            ).view(-1, 1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+
+        print(f"\nGuidance Method: {self.guidance_method.value}")
+        print(f"Weights: {weights.squeeze().tolist()}")
+        print(f"β range: [{self.args.beta_min}, {self.args.beta_max}]")
+        print(f"Acceptance criterion: {'ENABLED' if self.args.use_acceptance else 'DISABLED'}")
+
+        with torch.no_grad():
+            for step in tqdm(range(self.args.optimization_steps), desc="Guided Generation"):
+                improved_indices = []
+
+                progress = step / max(self.args.optimization_steps - 1, 1)
+                beta = self.args.beta_min + (self.args.beta_max - self.args.beta_min) * progress
+
+                if self.args.position_strategy == "random":
+                    mut_idx = random.randint(1, self.args.gen_len - 2)
+                else:
+                    mut_idx = (step % (self.args.gen_len - 2)) + 1
+
+                time_t = torch.full(
+                    (self.args.num_samples,),
+                    step / self.args.optimization_steps,
+                    device=self.device
+                )
+
+                with self.profiler.section("loop.model_forward"):
+                    logits = self.model(x, time_t)
+                    probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+
+                pos_probs[torch.arange(self.args.num_samples), x[:, mut_idx]] = 0
+                for inv_tok in self.invalid_tokens:
+                    pos_probs[:, inv_tok] = 0
+
+                with self.profiler.section("loop.top_p_sort"):
+                    sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+
+                with self.profiler.section("loop.current_scores"):
+                    current_scores = self._get_scores(x)
+                    current_aggregated = self._compute_aggregated_score(current_scores, weights)
+
+                if step == 0:
+                    print(f"\nInitial Scores: {current_scores.cpu().numpy()}")
+                    if self.tracker and self.conjugation_obj:
+                        detailed = self.conjugation_obj.get_last_detailed_results()
+                        self.tracker.log_step(step, current_scores[0], detailed)
+                        self.tracker.print_step_summary(
+                            step, current_scores[0], detailed, self.guidance_method.value
+                        )
+
+                final_tokens = []
+                for i in range(self.args.num_samples):
+                    sample_probs = sorted_probs[i]
+                    sample_indices = sorted_indices[i]
+
+                    cumsum = torch.cumsum(sample_probs, dim=0)
+                    mask = cumsum <= self.args.top_p
+                    mask[0] = True
+                    candidates = sample_indices[mask]
+
+                    valid_candidate_mask = ~torch.isin(
+                        candidates, torch.tensor(self.invalid_tokens, device=self.device)
+                    )
+                    candidates = candidates[valid_candidate_mask]
+
+                    current_sample_score = current_aggregated[i].item()
+
+                    if beta > 0 or self.guidance_method != GuidanceMethod.EXPONENTIAL_TILT:
+                        with self.profiler.section(f"loop.guidance.{self.guidance_method.value}"):
+                            if self.guidance_method == GuidanceMethod.EXPONENTIAL_TILT:
+                                token, _ = self._exponential_tilt_guidance(
+                                    x, pos_probs, candidates, mut_idx, i, beta, weights,
+                                    current_score=current_sample_score
+                                )
+                            elif self.guidance_method == GuidanceMethod.CLASSIFIER_GUIDANCE:
+                                token, _ = self._classifier_guidance(
+                                    x, logits, mut_idx, i, self.args.cg_weight, weights
+                                )
+                            elif self.guidance_method == GuidanceMethod.BEST_OF_K:
+                                token, _ = self._best_of_k_guidance(
+                                    x, pos_probs, candidates, mut_idx, i, self.args.bok_k, weights
+                                )
+                            elif self.guidance_method == GuidanceMethod.HYBRID:
+                                token, _ = self._hybrid_guidance(
+                                    x, pos_probs, logits, candidates, mut_idx, i, step, beta, weights
+                                )
+                            else:
+                                raise ValueError(f"Unknown guidance: {self.guidance_method}")
+                    else:
+                        if len(candidates) == 0:
+                            token = x[i, mut_idx].item()
+                        else:
+                            cprobs = pos_probs[i, candidates]
+                            cprobs = cprobs / (cprobs.sum() + 1e-9)
+                            idx = torch.multinomial(cprobs, 1).item()
+                            token = candidates[idx].item()
+
+                    final_tokens.append(token)
+
+                x[:, mut_idx] = torch.tensor(final_tokens, device=self.device)
+
+                with self.profiler.section("loop.new_scores"):
+                    new_scores = self._get_scores(x)
+                    new_aggregated = self._compute_aggregated_score(new_scores, weights)
+                for i in range(self.args.num_samples):
+                    if new_aggregated[i] > current_aggregated[i]:
+                        improved_indices.append(i)
+
+                if improved_indices:
+                    for i in improved_indices:
+                        prev_weighted = (weights.squeeze(1) * current_scores[:, i]).sum()
+                        new_weighted = (weights.squeeze(1) * new_scores[:, i]).sum()
+                        print(f"Previous Weighted Sum: {prev_weighted}")
+                        print(f"Previous Scores: {current_scores[:, i]}")
+                        print(f"New Weighted Sum: {new_weighted}")
+                        print(f"New Scores: {new_scores[:, i]}")
+                    print(self.tokenizer.batch_decode(x[improved_indices]))
+
+                if self.tracker and self.conjugation_obj and \
+                   (step % 4 == 0 or step == self.args.optimization_steps - 1):
+                    detailed = self.conjugation_obj.get_last_detailed_results()
+                    self.tracker.log_step(
+                        step, new_scores[0], detailed, improved_indices,
+                        guidance_info={'method': self.guidance_method.value, 'beta': beta}
+                    )
+                    if step == self.args.optimization_steps - 1:
+                        self.tracker.print_step_summary(
+                            step, new_scores[0], detailed, self.guidance_method.value
+                        )
+
+                if (self.profiler.enabled
+                    and self.args.profile_every > 0
+                    and (step + 1) % self.args.profile_every == 0):
+                    self.profiler.print_summary(header=f"(through step {step + 1})")
+
+        self.profiler.print_summary(header="(end of generate)")
+        return x
+
+
+# =============================================================================
+# Main
+# =============================================================================
+
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    print(f"Tokenizer type: {args.tokenizer_type}")
+    print(f"Guidance method: {args.guidance_method}")
+
+    if args.tokenizer_type == 'amino':
+        print("Using Amino Acid tokenizer")
+        tokenizer = AminoTokenizer()
+        smiles_tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+    else:
+        print("Using SMILES tokenizer")
+        tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+        smiles_tokenizer = tokenizer
+
+    conjugation_model = Conjugation(smiles_tokenizer, device=device, builder_type=args.builder_type)
+    # Binding affinity: AA-input or SMILES-input depending on --affinity_type.
+    affinity_model = make_affinity(
+        args.affinity_type,
+        target_protein_seq=args.target,
+        device=device,
+        aa_ckpt_path=args.aa_affinity_ckpt,
+    )
+    print(f"Affinity scorer: {type(affinity_model).__name__} (input_type={getattr(affinity_model, 'input_type', '?')})")
+    OBJECTIVE_FUNCTIONS = [conjugation_model, affinity_model]
+
+    tracker = GenerationTracker(output_dir=args.output_dir)
+
+    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)
+        print(f"Random seed: {args.seed}")
+
+    print(f"Loading model from: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+
+    if args.tokenizer_type == 'amino':
+        raw_model_args = checkpoint['args']
+        if isinstance(raw_model_args, argparse.Namespace):
+            raw_model_args = vars(raw_model_args)
+        model_args = edict(raw_model_args)
+        model = MDLM(
+            vocab_size=model_args.vocab_size,
+            seq_len=model_args.seq_len,
+            model_dim=model_args.model_dim,
+            n_heads=model_args.n_heads,
+            n_layers=model_args.n_layers,
+        ).to(device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print(f"Amino model loaded. Vocab size: {model_args.vocab_size}")
+    else:
+        model_hparams = checkpoint["hyper_parameters"]["args"]
+        model = MDLMLightningModule.load_from_checkpoint(
+            args.checkpoint, args=model_hparams, tokenizer=tokenizer,
+            map_location=device, strict=False
+        )
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully.")
+
+    generator = CupOfMoCAGeneratorUAA(
+        model, device, OBJECTIVE_FUNCTIONS, args,
+        tokenizer=tokenizer, tokenizer_type=args.tokenizer_type, tracker=tracker
+    )
+
+    all_sequences = []
+    all_smiles = []
+    all_final_results = []
+
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'=' * 60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'=' * 60}")
+
+        generated_tokens = generator.generate()
+        final_scores = generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+
+        print("\n--- Batch Results ---")
+        print(f"Sequences: {sequence_str}")
+        print(f"Scores: {final_scores}")
+
+        all_sequences.extend(sequence_str)
+
+        if args.tokenizer_type == 'amino':
+            smiles_list = []
+            for seq in sequence_str:
+                three_letter = one_letter_to_three_letter(seq)
+                smiles = aa_to_smiles(three_letter, cyclic=args.cyclic)
+                smiles_list.append(smiles if isinstance(smiles, str) else 'C')
+            smiles_tokens = smiles_tokenizer(
+                smiles_list, return_tensors='pt', padding=True
+            ).to(device)
+            detailed_with_props = conjugation_model.get_detailed_scores_with_properties(
+                smiles_tokens["input_ids"]
+            )
+            all_smiles.extend(smiles_list)
+            tracker.log_final(sequence_str, detailed_with_props, smiles_list=smiles_list)
+        else:
+            detailed_with_props = conjugation_model.get_detailed_scores_with_properties(generated_tokens)
+            tracker.log_final(sequence_str, detailed_with_props)
+        all_final_results.extend(detailed_with_props)
+
+    tracker.print_final_summary()
+    tracker.save(filename_prefix=f'cupofmoca_uaa_{args.guidance_method}')
+
+    with open(args.output_file, 'w') as f:
+        for i, seq in enumerate(all_sequences):
+            if args.tokenizer_type == 'amino':
+                smiles = all_smiles[i] if i < len(all_smiles) else ''
+                f.write(f"{seq}\t{smiles}\n")
+            else:
+                f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+    print("\nGeneration complete.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="CupOfMoCA-style Guided Generation with UAA Foundation Model"
+    )
+
+    parser.add_argument("--checkpoint", type=str, required=True)
+    parser.add_argument("--target", type=str, required=True)
+
+    parser.add_argument("--num_samples", type=int, default=10)
+    parser.add_argument("--num_batches", type=int, default=10)
+    parser.add_argument("--gen_len", type=int, default=50)
+    parser.add_argument("--optimization_steps", type=int, default=128)
+
+    parser.add_argument("--guidance_method", type=str, default="exp_tilt",
+                        choices=["exp_tilt", "cg", "bok", "hybrid"])
+
+    parser.add_argument("--beta_min", type=float, default=1.0)
+    parser.add_argument("--beta_max", type=float, default=20.0)
+    parser.add_argument("--cg_weight", type=float, default=2.0)
+    parser.add_argument("--use_acceptance", action="store_true")
+    parser.add_argument("--bok_k", type=int, default=8)
+    parser.add_argument("--top_k", type=int, default=200)
+    parser.add_argument("--top_p", type=float, default=0.9)
+    parser.add_argument("--score_temperature", type=float, default=1.0)
+    parser.add_argument("--score_aggregation", type=str, default="weighted_sum",
+                        choices=["weighted_sum", "weighted_min"])
+
+    parser.add_argument("--weights", type=float, nargs='+', default=None)
+
+    parser.add_argument("--position_strategy", type=str, default="random",
+                        choices=["random", "sequential"])
+
+    parser.add_argument("--require_valid_x0", action="store_true")
+    parser.add_argument("--max_x0_retries", type=int, default=10)
+
+    parser.add_argument("--builder_type", type=str, default="AoaDau",
+                        choices=["AoaDau", "ValAlaDau", "ValCitDau", "GlutarylDau", "Dau",
+                                 "SSCpt", "VcMMAE", "DSPACpt", "MalHexCpt", "DSECCpt", "Cpt",
+                                 "DSPASN38", "DSEBSN38", "DSEPSN38", "SN38"])
+
+    parser.add_argument("--tokenizer_type", type=str, default="amino",
+                        choices=["smiles", "amino"])
+    parser.add_argument("--cyclic", action="store_true",
+                        help="Generate cyclic peptides (only for amino tokenizer).")
+    parser.add_argument("--smiles_vocab_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_vocab.txt")
+    parser.add_argument("--smiles_splits_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_splits.txt")
+
+    # AA-mode binding-affinity selection (this is the WT-AA / PepReDi script).
+    parser.add_argument("--affinity_type", type=str, default="smiles",
+                        choices=["aa", "smiles"],
+                        help="Which binding-affinity oracle to guide on: "
+                             "'aa' = AA-input UnpooledBindingPredictor "
+                             "(matches moPPIt's setup); 'smiles' = "
+                             "SMILES-input BindingAffinity (matches the "
+                             "SMILES CupOFMoCA tables).")
+    parser.add_argument("--aa_affinity_ckpt", type=str,
+                        default=DEFAULT_AA_AFFINITY_CKPT,
+                        help="Path to AA-input affinity predictor checkpoint.")
+
+    parser.add_argument("--output_file", type=str, default="./smiles.txt")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs")
+
+    parser.add_argument("--seed", type=int, default=None)
+
+    parser.add_argument("--profile", action="store_true",
+                        help="Enable per-section wall-clock profiling.")
+    parser.add_argument("--profile_every", type=int, default=0,
+                        help="Print profiler summary every N optimization steps "
+                             "(0 = only at end of each generate()).")
+
+    parser.add_argument("--cascade_k", type=int, default=0,
+                        help="Cascade scoring: score all proposals with cheap "
+                             "Conjugation pre-filter + BindingAffinity, then run "
+                             "the expensive PDC validation only on the top-K by "
+                             "affinity (among cheap-filter survivors). 0 disables "
+                             "(full scoring on all proposals, original behavior).")
+    parser.add_argument("--cascade_workers", type=int, default=1,
+                        help="Number of worker threads for parallelizing "
+                             "cascade.full_validate (CPU-bound RDKit work). "
+                             "1 = sequential (original behavior). Try 4-8 when "
+                             "--cascade_k is high. Ignored if --cascade_k=0.")
+
+    args = parser.parse_args()
+    main(args)

smiles/generation/cupofmoca_guided_uaa.py

@@ -0,0 +1,1196 @@
+"""
+CupOfMoCA-style Guided Generation with UAA Foundation Model
+
+Combines CupOfMoCA Algorithm 2 guidance strategies (exponential tilting,
+classifier guidance, best-of-K, hybrid) with the amino-acid (UAA) foundation
+model. Sequences are generated in AA token space, converted to SMILES for
+scoring against Conjugation (peptide+PDC validity) and BindingAffinity.
+"""
+
+import argparse
+import csv
+import json
+import os
+import random
+import sys
+import time
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from datetime import datetime
+from enum import Enum
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from easydict import EasyDict as edict
+from tqdm import tqdm
+
+from smiles.models.train import MDLMLightningModule
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+from smiles.evaluation.peptide_analyzer import PeptideAnalyzer
+from smiles.evaluation.pdc_analyzer import PDCAnalyzer
+from smiles.scoring.classifiers import Conjugation, BindingAffinity
+
+from third_party.SMILESToPeptide.peptide_to_smiles import aa_to_smiles
+from third_party.SMILESToPeptide.aminoacid_selective import one_letter_to_three_letter
+
+
+# =============================================================================
+# Amino Acid Foundation Model (DiT-based MDLM)
+# =============================================================================
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(1, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+
+    def forward(self, t):
+        return self.mlp(t.unsqueeze(-1))
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, hidden_size, n_heads):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = nn.MultiheadAttention(hidden_size, n_heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, 4 * hidden_size),
+            nn.GELU(),
+            nn.Linear(4 * hidden_size, hidden_size),
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True),
+        )
+
+    def forward(self, x, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
+            self.adaLN_modulation(c).chunk(6, dim=1)
+        x_norm1 = modulate(self.norm1(x), shift_msa, scale_msa)
+        attn_output, _ = self.attn(x_norm1, x_norm1, x_norm1)
+        x = x + gate_msa.unsqueeze(1) * attn_output
+        x_norm2 = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        mlp_output = self.mlp(x_norm2)
+        x = x + gate_mlp.unsqueeze(1) * mlp_output
+        return x
+
+
+class MDLM(nn.Module):
+    def __init__(self, vocab_size, seq_len, model_dim, n_heads, n_layers):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.seq_len = seq_len
+        self.model_dim = model_dim
+
+        self.token_embedder = nn.Embedding(vocab_size, model_dim)
+        self.pos_embedder = nn.Parameter(torch.randn(1, seq_len, model_dim))
+        self.time_embedder = TimestepEmbedder(model_dim)
+
+        self.transformer_blocks = nn.ModuleList([
+            DiTBlock(model_dim, n_heads) for _ in range(n_layers)
+        ])
+        self.final_norm = nn.LayerNorm(model_dim)
+        self.lm_head = nn.Linear(model_dim, vocab_size)
+
+    def forward(self, x, t):
+        seq_len = x.shape[1]
+        x_embed = self.token_embedder(x) + self.pos_embedder[:, :seq_len, :]
+        t_embed = self.time_embedder(t)
+        for block in self.transformer_blocks:
+            x_embed = block(x_embed, t_embed)
+        x_embed = self.final_norm(x_embed)
+        return self.lm_head(x_embed)
+
+
+class AminoTokenizer:
+    def __init__(self):
+        self.vocab = {
+            '<cls>': 0, '<pad>': 1, '<eos>': 2, '<unk>': 3,
+            'L': 4, 'A': 5, 'G': 6, 'V': 7, 'S': 8, 'E': 9,
+            'R': 10, 'T': 11, 'I': 12, 'D': 13, 'P': 14, 'K': 15,
+            'Q': 16, 'N': 17, 'F': 18, 'Y': 19, 'M': 20, 'H': 21,
+            'W': 22, 'C': 23,
+            'Ŷ': 24, 'Ĉ': 25, 'Ĝ': 26, 'Ė': 27, 'Ƙ': 28,
+            'Ṃ': 29, 'Ł': 30, 'Ċ': 31, 'Ṯ': 32,
+            'a': 33, 'c': 34, 'd': 35, 'e': 36, 'f': 37,
+            'g': 38, 'h': 39, 'i': 40, 'k': 41, 'l': 42,
+            'm': 43, 'n': 44, 'p': 45, 'q': 46, 'r': 47,
+            's': 48, 't': 49, 'v': 50, 'w': 51, 'y': 52,
+        }
+        self.id_to_token = {i: t for t, i in self.vocab.items()}
+        self.cls_token, self.pad_token = '<cls>', '<pad>'
+        self.eos_token, self.unk_token = '<eos>', '<unk>'
+        self.cls_id = self.vocab[self.cls_token]
+        self.pad_id = self.vocab[self.pad_token]
+        self.eos_id = self.vocab[self.eos_token]
+        self.unk_id = self.vocab[self.unk_token]
+        self.vocab_size = len(self.vocab)
+
+    def decode(self, ids, skip_special_tokens=True):
+        if isinstance(ids, torch.Tensor):
+            ids = ids.tolist()
+        tokens = []
+        for i in ids:
+            tok = self.id_to_token.get(int(i), self.unk_token)
+            if skip_special_tokens and tok in {self.cls_token, self.eos_token, self.pad_token}:
+                continue
+            tokens.append(tok)
+        return ''.join(tokens)
+
+    def batch_decode(self, batch_ids, skip_special_tokens=True):
+        return [self.decode(ids, skip_special_tokens=skip_special_tokens) for ids in batch_ids]
+
+
+# =============================================================================
+# Profiler
+# =============================================================================
+
+class Profiler:
+    """Aggregates wall-clock time per named section. Optionally syncs CUDA.
+
+    `reset()` clears all state (use between batches for self-contained summaries).
+    `print_summary()` shows cumulative totals since the last reset, plus a
+    delta column showing time since the previous `print_summary()` call.
+    """
+
+    def __init__(self, enabled=False, cuda_sync=True):
+        self.enabled = enabled
+        self.cuda_sync = cuda_sync and torch.cuda.is_available()
+        self.totals = defaultdict(float)
+        self.counts = defaultdict(int)
+        self._last_totals = defaultdict(float)
+        self._last_counts = defaultdict(int)
+
+    @contextmanager
+    def section(self, name):
+        if not self.enabled:
+            yield
+            return
+        if self.cuda_sync:
+            torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        try:
+            yield
+        finally:
+            if self.cuda_sync:
+                torch.cuda.synchronize()
+            self.totals[name] += time.perf_counter() - t0
+            self.counts[name] += 1
+
+    def reset(self):
+        self.totals.clear()
+        self.counts.clear()
+        self._last_totals.clear()
+        self._last_counts.clear()
+
+    def print_summary(self, header=""):
+        if not self.enabled or not self.totals:
+            return
+        cumulative_total = sum(self.totals.values())
+        delta_totals = {k: self.totals[k] - self._last_totals.get(k, 0.0)
+                        for k in self.totals}
+        delta_counts = {k: self.counts[k] - self._last_counts.get(k, 0)
+                        for k in self.counts}
+        interval_total = sum(delta_totals.values())
+
+        print(f"\n=== Profiler summary {header} ===")
+        print(f"  (cumulative since reset | interval since previous summary)")
+        print(f"  {'section':<34} "
+              f"{'calls':>6} {'cum(s)':>8} {'cum_avg':>8}  "
+              f"{'Δcalls':>6} {'Δ(s)':>7} {'Δavg':>7} {'%Δ':>5}")
+        for name, tot in sorted(self.totals.items(), key=lambda kv: -kv[1]):
+            n = self.counts[name]
+            cum_avg = (tot / n) * 1000.0 if n else 0.0
+            dt = delta_totals[name]
+            dn = delta_counts[name]
+            d_avg = (dt / dn) * 1000.0 if dn else 0.0
+            d_pct = (dt / interval_total) * 100.0 if interval_total > 0 else 0.0
+            print(f"  {name:<34} "
+                  f"{n:>6d} {tot:>8.2f} {cum_avg:>7.1f}ms  "
+                  f"{dn:>6d} {dt:>7.2f} {d_avg:>6.1f}ms {d_pct:>4.1f}%")
+        print(f"  {'-' * 34} {'-' * 6} {'-' * 8} {'-' * 8}  "
+              f"{'-' * 6} {'-' * 7} {'-' * 7} {'-' * 5}")
+        print(f"  {'TOTAL':<34} {'':>6} {cumulative_total:>8.2f} {'':>8}  "
+              f"{'':>6} {interval_total:>7.2f}")
+
+        self._last_totals = dict(self.totals)
+        self._last_counts = dict(self.counts)
+
+
+# =============================================================================
+# Guidance Method Enum
+# =============================================================================
+
+class GuidanceMethod(Enum):
+    EXPONENTIAL_TILT = "exp_tilt"
+    CLASSIFIER_GUIDANCE = "cg"
+    BEST_OF_K = "bok"
+    HYBRID = "hybrid"
+
+
+# =============================================================================
+# Generation Tracker
+# =============================================================================
+
+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, 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(sum(1 for s in scores if s > 0)) if hasattr(scores, '__iter__') else 0,
+            'improved_indices': improved_indices or [],
+            'guidance_info': guidance_info or {}
+        }
+        if detailed_results:
+            step_info['detailed'] = []
+            for i, result in enumerate(detailed_results):
+                step_info['detailed'].append({
+                    'sample_idx': i,
+                    'is_valid_peptide': result.get('is_valid_peptide', False),
+                    'has_conjugation_site': result.get('has_conjugation_site', False),
+                    'conjugation_success': result.get('conjugation_success', False),
+                    'pdc_valid': result.get('pdc_valid', False),
+                    'num_conjugates': len(result.get('conjugates', [])),
+                    'error': result.get('error')
+                })
+        self.step_logs.append(step_info)
+
+    def log_final(self, sequences, detailed_results_with_props, smiles_list=None):
+        for i, (seq, result) in enumerate(zip(sequences, detailed_results_with_props)):
+            final_info = {
+                'sample_idx': i,
+                'aa_sequence': seq,
+                'peptide_smiles': smiles_list[i] if smiles_list else seq,
+                'is_valid_peptide': result.get('is_valid_peptide', False),
+                'has_conjugation_site': result.get('has_conjugation_site', False),
+                'num_conjugates': len(result.get('conjugates', [])),
+                'num_valid_pdcs': len(result.get('valid_pdcs', [])),
+                'best_pdc': None,
+                'linker_used': result.get('linker_used', ''),
+                'all_conjugates': result.get('conjugates', []),
+                'properties': result.get('properties', {})
+            }
+            if result.get('best_pdc'):
+                final_info['best_pdc'] = {
+                    'smiles': result['best_pdc'].get('smiles', ''),
+                    'checks': result['best_pdc'].get('checks', {})
+                }
+            self.final_results.append(final_info)
+
+    def save(self, filename_prefix='generation'):
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_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', 'aa_sequence', 'peptide_smiles', 'is_valid_peptide',
+                'has_conjugation_site', 'linker_used', 'num_conjugates', 'num_valid_pdcs',
+                'best_pdc_smiles', 'mf', 'polararea', 'xlogp', 'heavycnt', 'rotbonds'
+            ])
+            for result in self.final_results:
+                props = result.get('properties', {})
+                writer.writerow([
+                    result['sample_idx'],
+                    result.get('aa_sequence', ''),
+                    result['peptide_smiles'],
+                    result['is_valid_peptide'],
+                    result['has_conjugation_site'],
+                    result.get('linker_used', ''),
+                    result['num_conjugates'],
+                    result['num_valid_pdcs'],
+                    result['best_pdc']['smiles'] if result.get('best_pdc') else '',
+                    props.get('mf', ''),
+                    props.get('polararea', ''),
+                    props.get('xlogp', ''),
+                    props.get('heavycnt', ''),
+                    props.get('rotbonds', '')
+                ])
+
+        pdc_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_pdcs.csv")
+        with open(pdc_path, 'w', newline='') as f:
+            writer = csv.writer(f)
+            writer.writerow(['sample_idx', 'aa_sequence', 'peptide_smiles', 'linker_used',
+                             'conjugate_idx', 'pdc_smiles'])
+            for result in self.final_results:
+                for j, conj in enumerate(result.get('all_conjugates', [])):
+                    writer.writerow([
+                        result['sample_idx'],
+                        result.get('aa_sequence', ''),
+                        result['peptide_smiles'],
+                        result.get('linker_used', ''),
+                        j,
+                        conj
+                    ])
+
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        print(f"  PDC conjugates: {pdc_path}")
+        return step_log_path, final_path, csv_path, pdc_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)
+        num_valid = sum(1 for s in scores_list if s > 0)
+        print(f"\n--- Step {step} Summary ({guidance_method or 'unknown'}) ---")
+        print(f"  Valid: {num_valid}/{len(scores_list)}")
+        if detailed_results:
+            print(f"  Breakdown:")
+            print(f"    Valid peptides: {sum(1 for r in detailed_results if r.get('is_valid_peptide'))}")
+            print(f"    Has conjugation site: {sum(1 for r in detailed_results if r.get('has_conjugation_site'))}")
+            print(f"    Conjugation success: {sum(1 for r in detailed_results if r.get('conjugation_success'))}")
+            print(f"    PDC valid: {sum(1 for r in detailed_results if r.get('pdc_valid'))}")
+
+    def print_final_summary(self):
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        total = len(self.final_results)
+        valid_pdcs = sum(1 for r in self.final_results if r.get('num_valid_pdcs', 0) > 0)
+        print(f"\nTotal sequences: {total}")
+        print(f"With valid PDCs: {valid_pdcs}")
+        for result in self.final_results:
+            status = "" if result.get('num_valid_pdcs', 0) > 0 else ""
+            print(f"\n  [{status}] Sample {result['sample_idx']}:")
+            print(f"      AA seq: {result.get('aa_sequence', '')}")
+            print(f"      Peptide SMILES: {result['peptide_smiles'][:60]}...")
+            print(f"      Valid peptide: {result['is_valid_peptide']}")
+            print(f"      Has conjugation site: {result['has_conjugation_site']}")
+            print(f"      Conjugates: {result['num_conjugates']}")
+            print(f"      Valid PDCs: {result['num_valid_pdcs']}")
+
+
+# =============================================================================
+# CupOfMoCA Generator (UAA-aware)
+# =============================================================================
+
+class CupOfMoCAGeneratorUAA:
+    """CupOfMoCA Algorithm 2 generator that operates in AA token space.
+
+    AA tokens are converted to SMILES on the fly for scoring with the
+    Conjugation and BindingAffinity objectives.
+    """
+
+    def __init__(self, model, device, objectives, args, tokenizer, tokenizer_type='amino',
+                 tracker=None):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.pdc_analyzer = PDCAnalyzer(device)
+        self.tracker = tracker
+        self.tokenizer = tokenizer
+        self.tokenizer_type = tokenizer_type
+        self.guidance_method = GuidanceMethod(args.guidance_method)
+        self.profiler = Profiler(enabled=getattr(args, 'profile', False))
+
+        cascade_workers = max(1, getattr(args, 'cascade_workers', 1))
+        self._cascade_pool = (
+            ThreadPoolExecutor(max_workers=cascade_workers)
+            if cascade_workers > 1 else None
+        )
+
+        if tokenizer_type == 'amino':
+            self.invalid_tokens = [0, 1, 2, 3, 38]  # specials + 'g'
+            self.smiles_tokenizer = SMILES_SPE_Tokenizer(
+                args.smiles_vocab_path, args.smiles_splits_path
+            )
+        else:
+            self.invalid_tokens = [0, 1, 2, 3, 4, 585, 586]
+            self.smiles_tokenizer = None
+
+        self.conjugation_obj = None
+        for obj in objectives:
+            if isinstance(obj, Conjugation):
+                self.conjugation_obj = obj
+                break
+
+    # ------------------------------------------------------------------
+    # x0 generation (Algorithm 2, Line 6)
+    # ------------------------------------------------------------------
+    def _generate_single_x0(self, n_steps=16, temperature=1.0):
+        self.model.eval()
+
+        if self.tokenizer_type == 'amino':
+            vocab_size = self.model.vocab_size
+            valid_tokens = [i for i in range(vocab_size) if i not in self.invalid_tokens]
+            valid_tokens = torch.tensor(valid_tokens, device=self.device)
+            rand_indices = torch.randint(
+                0, len(valid_tokens),
+                (self.args.num_samples, self.args.gen_len),
+                device=self.device
+            )
+            body = valid_tokens[rand_indices]
+            cls_col = torch.full((self.args.num_samples, 1), self.tokenizer.cls_id, device=self.device)
+            eos_col = torch.full((self.args.num_samples, 1), self.tokenizer.eos_id, device=self.device)
+            x = torch.cat([cls_col, body, eos_col], dim=1)
+        else:
+            vocab_size = self.model.tokenizer.vocab_size
+            x = torch.randint(
+                0, vocab_size,
+                (self.args.num_samples, self.args.gen_len),
+                device=self.device
+            )
+
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Initial Flow Matching", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i + 1]
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+
+                logits = self.model(x, t_tensor) / temperature
+                for inv_tok in self.invalid_tokens:
+                    if inv_tok < logits.shape[-1]:
+                        logits[:, :, inv_tok] = -1000
+
+                pred_x1 = torch.argmax(logits, dim=-1)
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                noise = torch.randint(0, vocab_size, x.shape, device=self.device)
+                x = torch.where(mask, noise, pred_x1)
+
+        return x
+
+    def _aa_to_smiles_batch(self, x_batch):
+        """Decode AA tokens, convert to SMILES via three-letter codes."""
+        with self.profiler.section("aa2smiles.decode"):
+            seqs = self.tokenizer.batch_decode(x_batch)
+        with self.profiler.section("aa2smiles.one_to_three"):
+            three_letter_seqs = [one_letter_to_three_letter(s) for s in seqs]
+        with self.profiler.section("aa2smiles.aa_to_smiles"):
+            with ThreadPoolExecutor() as pool:
+                smiles_seqs = list(pool.map(
+                    lambda s: aa_to_smiles(s, cyclic=self.args.cyclic),
+                    three_letter_seqs
+                ))
+        return seqs, smiles_seqs
+
+    def _check_conjugation_validity(self, x):
+        if self.tokenizer_type == 'amino':
+            _, smiles_seqs = self._aa_to_smiles_batch(x)
+            validities = [
+                self.peptide_analyzer.is_peptide(s) if isinstance(s, str) else False
+                for s in smiles_seqs
+            ]
+            return torch.tensor(validities, device=self.device)
+
+        if self.conjugation_obj is None:
+            sequences = self.tokenizer.batch_decode(x)
+            return torch.tensor(
+                [self.peptide_analyzer.is_peptide(seq) for seq in sequences],
+                device=self.device
+            )
+        return self.conjugation_obj(x) > 0
+
+    def generate_x0(self, n_steps=16, temperature=1.0,
+                    require_valid=False, max_retries=10):
+        print("Generating initial sequences (x0 from prior)...")
+        x = self._generate_single_x0(n_steps, temperature)
+
+        if require_valid:
+            print("Enforcing initial validity constraint...")
+            for retry in range(max_retries):
+                validity = self._check_conjugation_validity(x)
+                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 are valid!")
+                    break
+                print(f"  Retry {retry + 1}/{max_retries}: {num_valid}/{self.args.num_samples} valid")
+                invalid_indices = torch.where(~validity)[0]
+                original_num = self.args.num_samples
+                self.args.num_samples = num_invalid
+                new_x = self._generate_single_x0(n_steps, temperature)
+                self.args.num_samples = original_num
+                x[invalid_indices] = new_x
+
+            final_validity = self._check_conjugation_validity(x)
+            if final_validity.sum().item() < self.args.num_samples:
+                print(f"⚠ After {max_retries} retries, "
+                      f"{final_validity.sum().item()}/{self.args.num_samples} valid")
+
+        generated_sequences = self.tokenizer.batch_decode(x)
+        for seq in generated_sequences:
+            print(f"  {seq}")
+        return x
+
+    # ------------------------------------------------------------------
+    # Scoring
+    # ------------------------------------------------------------------
+    def _get_scores(self, x_batch):
+        with self.profiler.section("get_scores.total"):
+            scores = []
+
+            if self.tokenizer_type == 'amino':
+                _, smiles_seqs = self._aa_to_smiles_batch(x_batch)
+                with self.profiler.section("get_scores.valid_mask"):
+                    valid_mask = torch.tensor(
+                        [1 if isinstance(s, str) else 0 for s in smiles_seqs],
+                        device=self.device
+                    )
+                    for i in torch.where((1 - valid_mask) == 1)[0].tolist():
+                        smiles_seqs[i] = 'C'
+
+                with self.profiler.section("get_scores.smiles_tokenize"):
+                    smiles_batch = self.smiles_tokenizer(
+                        smiles_seqs, return_tensors='pt', padding=True
+                    ).to(self.device)
+                attention_mask = smiles_batch.get("attention_mask", None)
+
+                for obj_func in self.objectives:
+                    obj_name = type(obj_func).__name__
+                    with self.profiler.section(f"get_scores.obj.{obj_name}"):
+                        try:
+                            score = obj_func(
+                                smiles_batch["input_ids"], attention_mask=attention_mask
+                            ).to(self.device)
+                        except TypeError:
+                            score = obj_func(smiles_batch["input_ids"]).to(self.device)
+                    scores.append(score * valid_mask)
+            else:
+                for obj_func in self.objectives:
+                    obj_name = type(obj_func).__name__
+                    with self.profiler.section(f"get_scores.obj.{obj_name}"):
+                        scores.append(obj_func(x_batch.to(self.device)))
+
+            return torch.stack(scores, dim=0).to(self.device)
+
+    def _compute_aggregated_score(self, scores, weights):
+        if self.args.score_aggregation == "weighted_min":
+            return torch.min(weights * scores, dim=0).values
+        return (weights * scores).sum(dim=0)
+
+    # ------------------------------------------------------------------
+    # Cascade scoring (cheap filter + affinity on all → full PDC on top-K)
+    # ------------------------------------------------------------------
+    def _get_scores_cascaded(self, x_batch):
+        """Run cheap conjugation pre-filter + BindingAffinity on all candidates,
+        then run the expensive PDC validation only on the top `cascade_k` by
+        affinity (among those that pass the cheap filter).
+
+        Returns [num_objectives, batch] in the same layout as `_get_scores`.
+        Falls back to full scoring if Conjugation isn't an objective or
+        tokenizer_type != 'amino'.
+        """
+        if self.tokenizer_type != 'amino' or self.conjugation_obj is None:
+            return self._get_scores(x_batch)
+
+        with self.profiler.section("cascade.total"):
+            n = x_batch.shape[0]
+            cascade_k = self.args.cascade_k
+            conj_obj = self.conjugation_obj
+
+            _, smiles_seqs = self._aa_to_smiles_batch(x_batch)
+            valid_mask = torch.tensor(
+                [1.0 if isinstance(s, str) else 0.0 for s in smiles_seqs],
+                device=self.device
+            )
+            safe_smiles = [s if isinstance(s, str) else 'C' for s in smiles_seqs]
+
+            with self.profiler.section("cascade.cheap_filter"):
+                cheap_pass = []
+                for s, vm in zip(safe_smiles, valid_mask.tolist()):
+                    if vm == 0.0:
+                        cheap_pass.append(False)
+                        continue
+                    cheap_pass.append(
+                        conj_obj.is_valid_peptide(s)
+                        and conj_obj.has_conjugation_site(s)
+                    )
+                cheap_pass_t = torch.tensor(
+                    cheap_pass, device=self.device, dtype=torch.float
+                )
+
+            with self.profiler.section("cascade.smiles_tokenize"):
+                smiles_batch = self.smiles_tokenizer(
+                    safe_smiles, return_tensors='pt', padding=True
+                ).to(self.device)
+                attention_mask = smiles_batch.get("attention_mask", None)
+
+            aff_obj = next(
+                (o for o in self.objectives if not isinstance(o, Conjugation)),
+                None
+            )
+            if aff_obj is None:
+                aff_scores = torch.zeros(n, device=self.device)
+            else:
+                with self.profiler.section(f"cascade.obj.{type(aff_obj).__name__}"):
+                    try:
+                        aff_scores = aff_obj(
+                            smiles_batch["input_ids"], attention_mask=attention_mask
+                        ).to(self.device)
+                    except TypeError:
+                        aff_scores = aff_obj(smiles_batch["input_ids"]).to(self.device)
+                    aff_scores = aff_scores * valid_mask
+
+            eligible_idx = torch.where(cheap_pass_t > 0)[0]
+            if cascade_k > 0 and len(eligible_idx) > cascade_k:
+                topk_local = torch.topk(aff_scores[eligible_idx], k=cascade_k).indices
+                validate_idx = eligible_idx[topk_local]
+            else:
+                validate_idx = eligible_idx
+
+            conj_scores = torch.zeros(n, device=self.device)
+            with self.profiler.section("cascade.full_validate"):
+                for j in validate_idx.tolist():
+                    ok, conjs, _, _ = conj_obj.conjugate_peptide(
+                        safe_smiles[j], exclude_n_terminus=True
+                    )
+                    if not ok or not conjs:
+                        continue
+                    has_valid, _ = conj_obj.validate_conjugates(conjs)
+                    if has_valid:
+                        conj_scores[j] = 1.0
+
+            out = []
+            for obj in self.objectives:
+                if isinstance(obj, Conjugation):
+                    out.append(conj_scores)
+                else:
+                    out.append(aff_scores)
+            return torch.stack(out, dim=0).to(self.device)
+
+    # ------------------------------------------------------------------
+    # Guidance methods
+    # ------------------------------------------------------------------
+    def _exponential_tilt_guidance(self, x, pos_probs, candidates, mut_idx, sample_idx,
+                                   beta, weights, current_score=None):
+        num_candidates = len(candidates)
+        current_token = x[sample_idx, mut_idx].item()
+
+        if num_candidates == 0:
+            return current_token, {'method': 'exp_tilt', 'fallback': True, 'accepted': False}
+
+        if num_candidates > self.args.top_k:
+            candidates = candidates[:self.args.top_k]
+            num_candidates = self.args.top_k
+
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(num_candidates, 1)
+        x_proposals[:, mut_idx] = candidates
+
+        with self.profiler.section("exp_tilt.proposal_scores"):
+            if self.args.cascade_k > 0:
+                proposal_scores = self._get_scores_cascaded(x_proposals)
+            else:
+                proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+        validity_mask = (proposal_scores[0] > 0).float()
+        base_probs = pos_probs[sample_idx, candidates]
+
+        log_base_probs = torch.log(base_probs + 1e-10)
+        log_exp_scores = beta * aggregated_scores
+        log_validity = torch.log(validity_mask + 1e-10)
+        log_tilted = log_base_probs + log_exp_scores + log_validity
+
+        log_tilted_max = log_tilted.max()
+        fallback = False
+        if torch.isinf(log_tilted_max) or torch.isnan(log_tilted_max):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+        else:
+            log_tilted_stable = log_tilted - log_tilted_max
+            tilted_probs = torch.exp(log_tilted_stable)
+            prob_sum = tilted_probs.sum()
+            if prob_sum < 1e-10 or torch.isnan(prob_sum) or torch.isinf(prob_sum):
+                fallback = True
+                final_probs = base_probs / (base_probs.sum() + 1e-9)
+            else:
+                final_probs = tilted_probs / prob_sum
+
+        if torch.any(torch.isnan(final_probs)) or torch.any(torch.isinf(final_probs)):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+
+        final_probs = torch.clamp(final_probs, min=1e-10)
+        final_probs = final_probs / final_probs.sum()
+
+        selected_idx = torch.multinomial(final_probs, 1).item()
+        selected_token = candidates[selected_idx].item()
+        selected_score = aggregated_scores[selected_idx].item()
+
+        accepted = True
+        if self.args.use_acceptance and current_score is not None:
+            if selected_score < current_score:
+                selected_token = current_token
+                accepted = False
+
+        return selected_token, {
+            'method': 'exp_tilt', 'beta': beta, 'num_candidates': num_candidates,
+            'proposal_score': aggregated_scores[selected_idx].item(),
+            'selected_score': selected_score if accepted else current_score,
+            'fallback': fallback, 'validity_rate': validity_mask.mean().item(),
+            'accepted': accepted
+        }
+
+    def _classifier_guidance(self, x, logits, mut_idx, sample_idx,
+                             guidance_weight, weights):
+        base_logits = logits[sample_idx, mut_idx, :]
+        sorted_logits, sorted_indices = torch.sort(base_logits, descending=True)
+        candidates = sorted_indices[:self.args.top_k]
+
+        valid_mask = ~torch.isin(
+            candidates, torch.tensor(self.invalid_tokens, device=self.device)
+        )
+        candidates = candidates[valid_mask]
+
+        if len(candidates) == 0:
+            return x[sample_idx, mut_idx].item(), {'method': 'cg', 'fallback': True}
+
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(len(candidates), 1)
+        x_proposals[:, mut_idx] = candidates
+
+        with self.profiler.section("cg.proposal_scores"):
+            if self.args.cascade_k > 0:
+                proposal_scores = self._get_scores_cascaded(x_proposals)
+            else:
+                proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+
+        score_logits = aggregated_scores * self.args.score_temperature
+        base_candidate_logits = base_logits[candidates]
+        guided_logits = base_candidate_logits + guidance_weight * score_logits
+        guided_probs = F.softmax(guided_logits, dim=-1)
+
+        selected_idx = torch.multinomial(guided_probs, 1).item()
+        return candidates[selected_idx].item(), {
+            'method': 'cg', 'guidance_weight': guidance_weight,
+            'num_candidates': len(candidates),
+            'selected_score': aggregated_scores[selected_idx].item(), 'fallback': False
+        }
+
+    def _best_of_k_guidance(self, x, pos_probs, candidates, mut_idx, sample_idx,
+                            k, weights):
+        num_candidates = len(candidates)
+        if num_candidates == 0:
+            return x[sample_idx, mut_idx].item(), {'method': 'bok', 'fallback': True}
+
+        k = min(k, num_candidates)
+        candidate_probs = pos_probs[sample_idx, candidates]
+        candidate_probs = candidate_probs / (candidate_probs.sum() + 1e-9)
+
+        if k >= num_candidates:
+            sampled_indices = torch.arange(num_candidates, device=self.device)
+        else:
+            sampled_indices = torch.multinomial(candidate_probs, k, replacement=False)
+
+        sampled_tokens = candidates[sampled_indices]
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(len(sampled_tokens), 1)
+        x_proposals[:, mut_idx] = sampled_tokens
+
+        with self.profiler.section("bok.proposal_scores"):
+            if self.args.cascade_k > 0:
+                proposal_scores = self._get_scores_cascaded(x_proposals)
+            else:
+                proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+
+        best_idx = torch.argmax(aggregated_scores).item()
+        return sampled_tokens[best_idx].item(), {
+            'method': 'bok', 'k': k, 'num_candidates': num_candidates,
+            'selected_score': aggregated_scores[best_idx].item(),
+            'mean_score': aggregated_scores.mean().item(),
+            'max_score': aggregated_scores.max().item(), 'fallback': False
+        }
+
+    def _hybrid_guidance(self, x, pos_probs, logits, candidates, mut_idx, sample_idx,
+                         step, beta, weights):
+        progress = step / max(self.args.optimization_steps - 1, 1)
+        if progress < 0.3:
+            return self._best_of_k_guidance(
+                x, pos_probs, candidates, mut_idx, sample_idx,
+                k=self.args.bok_k, weights=weights
+            )
+        elif progress < 0.7:
+            return self._exponential_tilt_guidance(
+                x, pos_probs, candidates, mut_idx, sample_idx,
+                beta=beta, weights=weights
+            )
+        else:
+            return self._classifier_guidance(
+                x, logits, mut_idx, sample_idx,
+                guidance_weight=self.args.cg_weight, weights=weights
+            )
+
+    # ------------------------------------------------------------------
+    # Main loop
+    # ------------------------------------------------------------------
+    def generate(self):
+        self.profiler.reset()
+        x = self.generate_x0(
+            require_valid=self.args.require_valid_x0,
+            max_retries=self.args.max_x0_retries
+        )
+
+        if self.args.weights is None:
+            weights = torch.tensor(
+                [1.0] + [1.0 / (self.num_objectives - 1)] * (self.num_objectives - 1),
+                device=self.device
+            ).view(-1, 1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+
+        print(f"\nGuidance Method: {self.guidance_method.value}")
+        print(f"Weights: {weights.squeeze().tolist()}")
+        print(f"β range: [{self.args.beta_min}, {self.args.beta_max}]")
+        print(f"Acceptance criterion: {'ENABLED' if self.args.use_acceptance else 'DISABLED'}")
+
+        with torch.no_grad():
+            for step in tqdm(range(self.args.optimization_steps), desc="Guided Generation"):
+                improved_indices = []
+
+                progress = step / max(self.args.optimization_steps - 1, 1)
+                beta = self.args.beta_min + (self.args.beta_max - self.args.beta_min) * progress
+
+                if self.args.position_strategy == "random":
+                    mut_idx = random.randint(1, self.args.gen_len - 2)
+                else:
+                    mut_idx = (step % (self.args.gen_len - 2)) + 1
+
+                time_t = torch.full(
+                    (self.args.num_samples,),
+                    step / self.args.optimization_steps,
+                    device=self.device
+                )
+
+                with self.profiler.section("loop.model_forward"):
+                    logits = self.model(x, time_t)
+                    probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+
+                pos_probs[torch.arange(self.args.num_samples), x[:, mut_idx]] = 0
+                for inv_tok in self.invalid_tokens:
+                    pos_probs[:, inv_tok] = 0
+
+                with self.profiler.section("loop.top_p_sort"):
+                    sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+
+                with self.profiler.section("loop.current_scores"):
+                    current_scores = self._get_scores(x)
+                    current_aggregated = self._compute_aggregated_score(current_scores, weights)
+
+                if step == 0:
+                    print(f"\nInitial Scores: {current_scores.cpu().numpy()}")
+                    if self.tracker and self.conjugation_obj:
+                        detailed = self.conjugation_obj.get_last_detailed_results()
+                        self.tracker.log_step(step, current_scores[0], detailed)
+                        self.tracker.print_step_summary(
+                            step, current_scores[0], detailed, self.guidance_method.value
+                        )
+
+                final_tokens = []
+                for i in range(self.args.num_samples):
+                    sample_probs = sorted_probs[i]
+                    sample_indices = sorted_indices[i]
+
+                    cumsum = torch.cumsum(sample_probs, dim=0)
+                    mask = cumsum <= self.args.top_p
+                    mask[0] = True
+                    candidates = sample_indices[mask]
+
+                    valid_candidate_mask = ~torch.isin(
+                        candidates, torch.tensor(self.invalid_tokens, device=self.device)
+                    )
+                    candidates = candidates[valid_candidate_mask]
+
+                    current_sample_score = current_aggregated[i].item()
+
+                    if beta > 0 or self.guidance_method != GuidanceMethod.EXPONENTIAL_TILT:
+                        with self.profiler.section(f"loop.guidance.{self.guidance_method.value}"):
+                            if self.guidance_method == GuidanceMethod.EXPONENTIAL_TILT:
+                                token, _ = self._exponential_tilt_guidance(
+                                    x, pos_probs, candidates, mut_idx, i, beta, weights,
+                                    current_score=current_sample_score
+                                )
+                            elif self.guidance_method == GuidanceMethod.CLASSIFIER_GUIDANCE:
+                                token, _ = self._classifier_guidance(
+                                    x, logits, mut_idx, i, self.args.cg_weight, weights
+                                )
+                            elif self.guidance_method == GuidanceMethod.BEST_OF_K:
+                                token, _ = self._best_of_k_guidance(
+                                    x, pos_probs, candidates, mut_idx, i, self.args.bok_k, weights
+                                )
+                            elif self.guidance_method == GuidanceMethod.HYBRID:
+                                token, _ = self._hybrid_guidance(
+                                    x, pos_probs, logits, candidates, mut_idx, i, step, beta, weights
+                                )
+                            else:
+                                raise ValueError(f"Unknown guidance: {self.guidance_method}")
+                    else:
+                        if len(candidates) == 0:
+                            token = x[i, mut_idx].item()
+                        else:
+                            cprobs = pos_probs[i, candidates]
+                            cprobs = cprobs / (cprobs.sum() + 1e-9)
+                            idx = torch.multinomial(cprobs, 1).item()
+                            token = candidates[idx].item()
+
+                    final_tokens.append(token)
+
+                x[:, mut_idx] = torch.tensor(final_tokens, device=self.device)
+
+                with self.profiler.section("loop.new_scores"):
+                    new_scores = self._get_scores(x)
+                    new_aggregated = self._compute_aggregated_score(new_scores, weights)
+                for i in range(self.args.num_samples):
+                    if new_aggregated[i] > current_aggregated[i]:
+                        improved_indices.append(i)
+
+                if improved_indices:
+                    for i in improved_indices:
+                        prev_weighted = (weights.squeeze(1) * current_scores[:, i]).sum()
+                        new_weighted = (weights.squeeze(1) * new_scores[:, i]).sum()
+                        print(f"Previous Weighted Sum: {prev_weighted}")
+                        print(f"Previous Scores: {current_scores[:, i]}")
+                        print(f"New Weighted Sum: {new_weighted}")
+                        print(f"New Scores: {new_scores[:, i]}")
+                    print(self.tokenizer.batch_decode(x[improved_indices]))
+
+                if self.tracker and self.conjugation_obj and \
+                   (step % 4 == 0 or step == self.args.optimization_steps - 1):
+                    detailed = self.conjugation_obj.get_last_detailed_results()
+                    self.tracker.log_step(
+                        step, new_scores[0], detailed, improved_indices,
+                        guidance_info={'method': self.guidance_method.value, 'beta': beta}
+                    )
+                    if step == self.args.optimization_steps - 1:
+                        self.tracker.print_step_summary(
+                            step, new_scores[0], detailed, self.guidance_method.value
+                        )
+
+                if (self.profiler.enabled
+                    and self.args.profile_every > 0
+                    and (step + 1) % self.args.profile_every == 0):
+                    self.profiler.print_summary(header=f"(through step {step + 1})")
+
+        self.profiler.print_summary(header="(end of generate)")
+        return x
+
+
+# =============================================================================
+# Main
+# =============================================================================
+
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    print(f"Tokenizer type: {args.tokenizer_type}")
+    print(f"Guidance method: {args.guidance_method}")
+
+    if args.tokenizer_type == 'amino':
+        print("Using Amino Acid tokenizer")
+        tokenizer = AminoTokenizer()
+        smiles_tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+    else:
+        print("Using SMILES tokenizer")
+        tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+        smiles_tokenizer = tokenizer
+
+    conjugation_model = Conjugation(smiles_tokenizer, device=device, builder_type=args.builder_type)
+    affinity_model = BindingAffinity(args.target, device)
+    OBJECTIVE_FUNCTIONS = [conjugation_model, affinity_model]
+
+    tracker = GenerationTracker(output_dir=args.output_dir)
+
+    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)
+        print(f"Random seed: {args.seed}")
+
+    print(f"Loading model from: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+
+    if args.tokenizer_type == 'amino':
+        raw_model_args = checkpoint['args']
+        if isinstance(raw_model_args, argparse.Namespace):
+            raw_model_args = vars(raw_model_args)
+        model_args = edict(raw_model_args)
+        model = MDLM(
+            vocab_size=model_args.vocab_size,
+            seq_len=model_args.seq_len,
+            model_dim=model_args.model_dim,
+            n_heads=model_args.n_heads,
+            n_layers=model_args.n_layers,
+        ).to(device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print(f"Amino model loaded. Vocab size: {model_args.vocab_size}")
+    else:
+        model_hparams = checkpoint["hyper_parameters"]["args"]
+        model = MDLMLightningModule.load_from_checkpoint(
+            args.checkpoint, args=model_hparams, tokenizer=tokenizer,
+            map_location=device, strict=False
+        )
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully.")
+
+    generator = CupOfMoCAGeneratorUAA(
+        model, device, OBJECTIVE_FUNCTIONS, args,
+        tokenizer=tokenizer, tokenizer_type=args.tokenizer_type, tracker=tracker
+    )
+
+    all_sequences = []
+    all_smiles = []
+    all_final_results = []
+
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'=' * 60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'=' * 60}")
+
+        generated_tokens = generator.generate()
+        final_scores = generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+
+        print("\n--- Batch Results ---")
+        print(f"Sequences: {sequence_str}")
+        print(f"Scores: {final_scores}")
+
+        all_sequences.extend(sequence_str)
+
+        if args.tokenizer_type == 'amino':
+            smiles_list = []
+            for seq in sequence_str:
+                three_letter = one_letter_to_three_letter(seq)
+                smiles = aa_to_smiles(three_letter, cyclic=args.cyclic)
+                smiles_list.append(smiles if isinstance(smiles, str) else 'C')
+            smiles_tokens = smiles_tokenizer(
+                smiles_list, return_tensors='pt', padding=True
+            ).to(device)
+            detailed_with_props = conjugation_model.get_detailed_scores_with_properties(
+                smiles_tokens["input_ids"]
+            )
+            all_smiles.extend(smiles_list)
+            tracker.log_final(sequence_str, detailed_with_props, smiles_list=smiles_list)
+        else:
+            detailed_with_props = conjugation_model.get_detailed_scores_with_properties(generated_tokens)
+            tracker.log_final(sequence_str, detailed_with_props)
+        all_final_results.extend(detailed_with_props)
+
+    tracker.print_final_summary()
+    tracker.save(filename_prefix=f'cupofmoca_uaa_{args.guidance_method}')
+
+    with open(args.output_file, 'w') as f:
+        for i, seq in enumerate(all_sequences):
+            if args.tokenizer_type == 'amino':
+                smiles = all_smiles[i] if i < len(all_smiles) else ''
+                f.write(f"{seq}\t{smiles}\n")
+            else:
+                f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+    print("\nGeneration complete.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="CupOfMoCA-style Guided Generation with UAA Foundation Model"
+    )
+
+    parser.add_argument("--checkpoint", type=str, required=True)
+    parser.add_argument("--target", type=str, required=True)
+
+    parser.add_argument("--num_samples", type=int, default=10)
+    parser.add_argument("--num_batches", type=int, default=10)
+    parser.add_argument("--gen_len", type=int, default=50)
+    parser.add_argument("--optimization_steps", type=int, default=128)
+
+    parser.add_argument("--guidance_method", type=str, default="exp_tilt",
+                        choices=["exp_tilt", "cg", "bok", "hybrid"])
+
+    parser.add_argument("--beta_min", type=float, default=1.0)
+    parser.add_argument("--beta_max", type=float, default=20.0)
+    parser.add_argument("--cg_weight", type=float, default=2.0)
+    parser.add_argument("--use_acceptance", action="store_true")
+    parser.add_argument("--bok_k", type=int, default=8)
+    parser.add_argument("--top_k", type=int, default=200)
+    parser.add_argument("--top_p", type=float, default=0.9)
+    parser.add_argument("--score_temperature", type=float, default=1.0)
+    parser.add_argument("--score_aggregation", type=str, default="weighted_sum",
+                        choices=["weighted_sum", "weighted_min"])
+
+    parser.add_argument("--weights", type=float, nargs='+', default=None)
+
+    parser.add_argument("--position_strategy", type=str, default="random",
+                        choices=["random", "sequential"])
+
+    parser.add_argument("--require_valid_x0", action="store_true")
+    parser.add_argument("--max_x0_retries", type=int, default=10)
+
+    parser.add_argument("--builder_type", type=str, default="AoaDau",
+                        choices=["AoaDau", "ValAlaDau", "ValCitDau", "GlutarylDau", "Dau",
+                                 "SSCpt", "VcMMAE", "DSPACpt", "MalHexCpt", "DSECCpt", "Cpt",
+                                 "DSPASN38", "DSEBSN38", "DSEPSN38", "SN38"])
+
+    parser.add_argument("--tokenizer_type", type=str, default="amino",
+                        choices=["smiles", "amino"])
+    parser.add_argument("--cyclic", action="store_true",
+                        help="Generate cyclic peptides (only for amino tokenizer).")
+    parser.add_argument("--smiles_vocab_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_vocab.txt")
+    parser.add_argument("--smiles_splits_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_splits.txt")
+
+    parser.add_argument("--output_file", type=str, default="./smiles.txt")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs")
+
+    parser.add_argument("--seed", type=int, default=None)
+
+    parser.add_argument("--profile", action="store_true",
+                        help="Enable per-section wall-clock profiling.")
+    parser.add_argument("--profile_every", type=int, default=0,
+                        help="Print profiler summary every N optimization steps "
+                             "(0 = only at end of each generate()).")
+
+    parser.add_argument("--cascade_k", type=int, default=0,
+                        help="Cascade scoring: score all proposals with cheap "
+                             "Conjugation pre-filter + BindingAffinity, then run "
+                             "the expensive PDC validation only on the top-K by "
+                             "affinity (among cheap-filter survivors). 0 disables "
+                             "(full scoring on all proposals, original behavior).")
+    parser.add_argument("--cascade_workers", type=int, default=1,
+                        help="Number of worker threads for parallelizing "
+                             "cascade.full_validate (CPU-bound RDKit work). "
+                             "1 = sequential (original behavior). Try 4-8 when "
+                             "--cascade_k is high. Ignored if --cascade_k=0.")
+
+    args = parser.parse_args()
+    main(args)

smiles/generation/cupofmoca_staged.py

@@ -0,0 +1,584 @@
+"""
+CupOfMoCA-style Staged Generation with Exponential Tilting
+
+This implements Algorithm 2 (Sampling with CupOfMoCA) with exponential tilting for staged generation.
+Unlike the guided version, this uses:
+- Peptide analyzer (not conjugation) as validity checker
+- No PDC building during generation
+- Simpler objectives focused on peptide validity and binding affinity
+
+Exponential Tilting: p'(z') ∝ p(z') · exp(β · S(z')) · 1_valid(z')
+"""
+
+import argparse
+import os
+import json
+import csv
+import random
+from collections import Counter
+from datetime import datetime
+from typing import List, Optional, Tuple, Dict, Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+
+# Import from the original codebase
+from smiles.models.train import MDLMLightningModule, PeptideAnalyzer
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+from smiles.scoring.classifiers import Analyzer, BindingAffinity
+
+
+class GenerationTracker:
+    """Tracks detailed validity information during generation"""
+    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, improved_indices=None, guidance_info: Optional[Dict] = None):
+        """Log information for a generation step"""
+        step_info = {
+            'step': step,
+            'scores': scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else scores,
+            'num_valid': int(sum(1 for s in scores if s > 0)) 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, sequences, scores):
+        """Log final results"""
+        for i, (seq, score_vec) in enumerate(zip(sequences, scores.T)):
+            final_info = {
+                'sample_idx': i,
+                'peptide_smiles': seq,
+                'is_valid_peptide': bool(score_vec[0] > 0),
+                'binding_score': float(score_vec[1]) if len(score_vec) > 1 else None,
+            }
+            self.final_results.append(final_info)
+    
+    def save(self, filename_prefix='generation'):
+        """Save all logs to files"""
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_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', 'peptide_smiles', 'is_valid_peptide', 'binding_score'])
+            for result in self.final_results:
+                writer.writerow([
+                    result['sample_idx'],
+                    result['peptide_smiles'],
+                    result['is_valid_peptide'],
+                    result.get('binding_score', '')
+                ])
+        
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        
+        return step_log_path, final_path, csv_path
+    
+    def print_step_summary(self, step, scores):
+        """Print a summary for the current step"""
+        if isinstance(scores, torch.Tensor):
+            scores_list = scores.cpu().numpy().tolist()
+        else:
+            scores_list = list(scores)
+        
+        num_valid = sum(1 for s in scores_list if s > 0)
+        
+        print(f"\n--- Step {step} Summary ---")
+        print(f"  Valid peptides: {num_valid}/{len(scores_list)}")
+    
+    def print_final_summary(self):
+        """Print final summary"""
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        
+        total = len(self.final_results)
+        valid_peptides = sum(1 for r in self.final_results if r.get('is_valid_peptide', False))
+        
+        print(f"\nTotal sequences: {total}")
+        print(f"Valid peptides: {valid_peptides}")
+        
+        if any(r.get('binding_score') is not None for r in self.final_results):
+            avg_binding = np.mean([r['binding_score'] for r in self.final_results if r.get('binding_score') is not None])
+            print(f"Average binding score: {avg_binding:.3f}")
+
+
+class CupOfMoCAStaged:
+    """
+    CupOfMoCA-style staged generator with exponential tilting.
+    
+    Uses peptide analyzer instead of conjugation model.
+    Focuses on generating valid peptides with good binding affinity.
+    """
+    
+    def __init__(self, model, device, objectives, args, tracker=None):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.invalid_tokens = [0, 1, 2, 3, 4, 585, 586]  # Special tokens to avoid
+        self.tracker = tracker
+        
+        # Find the Analyzer objective for validity checking
+        self.analyzer_obj = None
+        for obj in objectives:
+            if isinstance(obj, Analyzer):
+                self.analyzer_obj = obj
+                break
+    
+    # =========================================================================
+    # Initial Generation (x0 from prior)
+    # =========================================================================
+    
+    def generate_x0(self, n_steps: int = 16, temperature: float = 1.0) -> torch.Tensor:
+        """
+        Generate initial sequences using flow matching.
+        This follows the flow from noise to data without guidance.
+        """
+        print("Generating initial sequences (x0 from prior)...")
+        self.model.eval()
+        
+        # Start with random tokens (noise at t=0)
+        x = torch.randint(
+            0, self.model.tokenizer.vocab_size,
+            (self.args.num_samples, self.args.gen_len),
+            device=self.device
+        )
+        
+        # Time schedule
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+        
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Initial Flow Matching", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i+1]
+                
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+                
+                # Get model's prediction
+                logits = self.model(x, t_tensor)
+                logits = logits / temperature
+                logits[:, :, 586] = -1000  # Avoid padding token
+                
+                pred_x1 = torch.argmax(logits, dim=-1)
+                
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+                
+                # Construct next state with noise
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                noise = torch.randint(0, self.model.tokenizer.vocab_size, x.shape, device=self.device)
+                x = torch.where(mask, noise, pred_x1)
+        
+        generated_sequences = self.model.tokenizer.batch_decode(x)
+        
+        # Analyze the validity of the generated sequences
+        validities = []
+        for seq in generated_sequences:
+            validities.append(self.peptide_analyzer.is_peptide(seq))
+            print(seq)
+        print(f"Initial Sequence Validity: {validities}")
+        
+        return x
+    
+    # =========================================================================
+    # Score Functions
+    # =========================================================================
+    
+    def _get_scores(self, x_batch: torch.Tensor) -> torch.Tensor:
+        """Calculate scores for all objectives"""
+        scores = []
+        for obj_func in self.objectives:
+            scores.append(obj_func(x_batch.to(self.device)))
+        return torch.stack(scores, dim=0).to(self.device)
+    
+    def _compute_aggregated_score(self, scores: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
+        """Compute weighted aggregated score for guidance"""
+        if self.args.score_aggregation == "weighted_sum":
+            return (weights * scores).sum(dim=0)
+        elif self.args.score_aggregation == "weighted_min":
+            return torch.min(weights * scores, dim=0).values
+        else:
+            return (weights * scores).sum(dim=0)
+    
+    # =========================================================================
+    # Exponential Tilting Guidance
+    # =========================================================================
+    
+    def _exponential_tilt_guidance(
+        self,
+        x: torch.Tensor,
+        pos_probs: torch.Tensor,
+        candidates: torch.Tensor,
+        mut_idx: int,
+        sample_idx: int,
+        beta: float,
+        weights: torch.Tensor,
+        current_score: float = None
+    ) -> Tuple[int, Dict]:
+        """
+        Exponential Tilting Guidance with acceptance criterion
+        
+        p'(z') ∝ p(z') · exp(β · S(z')) · 1_valid(z')
+        
+        Only accept if new_score >= current_score (prevents regression)
+        """
+        num_candidates = len(candidates)
+        current_token = x[sample_idx, mut_idx].item()
+        
+        if num_candidates == 0:
+            return current_token, {'method': 'exp_tilt', 'fallback': True, 'accepted': False}
+        
+        # Limit candidates for efficiency
+        if num_candidates > self.args.top_k:
+            candidates = candidates[:self.args.top_k]
+            num_candidates = self.args.top_k
+        
+        # Create batch of proposals
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(num_candidates, 1)
+        x_proposals[:, mut_idx] = candidates
+        
+        # Compute scores for all candidates
+        proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+        
+        # Validity mask (using first objective - peptide analyzer)
+        validity_mask = (proposal_scores[0] > 0).float()
+        
+        # Get base probabilities
+        base_probs = pos_probs[sample_idx, candidates]
+        
+        # Apply exponential tilting in log-space for numerical stability
+        log_base_probs = torch.log(base_probs + 1e-10)
+        log_exp_scores = beta * aggregated_scores
+        log_validity = torch.log(validity_mask + 1e-10)
+        
+        log_tilted = log_base_probs + log_exp_scores + log_validity
+        
+        # Numerical stability
+        log_tilted_max = log_tilted.max()
+        fallback = False
+        
+        if torch.isinf(log_tilted_max) or torch.isnan(log_tilted_max):
+            # Fallback to base distribution
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+        else:
+            log_tilted_stable = log_tilted - log_tilted_max
+            tilted_probs = torch.exp(log_tilted_stable)
+            
+            prob_sum = tilted_probs.sum()
+            if prob_sum < 1e-10 or torch.isnan(prob_sum) or torch.isinf(prob_sum):
+                fallback = True
+                final_probs = base_probs / (base_probs.sum() + 1e-9)
+            else:
+                final_probs = tilted_probs / prob_sum
+        
+        # Safety check
+        if torch.any(torch.isnan(final_probs)) or torch.any(torch.isinf(final_probs)):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+        
+        # Ensure valid probabilities
+        final_probs = torch.clamp(final_probs, min=1e-10)
+        final_probs = final_probs / final_probs.sum()
+        
+        # Sample from guided distribution
+        selected_idx = torch.multinomial(final_probs, 1).item()
+        selected_token = candidates[selected_idx].item()
+        selected_score = aggregated_scores[selected_idx].item()
+        
+        # ACCEPTANCE CRITERION: Only accept if score improves or stays same
+        accepted = True
+        if self.args.use_acceptance and current_score is not None:
+            if selected_score < current_score:
+                selected_token = current_token
+                selected_score = current_score
+                accepted = False
+        
+        info = {
+            'method': 'exp_tilt',
+            'beta': beta,
+            'num_candidates': num_candidates,
+            'proposal_score': aggregated_scores[selected_idx].item(),
+            'selected_score': selected_score,
+            'fallback': fallback,
+            'validity_rate': validity_mask.mean().item(),
+            'accepted': accepted
+        }
+        
+        return selected_token, info
+    
+    # =========================================================================
+    # Main Generation Loop
+    # =========================================================================
+    
+    def generate(self) -> torch.Tensor:
+        """Main generation loop with exponential tilting guidance"""
+        # Generate initial x0
+        x = self.generate_x0()
+        
+        # Setup weights
+        if self.args.weights is None:
+            weights = torch.tensor(
+                [1.0] + [1.0 / (self.num_objectives - 1)] * (self.num_objectives - 1),
+                device=self.device
+            ).view(-1, 1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+        
+        print(f"\nGuidance Method: Exponential Tilting")
+        print(f"Weights: {weights.squeeze().tolist()}")
+        print(f"β range: [{self.args.beta_min}, {self.args.beta_max}]")
+        print(f"Acceptance criterion: {'ENABLED' if self.args.use_acceptance else 'DISABLED'}")
+        
+        # Main optimization loop
+        with torch.no_grad():
+            for step in tqdm(range(self.args.optimization_steps), desc="Guided Generation"):
+                # Anneal guidance strength
+                beta = self.args.beta_min + (self.args.beta_max - self.args.beta_min) * (
+                    step / max(self.args.optimization_steps - 1, 1)
+                )
+                
+                # Select mutation position
+                if self.args.position_strategy == "random":
+                    mut_idx = random.randint(1, self.args.gen_len - 2)
+                else:  # sequential
+                    mut_idx = (step % (self.args.gen_len - 2)) + 1
+                
+                # Determine generation timestep
+                generation_step = step % self.args.optimization_steps
+                time_t = torch.full(
+                    (self.args.num_samples,),
+                    (generation_step / self.args.optimization_steps),
+                    device=self.device
+                )
+                
+                # Get proposal distribution from model
+                logits = self.model(x, time_t)
+                probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+                
+                # Zero out current token and invalid tokens
+                pos_probs[torch.arange(self.args.num_samples), x[:, mut_idx]] = 0
+                for invalid_tok in self.invalid_tokens:
+                    pos_probs[:, invalid_tok] = 0
+                
+                # Top-p filtering
+                sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                remove_mask = cumulative_probs > self.args.top_p
+                remove_mask[..., 1:] = remove_mask[..., :-1].clone()
+                remove_mask[..., 0] = 0
+                
+                # Get current scores
+                current_scores = self._get_scores(x)
+                current_aggregated = self._compute_aggregated_score(current_scores, weights)
+                
+                if step == 0:
+                    print(f"\nInitial Scores: {current_scores.cpu().numpy()}")
+                
+                # Apply guidance for each sample
+                improved_count = 0
+                for i in range(self.args.num_samples):
+                    sample_mask = remove_mask[i]
+                    candidates = sorted_indices[i, ~sample_mask]
+                    
+                    # Filter invalid tokens from candidates
+                    valid_candidate_mask = ~torch.isin(
+                        candidates,
+                        torch.tensor(self.invalid_tokens, device=self.device)
+                    )
+                    candidates = candidates[valid_candidate_mask]
+                    
+                    if len(candidates) == 0:
+                        continue
+                    
+                    # Apply exponential tilting
+                    selected_token, guidance_info = self._exponential_tilt_guidance(
+                        x, pos_probs, candidates, mut_idx, i, beta, weights,
+                        current_score=current_aggregated[i].item()
+                    )
+                    
+                    if guidance_info['accepted']:
+                        improved_count += 1
+                    
+                    x[i, mut_idx] = selected_token
+                
+                # Log progress
+                if step % 10 == 0 or improved_count > 0:
+                    new_scores = self._get_scores(x)
+                    if self.tracker:
+                        self.tracker.log_step(step, new_scores[0])
+                    
+                    if step % 20 == 0:
+                        print(f"\nStep {step}: {improved_count}/{self.args.num_samples} improved")
+                        print(f"  β = {beta:.2f}")
+                        print(f"  Valid peptides: {(new_scores[0] > 0).sum().item()}/{self.args.num_samples}")
+        
+        return x
+
+
+# =============================================================================
+# Main Execution
+# =============================================================================
+
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    tokenizer = SMILES_SPE_Tokenizer(
+        './smiles/smiles_tokenizer/new_vocab.txt',
+        './smiles/smiles_tokenizer/new_splits.txt'
+    )
+    
+    # Initialize objectives
+    analyzer_model = Analyzer(tokenizer)
+    affinity_model = BindingAffinity(args.target, device)
+    OBJECTIVE_FUNCTIONS = [analyzer_model, affinity_model]
+    
+    # Initialize tracker
+    tracker = GenerationTracker(output_dir=args.output_dir)
+    
+    # Set random seeds
+    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)
+        print(f"Random seed: {args.seed}")
+    
+    # Load model
+    print(f"Loading model from: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+    model_hparams = checkpoint["hyper_parameters"]["args"]
+    
+    model = MDLMLightningModule.load_from_checkpoint(
+        args.checkpoint,
+        args=model_hparams,
+        tokenizer=tokenizer,
+        map_location=device,
+        strict=False
+    )
+    model.to(device)
+    print("Model loaded successfully.")
+    
+    # Create generator
+    generator = CupOfMoCAStaged(model, device, OBJECTIVE_FUNCTIONS, args, tracker=tracker)
+    
+    all_sequences = []
+    all_scores = []
+    
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'=' * 60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'=' * 60}")
+        
+        generated_tokens = generator.generate()
+        final_scores = generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+        
+        print("\n--- Batch Results ---")
+        for i, seq in enumerate(sequence_str):
+            print(f"  {i}: {seq}")
+        print(f"Scores: {final_scores}")
+        
+        all_sequences.extend(sequence_str)
+        all_scores.append(final_scores)
+        
+        tracker.log_final(sequence_str, final_scores)
+    
+    # Save results
+    tracker.print_final_summary()
+    tracker.save(filename_prefix='cupofmoca_staged')
+    
+    with open(args.output_file, 'w') as f:
+        for seq in all_sequences:
+            f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+    print("\nGeneration complete.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="CupOfMoCA-style Staged Generation with Exponential Tilting"
+    )
+    
+    # Model and data
+    parser.add_argument("--checkpoint", type=str, required=True,
+                        help="Path to trained model checkpoint")
+    parser.add_argument("--target", type=str, required=True,
+                        help="Target protein sequence for binding affinity")
+    
+    # Generation settings
+    parser.add_argument("--num_samples", type=int, default=10,
+                        help="Number of samples per batch")
+    parser.add_argument("--num_batches", type=int, default=10,
+                        help="Number of batches to generate")
+    parser.add_argument("--gen_len", type=int, default=50,
+                        help="Length of sequences to generate")
+    parser.add_argument("--optimization_steps", type=int, default=128,
+                        help="Number of guided optimization steps")
+    
+    # Guidance parameters
+    parser.add_argument("--beta_min", type=float, default=1.0,
+                        help="Minimum guidance strength β")
+    parser.add_argument("--beta_max", type=float, default=20.0,
+                        help="Maximum guidance strength β")
+    parser.add_argument("--use_acceptance", action="store_true",
+                        help="Use acceptance criterion: only accept if new_score >= current_score")
+    parser.add_argument("--top_k", type=int, default=200,
+                        help="Top-K candidates to consider for guidance")
+    parser.add_argument("--top_p", type=float, default=0.9,
+                        help="Top-p threshold for candidate filtering")
+    parser.add_argument("--score_aggregation", type=str, default="weighted_sum",
+                        choices=["weighted_sum", "weighted_min"],
+                        help="How to aggregate multi-objective scores")
+    
+    # Weights
+    parser.add_argument("--weights", type=float, nargs='+', default=None,
+                        help="Weights for objectives (e.g., 1.0 1.0)")
+    
+    # Position selection
+    parser.add_argument("--position_strategy", type=str, default="random",
+                        choices=["random", "sequential"],
+                        help="Strategy for selecting mutation positions")
+    
+    # Output
+    parser.add_argument("--output_file", type=str, default="./smiles.txt",
+                        help="File to save generated sequences")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs",
+                        help="Directory for detailed logs")
+    
+    # Reproducibility
+    parser.add_argument("--seed", type=int, default=None,
+                        help="Random seed")
+    
+    args = parser.parse_args()
+    main(args)

smiles/generation/cupofmoca_staged_aa.py

@@ -0,0 +1,724 @@
+"""CupOfMoCA-style staged generation with the natural-AA foundation model."""
+
+import argparse
+import csv
+import json
+import os
+import random
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from easydict import EasyDict as edict
+from tqdm import tqdm
+
+from smiles.models.train import MDLMLightningModule, PeptideAnalyzer
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+from smiles.scoring.classifiers import Analyzer, BindingAffinity
+
+# AA-mode shared utilities
+from smiles.models.aa_helpers import (
+    ESM2AminoTokenizer as _AAEsmTokenizer,
+    MDLM_PepReDi as _AAMDLM,
+    make_affinity,
+    DEFAULT_AA_AFFINITY_CKPT,
+    INVALID_AA_TOKEN_IDS as _AA_INVALID_IDS,
+)
+
+from third_party.SMILESToPeptide.peptide_to_smiles import aa_to_smiles
+from third_party.SMILESToPeptide.aminoacid_selective import one_letter_to_three_letter
+
+
+# =============================================================================
+# Amino Acid Foundation Model (DiT-based MDLM)
+# =============================================================================
+
+# --- AA-mode Model + Tokenizer (imported from aa_helpers) ---
+MDLM = _AAMDLM
+AminoTokenizer = _AAEsmTokenizer
+
+
+# =============================================================================
+# Generation Tracker
+# =============================================================================
+
+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, scores, 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(sum(1 for s in scores if s > 0)) 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, sequences, scores, smiles_list=None):
+        for i, (seq, score_vec) in enumerate(zip(sequences, scores.T)):
+            self.final_results.append({
+                'sample_idx': i,
+                'aa_sequence': seq,
+                'peptide_smiles': smiles_list[i] if smiles_list else seq,
+                'is_valid_peptide': bool(score_vec[0] > 0),
+                'binding_score': float(score_vec[1]) if len(score_vec) > 1 else None,
+            })
+
+    def save(self, filename_prefix='generation'):
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_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', 'aa_sequence', 'peptide_smiles',
+                             'is_valid_peptide', 'binding_score'])
+            for r in self.final_results:
+                writer.writerow([
+                    r['sample_idx'], r.get('aa_sequence', ''), r['peptide_smiles'],
+                    r['is_valid_peptide'], r.get('binding_score', '')
+                ])
+
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        return step_log_path, final_path, csv_path
+
+    def print_step_summary(self, step, scores):
+        scores_list = scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else list(scores)
+        num_valid = sum(1 for s in scores_list if s > 0)
+        print(f"\n--- Step {step} Summary ---")
+        print(f"  Valid peptides: {num_valid}/{len(scores_list)}")
+
+    def print_final_summary(self):
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        total = len(self.final_results)
+        valid = sum(1 for r in self.final_results if r.get('is_valid_peptide', False))
+        print(f"\nTotal sequences: {total}")
+        print(f"Valid peptides: {valid}")
+        if any(r.get('binding_score') is not None for r in self.final_results):
+            avg = np.mean([r['binding_score'] for r in self.final_results
+                           if r.get('binding_score') is not None])
+            print(f"Average binding score: {avg:.3f}")
+
+
+# =============================================================================
+# CupOfMoCA Staged Generator (UAA-aware)
+# =============================================================================
+
+class CupOfMoCAStagedUAA:
+    """Exponential-tilting staged generator that operates in AA token space."""
+
+    def __init__(self, model, device, objectives, args, tokenizer, tokenizer_type='amino',
+                 tracker=None):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.tracker = tracker
+        self.tokenizer = tokenizer
+        self.tokenizer_type = tokenizer_type
+
+        if tokenizer_type == 'amino':
+            # PepReDi natural-AA vocab is 24 tokens; 0-3 = cls/pad/eos/unk,
+            # 4-23 = canonical AAs. Block specials only. (UAA used id 38
+            # for D-AA 'g'; that's out of PepReDi vocab range.)
+            self.invalid_tokens = [0, 1, 2, 3]
+            self.smiles_tokenizer = SMILES_SPE_Tokenizer(
+                args.smiles_vocab_path, args.smiles_splits_path
+            )
+        else:
+            self.invalid_tokens = [0, 1, 2, 3, 4, 585, 586]
+            self.smiles_tokenizer = None
+
+        self.analyzer_obj = None
+        for obj in objectives:
+            if isinstance(obj, Analyzer):
+                self.analyzer_obj = obj
+                break
+
+    # ------------------------------------------------------------------
+    # x0 generation
+    # ------------------------------------------------------------------
+    def generate_x0(self, n_steps=16, temperature=1.0):
+        print("Generating initial sequences (x0 from prior)...")
+        self.model.eval()
+
+        if self.tokenizer_type == 'amino':
+            vocab_size = self.model.vocab_size
+            valid_tokens = [i for i in range(vocab_size) if i not in self.invalid_tokens]
+            valid_tokens = torch.tensor(valid_tokens, device=self.device)
+            rand_indices = torch.randint(
+                0, len(valid_tokens),
+                (self.args.num_samples, self.args.gen_len),
+                device=self.device
+            )
+            body = valid_tokens[rand_indices]
+            cls_col = torch.full((self.args.num_samples, 1), self.tokenizer.cls_id, device=self.device)
+            eos_col = torch.full((self.args.num_samples, 1), self.tokenizer.eos_id, device=self.device)
+            x = torch.cat([cls_col, body, eos_col], dim=1)
+        else:
+            vocab_size = self.model.tokenizer.vocab_size
+            x = torch.randint(0, vocab_size,
+                              (self.args.num_samples, self.args.gen_len),
+                              device=self.device)
+
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Initial Flow Matching", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i + 1]
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+
+                logits = self.model(x, t_tensor) / temperature
+                for inv_tok in self.invalid_tokens:
+                    if inv_tok < logits.shape[-1]:
+                        logits[:, :, inv_tok] = -1000
+
+                pred_x1 = torch.argmax(logits, dim=-1)
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                noise = torch.randint(0, vocab_size, x.shape, device=self.device)
+                x = torch.where(mask, noise, pred_x1)
+
+        generated_sequences = self.tokenizer.batch_decode(x)
+        for seq in generated_sequences:
+            print(f"  {seq}")
+        return x
+
+    # ------------------------------------------------------------------
+    # Scoring
+    # ------------------------------------------------------------------
+    def _aa_to_smiles_batch(self, x_batch):
+        seqs = self.tokenizer.batch_decode(x_batch)
+        three_letter_seqs = [one_letter_to_three_letter(s) for s in seqs]
+        with ThreadPoolExecutor() as pool:
+            smiles_seqs = list(pool.map(
+                lambda s: aa_to_smiles(s, cyclic=self.args.cyclic),
+                three_letter_seqs
+            ))
+        return seqs, smiles_seqs
+
+    def _get_scores(self, x_batch):
+        scores = []
+
+        if self.tokenizer_type == 'amino':
+            _, smiles_seqs = self._aa_to_smiles_batch(x_batch)
+            valid_mask = torch.tensor(
+                [1 if isinstance(s, str) else 0 for s in smiles_seqs],
+                device=self.device
+            )
+            for i in torch.where((1 - valid_mask) == 1)[0].tolist():
+                smiles_seqs[i] = 'C'
+
+            smiles_batch = self.smiles_tokenizer(
+                smiles_seqs, return_tensors='pt', padding=True
+            ).to(self.device)
+            attention_mask = smiles_batch.get("attention_mask", None)
+
+            for obj_func in self.objectives:
+                if getattr(obj_func, "input_type", None) == "aa":
+                    # AA-input objective: score raw AA tokens directly.
+                    score = obj_func(x_batch.to(self.device)).to(self.device)
+                    scores.append(score)
+                else:
+                    try:
+                        score = obj_func(
+                            smiles_batch["input_ids"], attention_mask=attention_mask
+                        ).to(self.device)
+                    except TypeError:
+                        score = obj_func(smiles_batch["input_ids"]).to(self.device)
+                    scores.append(score * valid_mask)
+        else:
+            for obj_func in self.objectives:
+                scores.append(obj_func(x_batch.to(self.device)))
+
+        return torch.stack(scores, dim=0).to(self.device)
+
+    def _compute_aggregated_score(self, scores, weights):
+        if self.args.score_aggregation == "weighted_min":
+            return torch.min(weights * scores, dim=0).values
+        return (weights * scores).sum(dim=0)
+
+    # ------------------------------------------------------------------
+    # Exponential tilting
+    # ------------------------------------------------------------------
+    def _exponential_tilt_guidance(self, x, pos_probs, candidates, mut_idx, sample_idx,
+                                   beta, weights, current_score=None, print_tilted=False,
+                                   step=None):
+        num_candidates = len(candidates)
+        current_token = x[sample_idx, mut_idx].item()
+
+        if num_candidates == 0:
+            return current_token, {'method': 'exp_tilt', 'fallback': True, 'accepted': False}
+
+        if num_candidates > self.args.top_k:
+            candidates = candidates[:self.args.top_k]
+            num_candidates = self.args.top_k
+
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(num_candidates, 1)
+        x_proposals[:, mut_idx] = candidates
+
+        proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+        validity_mask = (proposal_scores[0] > 0).float()
+        base_probs = pos_probs[sample_idx, candidates]
+
+        log_base_probs = torch.log(base_probs + 1e-10)
+        log_exp_scores = beta * aggregated_scores
+        log_validity = torch.log(validity_mask + 1e-10)
+        log_tilted = log_base_probs + log_exp_scores + log_validity
+
+        log_tilted_max = log_tilted.max()
+        fallback = False
+        if torch.isinf(log_tilted_max) or torch.isnan(log_tilted_max):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+        else:
+            log_tilted_stable = log_tilted - log_tilted_max
+            tilted_probs = torch.exp(log_tilted_stable)
+            prob_sum = tilted_probs.sum()
+            if prob_sum < 1e-10 or torch.isnan(prob_sum) or torch.isinf(prob_sum):
+                fallback = True
+                final_probs = base_probs / (base_probs.sum() + 1e-9)
+            else:
+                final_probs = tilted_probs / prob_sum
+
+        if torch.any(torch.isnan(final_probs)) or torch.any(torch.isinf(final_probs)):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+
+        final_probs = torch.clamp(final_probs, min=1e-10)
+        final_probs = final_probs / final_probs.sum()
+
+        if print_tilted:
+            id_to_token = self.tokenizer.id_to_token
+            topk = min(self.args.print_logits_topk, num_candidates)
+            top_vals, top_local = torch.topk(final_probs, k=topk)
+            print(f"\n[tilted] step={step} sample={sample_idx} mut_idx={mut_idx} "
+                  f"beta={beta:.2f} fallback={fallback}")
+            print(f"  Top-{topk} candidates by TILTED prob "
+                  f"(shows: base p, agg score S, β·S, valid, tilted p):")
+            for tv, li in zip(top_vals.tolist(), top_local.tolist()):
+                tid = candidates[li].item()
+                tok = id_to_token.get(int(tid), '?')
+                bp = base_probs[li].item()
+                sc = aggregated_scores[li].item()
+                bs = (beta * aggregated_scores[li]).item()
+                vd = validity_mask[li].item()
+                print(f"    id={tid:>3} tok={tok!r:>6}  base_p={bp:.3e}  "
+                      f"S={sc:+.3f}  β·S={bs:+.3f}  valid={vd:.0f}  tilted_p={tv:.3e}")
+            d_local = [li for li in range(num_candidates)
+                       if 33 <= candidates[li].item() <= 52
+                       and candidates[li].item() not in self.invalid_tokens]
+            if d_local:
+                print("  D-AA lowercase candidates (if any survived top-p):")
+                for li in d_local:
+                    tid = candidates[li].item()
+                    tok = id_to_token.get(int(tid), '?')
+                    print(f"    id={tid:>3} tok={tok!r:>4}  base_p={base_probs[li]:.3e}  "
+                          f"S={aggregated_scores[li]:+.3f}  "
+                          f"valid={validity_mask[li]:.0f}  tilted_p={final_probs[li]:.3e}")
+            else:
+                print("  D-AA lowercase candidates: NONE made it into the top-p pool")
+
+        selected_idx = torch.multinomial(final_probs, 1).item()
+        selected_token = candidates[selected_idx].item()
+        selected_score = aggregated_scores[selected_idx].item()
+
+        accepted = True
+        if self.args.use_acceptance and current_score is not None:
+            if selected_score < current_score:
+                selected_token = current_token
+                selected_score = current_score
+                accepted = False
+
+        return selected_token, {
+            'method': 'exp_tilt', 'beta': beta, 'num_candidates': num_candidates,
+            'proposal_score': aggregated_scores[selected_idx].item(),
+            'selected_score': selected_score, 'fallback': fallback,
+            'validity_rate': validity_mask.mean().item(), 'accepted': accepted
+        }
+
+    # ------------------------------------------------------------------
+    # Logit inspection
+    # ------------------------------------------------------------------
+    def _print_logit_distribution(self, raw_probs, masked_probs, mut_idx, step, sample_idx):
+        """Print the model's per-position probability distribution for inspection.
+
+        raw_probs:    [vocab_size] probabilities BEFORE masking current/invalid tokens
+        masked_probs: [vocab_size] probabilities AFTER masking
+        """
+        topk = self.args.print_logits_topk
+        id_to_token = self.tokenizer.id_to_token
+
+        # Token-category index ranges. Clip to model.vocab_size so we don't
+        # index out of bounds when the model uses a smaller vocab than the
+        # UAA tokenizer (e.g. PepReDi natural-AA vocab = 24).
+        V = int(raw_probs.shape[0])
+        cat_ranges = {
+            'special'   : [i for i in range(0, 4)   if i < V],
+            'L-AA upper': [i for i in range(4, 24)  if i < V],
+            'non-canon' : [i for i in range(24, 33) if i < V],
+            'D-AA lower': [i for i in range(33, 53) if i < V],
+        }
+
+        print(f"\n[logits] step={step} sample={sample_idx} mut_idx={mut_idx}")
+
+        # Top-k from raw probs
+        top_vals, top_ids = torch.topk(raw_probs, k=min(topk, raw_probs.shape[0]))
+        print(f"  Top-{topk} (raw model probs, pre-mask):")
+        for v, tid in zip(top_vals.tolist(), top_ids.tolist()):
+            tok = id_to_token.get(int(tid), '?')
+            print(f"    id={tid:>3} tok={tok!r:>6}  p={v:.4e}")
+
+        # Per-category mass
+        print("  Probability mass by category (raw / masked):")
+        for name, ids in cat_ranges.items():
+            ids_t = torch.tensor(ids, device=raw_probs.device)
+            raw_mass = raw_probs[ids_t].sum().item()
+            msk_mass = masked_probs[ids_t].sum().item()
+            print(f"    {name:>10}: raw={raw_mass:.4e}  masked={msk_mass:.4e}")
+
+        # Specifically: per-token probability for D-AA lowercase (excluding 'g'=38)
+        d_ids = [i for i in cat_ranges['D-AA lower'] if i not in self.invalid_tokens]
+        d_ids_t = torch.tensor(d_ids, device=raw_probs.device)
+        d_vals_raw = raw_probs[d_ids_t]
+        d_vals_msk = masked_probs[d_ids_t]
+        print("  D-AA lowercase per-token probs (raw / masked):")
+        for tid, vr, vm in zip(d_ids, d_vals_raw.tolist(), d_vals_msk.tolist()):
+            tok = id_to_token.get(int(tid), '?')
+            print(f"    id={tid:>3} tok={tok!r:>4}  raw={vr:.4e}  masked={vm:.4e}")
+
+    # ------------------------------------------------------------------
+    # Main loop
+    # ------------------------------------------------------------------
+    def generate(self):
+        x = self.generate_x0()
+
+        if self.args.weights is None:
+            weights = torch.tensor(
+                [1.0] + [1.0 / (self.num_objectives - 1)] * (self.num_objectives - 1),
+                device=self.device
+            ).view(-1, 1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+
+        print(f"\nGuidance Method: Exponential Tilting")
+        print(f"Weights: {weights.squeeze().tolist()}")
+        print(f"β range: [{self.args.beta_min}, {self.args.beta_max}]")
+        print(f"Acceptance criterion: {'ENABLED' if self.args.use_acceptance else 'DISABLED'}")
+
+        with torch.no_grad():
+            for step in tqdm(range(self.args.optimization_steps), desc="Guided Generation"):
+                beta = self.args.beta_min + (self.args.beta_max - self.args.beta_min) * (
+                    step / max(self.args.optimization_steps - 1, 1)
+                )
+
+                if self.args.position_strategy == "random":
+                    mut_idx = random.randint(1, self.args.gen_len - 2)
+                else:
+                    mut_idx = (step % (self.args.gen_len - 2)) + 1
+
+                generation_step = step % self.args.optimization_steps
+                time_t = torch.full(
+                    (self.args.num_samples,),
+                    generation_step / self.args.optimization_steps,
+                    device=self.device
+                )
+
+                logits = self.model(x, time_t)
+                probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+                raw_pos_probs = pos_probs.clone()
+
+                pos_probs[torch.arange(self.args.num_samples), x[:, mut_idx]] = 0
+                for inv_tok in self.invalid_tokens:
+                    pos_probs[:, inv_tok] = 0
+
+                should_print_logits = (
+                    self.args.print_logits
+                    and (step % max(self.args.print_logits_every, 1) == 0)
+                )
+                print_sample_idx = min(self.args.print_logits_sample,
+                                       self.args.num_samples - 1)
+                if should_print_logits:
+                    self._print_logit_distribution(
+                        raw_pos_probs[print_sample_idx], pos_probs[print_sample_idx],
+                        mut_idx, step, print_sample_idx
+                    )
+
+                sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                remove_mask = cumulative_probs > self.args.top_p
+                remove_mask[..., 1:] = remove_mask[..., :-1].clone()
+                remove_mask[..., 0] = 0
+
+                current_scores = self._get_scores(x)
+                current_aggregated = self._compute_aggregated_score(current_scores, weights)
+
+                if step == 0:
+                    print(f"\nInitial Scores: {current_scores.cpu().numpy()}")
+
+                improved = False
+                for i in range(self.args.num_samples):
+                    sample_mask = remove_mask[i]
+                    candidates = sorted_indices[i, ~sample_mask]
+
+                    valid_candidate_mask = ~torch.isin(
+                        candidates, torch.tensor(self.invalid_tokens, device=self.device)
+                    )
+                    candidates = candidates[valid_candidate_mask]
+
+                    if len(candidates) == 0:
+                        continue
+
+                    prev_token = x[i, mut_idx].item()
+                    prev_sample_scores = current_scores[:, i].clone()
+                    prev_weighted_sum = (weights.squeeze(1) * prev_sample_scores).sum()
+
+                    selected_token, info = self._exponential_tilt_guidance(
+                        x, pos_probs, candidates, mut_idx, i, beta, weights,
+                        current_score=current_aggregated[i].item(),
+                        print_tilted=(should_print_logits and i == print_sample_idx),
+                        step=step,
+                    )
+
+                    if selected_token == prev_token:
+                        continue
+
+                    x[i, mut_idx] = selected_token
+
+                    new_sample_scores = self._get_scores(x[i].unsqueeze(0))[:, 0]
+                    new_weighted_sum = (weights.squeeze(1) * new_sample_scores).sum()
+
+                    if new_weighted_sum >= prev_weighted_sum:
+                        improved = True
+                        print(f"Previous Weighted Sum: {prev_weighted_sum}")
+                        print(f"Previous Scores: {prev_sample_scores}")
+                        print(f"New Weighted Sum: {new_weighted_sum}")
+                        print(f"New Scores: {new_sample_scores}")
+                        current_scores[:, i] = new_sample_scores
+                        current_aggregated[i] = self._compute_aggregated_score(
+                            new_sample_scores.unsqueeze(1), weights
+                        )[0]
+                    else:
+                        x[i, mut_idx] = prev_token
+
+                if improved:
+                    print(self.tokenizer.batch_decode(x))
+                    if self.tracker:
+                        self.tracker.log_step(
+                            step, current_scores[0],
+                            guidance_info={'method': 'exp_tilt', 'beta': beta}
+                        )
+
+        return x
+
+
+# =============================================================================
+# Main
+# =============================================================================
+
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    print(f"Tokenizer type: {args.tokenizer_type}")
+
+    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)
+        print(f"Random seed: {args.seed}")
+
+    if args.tokenizer_type == 'amino':
+        print("Using Amino Acid tokenizer")
+        tokenizer = AminoTokenizer()
+        smiles_tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+    else:
+        print("Using SMILES tokenizer")
+        tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+        smiles_tokenizer = tokenizer
+
+    analyzer_model = Analyzer(smiles_tokenizer)
+    # Binding affinity: AA-input or SMILES-input depending on --affinity_type.
+    affinity_model = make_affinity(
+        args.affinity_type,
+        target_protein_seq=args.target,
+        device=device,
+        aa_ckpt_path=args.aa_affinity_ckpt,
+    )
+    print(f"Affinity scorer: {type(affinity_model).__name__} (input_type={getattr(affinity_model, 'input_type', '?')})")
+    OBJECTIVE_FUNCTIONS = [analyzer_model, affinity_model]
+
+    tracker = GenerationTracker(output_dir=args.output_dir)
+
+    print(f"Loading model from: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+
+    if args.tokenizer_type == 'amino':
+        raw_model_args = checkpoint['args']
+        if isinstance(raw_model_args, argparse.Namespace):
+            raw_model_args = vars(raw_model_args)
+        model_args = edict(raw_model_args)
+        model = MDLM(
+            vocab_size=model_args.vocab_size,
+            seq_len=model_args.seq_len,
+            model_dim=model_args.model_dim,
+            n_heads=model_args.n_heads,
+            n_layers=model_args.n_layers,
+        ).to(device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print(f"Amino model loaded. Vocab size: {model_args.vocab_size}")
+    else:
+        model_hparams = checkpoint["hyper_parameters"]["args"]
+        model = MDLMLightningModule.load_from_checkpoint(
+            args.checkpoint, args=model_hparams, tokenizer=tokenizer,
+            map_location=device, strict=False
+        )
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully.")
+
+    generator = CupOfMoCAStagedUAA(
+        model, device, OBJECTIVE_FUNCTIONS, args,
+        tokenizer=tokenizer, tokenizer_type=args.tokenizer_type, tracker=tracker
+    )
+
+    all_sequences = []
+    all_smiles = []
+    all_scores = []
+
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'=' * 60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'=' * 60}")
+
+        generated_tokens = generator.generate()
+        final_scores = generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+
+        print("\n--- Batch Results ---")
+        for i, seq in enumerate(sequence_str):
+            print(f"  {i}: {seq}")
+        print(f"Scores: {final_scores}")
+
+        all_sequences.extend(sequence_str)
+        all_scores.append(final_scores)
+
+        if args.tokenizer_type == 'amino':
+            smiles_list = []
+            for seq in sequence_str:
+                three_letter = one_letter_to_three_letter(seq)
+                smiles = aa_to_smiles(three_letter, cyclic=args.cyclic)
+                smiles_list.append(smiles if isinstance(smiles, str) else 'INVALID_SMILES')
+            all_smiles.extend(smiles_list)
+            tracker.log_final(sequence_str, final_scores, smiles_list=smiles_list)
+        else:
+            tracker.log_final(sequence_str, final_scores)
+
+    tracker.print_final_summary()
+    tracker.save(filename_prefix='cupofmoca_staged_uaa')
+
+    with open(args.output_file, 'w') as f:
+        for i, seq in enumerate(all_sequences):
+            if args.tokenizer_type == 'amino':
+                smiles = all_smiles[i] if i < len(all_smiles) else ''
+                f.write(f"{seq}\t{smiles}\n")
+            else:
+                f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+    print("\nGeneration complete.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="CupOfMoCA-style Staged Generation with UAA Foundation Model"
+    )
+
+    parser.add_argument("--checkpoint", type=str, required=True)
+    parser.add_argument("--target", type=str, required=True)
+
+    parser.add_argument("--num_samples", type=int, default=10)
+    parser.add_argument("--num_batches", type=int, default=10)
+    parser.add_argument("--gen_len", type=int, default=50)
+    parser.add_argument("--optimization_steps", type=int, default=128)
+
+    parser.add_argument("--beta_min", type=float, default=1.0)
+    parser.add_argument("--beta_max", type=float, default=20.0)
+    parser.add_argument("--use_acceptance", action="store_true")
+    parser.add_argument("--top_k", type=int, default=200)
+    parser.add_argument("--top_p", type=float, default=0.9)
+    parser.add_argument("--score_aggregation", type=str, default="weighted_sum",
+                        choices=["weighted_sum", "weighted_min"])
+
+    parser.add_argument("--weights", type=float, nargs='+', default=None)
+
+    parser.add_argument("--position_strategy", type=str, default="random",
+                        choices=["random", "sequential"])
+
+    parser.add_argument("--tokenizer_type", type=str, default="amino",
+                        choices=["smiles", "amino"])
+    parser.add_argument("--cyclic", action="store_true",
+                        help="Generate cyclic peptides (only for amino tokenizer).")
+    parser.add_argument("--smiles_vocab_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_vocab.txt")
+    parser.add_argument("--smiles_splits_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_splits.txt")
+
+    # AA-mode binding-affinity selection (this is the WT-AA / PepReDi script).
+    parser.add_argument("--affinity_type", type=str, default="smiles",
+                        choices=["aa", "smiles"],
+                        help="Which binding-affinity oracle to use: "
+                             "'aa' = AA-input UnpooledBindingPredictor, "
+                             "'smiles' = SMILES-input BindingAffinity.")
+    parser.add_argument("--aa_affinity_ckpt", type=str,
+                        default=DEFAULT_AA_AFFINITY_CKPT,
+                        help="Path to AA-input affinity predictor checkpoint.")
+
+    parser.add_argument("--output_file", type=str, default="./smiles.txt")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs")
+
+    parser.add_argument("--seed", type=int, default=None)
+
+    parser.add_argument("--print_logits", action="store_true",
+                        help="Print per-step probability distribution at the mutation position.")
+    parser.add_argument("--print_logits_every", type=int, default=1,
+                        help="Print logits every N optimization steps (requires --print_logits).")
+    parser.add_argument("--print_logits_topk", type=int, default=20,
+                        help="Number of top tokens to show per logit printout.")
+    parser.add_argument("--print_logits_sample", type=int, default=0,
+                        help="Which sample index to print logits for (default: 0).")
+
+    args = parser.parse_args()
+    main(args)

smiles/generation/cupofmoca_staged_uaa.py

@@ -0,0 +1,797 @@
+"""CupOfMoCA-style staged generation with the UAA foundation model."""
+
+import argparse
+import csv
+import json
+import os
+import random
+import sys
+from concurrent.futures import ThreadPoolExecutor
+from datetime import datetime
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from easydict import EasyDict as edict
+from tqdm import tqdm
+
+from smiles.models.train import MDLMLightningModule, PeptideAnalyzer
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+from smiles.scoring.classifiers import Analyzer, BindingAffinity
+
+from third_party.SMILESToPeptide.peptide_to_smiles import aa_to_smiles
+from third_party.SMILESToPeptide.aminoacid_selective import one_letter_to_three_letter
+
+
+# =============================================================================
+# Amino Acid Foundation Model (DiT-based MDLM)
+# =============================================================================
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(1, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+
+    def forward(self, t):
+        return self.mlp(t.unsqueeze(-1))
+
+
+class DiTBlock(nn.Module):
+    def __init__(self, hidden_size, n_heads):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = nn.MultiheadAttention(hidden_size, n_heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, 4 * hidden_size),
+            nn.GELU(),
+            nn.Linear(4 * hidden_size, hidden_size),
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True),
+        )
+
+    def forward(self, x, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = \
+            self.adaLN_modulation(c).chunk(6, dim=1)
+        x_norm1 = modulate(self.norm1(x), shift_msa, scale_msa)
+        attn_output, _ = self.attn(x_norm1, x_norm1, x_norm1)
+        x = x + gate_msa.unsqueeze(1) * attn_output
+        x_norm2 = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        mlp_output = self.mlp(x_norm2)
+        x = x + gate_mlp.unsqueeze(1) * mlp_output
+        return x
+
+
+class MDLM(nn.Module):
+    def __init__(self, vocab_size, seq_len, model_dim, n_heads, n_layers):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.seq_len = seq_len
+        self.model_dim = model_dim
+
+        self.token_embedder = nn.Embedding(vocab_size, model_dim)
+        self.pos_embedder = nn.Parameter(torch.randn(1, seq_len, model_dim))
+        self.time_embedder = TimestepEmbedder(model_dim)
+
+        self.transformer_blocks = nn.ModuleList([
+            DiTBlock(model_dim, n_heads) for _ in range(n_layers)
+        ])
+        self.final_norm = nn.LayerNorm(model_dim)
+        self.lm_head = nn.Linear(model_dim, vocab_size)
+
+    def forward(self, x, t):
+        seq_len = x.shape[1]
+        x_embed = self.token_embedder(x) + self.pos_embedder[:, :seq_len, :]
+        t_embed = self.time_embedder(t)
+        for block in self.transformer_blocks:
+            x_embed = block(x_embed, t_embed)
+        x_embed = self.final_norm(x_embed)
+        return self.lm_head(x_embed)
+
+
+class AminoTokenizer:
+    def __init__(self):
+        self.vocab = {
+            '<cls>': 0, '<pad>': 1, '<eos>': 2, '<unk>': 3,
+            'L': 4, 'A': 5, 'G': 6, 'V': 7, 'S': 8, 'E': 9,
+            'R': 10, 'T': 11, 'I': 12, 'D': 13, 'P': 14, 'K': 15,
+            'Q': 16, 'N': 17, 'F': 18, 'Y': 19, 'M': 20, 'H': 21,
+            'W': 22, 'C': 23,
+            'Ŷ': 24, 'Ĉ': 25, 'Ĝ': 26, 'Ė': 27, 'Ƙ': 28,
+            'Ṃ': 29, 'Ł': 30, 'Ċ': 31, 'Ṯ': 32,
+            'a': 33, 'c': 34, 'd': 35, 'e': 36, 'f': 37,
+            'g': 38, 'h': 39, 'i': 40, 'k': 41, 'l': 42,
+            'm': 43, 'n': 44, 'p': 45, 'q': 46, 'r': 47,
+            's': 48, 't': 49, 'v': 50, 'w': 51, 'y': 52,
+        }
+        self.id_to_token = {i: t for t, i in self.vocab.items()}
+        self.cls_token, self.pad_token = '<cls>', '<pad>'
+        self.eos_token, self.unk_token = '<eos>', '<unk>'
+        self.cls_id = self.vocab[self.cls_token]
+        self.pad_id = self.vocab[self.pad_token]
+        self.eos_id = self.vocab[self.eos_token]
+        self.unk_id = self.vocab[self.unk_token]
+        self.vocab_size = len(self.vocab)
+
+    def decode(self, ids, skip_special_tokens=True):
+        if isinstance(ids, torch.Tensor):
+            ids = ids.tolist()
+        tokens = []
+        for i in ids:
+            tok = self.id_to_token.get(int(i), self.unk_token)
+            if skip_special_tokens and tok in {self.cls_token, self.eos_token, self.pad_token}:
+                continue
+            tokens.append(tok)
+        return ''.join(tokens)
+
+    def batch_decode(self, batch_ids, skip_special_tokens=True):
+        return [self.decode(ids, skip_special_tokens=skip_special_tokens) for ids in batch_ids]
+
+
+# =============================================================================
+# Generation Tracker
+# =============================================================================
+
+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, scores, 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(sum(1 for s in scores if s > 0)) 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, sequences, scores, smiles_list=None):
+        for i, (seq, score_vec) in enumerate(zip(sequences, scores.T)):
+            self.final_results.append({
+                'sample_idx': i,
+                'aa_sequence': seq,
+                'peptide_smiles': smiles_list[i] if smiles_list else seq,
+                'is_valid_peptide': bool(score_vec[0] > 0),
+                'binding_score': float(score_vec[1]) if len(score_vec) > 1 else None,
+            })
+
+    def save(self, filename_prefix='generation'):
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_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', 'aa_sequence', 'peptide_smiles',
+                             'is_valid_peptide', 'binding_score'])
+            for r in self.final_results:
+                writer.writerow([
+                    r['sample_idx'], r.get('aa_sequence', ''), r['peptide_smiles'],
+                    r['is_valid_peptide'], r.get('binding_score', '')
+                ])
+
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        return step_log_path, final_path, csv_path
+
+    def print_step_summary(self, step, scores):
+        scores_list = scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else list(scores)
+        num_valid = sum(1 for s in scores_list if s > 0)
+        print(f"\n--- Step {step} Summary ---")
+        print(f"  Valid peptides: {num_valid}/{len(scores_list)}")
+
+    def print_final_summary(self):
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        total = len(self.final_results)
+        valid = sum(1 for r in self.final_results if r.get('is_valid_peptide', False))
+        print(f"\nTotal sequences: {total}")
+        print(f"Valid peptides: {valid}")
+        if any(r.get('binding_score') is not None for r in self.final_results):
+            avg = np.mean([r['binding_score'] for r in self.final_results
+                           if r.get('binding_score') is not None])
+            print(f"Average binding score: {avg:.3f}")
+
+
+# =============================================================================
+# CupOfMoCA Staged Generator (UAA-aware)
+# =============================================================================
+
+class CupOfMoCAStagedUAA:
+    """Exponential-tilting staged generator that operates in AA token space."""
+
+    def __init__(self, model, device, objectives, args, tokenizer, tokenizer_type='amino',
+                 tracker=None):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.tracker = tracker
+        self.tokenizer = tokenizer
+        self.tokenizer_type = tokenizer_type
+
+        if tokenizer_type == 'amino':
+            self.invalid_tokens = [0, 1, 2, 3, 38]  # specials + 'g'
+            self.smiles_tokenizer = SMILES_SPE_Tokenizer(
+                args.smiles_vocab_path, args.smiles_splits_path
+            )
+        else:
+            self.invalid_tokens = [0, 1, 2, 3, 4, 585, 586]
+            self.smiles_tokenizer = None
+
+        self.analyzer_obj = None
+        for obj in objectives:
+            if isinstance(obj, Analyzer):
+                self.analyzer_obj = obj
+                break
+
+    # ------------------------------------------------------------------
+    # x0 generation
+    # ------------------------------------------------------------------
+    def generate_x0(self, n_steps=16, temperature=1.0):
+        print("Generating initial sequences (x0 from prior)...")
+        self.model.eval()
+
+        if self.tokenizer_type == 'amino':
+            vocab_size = self.model.vocab_size
+            valid_tokens = [i for i in range(vocab_size) if i not in self.invalid_tokens]
+            valid_tokens = torch.tensor(valid_tokens, device=self.device)
+            rand_indices = torch.randint(
+                0, len(valid_tokens),
+                (self.args.num_samples, self.args.gen_len),
+                device=self.device
+            )
+            body = valid_tokens[rand_indices]
+            cls_col = torch.full((self.args.num_samples, 1), self.tokenizer.cls_id, device=self.device)
+            eos_col = torch.full((self.args.num_samples, 1), self.tokenizer.eos_id, device=self.device)
+            x = torch.cat([cls_col, body, eos_col], dim=1)
+        else:
+            vocab_size = self.model.tokenizer.vocab_size
+            x = torch.randint(0, vocab_size,
+                              (self.args.num_samples, self.args.gen_len),
+                              device=self.device)
+
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Initial Flow Matching", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i + 1]
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+
+                logits = self.model(x, t_tensor) / temperature
+                for inv_tok in self.invalid_tokens:
+                    if inv_tok < logits.shape[-1]:
+                        logits[:, :, inv_tok] = -1000
+
+                pred_x1 = torch.argmax(logits, dim=-1)
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                noise = torch.randint(0, vocab_size, x.shape, device=self.device)
+                x = torch.where(mask, noise, pred_x1)
+
+        generated_sequences = self.tokenizer.batch_decode(x)
+        for seq in generated_sequences:
+            print(f"  {seq}")
+        return x
+
+    # ------------------------------------------------------------------
+    # Scoring
+    # ------------------------------------------------------------------
+    def _aa_to_smiles_batch(self, x_batch):
+        seqs = self.tokenizer.batch_decode(x_batch)
+        three_letter_seqs = [one_letter_to_three_letter(s) for s in seqs]
+        with ThreadPoolExecutor() as pool:
+            smiles_seqs = list(pool.map(
+                lambda s: aa_to_smiles(s, cyclic=self.args.cyclic),
+                three_letter_seqs
+            ))
+        return seqs, smiles_seqs
+
+    def _get_scores(self, x_batch):
+        scores = []
+
+        if self.tokenizer_type == 'amino':
+            _, smiles_seqs = self._aa_to_smiles_batch(x_batch)
+            valid_mask = torch.tensor(
+                [1 if isinstance(s, str) else 0 for s in smiles_seqs],
+                device=self.device
+            )
+            for i in torch.where((1 - valid_mask) == 1)[0].tolist():
+                smiles_seqs[i] = 'C'
+
+            smiles_batch = self.smiles_tokenizer(
+                smiles_seqs, return_tensors='pt', padding=True
+            ).to(self.device)
+            attention_mask = smiles_batch.get("attention_mask", None)
+
+            for obj_func in self.objectives:
+                try:
+                    score = obj_func(
+                        smiles_batch["input_ids"], attention_mask=attention_mask
+                    ).to(self.device)
+                except TypeError:
+                    score = obj_func(smiles_batch["input_ids"]).to(self.device)
+                scores.append(score * valid_mask)
+        else:
+            for obj_func in self.objectives:
+                scores.append(obj_func(x_batch.to(self.device)))
+
+        return torch.stack(scores, dim=0).to(self.device)
+
+    def _compute_aggregated_score(self, scores, weights):
+        if self.args.score_aggregation == "weighted_min":
+            return torch.min(weights * scores, dim=0).values
+        return (weights * scores).sum(dim=0)
+
+    # ------------------------------------------------------------------
+    # Exponential tilting
+    # ------------------------------------------------------------------
+    def _exponential_tilt_guidance(self, x, pos_probs, candidates, mut_idx, sample_idx,
+                                   beta, weights, current_score=None, print_tilted=False,
+                                   step=None):
+        num_candidates = len(candidates)
+        current_token = x[sample_idx, mut_idx].item()
+
+        if num_candidates == 0:
+            return current_token, {'method': 'exp_tilt', 'fallback': True, 'accepted': False}
+
+        if num_candidates > self.args.top_k:
+            candidates = candidates[:self.args.top_k]
+            num_candidates = self.args.top_k
+
+        x_proposals = x[sample_idx].unsqueeze(0).repeat(num_candidates, 1)
+        x_proposals[:, mut_idx] = candidates
+
+        proposal_scores = self._get_scores(x_proposals)
+        aggregated_scores = self._compute_aggregated_score(proposal_scores, weights)
+        validity_mask = (proposal_scores[0] > 0).float()
+        base_probs = pos_probs[sample_idx, candidates]
+
+        log_base_probs = torch.log(base_probs + 1e-10)
+        log_exp_scores = beta * aggregated_scores
+        log_validity = torch.log(validity_mask + 1e-10)
+        log_tilted = log_base_probs + log_exp_scores + log_validity
+
+        log_tilted_max = log_tilted.max()
+        fallback = False
+        if torch.isinf(log_tilted_max) or torch.isnan(log_tilted_max):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+        else:
+            log_tilted_stable = log_tilted - log_tilted_max
+            tilted_probs = torch.exp(log_tilted_stable)
+            prob_sum = tilted_probs.sum()
+            if prob_sum < 1e-10 or torch.isnan(prob_sum) or torch.isinf(prob_sum):
+                fallback = True
+                final_probs = base_probs / (base_probs.sum() + 1e-9)
+            else:
+                final_probs = tilted_probs / prob_sum
+
+        if torch.any(torch.isnan(final_probs)) or torch.any(torch.isinf(final_probs)):
+            fallback = True
+            final_probs = base_probs / (base_probs.sum() + 1e-9)
+
+        final_probs = torch.clamp(final_probs, min=1e-10)
+        final_probs = final_probs / final_probs.sum()
+
+        if print_tilted:
+            id_to_token = self.tokenizer.id_to_token
+            topk = min(self.args.print_logits_topk, num_candidates)
+            top_vals, top_local = torch.topk(final_probs, k=topk)
+            print(f"\n[tilted] step={step} sample={sample_idx} mut_idx={mut_idx} "
+                  f"beta={beta:.2f} fallback={fallback}")
+            print(f"  Top-{topk} candidates by TILTED prob "
+                  f"(shows: base p, agg score S, β·S, valid, tilted p):")
+            for tv, li in zip(top_vals.tolist(), top_local.tolist()):
+                tid = candidates[li].item()
+                tok = id_to_token.get(int(tid), '?')
+                bp = base_probs[li].item()
+                sc = aggregated_scores[li].item()
+                bs = (beta * aggregated_scores[li]).item()
+                vd = validity_mask[li].item()
+                print(f"    id={tid:>3} tok={tok!r:>6}  base_p={bp:.3e}  "
+                      f"S={sc:+.3f}  β·S={bs:+.3f}  valid={vd:.0f}  tilted_p={tv:.3e}")
+            d_local = [li for li in range(num_candidates)
+                       if 33 <= candidates[li].item() <= 52
+                       and candidates[li].item() not in self.invalid_tokens]
+            if d_local:
+                print("  D-AA lowercase candidates (if any survived top-p):")
+                for li in d_local:
+                    tid = candidates[li].item()
+                    tok = id_to_token.get(int(tid), '?')
+                    print(f"    id={tid:>3} tok={tok!r:>4}  base_p={base_probs[li]:.3e}  "
+                          f"S={aggregated_scores[li]:+.3f}  "
+                          f"valid={validity_mask[li]:.0f}  tilted_p={final_probs[li]:.3e}")
+            else:
+                print("  D-AA lowercase candidates: NONE made it into the top-p pool")
+
+        selected_idx = torch.multinomial(final_probs, 1).item()
+        selected_token = candidates[selected_idx].item()
+        selected_score = aggregated_scores[selected_idx].item()
+
+        accepted = True
+        if self.args.use_acceptance and current_score is not None:
+            if selected_score < current_score:
+                selected_token = current_token
+                selected_score = current_score
+                accepted = False
+
+        return selected_token, {
+            'method': 'exp_tilt', 'beta': beta, 'num_candidates': num_candidates,
+            'proposal_score': aggregated_scores[selected_idx].item(),
+            'selected_score': selected_score, 'fallback': fallback,
+            'validity_rate': validity_mask.mean().item(), 'accepted': accepted
+        }
+
+    # ------------------------------------------------------------------
+    # Logit inspection
+    # ------------------------------------------------------------------
+    def _print_logit_distribution(self, raw_probs, masked_probs, mut_idx, step, sample_idx):
+        """Print the model's per-position probability distribution for inspection.
+
+        raw_probs:    [vocab_size] probabilities BEFORE masking current/invalid tokens
+        masked_probs: [vocab_size] probabilities AFTER masking
+        """
+        topk = self.args.print_logits_topk
+        id_to_token = self.tokenizer.id_to_token
+
+        # Token-category index ranges (AminoTokenizer):
+        #   specials  : 0-3
+        #   L-AA upper: 4-23
+        #   non-canon : 24-32
+        #   D-AA lower: 33-52  (38='g' is invalid)
+        cat_ranges = {
+            'special'   : list(range(0, 4)),
+            'L-AA upper': list(range(4, 24)),
+            'non-canon' : list(range(24, 33)),
+            'D-AA lower': list(range(33, 53)),
+        }
+
+        print(f"\n[logits] step={step} sample={sample_idx} mut_idx={mut_idx}")
+
+        # Top-k from raw probs
+        top_vals, top_ids = torch.topk(raw_probs, k=min(topk, raw_probs.shape[0]))
+        print(f"  Top-{topk} (raw model probs, pre-mask):")
+        for v, tid in zip(top_vals.tolist(), top_ids.tolist()):
+            tok = id_to_token.get(int(tid), '?')
+            print(f"    id={tid:>3} tok={tok!r:>6}  p={v:.4e}")
+
+        # Per-category mass
+        print("  Probability mass by category (raw / masked):")
+        for name, ids in cat_ranges.items():
+            ids_t = torch.tensor(ids, device=raw_probs.device)
+            raw_mass = raw_probs[ids_t].sum().item()
+            msk_mass = masked_probs[ids_t].sum().item()
+            print(f"    {name:>10}: raw={raw_mass:.4e}  masked={msk_mass:.4e}")
+
+        # Specifically: per-token probability for D-AA lowercase (excluding 'g'=38)
+        d_ids = [i for i in cat_ranges['D-AA lower'] if i not in self.invalid_tokens]
+        d_ids_t = torch.tensor(d_ids, device=raw_probs.device)
+        d_vals_raw = raw_probs[d_ids_t]
+        d_vals_msk = masked_probs[d_ids_t]
+        print("  D-AA lowercase per-token probs (raw / masked):")
+        for tid, vr, vm in zip(d_ids, d_vals_raw.tolist(), d_vals_msk.tolist()):
+            tok = id_to_token.get(int(tid), '?')
+            print(f"    id={tid:>3} tok={tok!r:>4}  raw={vr:.4e}  masked={vm:.4e}")
+
+    # ------------------------------------------------------------------
+    # Main loop
+    # ------------------------------------------------------------------
+    def generate(self):
+        x = self.generate_x0()
+
+        if self.args.weights is None:
+            weights = torch.tensor(
+                [1.0] + [1.0 / (self.num_objectives - 1)] * (self.num_objectives - 1),
+                device=self.device
+            ).view(-1, 1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+
+        print(f"\nGuidance Method: Exponential Tilting")
+        print(f"Weights: {weights.squeeze().tolist()}")
+        print(f"β range: [{self.args.beta_min}, {self.args.beta_max}]")
+        print(f"Acceptance criterion: {'ENABLED' if self.args.use_acceptance else 'DISABLED'}")
+
+        with torch.no_grad():
+            for step in tqdm(range(self.args.optimization_steps), desc="Guided Generation"):
+                beta = self.args.beta_min + (self.args.beta_max - self.args.beta_min) * (
+                    step / max(self.args.optimization_steps - 1, 1)
+                )
+
+                if self.args.position_strategy == "random":
+                    mut_idx = random.randint(1, self.args.gen_len - 2)
+                else:
+                    mut_idx = (step % (self.args.gen_len - 2)) + 1
+
+                generation_step = step % self.args.optimization_steps
+                time_t = torch.full(
+                    (self.args.num_samples,),
+                    generation_step / self.args.optimization_steps,
+                    device=self.device
+                )
+
+                logits = self.model(x, time_t)
+                probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+                raw_pos_probs = pos_probs.clone()
+
+                pos_probs[torch.arange(self.args.num_samples), x[:, mut_idx]] = 0
+                for inv_tok in self.invalid_tokens:
+                    pos_probs[:, inv_tok] = 0
+
+                should_print_logits = (
+                    self.args.print_logits
+                    and (step % max(self.args.print_logits_every, 1) == 0)
+                )
+                print_sample_idx = min(self.args.print_logits_sample,
+                                       self.args.num_samples - 1)
+                if should_print_logits:
+                    self._print_logit_distribution(
+                        raw_pos_probs[print_sample_idx], pos_probs[print_sample_idx],
+                        mut_idx, step, print_sample_idx
+                    )
+
+                sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                remove_mask = cumulative_probs > self.args.top_p
+                remove_mask[..., 1:] = remove_mask[..., :-1].clone()
+                remove_mask[..., 0] = 0
+
+                current_scores = self._get_scores(x)
+                current_aggregated = self._compute_aggregated_score(current_scores, weights)
+
+                if step == 0:
+                    print(f"\nInitial Scores: {current_scores.cpu().numpy()}")
+
+                improved = False
+                for i in range(self.args.num_samples):
+                    sample_mask = remove_mask[i]
+                    candidates = sorted_indices[i, ~sample_mask]
+
+                    valid_candidate_mask = ~torch.isin(
+                        candidates, torch.tensor(self.invalid_tokens, device=self.device)
+                    )
+                    candidates = candidates[valid_candidate_mask]
+
+                    if len(candidates) == 0:
+                        continue
+
+                    prev_token = x[i, mut_idx].item()
+                    prev_sample_scores = current_scores[:, i].clone()
+                    prev_weighted_sum = (weights.squeeze(1) * prev_sample_scores).sum()
+
+                    selected_token, info = self._exponential_tilt_guidance(
+                        x, pos_probs, candidates, mut_idx, i, beta, weights,
+                        current_score=current_aggregated[i].item(),
+                        print_tilted=(should_print_logits and i == print_sample_idx),
+                        step=step,
+                    )
+
+                    if selected_token == prev_token:
+                        continue
+
+                    x[i, mut_idx] = selected_token
+
+                    new_sample_scores = self._get_scores(x[i].unsqueeze(0))[:, 0]
+                    new_weighted_sum = (weights.squeeze(1) * new_sample_scores).sum()
+
+                    if new_weighted_sum >= prev_weighted_sum:
+                        improved = True
+                        print(f"Previous Weighted Sum: {prev_weighted_sum}")
+                        print(f"Previous Scores: {prev_sample_scores}")
+                        print(f"New Weighted Sum: {new_weighted_sum}")
+                        print(f"New Scores: {new_sample_scores}")
+                        current_scores[:, i] = new_sample_scores
+                        current_aggregated[i] = self._compute_aggregated_score(
+                            new_sample_scores.unsqueeze(1), weights
+                        )[0]
+                    else:
+                        x[i, mut_idx] = prev_token
+
+                if improved:
+                    print(self.tokenizer.batch_decode(x))
+                    if self.tracker:
+                        self.tracker.log_step(
+                            step, current_scores[0],
+                            guidance_info={'method': 'exp_tilt', 'beta': beta}
+                        )
+
+        return x
+
+
+# =============================================================================
+# Main
+# =============================================================================
+
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    print(f"Tokenizer type: {args.tokenizer_type}")
+
+    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)
+        print(f"Random seed: {args.seed}")
+
+    if args.tokenizer_type == 'amino':
+        print("Using Amino Acid tokenizer")
+        tokenizer = AminoTokenizer()
+        smiles_tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+    else:
+        print("Using SMILES tokenizer")
+        tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+        smiles_tokenizer = tokenizer
+
+    analyzer_model = Analyzer(smiles_tokenizer)
+    affinity_model = BindingAffinity(args.target, device)
+    OBJECTIVE_FUNCTIONS = [analyzer_model, affinity_model]
+
+    tracker = GenerationTracker(output_dir=args.output_dir)
+
+    print(f"Loading model from: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+
+    if args.tokenizer_type == 'amino':
+        raw_model_args = checkpoint['args']
+        if isinstance(raw_model_args, argparse.Namespace):
+            raw_model_args = vars(raw_model_args)
+        model_args = edict(raw_model_args)
+        model = MDLM(
+            vocab_size=model_args.vocab_size,
+            seq_len=model_args.seq_len,
+            model_dim=model_args.model_dim,
+            n_heads=model_args.n_heads,
+            n_layers=model_args.n_layers,
+        ).to(device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print(f"Amino model loaded. Vocab size: {model_args.vocab_size}")
+    else:
+        model_hparams = checkpoint["hyper_parameters"]["args"]
+        model = MDLMLightningModule.load_from_checkpoint(
+            args.checkpoint, args=model_hparams, tokenizer=tokenizer,
+            map_location=device, strict=False
+        )
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully.")
+
+    generator = CupOfMoCAStagedUAA(
+        model, device, OBJECTIVE_FUNCTIONS, args,
+        tokenizer=tokenizer, tokenizer_type=args.tokenizer_type, tracker=tracker
+    )
+
+    all_sequences = []
+    all_smiles = []
+    all_scores = []
+
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'=' * 60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'=' * 60}")
+
+        generated_tokens = generator.generate()
+        final_scores = generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+
+        print("\n--- Batch Results ---")
+        for i, seq in enumerate(sequence_str):
+            print(f"  {i}: {seq}")
+        print(f"Scores: {final_scores}")
+
+        all_sequences.extend(sequence_str)
+        all_scores.append(final_scores)
+
+        if args.tokenizer_type == 'amino':
+            smiles_list = []
+            for seq in sequence_str:
+                three_letter = one_letter_to_three_letter(seq)
+                smiles = aa_to_smiles(three_letter, cyclic=args.cyclic)
+                smiles_list.append(smiles if isinstance(smiles, str) else 'INVALID_SMILES')
+            all_smiles.extend(smiles_list)
+            tracker.log_final(sequence_str, final_scores, smiles_list=smiles_list)
+        else:
+            tracker.log_final(sequence_str, final_scores)
+
+    tracker.print_final_summary()
+    tracker.save(filename_prefix='cupofmoca_staged_uaa')
+
+    with open(args.output_file, 'w') as f:
+        for i, seq in enumerate(all_sequences):
+            if args.tokenizer_type == 'amino':
+                smiles = all_smiles[i] if i < len(all_smiles) else ''
+                f.write(f"{seq}\t{smiles}\n")
+            else:
+                f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+    print("\nGeneration complete.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="CupOfMoCA-style Staged Generation with UAA Foundation Model"
+    )
+
+    parser.add_argument("--checkpoint", type=str, required=True)
+    parser.add_argument("--target", type=str, required=True)
+
+    parser.add_argument("--num_samples", type=int, default=10)
+    parser.add_argument("--num_batches", type=int, default=10)
+    parser.add_argument("--gen_len", type=int, default=50)
+    parser.add_argument("--optimization_steps", type=int, default=128)
+
+    parser.add_argument("--beta_min", type=float, default=1.0)
+    parser.add_argument("--beta_max", type=float, default=20.0)
+    parser.add_argument("--use_acceptance", action="store_true")
+    parser.add_argument("--top_k", type=int, default=200)
+    parser.add_argument("--top_p", type=float, default=0.9)
+    parser.add_argument("--score_aggregation", type=str, default="weighted_sum",
+                        choices=["weighted_sum", "weighted_min"])
+
+    parser.add_argument("--weights", type=float, nargs='+', default=None)
+
+    parser.add_argument("--position_strategy", type=str, default="random",
+                        choices=["random", "sequential"])
+
+    parser.add_argument("--tokenizer_type", type=str, default="amino",
+                        choices=["smiles", "amino"])
+    parser.add_argument("--cyclic", action="store_true",
+                        help="Generate cyclic peptides (only for amino tokenizer).")
+    parser.add_argument("--smiles_vocab_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_vocab.txt")
+    parser.add_argument("--smiles_splits_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_splits.txt")
+
+    parser.add_argument("--output_file", type=str, default="./smiles.txt")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs")
+
+    parser.add_argument("--seed", type=int, default=None)
+
+    parser.add_argument("--print_logits", action="store_true",
+                        help="Print per-step probability distribution at the mutation position.")
+    parser.add_argument("--print_logits_every", type=int, default=1,
+                        help="Print logits every N optimization steps (requires --print_logits).")
+    parser.add_argument("--print_logits_topk", type=int, default=20,
+                        help="Number of top tokens to show per logit printout.")
+    parser.add_argument("--print_logits_sample", type=int, default=0,
+                        help="Which sample index to print logits for (default: 0).")
+
+    args = parser.parse_args()
+    main(args)

smiles/generation/guided.py

@@ -0,0 +1,651 @@
+import argparse
+import re
+import random
+from collections import Counter
+import csv
+import json
+import os
+from datetime import datetime
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+from smiles.models.train import MDLMLightningModule
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+import numpy as np
+
+from smiles.evaluation.peptide_analyzer import PeptideAnalyzer
+from smiles.evaluation.pdc_analyzer import PDCAnalyzer
+from smiles.scoring.classifiers import Conjugation, BindingAffinity
+
+
+class GenerationTracker:
+    """
+    Tracks detailed validity information during generation
+    """
+    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, scores, detailed_results=None, improved_indices=None):
+        """Log information for a generation step"""
+        step_info = {
+            'step': step,
+            'scores': scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else scores,
+            'num_valid': int(sum(1 for s in scores if s > 0)) if hasattr(scores, '__iter__') else 0,
+            'improved_indices': improved_indices or []
+        }
+        
+        if detailed_results:
+            step_info['detailed'] = []
+            for i, result in enumerate(detailed_results):
+                step_info['detailed'].append({
+                    'sample_idx': i,
+                    'is_valid_peptide': result.get('is_valid_peptide', False),
+                    'has_conjugation_site': result.get('has_conjugation_site', False),
+                    'conjugation_success': result.get('conjugation_success', False),
+                    'pdc_valid': result.get('pdc_valid', False),
+                    'num_conjugates': len(result.get('conjugates', [])),
+                    'error': result.get('error')
+                })
+        
+        self.step_logs.append(step_info)
+        
+    def log_final(self, sequences, detailed_results_with_props):
+        """Log final results with properties"""
+        for i, (seq, result) in enumerate(zip(sequences, detailed_results_with_props)):
+            final_info = {
+                'sample_idx': i,
+                'peptide_smiles': seq,
+                'is_valid_peptide': result.get('is_valid_peptide', False),
+                'has_conjugation_site': result.get('has_conjugation_site', False),
+                'num_conjugates': len(result.get('conjugates', [])),
+                'num_valid_pdcs': len(result.get('valid_pdcs', [])),
+                'best_pdc': None,
+                'properties': result.get('properties', {})
+            }
+            
+            if result.get('best_pdc'):
+                final_info['best_pdc'] = {
+                    'smiles': result['best_pdc'].get('smiles', ''),
+                    'checks': result['best_pdc'].get('checks', {})
+                }
+            
+            self.final_results.append(final_info)
+    
+    def save(self, filename_prefix='generation'):
+        """Save all logs to files"""
+        # Save step logs
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_path, 'w') as f:
+            json.dump(self.step_logs, f, indent=2, default=str)
+        
+        # Save final results
+        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)
+        
+        # Save summary CSV
+        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', 'peptide_smiles', 'is_valid_peptide', 'has_conjugation_site',
+                'num_conjugates', 'num_valid_pdcs', 'best_pdc_smiles',
+                'mf', 'polararea', 'xlogp', 'heavycnt', 'rotbonds'
+            ])
+            for result in self.final_results:
+                props = result.get('properties', {})
+                writer.writerow([
+                    result['sample_idx'],
+                    result['peptide_smiles'],
+                    result['is_valid_peptide'],
+                    result['has_conjugation_site'],
+                    result['num_conjugates'],
+                    result['num_valid_pdcs'],
+                    result['best_pdc']['smiles'] if result.get('best_pdc') else '',
+                    props.get('mf', ''),
+                    props.get('polararea', ''),
+                    props.get('xlogp', ''),
+                    props.get('heavycnt', ''),
+                    props.get('rotbonds', '')
+                ])
+        
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        
+        return step_log_path, final_path, csv_path
+    
+    def print_step_summary(self, step, scores, detailed_results=None):
+        """Print a summary for the current step"""
+        if isinstance(scores, torch.Tensor):
+            scores_list = scores.cpu().numpy().tolist()
+        else:
+            scores_list = list(scores)
+        
+        num_valid = sum(1 for s in scores_list if s > 0)
+        
+        print(f"\n--- Step {step} Summary ---")
+        print(f"  Valid: {num_valid}/{len(scores_list)}")
+        
+        if detailed_results:
+            valid_peptide = sum(1 for r in detailed_results if r.get('is_valid_peptide'))
+            has_site = sum(1 for r in detailed_results if r.get('has_conjugation_site'))
+            conj_success = sum(1 for r in detailed_results if r.get('conjugation_success'))
+            pdc_valid = sum(1 for r in detailed_results if r.get('pdc_valid'))
+            
+            print(f"  Breakdown:")
+            print(f"    Valid peptides: {valid_peptide}")
+            print(f"    Has conjugation site: {has_site}")
+            print(f"    Conjugation success: {conj_success}")
+            print(f"    PDC valid: {pdc_valid}")
+    
+    def print_final_summary(self):
+        """Print final summary"""
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        
+        total = len(self.final_results)
+        valid_pdcs = sum(1 for r in self.final_results if r.get('num_valid_pdcs', 0) > 0)
+        
+        print(f"\nTotal sequences: {total}")
+        print(f"With valid PDCs: {valid_pdcs}")
+        
+        print("\nDetailed results:")
+        for result in self.final_results:
+            status = "" if result.get('num_valid_pdcs', 0) > 0 else ""
+            print(f"\n  [{status}] Sample {result['sample_idx']}:")
+            print(f"      Peptide: {result['peptide_smiles'][:60]}...")
+            print(f"      Valid peptide: {result['is_valid_peptide']}")
+            print(f"      Has conjugation site: {result['has_conjugation_site']}")            
+            print(f"      Conjugates: {result['num_conjugates']}")
+            print(f"      Valid PDCs: {result['num_valid_pdcs']}")
+            
+            if result.get('properties'):
+                props = result['properties']
+                print(f"      Properties:")
+                for k, v in props.items():
+                    print(f"        {k}: {v}")
+
+
+def peptide_bond_mask(smiles_list):
+    # Initialize the batch mask
+    batch_size = len(smiles_list)
+    max_seq_length = 1035 #max(len(smiles) for smiles in smiles_list)  # Find the longest SMILES
+    mask = torch.zeros((batch_size, max_seq_length), dtype=torch.int)  # Mask filled with zeros
+
+    bond_patterns = [
+        (r'OC\(=O\)', 'ester'),
+        (r'N\(C\)C\(=O\)', 'n_methyl'),
+        (r'N[12]C\(=O\)', 'peptide'),  # Pro peptide bonds
+        (r'NC\(=O\)', 'peptide'),  # Regular peptide bonds
+        (r'C\(=O\)N\(C\)', 'n_methyl'),
+        (r'C\(=O\)N[12]?', 'peptide')
+    ]
+
+    for batch_idx, smiles in enumerate(smiles_list):
+        positions = []
+        used = set()
+
+        # Identify bonds
+        for pattern, bond_type in bond_patterns:
+            for match in re.finditer(pattern, smiles):
+                if not any(p in range(match.start(), match.end()) for p in used):
+                    positions.append({
+                        'start': match.start(),
+                        'end': match.end(),
+                        'type': bond_type,
+                        'pattern': match.group()
+                    })
+                    used.update(range(match.start(), match.end()))
+
+        # Update the mask for the current SMILES
+        for pos in positions:
+            mask[batch_idx, pos['start']:pos['end']] = 1
+
+    return mask
+
+def peptide_token_mask(smiles_list, token_lists):
+    # Initialize the batch mask
+    batch_size = len(smiles_list)
+    token_seq_length = max(len(tokens) for tokens in token_lists)  # Find the longest tokenized sequence
+    tokenized_masks = torch.zeros((batch_size, token_seq_length), dtype=torch.int)  # Mask filled with zeros
+    atomwise_masks = peptide_bond_mask(smiles_list)
+
+
+    for batch_idx, atomwise_mask in enumerate(atomwise_masks):
+        token_seq = token_lists[batch_idx]
+        atom_idx = 0
+        
+        for token_idx, token in enumerate(token_seq):
+            if token_idx != 0 and token_idx != len(token_seq) - 1:
+                if torch.sum(atomwise_mask[atom_idx:atom_idx+len(token)]) >= 1:
+                    tokenized_masks[batch_idx][token_idx] = 1
+                atom_idx += len(token)
+    
+    return tokenized_masks
+
+class MOGGenerator:
+    def __init__(self, model, device, objectives, args, tracker=None):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.pdc_analyzer = PDCAnalyzer(device)
+        self.invalid = [0, 1, 2, 3, 4, 585, 586]
+        self.tracker = tracker
+        
+        # Find the Conjugation objective for detailed tracking
+        self.conjugation_obj = None
+        for obj in objectives:
+            if isinstance(obj, Conjugation):
+                self.conjugation_obj = obj
+                break
+
+    def _generate_single_x0(self, n_steps=16, temperature=1.0):
+        """Generate a single batch of x0 sequences."""
+        self.model.eval()
+
+        # 1. Start with a tensor of random tokens (pure noise at t=0)
+        x = torch.randint(
+            0,
+            self.model.tokenizer.vocab_size,
+            (self.args.num_samples, self.args.gen_len),
+            device=self.device
+        )
+
+        # 2. Define the time schedule for the forward process (0.0 to 1.0)
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+
+        # 3. Iteratively follow the flow from noise to data
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Flow Matching Steps", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i+1]
+
+                # Prepare the current timestep tensor for the model
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+
+                # Get the model's prediction for the final clean sequence (at t=1)
+                logits = self.model(x, t_tensor) 
+                logits = logits / temperature
+                logits[:, :, 586] = -1000
+
+                pred_x1 = torch.argmax(logits, dim=-1)
+
+                # On the last step, the result is the final prediction
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+
+                # --- Construct the next state x_{t_next} ---
+                # The probability of a token being noise at time t_next is (1 - t_next).
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                
+                # Generate new random tokens for the noise positions
+                noise = torch.randint(
+                    0,
+                    self.model.tokenizer.vocab_size,
+                    x.shape,
+                    device=self.device
+                )
+
+                # Combine the final prediction with noise to form the next intermediate state
+                x = torch.where(mask, noise, pred_x1)
+
+        return x
+
+    def _check_conjugation_validity(self, x):
+        """
+        Check if sequences can be conjugated.
+        Returns a boolean tensor indicating validity for each sample.
+        """
+        if self.conjugation_obj is None:
+            # Fall back to peptide validity check
+            sequences = self.model.tokenizer.batch_decode(x)
+            return torch.tensor([self.peptide_analyzer.is_peptide(seq) for seq in sequences], device=self.device)
+        
+        # Use conjugation objective - score > 0 means valid
+        scores = self.conjugation_obj(x)
+        return scores > 0
+
+    def generate_x0(self, n_steps=16, temperature=1.0, require_valid=False, max_retries=10):
+        """
+        Generate initial x0 sequences.
+        
+        Args:
+            n_steps: Number of flow matching steps
+            temperature: Sampling temperature
+            require_valid: If True, regenerate invalid samples until all are valid conjugatable peptides
+            max_retries: Maximum number of regeneration attempts per sample
+        """
+        print("Starting initial SMILES generation...")
+        
+        x = self._generate_single_x0(n_steps, temperature)
+        
+        if require_valid:
+            print("Checking conjugation validity and regenerating invalid samples...")
+            
+            for retry in range(max_retries):
+                validity = self._check_conjugation_validity(x)
+                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 are valid conjugatable peptides!")
+                    break
+                
+                print(f"Retry {retry + 1}/{max_retries}: {num_valid}/{self.args.num_samples} valid, regenerating {num_invalid} invalid samples...")
+                
+                # Get indices of invalid samples
+                invalid_indices = torch.where(~validity)[0]
+                
+                # Generate new candidates for invalid samples
+                # Temporarily adjust num_samples to generate replacements
+                original_num_samples = self.args.num_samples
+                self.args.num_samples = num_invalid
+                
+                new_x = self._generate_single_x0(n_steps, temperature)
+                
+                self.args.num_samples = original_num_samples
+                
+                # Replace invalid samples with new ones
+                x[invalid_indices] = new_x
+            
+            # Final validity check
+            final_validity = self._check_conjugation_validity(x)
+            num_final_valid = final_validity.sum().item()
+            if num_final_valid < self.args.num_samples:
+                print(f"Warning: After {max_retries} retries, only {num_final_valid}/{self.args.num_samples} samples are valid.")
+            else:
+                print(f"Success: All {self.args.num_samples} samples are valid conjugatable peptides.")
+        
+        # 4. Decode the final token IDs into SMILES strings
+        generated_sequences = self.model.tokenizer.batch_decode(x)
+        
+        # 5. Analyze the validity of the generated sequences
+        validities = []
+        for seq in generated_sequences:
+            validities.append(self.peptide_analyzer.is_peptide(seq))
+            print(seq)
+        print(f"Initial Sequence Validity: {validities}")
+        return x
+
+    def _get_scores(self, x_batch):
+        """Calculates the normalized scores for a batch of sequences."""
+        scores = []
+        for obj_func in self.objectives:
+            scores.append(obj_func(x_batch.to(self.device)))
+        return torch.stack(scores, dim=0).to(self.device)
+
+    def _barker_g(self, u):
+        """Barker balancing function."""
+        return u / (1 + u)
+    
+    def validity(self, x):
+        sampled_sequences = self.model.tokenizer.batch_decode(x)
+        
+        return [1.0 if self.peptide_analyzer.is_peptide(seq) else 0.0 for seq in sampled_sequences]
+
+    def generate(self):
+        """Main generation loop with tracking."""
+        shape = (self.args.num_samples, self.args.gen_len)
+        x = self.generate_x0(
+            require_valid=self.args.require_valid_x0,
+            max_retries=self.args.max_x0_retries
+        )
+        print(x)
+
+        if args.weights is None:
+            # The first weight is for conjugation (peptide + PDC validity)
+            # We need to ensure the SMILES sequences are valid peptides with valid PDCs
+            weights = torch.tensor([1] + [1/(self.num_objectives-1)] * (self.num_objectives-1), device=self.device).view(-1,1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+        print(f"Weights: {weights}")
+
+        with torch.no_grad():
+            for t in tqdm(range(self.args.optimization_steps), desc="MOG Generation"):
+                improved = False
+                improved_indices = []
+                
+                # Anneal guidance strength
+                eta_t = self.args.eta_min + (self.args.eta_max - self.args.eta_min) * (t / (self.args.optimization_steps - 1))
+                # Choose a random position to mutate
+                mut_idx = random.randint(1, self.args.gen_len-2)
+
+                # Determine the generation timestep
+                # We cycle through the timesteps to ensure all are visited
+                generation_step = t % self.args.optimization_steps
+                time_t = torch.full((self.args.num_samples,), (generation_step / self.args.optimization_steps), device=self.device)
+                
+                # Get proposal distribution from ReDi model for the chosen position
+                logits = self.model(x, time_t)
+                probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+                pos_probs[:, x[:, mut_idx]] = 0   # We don't evalute the same token
+                pos_probs[:, self.invalid] = 0
+
+                # Prune candidate vocabulary using top-p sampling
+                sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                remove_mask = cumulative_probs > self.args.top_p
+                remove_mask[..., 1:] = remove_mask[..., :-1].clone()
+                remove_mask[..., 0] = 0
+
+                # Get the set of candidate tokens for each sample in the batch
+                candidate_tokens_list = []
+                for i in range(self.args.num_samples):
+                    sample_mask = remove_mask[i]
+                    candidates = sorted_indices[i, ~sample_mask]
+                    candidate_tokens_list.append(candidates)
+                
+                # Get current scores
+                current_scores = self._get_scores(x)
+                w_current = torch.exp(eta_t * torch.min(weights * current_scores, dim=0).values)
+
+                if t == 0:
+                    print(f"Initial Scores: {current_scores}")
+                    # Log initial step with detailed results
+                    if self.tracker and self.conjugation_obj:
+                        detailed = self.conjugation_obj.get_last_detailed_results()
+                        self.tracker.log_step(t, current_scores[0], detailed)
+                        self.tracker.print_step_summary(t, current_scores[0], detailed)
+
+                # Evaluate all candidate tokens for each sample
+                final_proposal_tokens = []
+                for i in range(self.args.num_samples):
+                    candidates = candidate_tokens_list[i]
+                    candidates = torch.tensor([token for token in candidates if token not in self.invalid], device=candidates.device)
+                    if len(candidates) >= 200:
+                        candidates = candidates[:200]
+                    num_candidates = len(candidates)
+                    
+                    # Create a batch of proposed sequences for the current sample
+                    x_prop_batch = x[i].repeat(num_candidates, 1)
+                    x_prop_batch[:, mut_idx] = candidates
+                    
+                    # Evaluate all proposals
+                    proposal_scores = self._get_scores(x_prop_batch)
+                    proposal_s_omega = torch.min(weights * proposal_scores, dim=0).values
+                    w_proposal = torch.exp(eta_t * proposal_s_omega)
+                    
+                    # Get ReDi probabilities for the candidates
+                    redi_probs = pos_probs[i, candidates]
+
+                    # Calculate unnormalized guided probabilities
+                    tilde_q = redi_probs * self._barker_g(w_proposal / w_current[i])
+                    
+                    # Normalize and sample the final token
+                    final_probs = tilde_q / (torch.sum(tilde_q) + 1e-9)
+
+                    index = torch.multinomial(final_probs, 1).item()
+
+                    if (random.uniform(0,1) > 1 and proposal_scores[:,index][0] == 1) or torch.sum(weights.squeeze(1) * proposal_scores[:, index]) >= torch.sum(weights.squeeze(1) * current_scores[:,i]):
+                        final_token = candidates[index]
+                        print(f"Previous Weighted Sum: {torch.sum(weights.squeeze(1) * current_scores[:,i])}")
+                        print(f"Previous Scores: {current_scores[:,i]}")
+
+                        print(f"New Weighted Sum: {torch.sum(weights.squeeze(1) * proposal_scores[:, index])}")
+                        print(f"New Scores: {proposal_scores[:,index]}")
+                        improved = True
+                        improved_indices.append(i)
+                    
+                    else:
+                        final_token = x[i][mut_idx]
+                    
+                    final_proposal_tokens.append(final_token)
+
+                # Update the sequences with the chosen tokens
+                x[torch.arange(self.args.num_samples), mut_idx] = torch.stack(final_proposal_tokens)
+                if improved:
+                    print(self.model.tokenizer.batch_decode(x))
+                
+                # Log step with tracking (every 4 steps to reduce overhead)
+                if self.tracker and self.conjugation_obj and (t % 4 == 0 or t == self.args.optimization_steps - 1):
+                    # Re-evaluate for detailed logging
+                    step_scores = self._get_scores(x)
+                    detailed = self.conjugation_obj.get_last_detailed_results()
+                    self.tracker.log_step(t, step_scores[0], detailed, improved_indices)
+                    
+                    # Print summary for final step
+                    if t == self.args.optimization_steps - 1:
+                        self.tracker.print_step_summary(t, step_scores[0], detailed)
+
+        return x
+
+
+# --- Main Execution ---
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    target = args.target
+    tokenizer = SMILES_SPE_Tokenizer('./smiles/smiles_tokenizer/new_vocab.txt', './smiles/smiles_tokenizer/new_splits.txt')
+    
+    # Initialize Conjugation model with specified builder type
+    conjugation_model = Conjugation(tokenizer, device=device, builder_type=args.builder_type)
+    # hemolysis_model = Hemolysis(device)
+    # nonfouling_model = Nonfouling(device)
+    # solubility_model = Solubility(device)
+    # permeability_model = Permeability(device)
+    affinity_model = BindingAffinity(target, device)
+
+    # List of all objective functions
+    OBJECTIVE_FUNCTIONS = [conjugation_model, affinity_model]
+    
+    # Initialize tracker for logging
+    tracker = GenerationTracker(output_dir=args.output_dir)
+
+    # --- Set random seeds for reproducibility ---
+    if args.seed is not None:
+        import random
+        import numpy as np
+        random.seed(args.seed)
+        np.random.seed(args.seed)
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(args.seed)
+            torch.cuda.manual_seed_all(args.seed)
+        print(f"Random seed set to: {args.seed}")
+
+    # --- Load Model ---
+    print(f"Loading model from checkpoint: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+    model_hparams = checkpoint["hyper_parameters"]["args"]
+    
+    model = MDLMLightningModule.load_from_checkpoint(
+        args.checkpoint,
+        args=model_hparams,
+        tokenizer=tokenizer,
+        map_location=device,
+        strict=False
+    )
+    model.to(device)
+    print("Model loaded successfully.")
+
+    mog_generator = MOGGenerator(model, device, OBJECTIVE_FUNCTIONS, args, tracker=tracker)
+    
+    all_sequences = []
+    all_final_results = []
+    
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'='*60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'='*60}")
+        
+        generated_tokens = mog_generator.generate()
+        final_scores = mog_generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+
+        print("\n--- Batch Results ---")
+        print(f"Sequences: {sequence_str}")
+        print(f"Scores: {final_scores}")
+        
+        all_sequences.extend(sequence_str)
+        
+        # Get detailed results with properties for final sequences
+        detailed_with_props = conjugation_model.get_detailed_scores_with_properties(generated_tokens)
+        tracker.log_final(sequence_str, detailed_with_props)
+        all_final_results.extend(detailed_with_props)
+    
+    # Print and save final summary
+    tracker.print_final_summary()
+    tracker.save(filename_prefix='pdc_generation')
+    
+    # Also save sequences to simple text file
+    with open(args.output_file, 'w') as f:
+        for seq in all_sequences:
+            f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+
+    print("\nGeneration complete.")
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Multi-Objective Generation with LBP-MOG-ReDi (Single Mutation).")
+
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to the trained ReDi model checkpoint.")
+    parser.add_argument("--num_samples", type=int, default=10, help="Number of samples to generate.")
+    parser.add_argument("--num_batches", type=int, default=10, help="Number of batches to generate.")
+    parser.add_argument("--output_file", type=str, default="./smiles.txt", help="File to save the generated sequences.")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs", help="Directory to save detailed generation logs.")
+    parser.add_argument("--gen_len", type=int, default=50, help="Length of the sequences to generate.")
+    parser.add_argument("--optimization_steps", type=int, default=16, help="Number of optimization steps per generation.")
+    parser.add_argument("--weights", type=float, nargs='+', required=False, help="Weights for the objectives (e.g., 0.5 0.5).")
+    parser.add_argument("--eta_min", type=float, default=1.0, help="Minimum guidance strength for annealing.")
+    parser.add_argument("--eta_max", type=float, default=20.0, help="Maximum guidance strength for annealing.")
+    parser.add_argument("--top_p", type=float, default=0.9, help="Top-p for pruning candidate tokens.")
+    
+    # PDC builder
+    parser.add_argument("--builder_type", type=str, default="AoaDau",
+                        choices=["AoaDau", "ValAlaDau", "ValCitDau", "GlutarylDau", "Dau",
+                                 "SSCpt", "VcMMAE", "DSPACpt", "MalHexCpt", "DSECCpt", "Cpt",
+                                 "DSPASN38", "DSEBSN38", "DSEPSN38", "SN38"],
+                        help="PDC builder type for conjugation. "
+                             "'Dau' samples linker by GNRHR+Dau frequency. "
+                             "'Cpt' samples linker by ERBB2+Cpt frequency. "
+                             "'SN38' samples linker by NRP1+SN38 frequency.")
+    
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.")
+    parser.add_argument("--require_valid_x0", action="store_true",
+                        help="If set, regenerate x0 samples until they are valid conjugatable peptides.")
+    parser.add_argument("--max_x0_retries", type=int, default=10,
+                        help="Maximum retries when require_valid_x0 is enabled.")
+
+    parser.add_argument("--target", type=str, required=True, help="Target protein sequence for binding affinity.")
+    args = parser.parse_args()
+    main(args)

smiles/generation/guided_aa.py

@@ -0,0 +1,872 @@
+"""Natural-AA variant of guided_uaa.py."""
+import argparse
+import re
+import random
+import sys
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+import csv
+import json
+import os
+from datetime import datetime
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+from smiles.models.train import MDLMLightningModule
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+import numpy as np
+from easydict import EasyDict as edict
+
+from smiles.evaluation.peptide_analyzer import PeptideAnalyzer
+from smiles.evaluation.pdc_analyzer import PDCAnalyzer
+from smiles.scoring.classifiers import Conjugation, BindingAffinity
+
+from smiles.models.aa_helpers import (
+    ESM2AminoTokenizer as _AAEsmTokenizer,
+    MDLM_PepReDi as _AAMDLM,
+    make_affinity,
+    DEFAULT_AA_AFFINITY_CKPT,
+    INVALID_AA_TOKEN_IDS as _AA_INVALID_IDS,
+)
+
+from third_party.SMILESToPeptide.peptide_to_smiles import aa_to_smiles
+from third_party.SMILESToPeptide.aminoacid_selective import one_letter_to_three_letter
+
+
+MDLM = _AAMDLM
+AminoTokenizer = _AAEsmTokenizer
+
+
+class GenerationTracker:
+    """
+    Tracks detailed validity information during generation
+    """
+    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, scores, detailed_results=None, improved_indices=None):
+        """Log information for a generation step"""
+        step_info = {
+            'step': step,
+            'scores': scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else scores,
+            'num_valid': int(sum(1 for s in scores if s > 0)) if hasattr(scores, '__iter__') else 0,
+            'improved_indices': improved_indices or []
+        }
+        
+        if detailed_results:
+            step_info['detailed'] = []
+            for i, result in enumerate(detailed_results):
+                step_info['detailed'].append({
+                    'sample_idx': i,
+                    'is_valid_peptide': result.get('is_valid_peptide', False),
+                    'has_conjugation_site': result.get('has_conjugation_site', False),
+                    'conjugation_success': result.get('conjugation_success', False),
+                    'pdc_valid': result.get('pdc_valid', False),
+                    'num_conjugates': len(result.get('conjugates', [])),
+                    'error': result.get('error')
+                })
+        
+        self.step_logs.append(step_info)
+        
+    def log_final(self, sequences, detailed_results_with_props):
+        """Log final results with properties"""
+        for i, (seq, result) in enumerate(zip(sequences, detailed_results_with_props)):
+            final_info = {
+                'sample_idx': i,
+                'peptide_smiles': seq,
+                'is_valid_peptide': result.get('is_valid_peptide', False),
+                'has_conjugation_site': result.get('has_conjugation_site', False),
+                'num_conjugates': len(result.get('conjugates', [])),
+                'num_valid_pdcs': len(result.get('valid_pdcs', [])),
+                'best_pdc': None,
+                'properties': result.get('properties', {})
+            }
+            
+            if result.get('best_pdc'):
+                final_info['best_pdc'] = {
+                    'smiles': result['best_pdc'].get('smiles', ''),
+                    'checks': result['best_pdc'].get('checks', {})
+                }
+            
+            self.final_results.append(final_info)
+    
+    def save(self, filename_prefix='generation'):
+        """Save all logs to files"""
+        # Save step logs
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_path, 'w') as f:
+            json.dump(self.step_logs, f, indent=2, default=str)
+        
+        # Save final results
+        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)
+        
+        # Save summary CSV
+        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', 'peptide_smiles', 'is_valid_peptide', 'has_conjugation_site',
+                'num_conjugates', 'num_valid_pdcs', 'best_pdc_smiles',
+                'mf', 'polararea', 'xlogp', 'heavycnt', 'rotbonds'
+            ])
+            for result in self.final_results:
+                props = result.get('properties', {})
+                writer.writerow([
+                    result['sample_idx'],
+                    result['peptide_smiles'],
+                    result['is_valid_peptide'],
+                    result['has_conjugation_site'],
+                    result['num_conjugates'],
+                    result['num_valid_pdcs'],
+                    result['best_pdc']['smiles'] if result.get('best_pdc') else '',
+                    props.get('mf', ''),
+                    props.get('polararea', ''),
+                    props.get('xlogp', ''),
+                    props.get('heavycnt', ''),
+                    props.get('rotbonds', '')
+                ])
+        
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        
+        return step_log_path, final_path, csv_path
+    
+    def print_step_summary(self, step, scores, detailed_results=None):
+        """Print a summary for the current step"""
+        if isinstance(scores, torch.Tensor):
+            scores_list = scores.cpu().numpy().tolist()
+        else:
+            scores_list = list(scores)
+        
+        num_valid = sum(1 for s in scores_list if s > 0)
+        
+        print(f"\n--- Step {step} Summary ---")
+        print(f"  Valid: {num_valid}/{len(scores_list)}")
+        
+        if detailed_results:
+            valid_peptide = sum(1 for r in detailed_results if r.get('is_valid_peptide'))
+            has_site = sum(1 for r in detailed_results if r.get('has_conjugation_site'))
+            conj_success = sum(1 for r in detailed_results if r.get('conjugation_success'))
+            pdc_valid = sum(1 for r in detailed_results if r.get('pdc_valid'))
+            
+            print(f"  Breakdown:")
+            print(f"    Valid peptides: {valid_peptide}")
+            print(f"    Has conjugation site: {has_site}")
+            print(f"    Conjugation success: {conj_success}")
+            print(f"    PDC valid: {pdc_valid}")
+    
+    def print_final_summary(self):
+        """Print final summary"""
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        
+        total = len(self.final_results)
+        valid_pdcs = sum(1 for r in self.final_results if r.get('num_valid_pdcs', 0) > 0)
+        
+        print(f"\nTotal sequences: {total}")
+        print(f"With valid PDCs: {valid_pdcs}")
+        
+        print("\nDetailed results:")
+        for result in self.final_results:
+            status = "" if result.get('num_valid_pdcs', 0) > 0 else ""
+            print(f"\n  [{status}] Sample {result['sample_idx']}:")
+            print(f"      Peptide: {result['peptide_smiles'][:60]}...")
+            print(f"      Valid peptide: {result['is_valid_peptide']}")
+            print(f"      Has conjugation site: {result['has_conjugation_site']}")            
+            print(f"      Conjugates: {result['num_conjugates']}")
+            print(f"      Valid PDCs: {result['num_valid_pdcs']}")
+            
+            if result.get('properties'):
+                props = result['properties']
+                print(f"      Properties:")
+                for k, v in props.items():
+                    print(f"        {k}: {v}")
+
+
+def peptide_bond_mask(smiles_list):
+    # Initialize the batch mask
+    batch_size = len(smiles_list)
+    max_seq_length = 1035 #max(len(smiles) for smiles in smiles_list)  # Find the longest SMILES
+    mask = torch.zeros((batch_size, max_seq_length), dtype=torch.int)  # Mask filled with zeros
+
+    bond_patterns = [
+        (r'OC\(=O\)', 'ester'),
+        (r'N\(C\)C\(=O\)', 'n_methyl'),
+        (r'N[12]C\(=O\)', 'peptide'),  # Pro peptide bonds
+        (r'NC\(=O\)', 'peptide'),  # Regular peptide bonds
+        (r'C\(=O\)N\(C\)', 'n_methyl'),
+        (r'C\(=O\)N[12]?', 'peptide')
+    ]
+
+    for batch_idx, smiles in enumerate(smiles_list):
+        positions = []
+        used = set()
+
+        # Identify bonds
+        for pattern, bond_type in bond_patterns:
+            for match in re.finditer(pattern, smiles):
+                if not any(p in range(match.start(), match.end()) for p in used):
+                    positions.append({
+                        'start': match.start(),
+                        'end': match.end(),
+                        'type': bond_type,
+                        'pattern': match.group()
+                    })
+                    used.update(range(match.start(), match.end()))
+
+        # Update the mask for the current SMILES
+        for pos in positions:
+            mask[batch_idx, pos['start']:pos['end']] = 1
+
+    return mask
+
+def peptide_token_mask(smiles_list, token_lists):
+    # Initialize the batch mask
+    batch_size = len(smiles_list)
+    token_seq_length = max(len(tokens) for tokens in token_lists)  # Find the longest tokenized sequence
+    tokenized_masks = torch.zeros((batch_size, token_seq_length), dtype=torch.int)  # Mask filled with zeros
+    atomwise_masks = peptide_bond_mask(smiles_list)
+
+
+    for batch_idx, atomwise_mask in enumerate(atomwise_masks):
+        token_seq = token_lists[batch_idx]
+        atom_idx = 0
+        
+        for token_idx, token in enumerate(token_seq):
+            if token_idx != 0 and token_idx != len(token_seq) - 1:
+                if torch.sum(atomwise_mask[atom_idx:atom_idx+len(token)]) >= 1:
+                    tokenized_masks[batch_idx][token_idx] = 1
+                atom_idx += len(token)
+    
+    return tokenized_masks
+
+class MOGGenerator:
+    def __init__(self, model, device, objectives, args, tracker=None, tokenizer=None, tokenizer_type='smiles'):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.pdc_analyzer = PDCAnalyzer(device)
+        self.tracker = tracker
+        self.tokenizer = tokenizer
+        self.tokenizer_type = tokenizer_type
+        
+        # Set invalid tokens based on tokenizer type
+        if tokenizer_type == 'amino':
+            # PepReDi (natural-AA) vocab is 24 tokens: 0-3 are specials
+            # (cls/pad/eos/unk) and 4-23 are the 20 canonical AAs. The UAA
+            # script also blocked id 38 ('g' D-AA), but that's out of range
+            # for PepReDi and would crash the `pos_probs[:, self.invalid] = 0`
+            # masking. Block special tokens only.
+            self.invalid = [0, 1, 2, 3]
+            # Initialize SMILES tokenizer for AA-to-SMILES conversion
+            self.smiles_tokenizer = SMILES_SPE_Tokenizer(
+                args.smiles_vocab_path,
+                args.smiles_splits_path
+            )
+        else:
+            self.invalid = [0, 1, 2, 3, 4, 585, 586]  # SMILES special tokens
+            self.smiles_tokenizer = None
+        
+        # Find the Conjugation objective for detailed tracking
+        self.conjugation_obj = None
+        for obj in objectives:
+            if isinstance(obj, Conjugation):
+                self.conjugation_obj = obj
+                break
+
+    def _generate_single_x0(self, n_steps=16, temperature=1.0):
+        """Generate a single batch of x0 sequences."""
+        self.model.eval()
+
+        # Get vocab size based on tokenizer type
+        if self.tokenizer_type == 'amino':
+            vocab_size = self.model.vocab_size
+        else:
+            vocab_size = self.model.tokenizer.vocab_size
+
+        # 1. Start with a tensor of random tokens (pure noise at t=0)
+        # For amino tokenizer, exclude invalid tokens from initialization
+        if self.tokenizer_type == 'amino':
+            valid_tokens = [i for i in range(vocab_size) if i not in self.invalid]
+            valid_tokens = torch.tensor(valid_tokens, device=self.device)
+            rand_indices = torch.randint(0, len(valid_tokens), (self.args.num_samples, self.args.gen_len), device=self.device)
+            body = valid_tokens[rand_indices]
+            # Add CLS and EOS tokens
+            cls_col = torch.full((self.args.num_samples, 1), self.tokenizer.cls_id, device=self.device)
+            eos_col = torch.full((self.args.num_samples, 1), self.tokenizer.eos_id, device=self.device)
+            x = torch.cat([cls_col, body, eos_col], dim=1)
+        else:
+            x = torch.randint(
+                0,
+                vocab_size,
+                (self.args.num_samples, self.args.gen_len),
+                device=self.device
+            )
+
+        # 2. Define the time schedule for the forward process (0.0 to 1.0)
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+
+        # 3. Iteratively follow the flow from noise to data
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Flow Matching Steps", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i+1]
+
+                # Prepare the current timestep tensor for the model
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+
+                # Get the model's prediction for the final clean sequence (at t=1)
+                logits = self.model(x, t_tensor) 
+                logits = logits / temperature
+                
+                # Mask out invalid tokens
+                for inv_tok in self.invalid:
+                    if inv_tok < logits.shape[-1]:
+                        logits[:, :, inv_tok] = -1000
+
+                pred_x1 = torch.argmax(logits, dim=-1)
+
+                # On the last step, the result is the final prediction
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+
+                # --- Construct the next state x_{t_next} ---
+                # The probability of a token being noise at time t_next is (1 - t_next).
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                
+                # Generate new random tokens for the noise positions
+                noise = torch.randint(
+                    0,
+                    vocab_size,
+                    x.shape,
+                    device=self.device
+                )
+
+                # Combine the final prediction with noise to form the next intermediate state
+                x = torch.where(mask, noise, pred_x1)
+
+        return x
+
+    def _check_conjugation_validity(self, x):
+        """
+        Check if sequences can be conjugated.
+        Returns a boolean tensor indicating validity for each sample.
+        """
+        if self.tokenizer_type == 'amino':
+            # For amino tokenizer, convert to SMILES first
+            seqs = self.tokenizer.batch_decode(x)
+            three_letter_seqs = [one_letter_to_three_letter(seq) for seq in seqs]
+            smiles_seqs = [aa_to_smiles(three_letter_seq, cyclic=self.args.cyclic) 
+                          for three_letter_seq in three_letter_seqs]
+            
+            # Check if SMILES conversion succeeded and result is valid peptide
+            validities = []
+            for smiles in smiles_seqs:
+                if isinstance(smiles, str):
+                    validities.append(self.peptide_analyzer.is_peptide(smiles))
+                else:
+                    validities.append(False)
+            return torch.tensor(validities, device=self.device)
+        
+        if self.conjugation_obj is None:
+            # Fall back to peptide validity check
+            sequences = self.tokenizer.batch_decode(x)
+            return torch.tensor([self.peptide_analyzer.is_peptide(seq) for seq in sequences], device=self.device)
+        
+        # Use conjugation objective - score > 0 means valid
+        scores = self.conjugation_obj(x)
+        return scores > 0
+
+    def generate_x0(self, n_steps=16, temperature=1.0, require_valid=False, max_retries=10):
+        """
+        Generate initial x0 sequences.
+        
+        Args:
+            n_steps: Number of flow matching steps
+            temperature: Sampling temperature
+            require_valid: If True, regenerate invalid samples until all are valid conjugatable peptides
+            max_retries: Maximum number of regeneration attempts per sample
+        """
+        print("Starting initial SMILES generation...")
+        
+        x = self._generate_single_x0(n_steps, temperature)
+        
+        if require_valid:
+            print("Checking conjugation validity and regenerating invalid samples...")
+            
+            for retry in range(max_retries):
+                validity = self._check_conjugation_validity(x)
+                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 are valid conjugatable peptides!")
+                    break
+                
+                print(f"Retry {retry + 1}/{max_retries}: {num_valid}/{self.args.num_samples} valid, regenerating {num_invalid} invalid samples...")
+                
+                # Get indices of invalid samples
+                invalid_indices = torch.where(~validity)[0]
+                
+                # Generate new candidates for invalid samples
+                # Temporarily adjust num_samples to generate replacements
+                original_num_samples = self.args.num_samples
+                self.args.num_samples = num_invalid
+                
+                new_x = self._generate_single_x0(n_steps, temperature)
+                
+                self.args.num_samples = original_num_samples
+                
+                # Replace invalid samples with new ones
+                x[invalid_indices] = new_x
+            
+            # Final validity check
+            final_validity = self._check_conjugation_validity(x)
+            num_final_valid = final_validity.sum().item()
+            if num_final_valid < self.args.num_samples:
+                print(f"Warning: After {max_retries} retries, only {num_final_valid}/{self.args.num_samples} samples are valid.")
+            else:
+                print(f"Success: All {self.args.num_samples} samples are valid conjugatable peptides.")
+        
+        # 4. Decode the final token IDs into strings
+        generated_sequences = self.tokenizer.batch_decode(x)
+        
+        # 5. Analyze the validity of the generated sequences
+        validities = []
+        if self.tokenizer_type == 'amino':
+            # For amino mode, report conjugation-relevant validity after AA->SMILES conversion.
+            validity_tensor = self._check_conjugation_validity(x)
+            validities = [bool(v.item()) for v in validity_tensor]
+            for seq in generated_sequences:
+                print(seq)
+        else:
+            for seq in generated_sequences:
+                validities.append(self.peptide_analyzer.is_peptide(seq))
+                print(seq)
+        if self.tokenizer_type == 'amino':
+            print(f"Initial Peptide Validity (AA->SMILES): {validities}")
+        else:
+            print(f"Initial Sequence Validity: {validities}")
+        return x
+
+    def _get_scores(self, x_batch):
+        """Calculates the normalized scores for a batch of sequences."""
+        scores = []
+        
+        if self.tokenizer_type == 'amino':
+            # Convert AA tokens to SMILES for scoring
+            seqs = self.tokenizer.batch_decode(x_batch)
+            
+            # Convert one-letter AA codes to three-letter, then to SMILES
+            three_letter_seqs = [one_letter_to_three_letter(seq) for seq in seqs]
+            with ThreadPoolExecutor() as pool:
+                smiles_seqs = list(pool.map(
+                    lambda s: aa_to_smiles(s, cyclic=self.args.cyclic),
+                    three_letter_seqs
+                ))
+
+            # Track valid conversions
+            valid_mask = torch.tensor(
+                [1 if isinstance(seq, str) else 0 for seq in smiles_seqs], 
+                device=self.device
+            )
+            
+            # Replace invalid SMILES with dummy 'C' for tokenization
+            for i in torch.where((1 - valid_mask) == 1)[0].tolist():
+                smiles_seqs[i] = 'C'
+            
+            # Tokenize SMILES for objective functions
+            smiles_batch = self.smiles_tokenizer(
+                smiles_seqs, return_tensors='pt', padding=True
+            ).to(self.device)
+            attention_mask = smiles_batch.get("attention_mask", None)
+            
+            # Evaluate each objective. Objectives may declare `.input_type`:
+            #   'aa'     -> feed raw AA tokens (no SMILES conversion)
+            #   'smiles' -> feed SMILES tokens (default for the existing
+            #               Conjugation / SMILES BindingAffinity)
+            for obj_func in self.objectives:
+                if getattr(obj_func, "input_type", None) == "aa":
+                    # AA-input objective (e.g. UnpooledBindingPredictor):
+                    # score directly on the original AA tokens; do not multiply
+                    # by valid_mask (the AA->SMILES conversion outcome doesn't
+                    # affect this scorer).
+                    score = obj_func(x_batch.to(self.device)).to(self.device)
+                    scores.append(score)
+                else:
+                    try:
+                        score = obj_func(
+                            smiles_batch["input_ids"],
+                            attention_mask=attention_mask
+                        ).to(self.device)
+                    except TypeError:
+                        score = obj_func(smiles_batch["input_ids"]).to(self.device)
+                    scores.append(score * valid_mask)  # Zero out invalid conversions
+        else:
+            # Original SMILES-based scoring
+            for obj_func in self.objectives:
+                scores.append(obj_func(x_batch.to(self.device)))
+        
+        return torch.stack(scores, dim=0).to(self.device)
+
+    def _barker_g(self, u):
+        """Barker balancing function."""
+        return u / (1 + u)
+    
+    def validity(self, x):
+        sampled_sequences = self.tokenizer.batch_decode(x)
+        
+        if self.tokenizer_type == 'amino':
+            # For amino tokenizer, convert to SMILES first and check validity
+            three_letter_seqs = [one_letter_to_three_letter(seq) for seq in sampled_sequences]
+            smiles_seqs = [aa_to_smiles(three_letter_seq, cyclic=self.args.cyclic) 
+                          for three_letter_seq in three_letter_seqs]
+            return [1.0 if isinstance(seq, str) and self.peptide_analyzer.is_peptide(seq) else 0.0 
+                    for seq in smiles_seqs]
+        else:
+            return [1.0 if self.peptide_analyzer.is_peptide(seq) else 0.0 for seq in sampled_sequences]
+
+    def generate(self):
+        """Main generation loop with tracking."""
+        shape = (self.args.num_samples, self.args.gen_len)
+        x = self.generate_x0(
+            require_valid=self.args.require_valid_x0,
+            max_retries=self.args.max_x0_retries
+        )
+        print(x)
+
+        if self.args.weights is None:
+            # The first weight is for conjugation (peptide + PDC validity)
+            # We need to ensure the SMILES sequences are valid peptides with valid PDCs
+            weights = torch.tensor([1] + [1/(self.num_objectives-1)] * (self.num_objectives-1), device=self.device).view(-1,1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+        print(f"Weights: {weights}")
+
+        with torch.no_grad():
+            current_scores = self._get_scores(x)
+            for t in tqdm(range(self.args.optimization_steps), desc="MOG Generation"):
+                improved = False
+                improved_indices = []
+
+                # Anneal guidance strength
+                eta_t = self.args.eta_min + (self.args.eta_max - self.args.eta_min) * (t / (self.args.optimization_steps - 1))
+                # Choose a random position to mutate
+                mut_idx = random.randint(1, self.args.gen_len-2)
+
+                # Determine the generation timestep
+                # We cycle through the timesteps to ensure all are visited
+                generation_step = t % self.args.optimization_steps
+                time_t = torch.full((self.args.num_samples,), (generation_step / self.args.optimization_steps), device=self.device)
+                
+                # Get proposal distribution from ReDi model for the chosen position
+                logits = self.model(x, time_t)
+                probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+                pos_probs[:, x[:, mut_idx]] = 0   # We don't evalute the same token
+                pos_probs[:, self.invalid] = 0
+
+                # Prune candidate vocabulary using top-p sampling
+                sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                remove_mask = cumulative_probs > self.args.top_p
+                remove_mask[..., 1:] = remove_mask[..., :-1].clone()
+                remove_mask[..., 0] = 0
+
+                # Get the set of candidate tokens for each sample in the batch
+                candidate_tokens_list = []
+                for i in range(self.args.num_samples):
+                    sample_mask = remove_mask[i]
+                    candidates = sorted_indices[i, ~sample_mask]
+                    candidate_tokens_list.append(candidates)
+                
+                w_current = torch.exp(eta_t * torch.min(weights * current_scores, dim=0).values)
+
+                if t == 0:
+                    print(f"Initial Scores: {current_scores}")
+                    # Log initial step with detailed results
+                    if self.tracker and self.conjugation_obj:
+                        detailed = self.conjugation_obj.get_last_detailed_results()
+                        self.tracker.log_step(t, current_scores[0], detailed)
+                        self.tracker.print_step_summary(t, current_scores[0], detailed)
+
+                # Evaluate all candidate tokens for each sample
+                final_proposal_tokens = []
+                for i in range(self.args.num_samples):
+                    candidates = candidate_tokens_list[i]
+                    candidates = torch.tensor([token for token in candidates if token not in self.invalid], device=candidates.device)
+                    if len(candidates) >= 200:
+                        candidates = candidates[:200]
+                    num_candidates = len(candidates)
+                    
+                    # Create a batch of proposed sequences for the current sample
+                    x_prop_batch = x[i].repeat(num_candidates, 1)
+                    x_prop_batch[:, mut_idx] = candidates
+                    
+                    # Evaluate all proposals
+                    proposal_scores = self._get_scores(x_prop_batch)
+                    proposal_s_omega = torch.min(weights * proposal_scores, dim=0).values
+                    w_proposal = torch.exp(eta_t * proposal_s_omega)
+                    
+                    # Get ReDi probabilities for the candidates
+                    redi_probs = pos_probs[i, candidates]
+
+                    # Calculate unnormalized guided probabilities
+                    tilde_q = redi_probs * self._barker_g(w_proposal / w_current[i])
+                    
+                    # Normalize and sample the final token
+                    final_probs = tilde_q / (torch.sum(tilde_q) + 1e-9)
+
+                    index = torch.multinomial(final_probs, 1).item()
+
+                    if (random.uniform(0,1) > 1 and proposal_scores[:,index][0] == 1) or torch.sum(weights.squeeze(1) * proposal_scores[:, index]) >= torch.sum(weights.squeeze(1) * current_scores[:,i]):
+                        final_token = candidates[index]
+                        print(f"Previous Weighted Sum: {torch.sum(weights.squeeze(1) * current_scores[:,i])}")
+                        print(f"Previous Scores: {current_scores[:,i]}")
+
+                        print(f"New Weighted Sum: {torch.sum(weights.squeeze(1) * proposal_scores[:, index])}")
+                        print(f"New Scores: {proposal_scores[:,index]}")
+                        improved = True
+                        improved_indices.append(i)
+                    
+                    else:
+                        final_token = x[i][mut_idx]
+                    
+                    final_proposal_tokens.append(final_token)
+
+                # Update the sequences with the chosen tokens
+                x[torch.arange(self.args.num_samples), mut_idx] = torch.stack(final_proposal_tokens)
+                if improved:
+                    current_scores = self._get_scores(x)
+                    print(self.tokenizer.batch_decode(x))
+
+                # Log step with tracking (every 4 steps to reduce overhead)
+                if self.tracker and self.conjugation_obj and (t % 4 == 0 or t == self.args.optimization_steps - 1):
+                    detailed = self.conjugation_obj.get_last_detailed_results()
+                    self.tracker.log_step(t, current_scores[0], detailed, improved_indices)
+
+                    # Print summary for final step
+                    if t == self.args.optimization_steps - 1:
+                        self.tracker.print_step_summary(t, current_scores[0], detailed)
+
+        return x
+
+
+# --- Main Execution ---
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    target = args.target
+    
+    # --- Initialize tokenizer based on type ---
+    if args.tokenizer_type == 'amino':
+        print("Using Amino Acid tokenizer")
+        tokenizer = AminoTokenizer()
+        smiles_tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+    else:
+        print("Using SMILES tokenizer")
+        tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+        smiles_tokenizer = tokenizer
+    
+    # --- Initialize objective functions ---
+    # Conjugation always uses SMILES tokens (AA->SMILES->builder).
+    conjugation_model = Conjugation(smiles_tokenizer, device=device, builder_type=args.builder_type)
+    # Binding affinity: AA-input or SMILES-input, depending on --affinity_type.
+    # The returned object carries `.input_type` ('aa' or 'smiles') so the
+    # generator routes inputs correctly in _get_scores.
+    affinity_model = make_affinity(
+        args.affinity_type,
+        target_protein_seq=target,
+        device=device,
+        aa_ckpt_path=args.aa_affinity_ckpt,
+    )
+    print(f"Affinity scorer: {type(affinity_model).__name__} (input_type={getattr(affinity_model, 'input_type', '?')})")
+
+    # List of all objective functions
+    OBJECTIVE_FUNCTIONS = [conjugation_model, affinity_model]
+    
+    # Initialize tracker for logging
+    tracker = GenerationTracker(output_dir=args.output_dir)
+
+    # --- Set random seeds for reproducibility ---
+    if args.seed is not None:
+        import random
+        import numpy as np
+        random.seed(args.seed)
+        np.random.seed(args.seed)
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(args.seed)
+            torch.cuda.manual_seed_all(args.seed)
+        print(f"Random seed set to: {args.seed}")
+
+    # --- Load Model ---
+    print(f"Loading model from checkpoint: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+    
+    if args.tokenizer_type == 'amino':
+        # Vanilla PyTorch checkpoint for amino model
+        raw_model_args = checkpoint['args']
+        if isinstance(raw_model_args, argparse.Namespace):
+            raw_model_args = vars(raw_model_args)
+        model_args = edict(raw_model_args)
+        model = MDLM(
+            vocab_size=model_args.vocab_size,
+            seq_len=model_args.seq_len,
+            model_dim=model_args.model_dim,
+            n_heads=model_args.n_heads,
+            n_layers=model_args.n_layers
+        ).to(device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print(f"Amino model loaded. Vocab size: {model_args.vocab_size}")
+    else:
+        # PyTorch Lightning checkpoint for SMILES model
+        model_hparams = checkpoint["hyper_parameters"]["args"]
+        model = MDLMLightningModule.load_from_checkpoint(
+            args.checkpoint,
+            args=model_hparams,
+            tokenizer=tokenizer,
+            map_location=device,
+            strict=False
+        )
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully.")
+
+    mog_generator = MOGGenerator(
+        model, device, OBJECTIVE_FUNCTIONS, args, tracker=tracker,
+        tokenizer=tokenizer, tokenizer_type=args.tokenizer_type
+    )
+    
+    all_sequences = []
+    all_final_results = []
+    
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'='*60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'='*60}")
+        
+        generated_tokens = mog_generator.generate()
+        final_scores = mog_generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+
+        print("\n--- Batch Results ---")
+        print(f"Sequences: {sequence_str}")
+        print(f"Scores: {final_scores}")
+        
+        all_sequences.extend(sequence_str)
+        
+        # Get detailed results with properties for final sequences
+        # For amino tokenizer, we need to convert to SMILES first for detailed scoring
+        if args.tokenizer_type == 'amino':
+            # Convert AA to SMILES for detailed conjugation analysis
+            smiles_list = []
+            for seq in sequence_str:
+                three_letter = one_letter_to_three_letter(seq)
+                smiles = aa_to_smiles(three_letter, cyclic=args.cyclic)
+                smiles_list.append(smiles if isinstance(smiles, str) else 'C')
+            # Tokenize SMILES for conjugation model
+            smiles_tokens = smiles_tokenizer(smiles_list, return_tensors='pt', padding=True).to(device)
+            detailed_with_props = conjugation_model.get_detailed_scores_with_properties(smiles_tokens["input_ids"])
+        else:
+            detailed_with_props = conjugation_model.get_detailed_scores_with_properties(generated_tokens)
+        
+        tracker.log_final(sequence_str, detailed_with_props)
+        all_final_results.extend(detailed_with_props)
+    
+    # Print and save final summary
+    tracker.print_final_summary()
+    tracker.save(filename_prefix='pdc_generation')
+    
+    # Save sequences to text file
+    with open(args.output_file, 'w') as f:
+        for i, seq in enumerate(all_sequences):
+            if args.tokenizer_type == 'amino':
+                # Also save converted SMILES for amino sequences
+                three_letter = one_letter_to_three_letter(seq)
+                smiles = aa_to_smiles(three_letter, cyclic=args.cyclic)
+                if isinstance(smiles, str):
+                    f.write(f"{seq}\t{smiles}\n")
+                else:
+                    f.write(f"{seq}\tINVALID_SMILES\n")
+            else:
+                f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+
+    print("\nGeneration complete.")
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Multi-Objective Generation with LBP-MOG-ReDi (Single Mutation).")
+
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to the trained ReDi model checkpoint.")
+    parser.add_argument("--num_samples", type=int, default=10, help="Number of samples to generate.")
+    parser.add_argument("--num_batches", type=int, default=10, help="Number of batches to generate.")
+    parser.add_argument("--output_file", type=str, default="./smiles.txt", help="File to save the generated sequences.")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs", help="Directory to save detailed generation logs.")
+    parser.add_argument("--gen_len", type=int, default=50, help="Length of the sequences to generate.")
+    parser.add_argument("--optimization_steps", type=int, default=16, help="Number of optimization steps per generation.")
+    parser.add_argument("--weights", type=float, nargs='+', required=False, help="Weights for the objectives.")
+    parser.add_argument("--eta_min", type=float, default=1.0, help="Minimum guidance strength.")
+    parser.add_argument("--eta_max", type=float, default=20.0, help="Maximum guidance strength.")
+    parser.add_argument("--top_p", type=float, default=0.9, help="Top-p for pruning candidate tokens.")
+    parser.add_argument("--builder_type", type=str, default="AoaDau",
+                        choices=["AoaDau", "SSCpt", "VcMMAE", "DSPACpt", "MalHexCpt", 
+                                 "DSECCpt", "Cpt", "DSPASN38", "DSEBSN38", "DSEPSN38", "SN38"],
+                        help="PDC builder type for conjugation.")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.")
+    parser.add_argument("--require_valid_x0", action="store_true",
+                        help="If set, regenerate x0 samples until they are valid conjugatable peptides.")
+    parser.add_argument("--max_x0_retries", type=int, default=10,
+                        help="Maximum retries when require_valid_x0 is enabled.")
+
+    parser.add_argument("--target", type=str, required=True, help="Target protein sequence for binding affinity.")
+    
+    # Tokenizer and model type arguments
+    parser.add_argument("--tokenizer_type", type=str, default="smiles", choices=["smiles", "amino"],
+                        help="Type of tokenizer: 'smiles' for SMILES SPE tokenizer, 'amino' for amino acid tokenizer.")
+    parser.add_argument("--cyclic", action="store_true", 
+                        help="Generate cyclic peptides (only for amino tokenizer).")
+    parser.add_argument("--smiles_vocab_path", type=str, 
+                        default="./smiles/smiles_tokenizer/new_vocab.txt",
+                        help="Path to SMILES tokenizer vocabulary file.")
+    parser.add_argument("--smiles_splits_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_splits.txt",
+                        help="Path to SMILES tokenizer splits file.")
+
+    # AA-mode binding-affinity selection (this is the WT-AA / PepReDi script).
+    parser.add_argument("--affinity_type", type=str, default="smiles",
+                        choices=["aa", "smiles"],
+                        help="Which binding-affinity oracle to use during "
+                             "guidance: 'aa' = AA-input UnpooledBindingPredictor "
+                             "(matches moPPIt's setup), 'smiles' = SMILES-input "
+                             "BindingAffinity (matches the SMILES CupOFMoCA "
+                             "tables). Defaults to 'aa' since this script is "
+                             "the WT-AA pipeline.")
+    parser.add_argument("--aa_affinity_ckpt", type=str,
+                        default=DEFAULT_AA_AFFINITY_CKPT,
+                        help="Path to the AA-input affinity predictor "
+                             "checkpoint (binding_affinity_unpooled.pt).")
+
+    args = parser.parse_args()
+    main(args)

smiles/generation/guided_uaa.py

@@ -0,0 +1,981 @@
+import argparse
+import re
+import random
+import sys
+from collections import Counter
+from concurrent.futures import ThreadPoolExecutor
+import csv
+import json
+import os
+from datetime import datetime
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+from smiles.models.train import MDLMLightningModule
+from smiles.smiles_tokenizer.my_tokenizers import SMILES_SPE_Tokenizer
+import numpy as np
+from easydict import EasyDict as edict
+
+from smiles.evaluation.peptide_analyzer import PeptideAnalyzer
+from smiles.evaluation.pdc_analyzer import PDCAnalyzer
+from smiles.scoring.classifiers import Conjugation, BindingAffinity
+
+# Import AA-to-SMILES conversion utilities
+from third_party.SMILESToPeptide.peptide_to_smiles import aa_to_smiles
+from third_party.SMILESToPeptide.aminoacid_selective import one_letter_to_three_letter
+
+
+# --- Model Architecture for Amino Acid Model ---
+def modulate(x, shift, scale):
+    """Modulates the input tensor x with a shift and scale."""
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class TimestepEmbedder(nn.Module):
+    """Embeds a continuous scalar timestep t in [0, 1] into a vector representation."""
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(1, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+
+    def forward(self, t):
+        return self.mlp(t.unsqueeze(-1))
+
+
+class DiTBlock(nn.Module):
+    """A single block of the Diffusion Transformer."""
+    def __init__(self, hidden_size, n_heads):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.attn = nn.MultiheadAttention(hidden_size, n_heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, 4 * hidden_size),
+            nn.GELU(),
+            nn.Linear(4 * hidden_size, hidden_size)
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        )
+
+    def forward(self, x, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+        x_norm1 = modulate(self.norm1(x), shift_msa, scale_msa)
+        attn_output, _ = self.attn(x_norm1, x_norm1, x_norm1)
+        x = x + gate_msa.unsqueeze(1) * attn_output
+        x_norm2 = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        mlp_output = self.mlp(x_norm2)
+        x = x + gate_mlp.unsqueeze(1) * mlp_output
+        return x
+
+
+class MDLM(nn.Module):
+    """Masked Diffusion Language Model (MDLM) using a DiT backbone."""
+    def __init__(self, vocab_size, seq_len, model_dim, n_heads, n_layers):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.seq_len = seq_len
+        self.model_dim = model_dim
+
+        self.token_embedder = nn.Embedding(vocab_size, model_dim)
+        self.pos_embedder = nn.Parameter(torch.randn(1, seq_len, model_dim))
+        self.time_embedder = TimestepEmbedder(model_dim)
+
+        self.transformer_blocks = nn.ModuleList([
+            DiTBlock(model_dim, n_heads) for _ in range(n_layers)
+        ])
+
+        self.final_norm = nn.LayerNorm(model_dim)
+        self.lm_head = nn.Linear(model_dim, vocab_size)
+
+    def forward(self, x, t):
+        seq_len = x.shape[1]
+        x_embed = self.token_embedder(x) + self.pos_embedder[:, :seq_len, :]
+        t_embed = self.time_embedder(t)
+        for block in self.transformer_blocks:
+            x_embed = block(x_embed, t_embed)
+        x_embed = self.final_norm(x_embed)
+        logits = self.lm_head(x_embed)
+        return logits
+
+
+class AminoTokenizer:
+    """Tokenizer for amino acid sequences (natural and unnatural AAs)."""
+    def __init__(self):
+        self.vocab = {
+            '<cls>': 0, '<pad>': 1, '<eos>': 2, '<unk>': 3,
+            'L': 4, 'A': 5, 'G': 6, 'V': 7, 'S': 8, 'E': 9,
+            'R': 10, 'T': 11, 'I': 12, 'D': 13, 'P': 14, 'K': 15,
+            'Q': 16, 'N': 17, 'F': 18, 'Y': 19, 'M': 20, 'H': 21,
+            'W': 22, 'C': 23,
+            # Non-canonical amino acids
+            'Ŷ': 24, 'Ĉ': 25, 'Ĝ': 26, 'Ė': 27, 'Ƙ': 28,
+            'Ṃ': 29, 'Ł': 30, 'Ċ': 31, 'Ṯ': 32,
+            # Lowercase tokens (D-amino acids)
+            'a': 33, 'c': 34, 'd': 35, 'e': 36, 'f': 37,
+            'g': 38, 'h': 39, 'i': 40, 'k': 41, 'l': 42,
+            'm': 43, 'n': 44, 'p': 45, 'q': 46, 'r': 47,
+            's': 48, 't': 49, 'v': 50, 'w': 51, 'y': 52,
+        }
+        self.id_to_token = {i: t for t, i in self.vocab.items()}
+
+        self.cls_token = '<cls>'
+        self.pad_token = '<pad>'
+        self.eos_token = '<eos>'
+        self.unk_token = '<unk>'
+
+        self.cls_id = self.vocab[self.cls_token]
+        self.pad_id = self.vocab[self.pad_token]
+        self.eos_id = self.vocab[self.eos_token]
+        self.unk_id = self.vocab[self.unk_token]
+
+        self.vocab_size = len(self.vocab.keys())
+
+        self._ncaa_to_nat_token = {
+            'Ŷ': 'Y', 'Ĉ': 'C', 'Ċ': 'C', 'Ĝ': 'G', 'Ė': 'E',
+            'Ƙ': 'K', 'Ṃ': 'M', 'Ł': 'L', 'Ṯ': 'T',
+        }
+
+        id_map = torch.arange(self.vocab_size, dtype=torch.long)
+        for ncaa_tok, nat_tok in self._ncaa_to_nat_token.items():
+            if ncaa_tok in self.vocab and nat_tok in self.vocab:
+                id_map[self.vocab[ncaa_tok]] = self.vocab[nat_tok]
+        self._ncaa_to_nat_id_map = id_map
+
+    def encode(self, seq: str, add_special_tokens: bool = True):
+        token_ids = []
+        if add_special_tokens:
+            token_ids.append(self.cls_id)
+        for ch in seq:
+            token_ids.append(self.vocab.get(ch, self.unk_id))
+        if add_special_tokens:
+            token_ids.append(self.eos_id)
+        return token_ids
+
+    def batch_encode(self, seqs, add_special_tokens: bool = True):
+        return torch.tensor([self.encode(s, add_special_tokens=add_special_tokens) for s in seqs])
+
+    def decode(self, ids, skip_special_tokens: bool = True):
+        if isinstance(ids, torch.Tensor):
+            ids = ids.tolist()
+        tokens = []
+        for i in ids:
+            tok = self.id_to_token.get(int(i), self.unk_token)
+            if skip_special_tokens and tok in {self.cls_token, self.eos_token, self.pad_token}:
+                continue
+            tokens.append(tok)
+        return ''.join(tokens)
+
+    def batch_decode(self, batch_ids, skip_special_tokens: bool = True):
+        return [self.decode(ids, skip_special_tokens=skip_special_tokens) for ids in batch_ids]
+
+    def to_natural_ids(self, batch_ids: torch.Tensor) -> torch.Tensor:
+        """Map non-canonical AA token IDs to natural AA token IDs."""
+        if not torch.is_tensor(batch_ids):
+            batch_ids = torch.tensor(batch_ids, dtype=torch.long)
+        id_map = self._ncaa_to_nat_id_map.to(device=batch_ids.device)
+        out = batch_ids.clone()
+        valid = (out >= 0) & (out < id_map.numel())
+        out[valid] = id_map[out[valid]]
+        out[~valid] = self.unk_id
+        return out
+
+
+class GenerationTracker:
+    """
+    Tracks detailed validity information during generation
+    """
+    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, scores, detailed_results=None, improved_indices=None):
+        """Log information for a generation step"""
+        step_info = {
+            'step': step,
+            'scores': scores.cpu().numpy().tolist() if isinstance(scores, torch.Tensor) else scores,
+            'num_valid': int(sum(1 for s in scores if s > 0)) if hasattr(scores, '__iter__') else 0,
+            'improved_indices': improved_indices or []
+        }
+        
+        if detailed_results:
+            step_info['detailed'] = []
+            for i, result in enumerate(detailed_results):
+                step_info['detailed'].append({
+                    'sample_idx': i,
+                    'is_valid_peptide': result.get('is_valid_peptide', False),
+                    'has_conjugation_site': result.get('has_conjugation_site', False),
+                    'conjugation_success': result.get('conjugation_success', False),
+                    'pdc_valid': result.get('pdc_valid', False),
+                    'num_conjugates': len(result.get('conjugates', [])),
+                    'error': result.get('error')
+                })
+        
+        self.step_logs.append(step_info)
+        
+    def log_final(self, sequences, detailed_results_with_props):
+        """Log final results with properties"""
+        for i, (seq, result) in enumerate(zip(sequences, detailed_results_with_props)):
+            final_info = {
+                'sample_idx': i,
+                'peptide_smiles': seq,
+                'is_valid_peptide': result.get('is_valid_peptide', False),
+                'has_conjugation_site': result.get('has_conjugation_site', False),
+                'num_conjugates': len(result.get('conjugates', [])),
+                'num_valid_pdcs': len(result.get('valid_pdcs', [])),
+                'best_pdc': None,
+                'properties': result.get('properties', {})
+            }
+            
+            if result.get('best_pdc'):
+                final_info['best_pdc'] = {
+                    'smiles': result['best_pdc'].get('smiles', ''),
+                    'checks': result['best_pdc'].get('checks', {})
+                }
+            
+            self.final_results.append(final_info)
+    
+    def save(self, filename_prefix='generation'):
+        """Save all logs to files"""
+        # Save step logs
+        step_log_path = os.path.join(self.output_dir, f"{filename_prefix}_{self.run_id}_steps.json")
+        with open(step_log_path, 'w') as f:
+            json.dump(self.step_logs, f, indent=2, default=str)
+        
+        # Save final results
+        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)
+        
+        # Save summary CSV
+        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', 'peptide_smiles', 'is_valid_peptide', 'has_conjugation_site',
+                'num_conjugates', 'num_valid_pdcs', 'best_pdc_smiles',
+                'mf', 'polararea', 'xlogp', 'heavycnt', 'rotbonds'
+            ])
+            for result in self.final_results:
+                props = result.get('properties', {})
+                writer.writerow([
+                    result['sample_idx'],
+                    result['peptide_smiles'],
+                    result['is_valid_peptide'],
+                    result['has_conjugation_site'],
+                    result['num_conjugates'],
+                    result['num_valid_pdcs'],
+                    result['best_pdc']['smiles'] if result.get('best_pdc') else '',
+                    props.get('mf', ''),
+                    props.get('polararea', ''),
+                    props.get('xlogp', ''),
+                    props.get('heavycnt', ''),
+                    props.get('rotbonds', '')
+                ])
+        
+        print(f"\n=== Generation logs saved ===")
+        print(f"  Step logs: {step_log_path}")
+        print(f"  Final results: {final_path}")
+        print(f"  Summary CSV: {csv_path}")
+        
+        return step_log_path, final_path, csv_path
+    
+    def print_step_summary(self, step, scores, detailed_results=None):
+        """Print a summary for the current step"""
+        if isinstance(scores, torch.Tensor):
+            scores_list = scores.cpu().numpy().tolist()
+        else:
+            scores_list = list(scores)
+        
+        num_valid = sum(1 for s in scores_list if s > 0)
+        
+        print(f"\n--- Step {step} Summary ---")
+        print(f"  Valid: {num_valid}/{len(scores_list)}")
+        
+        if detailed_results:
+            valid_peptide = sum(1 for r in detailed_results if r.get('is_valid_peptide'))
+            has_site = sum(1 for r in detailed_results if r.get('has_conjugation_site'))
+            conj_success = sum(1 for r in detailed_results if r.get('conjugation_success'))
+            pdc_valid = sum(1 for r in detailed_results if r.get('pdc_valid'))
+            
+            print(f"  Breakdown:")
+            print(f"    Valid peptides: {valid_peptide}")
+            print(f"    Has conjugation site: {has_site}")
+            print(f"    Conjugation success: {conj_success}")
+            print(f"    PDC valid: {pdc_valid}")
+    
+    def print_final_summary(self):
+        """Print final summary"""
+        print("\n" + "=" * 60)
+        print("FINAL GENERATION SUMMARY")
+        print("=" * 60)
+        
+        total = len(self.final_results)
+        valid_pdcs = sum(1 for r in self.final_results if r.get('num_valid_pdcs', 0) > 0)
+        
+        print(f"\nTotal sequences: {total}")
+        print(f"With valid PDCs: {valid_pdcs}")
+        
+        print("\nDetailed results:")
+        for result in self.final_results:
+            status = "" if result.get('num_valid_pdcs', 0) > 0 else ""
+            print(f"\n  [{status}] Sample {result['sample_idx']}:")
+            print(f"      Peptide: {result['peptide_smiles'][:60]}...")
+            print(f"      Valid peptide: {result['is_valid_peptide']}")
+            print(f"      Has conjugation site: {result['has_conjugation_site']}")            
+            print(f"      Conjugates: {result['num_conjugates']}")
+            print(f"      Valid PDCs: {result['num_valid_pdcs']}")
+            
+            if result.get('properties'):
+                props = result['properties']
+                print(f"      Properties:")
+                for k, v in props.items():
+                    print(f"        {k}: {v}")
+
+
+def peptide_bond_mask(smiles_list):
+    # Initialize the batch mask
+    batch_size = len(smiles_list)
+    max_seq_length = 1035 #max(len(smiles) for smiles in smiles_list)  # Find the longest SMILES
+    mask = torch.zeros((batch_size, max_seq_length), dtype=torch.int)  # Mask filled with zeros
+
+    bond_patterns = [
+        (r'OC\(=O\)', 'ester'),
+        (r'N\(C\)C\(=O\)', 'n_methyl'),
+        (r'N[12]C\(=O\)', 'peptide'),  # Pro peptide bonds
+        (r'NC\(=O\)', 'peptide'),  # Regular peptide bonds
+        (r'C\(=O\)N\(C\)', 'n_methyl'),
+        (r'C\(=O\)N[12]?', 'peptide')
+    ]
+
+    for batch_idx, smiles in enumerate(smiles_list):
+        positions = []
+        used = set()
+
+        # Identify bonds
+        for pattern, bond_type in bond_patterns:
+            for match in re.finditer(pattern, smiles):
+                if not any(p in range(match.start(), match.end()) for p in used):
+                    positions.append({
+                        'start': match.start(),
+                        'end': match.end(),
+                        'type': bond_type,
+                        'pattern': match.group()
+                    })
+                    used.update(range(match.start(), match.end()))
+
+        # Update the mask for the current SMILES
+        for pos in positions:
+            mask[batch_idx, pos['start']:pos['end']] = 1
+
+    return mask
+
+def peptide_token_mask(smiles_list, token_lists):
+    # Initialize the batch mask
+    batch_size = len(smiles_list)
+    token_seq_length = max(len(tokens) for tokens in token_lists)  # Find the longest tokenized sequence
+    tokenized_masks = torch.zeros((batch_size, token_seq_length), dtype=torch.int)  # Mask filled with zeros
+    atomwise_masks = peptide_bond_mask(smiles_list)
+
+
+    for batch_idx, atomwise_mask in enumerate(atomwise_masks):
+        token_seq = token_lists[batch_idx]
+        atom_idx = 0
+        
+        for token_idx, token in enumerate(token_seq):
+            if token_idx != 0 and token_idx != len(token_seq) - 1:
+                if torch.sum(atomwise_mask[atom_idx:atom_idx+len(token)]) >= 1:
+                    tokenized_masks[batch_idx][token_idx] = 1
+                atom_idx += len(token)
+    
+    return tokenized_masks
+
+class MOGGenerator:
+    def __init__(self, model, device, objectives, args, tracker=None, tokenizer=None, tokenizer_type='smiles'):
+        self.model = model
+        self.device = device
+        self.objectives = objectives
+        self.args = args
+        self.num_objectives = len(objectives)
+        self.peptide_analyzer = PeptideAnalyzer()
+        self.pdc_analyzer = PDCAnalyzer(device)
+        self.tracker = tracker
+        self.tokenizer = tokenizer
+        self.tokenizer_type = tokenizer_type
+        
+        # Set invalid tokens based on tokenizer type
+        if tokenizer_type == 'amino':
+            self.invalid = [0, 1, 2, 3, 38]  # special tokens + 'g' for amino
+            # Initialize SMILES tokenizer for AA-to-SMILES conversion
+            self.smiles_tokenizer = SMILES_SPE_Tokenizer(
+                args.smiles_vocab_path,
+                args.smiles_splits_path
+            )
+        else:
+            self.invalid = [0, 1, 2, 3, 4, 585, 586]  # SMILES special tokens
+            self.smiles_tokenizer = None
+        
+        # Find the Conjugation objective for detailed tracking
+        self.conjugation_obj = None
+        for obj in objectives:
+            if isinstance(obj, Conjugation):
+                self.conjugation_obj = obj
+                break
+
+    def _generate_single_x0(self, n_steps=16, temperature=1.0):
+        """Generate a single batch of x0 sequences."""
+        self.model.eval()
+
+        # Get vocab size based on tokenizer type
+        if self.tokenizer_type == 'amino':
+            vocab_size = self.model.vocab_size
+        else:
+            vocab_size = self.model.tokenizer.vocab_size
+
+        # 1. Start with a tensor of random tokens (pure noise at t=0)
+        # For amino tokenizer, exclude invalid tokens from initialization
+        if self.tokenizer_type == 'amino':
+            valid_tokens = [i for i in range(vocab_size) if i not in self.invalid]
+            valid_tokens = torch.tensor(valid_tokens, device=self.device)
+            rand_indices = torch.randint(0, len(valid_tokens), (self.args.num_samples, self.args.gen_len), device=self.device)
+            body = valid_tokens[rand_indices]
+            # Add CLS and EOS tokens
+            cls_col = torch.full((self.args.num_samples, 1), self.tokenizer.cls_id, device=self.device)
+            eos_col = torch.full((self.args.num_samples, 1), self.tokenizer.eos_id, device=self.device)
+            x = torch.cat([cls_col, body, eos_col], dim=1)
+        else:
+            x = torch.randint(
+                0,
+                vocab_size,
+                (self.args.num_samples, self.args.gen_len),
+                device=self.device
+            )
+
+        # 2. Define the time schedule for the forward process (0.0 to 1.0)
+        time_steps = torch.linspace(0.0, 1.0, n_steps + 1, device=self.device)
+
+        # 3. Iteratively follow the flow from noise to data
+        with torch.no_grad():
+            for i in tqdm(range(n_steps), desc="Flow Matching Steps", leave=False):
+                t_curr = time_steps[i]
+                t_next = time_steps[i+1]
+
+                # Prepare the current timestep tensor for the model
+                t_tensor = torch.full((self.args.num_samples,), t_curr, device=self.device)
+
+                # Get the model's prediction for the final clean sequence (at t=1)
+                logits = self.model(x, t_tensor) 
+                logits = logits / temperature
+                
+                # Mask out invalid tokens
+                for inv_tok in self.invalid:
+                    if inv_tok < logits.shape[-1]:
+                        logits[:, :, inv_tok] = -1000
+
+                pred_x1 = torch.argmax(logits, dim=-1)
+
+                # On the last step, the result is the final prediction
+                if i == n_steps - 1:
+                    x = pred_x1
+                    break
+
+                # --- Construct the next state x_{t_next} ---
+                # The probability of a token being noise at time t_next is (1 - t_next).
+                noise_prob = 1.0 - t_next
+                mask = torch.rand(x.shape, device=self.device) < noise_prob
+                
+                # Generate new random tokens for the noise positions
+                noise = torch.randint(
+                    0,
+                    vocab_size,
+                    x.shape,
+                    device=self.device
+                )
+
+                # Combine the final prediction with noise to form the next intermediate state
+                x = torch.where(mask, noise, pred_x1)
+
+        return x
+
+    def _check_conjugation_validity(self, x):
+        """
+        Check if sequences can be conjugated.
+        Returns a boolean tensor indicating validity for each sample.
+        """
+        if self.tokenizer_type == 'amino':
+            # For amino tokenizer, convert to SMILES first
+            seqs = self.tokenizer.batch_decode(x)
+            three_letter_seqs = [one_letter_to_three_letter(seq) for seq in seqs]
+            smiles_seqs = [aa_to_smiles(three_letter_seq, cyclic=self.args.cyclic) 
+                          for three_letter_seq in three_letter_seqs]
+            
+            # Check if SMILES conversion succeeded and result is valid peptide
+            validities = []
+            for smiles in smiles_seqs:
+                if isinstance(smiles, str):
+                    validities.append(self.peptide_analyzer.is_peptide(smiles))
+                else:
+                    validities.append(False)
+            return torch.tensor(validities, device=self.device)
+        
+        if self.conjugation_obj is None:
+            # Fall back to peptide validity check
+            sequences = self.tokenizer.batch_decode(x)
+            return torch.tensor([self.peptide_analyzer.is_peptide(seq) for seq in sequences], device=self.device)
+        
+        # Use conjugation objective - score > 0 means valid
+        scores = self.conjugation_obj(x)
+        return scores > 0
+
+    def generate_x0(self, n_steps=16, temperature=1.0, require_valid=False, max_retries=10):
+        """
+        Generate initial x0 sequences.
+        
+        Args:
+            n_steps: Number of flow matching steps
+            temperature: Sampling temperature
+            require_valid: If True, regenerate invalid samples until all are valid conjugatable peptides
+            max_retries: Maximum number of regeneration attempts per sample
+        """
+        print("Starting initial SMILES generation...")
+        
+        x = self._generate_single_x0(n_steps, temperature)
+        
+        if require_valid:
+            print("Checking conjugation validity and regenerating invalid samples...")
+            
+            for retry in range(max_retries):
+                validity = self._check_conjugation_validity(x)
+                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 are valid conjugatable peptides!")
+                    break
+                
+                print(f"Retry {retry + 1}/{max_retries}: {num_valid}/{self.args.num_samples} valid, regenerating {num_invalid} invalid samples...")
+                
+                # Get indices of invalid samples
+                invalid_indices = torch.where(~validity)[0]
+                
+                # Generate new candidates for invalid samples
+                # Temporarily adjust num_samples to generate replacements
+                original_num_samples = self.args.num_samples
+                self.args.num_samples = num_invalid
+                
+                new_x = self._generate_single_x0(n_steps, temperature)
+                
+                self.args.num_samples = original_num_samples
+                
+                # Replace invalid samples with new ones
+                x[invalid_indices] = new_x
+            
+            # Final validity check
+            final_validity = self._check_conjugation_validity(x)
+            num_final_valid = final_validity.sum().item()
+            if num_final_valid < self.args.num_samples:
+                print(f"Warning: After {max_retries} retries, only {num_final_valid}/{self.args.num_samples} samples are valid.")
+            else:
+                print(f"Success: All {self.args.num_samples} samples are valid conjugatable peptides.")
+        
+        # 4. Decode the final token IDs into strings
+        generated_sequences = self.tokenizer.batch_decode(x)
+        
+        # 5. Analyze the validity of the generated sequences
+        validities = []
+        if self.tokenizer_type == 'amino':
+            # For amino mode, report conjugation-relevant validity after AA->SMILES conversion.
+            validity_tensor = self._check_conjugation_validity(x)
+            validities = [bool(v.item()) for v in validity_tensor]
+            for seq in generated_sequences:
+                print(seq)
+        else:
+            for seq in generated_sequences:
+                validities.append(self.peptide_analyzer.is_peptide(seq))
+                print(seq)
+        if self.tokenizer_type == 'amino':
+            print(f"Initial Peptide Validity (AA->SMILES): {validities}")
+        else:
+            print(f"Initial Sequence Validity: {validities}")
+        return x
+
+    def _get_scores(self, x_batch):
+        """Calculates the normalized scores for a batch of sequences."""
+        scores = []
+        
+        if self.tokenizer_type == 'amino':
+            # Convert AA tokens to SMILES for scoring
+            seqs = self.tokenizer.batch_decode(x_batch)
+            
+            # Convert one-letter AA codes to three-letter, then to SMILES
+            three_letter_seqs = [one_letter_to_three_letter(seq) for seq in seqs]
+            with ThreadPoolExecutor() as pool:
+                smiles_seqs = list(pool.map(
+                    lambda s: aa_to_smiles(s, cyclic=self.args.cyclic),
+                    three_letter_seqs
+                ))
+
+            # Track valid conversions
+            valid_mask = torch.tensor(
+                [1 if isinstance(seq, str) else 0 for seq in smiles_seqs], 
+                device=self.device
+            )
+            
+            # Replace invalid SMILES with dummy 'C' for tokenization
+            for i in torch.where((1 - valid_mask) == 1)[0].tolist():
+                smiles_seqs[i] = 'C'
+            
+            # Tokenize SMILES for objective functions
+            smiles_batch = self.smiles_tokenizer(
+                smiles_seqs, return_tensors='pt', padding=True
+            ).to(self.device)
+            attention_mask = smiles_batch.get("attention_mask", None)
+            
+            # Evaluate each objective function on SMILES tokens
+            for obj_func in self.objectives:
+                try:
+                    score = obj_func(
+                        smiles_batch["input_ids"],
+                        attention_mask=attention_mask
+                    ).to(self.device)
+                except TypeError:
+                    score = obj_func(smiles_batch["input_ids"]).to(self.device)
+                scores.append(score * valid_mask)  # Zero out invalid conversions
+        else:
+            # Original SMILES-based scoring
+            for obj_func in self.objectives:
+                scores.append(obj_func(x_batch.to(self.device)))
+        
+        return torch.stack(scores, dim=0).to(self.device)
+
+    def _barker_g(self, u):
+        """Barker balancing function."""
+        return u / (1 + u)
+    
+    def validity(self, x):
+        sampled_sequences = self.tokenizer.batch_decode(x)
+        
+        if self.tokenizer_type == 'amino':
+            # For amino tokenizer, convert to SMILES first and check validity
+            three_letter_seqs = [one_letter_to_three_letter(seq) for seq in sampled_sequences]
+            smiles_seqs = [aa_to_smiles(three_letter_seq, cyclic=self.args.cyclic) 
+                          for three_letter_seq in three_letter_seqs]
+            return [1.0 if isinstance(seq, str) and self.peptide_analyzer.is_peptide(seq) else 0.0 
+                    for seq in smiles_seqs]
+        else:
+            return [1.0 if self.peptide_analyzer.is_peptide(seq) else 0.0 for seq in sampled_sequences]
+
+    def generate(self):
+        """Main generation loop with tracking."""
+        shape = (self.args.num_samples, self.args.gen_len)
+        x = self.generate_x0(
+            require_valid=self.args.require_valid_x0,
+            max_retries=self.args.max_x0_retries
+        )
+        print(x)
+
+        if self.args.weights is None:
+            # The first weight is for conjugation (peptide + PDC validity)
+            # We need to ensure the SMILES sequences are valid peptides with valid PDCs
+            weights = torch.tensor([1] + [1/(self.num_objectives-1)] * (self.num_objectives-1), device=self.device).view(-1,1)
+        else:
+            weights = torch.tensor(self.args.weights, device=self.device).view(-1, 1)
+            if len(weights) != self.num_objectives:
+                raise ValueError("Number of weights must match number of objectives.")
+        print(f"Weights: {weights}")
+
+        with torch.no_grad():
+            current_scores = self._get_scores(x)
+            for t in tqdm(range(self.args.optimization_steps), desc="MOG Generation"):
+                improved = False
+                improved_indices = []
+
+                # Anneal guidance strength
+                eta_t = self.args.eta_min + (self.args.eta_max - self.args.eta_min) * (t / (self.args.optimization_steps - 1))
+                # Choose a random position to mutate
+                mut_idx = random.randint(1, self.args.gen_len-2)
+
+                # Determine the generation timestep
+                # We cycle through the timesteps to ensure all are visited
+                generation_step = t % self.args.optimization_steps
+                time_t = torch.full((self.args.num_samples,), (generation_step / self.args.optimization_steps), device=self.device)
+                
+                # Get proposal distribution from ReDi model for the chosen position
+                logits = self.model(x, time_t)
+                probs = F.softmax(logits, dim=-1)
+                pos_probs = probs[:, mut_idx, :]
+                pos_probs[:, x[:, mut_idx]] = 0   # We don't evalute the same token
+                pos_probs[:, self.invalid] = 0
+
+                # Prune candidate vocabulary using top-p sampling
+                sorted_probs, sorted_indices = torch.sort(pos_probs, descending=True)
+                cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                remove_mask = cumulative_probs > self.args.top_p
+                remove_mask[..., 1:] = remove_mask[..., :-1].clone()
+                remove_mask[..., 0] = 0
+
+                # Get the set of candidate tokens for each sample in the batch
+                candidate_tokens_list = []
+                for i in range(self.args.num_samples):
+                    sample_mask = remove_mask[i]
+                    candidates = sorted_indices[i, ~sample_mask]
+                    candidate_tokens_list.append(candidates)
+                
+                w_current = torch.exp(eta_t * torch.min(weights * current_scores, dim=0).values)
+
+                if t == 0:
+                    print(f"Initial Scores: {current_scores}")
+                    # Log initial step with detailed results
+                    if self.tracker and self.conjugation_obj:
+                        detailed = self.conjugation_obj.get_last_detailed_results()
+                        self.tracker.log_step(t, current_scores[0], detailed)
+                        self.tracker.print_step_summary(t, current_scores[0], detailed)
+
+                # Evaluate all candidate tokens for each sample
+                final_proposal_tokens = []
+                for i in range(self.args.num_samples):
+                    candidates = candidate_tokens_list[i]
+                    candidates = torch.tensor([token for token in candidates if token not in self.invalid], device=candidates.device)
+                    if len(candidates) >= 200:
+                        candidates = candidates[:200]
+                    num_candidates = len(candidates)
+                    
+                    # Create a batch of proposed sequences for the current sample
+                    x_prop_batch = x[i].repeat(num_candidates, 1)
+                    x_prop_batch[:, mut_idx] = candidates
+                    
+                    # Evaluate all proposals
+                    proposal_scores = self._get_scores(x_prop_batch)
+                    proposal_s_omega = torch.min(weights * proposal_scores, dim=0).values
+                    w_proposal = torch.exp(eta_t * proposal_s_omega)
+                    
+                    # Get ReDi probabilities for the candidates
+                    redi_probs = pos_probs[i, candidates]
+
+                    # Calculate unnormalized guided probabilities
+                    tilde_q = redi_probs * self._barker_g(w_proposal / w_current[i])
+                    
+                    # Normalize and sample the final token
+                    final_probs = tilde_q / (torch.sum(tilde_q) + 1e-9)
+
+                    index = torch.multinomial(final_probs, 1).item()
+
+                    if (random.uniform(0,1) > 1 and proposal_scores[:,index][0] == 1) or torch.sum(weights.squeeze(1) * proposal_scores[:, index]) >= torch.sum(weights.squeeze(1) * current_scores[:,i]):
+                        final_token = candidates[index]
+                        print(f"Previous Weighted Sum: {torch.sum(weights.squeeze(1) * current_scores[:,i])}")
+                        print(f"Previous Scores: {current_scores[:,i]}")
+
+                        print(f"New Weighted Sum: {torch.sum(weights.squeeze(1) * proposal_scores[:, index])}")
+                        print(f"New Scores: {proposal_scores[:,index]}")
+                        improved = True
+                        improved_indices.append(i)
+                    
+                    else:
+                        final_token = x[i][mut_idx]
+                    
+                    final_proposal_tokens.append(final_token)
+
+                # Update the sequences with the chosen tokens
+                x[torch.arange(self.args.num_samples), mut_idx] = torch.stack(final_proposal_tokens)
+                if improved:
+                    current_scores = self._get_scores(x)
+                    print(self.tokenizer.batch_decode(x))
+
+                # Log step with tracking (every 4 steps to reduce overhead)
+                if self.tracker and self.conjugation_obj and (t % 4 == 0 or t == self.args.optimization_steps - 1):
+                    detailed = self.conjugation_obj.get_last_detailed_results()
+                    self.tracker.log_step(t, current_scores[0], detailed, improved_indices)
+
+                    # Print summary for final step
+                    if t == self.args.optimization_steps - 1:
+                        self.tracker.print_step_summary(t, current_scores[0], detailed)
+
+        return x
+
+
+# --- Main Execution ---
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    target = args.target
+    
+    # --- Initialize tokenizer based on type ---
+    if args.tokenizer_type == 'amino':
+        print("Using Amino Acid tokenizer")
+        tokenizer = AminoTokenizer()
+        smiles_tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+    else:
+        print("Using SMILES tokenizer")
+        tokenizer = SMILES_SPE_Tokenizer(args.smiles_vocab_path, args.smiles_splits_path)
+        smiles_tokenizer = tokenizer
+    
+    # --- Initialize objective functions ---
+    # Use SMILES tokenizer for objectives (they expect SMILES tokens)
+    conjugation_model = Conjugation(smiles_tokenizer, device=device, builder_type=args.builder_type)
+    affinity_model = BindingAffinity(target, device)
+
+    # List of all objective functions
+    OBJECTIVE_FUNCTIONS = [conjugation_model, affinity_model]
+    
+    # Initialize tracker for logging
+    tracker = GenerationTracker(output_dir=args.output_dir)
+
+    # --- Set random seeds for reproducibility ---
+    if args.seed is not None:
+        import random
+        import numpy as np
+        random.seed(args.seed)
+        np.random.seed(args.seed)
+        torch.manual_seed(args.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed(args.seed)
+            torch.cuda.manual_seed_all(args.seed)
+        print(f"Random seed set to: {args.seed}")
+
+    # --- Load Model ---
+    print(f"Loading model from checkpoint: {args.checkpoint}")
+    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=False)
+    
+    if args.tokenizer_type == 'amino':
+        # Vanilla PyTorch checkpoint for amino model
+        raw_model_args = checkpoint['args']
+        if isinstance(raw_model_args, argparse.Namespace):
+            raw_model_args = vars(raw_model_args)
+        model_args = edict(raw_model_args)
+        model = MDLM(
+            vocab_size=model_args.vocab_size,
+            seq_len=model_args.seq_len,
+            model_dim=model_args.model_dim,
+            n_heads=model_args.n_heads,
+            n_layers=model_args.n_layers
+        ).to(device)
+        model.load_state_dict(checkpoint['model_state_dict'])
+        print(f"Amino model loaded. Vocab size: {model_args.vocab_size}")
+    else:
+        # PyTorch Lightning checkpoint for SMILES model
+        model_hparams = checkpoint["hyper_parameters"]["args"]
+        model = MDLMLightningModule.load_from_checkpoint(
+            args.checkpoint,
+            args=model_hparams,
+            tokenizer=tokenizer,
+            map_location=device,
+            strict=False
+        )
+    model.to(device)
+    model.eval()
+    print("Model loaded successfully.")
+
+    mog_generator = MOGGenerator(
+        model, device, OBJECTIVE_FUNCTIONS, args, tracker=tracker,
+        tokenizer=tokenizer, tokenizer_type=args.tokenizer_type
+    )
+    
+    all_sequences = []
+    all_final_results = []
+    
+    for batch_idx in range(args.num_batches):
+        print(f"\n{'='*60}")
+        print(f"BATCH {batch_idx + 1}/{args.num_batches}")
+        print(f"{'='*60}")
+        
+        generated_tokens = mog_generator.generate()
+        final_scores = mog_generator._get_scores(generated_tokens.detach()).detach().cpu().numpy()
+        sequence_str = tokenizer.batch_decode(generated_tokens)
+
+        print("\n--- Batch Results ---")
+        print(f"Sequences: {sequence_str}")
+        print(f"Scores: {final_scores}")
+        
+        all_sequences.extend(sequence_str)
+        
+        # Get detailed results with properties for final sequences
+        # For amino tokenizer, we need to convert to SMILES first for detailed scoring
+        if args.tokenizer_type == 'amino':
+            # Convert AA to SMILES for detailed conjugation analysis
+            smiles_list = []
+            for seq in sequence_str:
+                three_letter = one_letter_to_three_letter(seq)
+                smiles = aa_to_smiles(three_letter, cyclic=args.cyclic)
+                smiles_list.append(smiles if isinstance(smiles, str) else 'C')
+            # Tokenize SMILES for conjugation model
+            smiles_tokens = smiles_tokenizer(smiles_list, return_tensors='pt', padding=True).to(device)
+            detailed_with_props = conjugation_model.get_detailed_scores_with_properties(smiles_tokens["input_ids"])
+        else:
+            detailed_with_props = conjugation_model.get_detailed_scores_with_properties(generated_tokens)
+        
+        tracker.log_final(sequence_str, detailed_with_props)
+        all_final_results.extend(detailed_with_props)
+    
+    # Print and save final summary
+    tracker.print_final_summary()
+    tracker.save(filename_prefix='pdc_generation')
+    
+    # Save sequences to text file
+    with open(args.output_file, 'w') as f:
+        for i, seq in enumerate(all_sequences):
+            if args.tokenizer_type == 'amino':
+                # Also save converted SMILES for amino sequences
+                three_letter = one_letter_to_three_letter(seq)
+                smiles = aa_to_smiles(three_letter, cyclic=args.cyclic)
+                if isinstance(smiles, str):
+                    f.write(f"{seq}\t{smiles}\n")
+                else:
+                    f.write(f"{seq}\tINVALID_SMILES\n")
+            else:
+                f.write(seq + '\n')
+    print(f"\nSequences saved to: {args.output_file}")
+
+    print("\nGeneration complete.")
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Multi-Objective Generation with LBP-MOG-ReDi (Single Mutation).")
+
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to the trained ReDi model checkpoint.")
+    parser.add_argument("--num_samples", type=int, default=10, help="Number of samples to generate.")
+    parser.add_argument("--num_batches", type=int, default=10, help="Number of batches to generate.")
+    parser.add_argument("--output_file", type=str, default="./smiles.txt", help="File to save the generated sequences.")
+    parser.add_argument("--output_dir", type=str, default="./generation_logs", help="Directory to save detailed generation logs.")
+    parser.add_argument("--gen_len", type=int, default=50, help="Length of the sequences to generate.")
+    parser.add_argument("--optimization_steps", type=int, default=16, help="Number of optimization steps per generation.")
+    parser.add_argument("--weights", type=float, nargs='+', required=False, help="Weights for the objectives.")
+    parser.add_argument("--eta_min", type=float, default=1.0, help="Minimum guidance strength.")
+    parser.add_argument("--eta_max", type=float, default=20.0, help="Maximum guidance strength.")
+    parser.add_argument("--top_p", type=float, default=0.9, help="Top-p for pruning candidate tokens.")
+    parser.add_argument("--builder_type", type=str, default="AoaDau",
+                        choices=["AoaDau", "SSCpt", "VcMMAE", "DSPACpt", "MalHexCpt", 
+                                 "DSECCpt", "Cpt", "DSPASN38", "DSEBSN38", "DSEPSN38", "SN38"],
+                        help="PDC builder type for conjugation.")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.")
+    parser.add_argument("--require_valid_x0", action="store_true",
+                        help="If set, regenerate x0 samples until they are valid conjugatable peptides.")
+    parser.add_argument("--max_x0_retries", type=int, default=10,
+                        help="Maximum retries when require_valid_x0 is enabled.")
+
+    parser.add_argument("--target", type=str, required=True, help="Target protein sequence for binding affinity.")
+    
+    # Tokenizer and model type arguments
+    parser.add_argument("--tokenizer_type", type=str, default="smiles", choices=["smiles", "amino"],
+                        help="Type of tokenizer: 'smiles' for SMILES SPE tokenizer, 'amino' for amino acid tokenizer.")
+    parser.add_argument("--cyclic", action="store_true", 
+                        help="Generate cyclic peptides (only for amino tokenizer).")
+    parser.add_argument("--smiles_vocab_path", type=str, 
+                        default="./smiles/smiles_tokenizer/new_vocab.txt",
+                        help="Path to SMILES tokenizer vocabulary file.")
+    parser.add_argument("--smiles_splits_path", type=str,
+                        default="./smiles/smiles_tokenizer/new_splits.txt",
+                        help="Path to SMILES tokenizer splits file.")
+    
+    args = parser.parse_args()
+    main(args)