standalone_show_generated.py

"""
standalone_show_generated.py
--------------------------------
Runs Gen_PartialSMILES2.py with a user scaffold SMILES, loads the
generated_molecules.csv file, displays a 2D grid of molecules, and
creates small icon images for each generated molecule.

Usage:
    python standalone_show_generated.py "CC(N)C(=O)O"
"""

import subprocess
import sys
import os
import pandas as pd
from rdkit import Chem
from rdkit.Chem import Draw


# ---------------------------------------------------------
# 1. Ensure scaffold has attachment point [*]
# ---------------------------------------------------------
def ensure_attachment_point(smiles):
    """
    Ensures the scaffold contains the required [*] wildcard.
    If missing, add [*] to the front.
    """
    if "[*]" in smiles:
        return smiles
    return "[*]" + smiles


# ---------------------------------------------------------
# 2. Run the generator script (Gen_PartialSMILES2.py)
# ---------------------------------------------------------
def run_generator(scaffold_smiles, n_to_gen=16):
    """
    Runs the molecule generator and produces generated_molecules.csv.
    """
    scaffold_smiles = ensure_attachment_point(scaffold_smiles)

    cmd = [
        sys.executable,
        "Gen_PartialSMILES2.py",
        "--scaffold", scaffold_smiles,
        "--n_to_gen", str(n_to_gen)
    ]

    print(f"\nRunning generator with scaffold: {scaffold_smiles}")
    print("Command:", " ".join(cmd), "\n")

    subprocess.run(cmd, check=True)

    if not os.path.exists("generated_molecules.csv"):
        raise FileNotFoundError("generated_molecules.csv not found after generator run.")

    print("Generation complete!\n")


# ---------------------------------------------------------
# 3. Render full RDKit grid image & save as PNG
# ---------------------------------------------------------
def load_and_draw_generated(csv_path="generated_molecules.csv"):
    """
    Reads generated_molecules.csv and displays a 2D RDKit structure grid.
    Saves the grid as generated_molecules.png.
    """
    df = pd.read_csv(csv_path)
    smiles_list = df["smiles"].tolist()

    mols = []
    legends = []

    for smi in smiles_list:
        mol = Chem.MolFromSmiles(smi)
        if mol:
            mols.append(mol)
            legends.append(smi)

    if not mols:
        print("No valid molecules found in CSV.")
        return None, []

    print(f"Rendering {len(mols)} molecules into grid...")

    img = Draw.MolsToGridImage(
        mols,
        molsPerRow=4,
        subImgSize=(250, 250),
        legends=legends
    )

    img.show()
    img.save("generated_molecules.png")
    print("Saved structure grid → generated_molecules.png\n")

    return mols, smiles_list


# ---------------------------------------------------------
# 4. Create 150×150 PNG icons from each molecule
# ---------------------------------------------------------
def generate_icons(mols, smiles_list, save_folder="icons"):
    """
    Converts each RDKit Mol into a small 150×150 PNG icon.
    Saves them in icons/, returns PIL images and labels.
    """
    os.makedirs(save_folder, exist_ok=True)

    print("Creating icons...")

    icons = []
    labels = []

    for i, mol in enumerate(mols):
        img = Draw.MolToImage(mol, size=(150, 150), kekulize=True)

        icon_path = os.path.join(save_folder, f"icon_{i+1}.png")
        img.save(icon_path)

        icons.append(img)
        labels.append(smiles_list[i])

    print(f"Saved {len(icons)} icons → {save_folder}/")

    return icons, labels


# ---------------------------------------------------------
# 5. Main execution
# ---------------------------------------------------------
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python standalone_show_generated.py \"SMILES_STRING\"")
        sys.exit(1)

    scaffold = sys.argv[1]

    # Step 1: Run generator
    run_generator(scaffold, n_to_gen=16)

    # Step 2: Show big RDKit grid
    mols, smiles_out = load_and_draw_generated()

    # Step 3: Create icon images
    if mols:
        icons, labels = generate_icons(mols, smiles_out)
        print("\nIcons ready! You can now use them in your Gradio grid.")