import time
start_time = time.time()

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import pandas as pd
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit import DataStructs
from rdkit.Chem import Descriptors
from rdkit.Chem import Draw
from rdkit.Chem import RDConfig
import os
import sys
sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
import sascorer

# Function to calculate largest ring size
def get_largest_ring_size(mol):
    cycle_list = mol.GetRingInfo().AtomRings()
    if cycle_list:
        cycle_length = max([len(j) for j in cycle_list])
    else:
        cycle_length = 0
    return cycle_length

# Function to calculate penalized logP
def plogp(smile):
    if smile:
        mol = Chem.MolFromSmiles(smile)
        if mol:
            log_p = Descriptors.MolLogP(mol)
            sas_score = sascorer.calculateScore(mol)
            largest_ring_size = get_largest_ring_size(mol)
            cycle_score = max(largest_ring_size - 6, 0)
            if log_p and sas_score and largest_ring_size:
                p_logp = log_p - sas_score - cycle_score
                return p_logp
    return -100

# Function to calculate Tanimoto similarity
def sim(input_smile, output_smile):
    if input_smile and output_smile:
        input_mol = Chem.MolFromSmiles(input_smile)
        output_mol = Chem.MolFromSmiles(output_smile)
        if input_mol and output_mol:
            input_fp = AllChem.GetMorganFingerprint(input_mol, 2)
            output_fp = AllChem.GetMorganFingerprint(output_mol, 2)
            return DataStructs.TanimotoSimilarity(input_fp, output_fp)
    return None

def optimize_molecule(input_smiles):
    tokenizer = AutoTokenizer.from_pretrained("zjunlp/MolGen-large-opt")
    model = AutoModelForSeq2SeqLM.from_pretrained("zjunlp/MolGen-large-opt")

    mol_input = Chem.MolFromSmiles(input_smiles)
    if not mol_input:
        raise ValueError("Invalid input SMILES")

    sf_input = tokenizer(input_smiles, return_tensors="pt")
    molecules = model.generate(
        input_ids=sf_input["input_ids"],
        attention_mask=sf_input["attention_mask"],
        do_sample=True,
        max_length=100,
        min_length=5,
        top_k=30,
        top_p=1,
        num_return_sequences=10
    )
    smiles_outputs = list(set(tokenizer.decode(g, skip_special_tokens=True).replace(" ", "") for g in molecules))

    # Initialize lists to hold valid results
    valid_smiles = []
    valid_improvements = []
    valid_similarities = []

    # Calculate properties and filter valid molecules
    input_plogp = plogp(input_smiles)
    for sm in smiles_outputs:
        mol = Chem.MolFromSmiles(sm)
        if mol:
            p_logp = plogp(sm)
            if p_logp != -100:  # Ensure plogp calculation was successful
                plogp_improvement = p_logp - input_plogp
                similarity = sim(input_smiles, sm)
                if similarity and plogp_improvement > 0 and similarity > 0.4:
                    valid_smiles.append(sm)
                    valid_improvements.append(plogp_improvement)
                    valid_similarities.append(similarity)
        else:
            print(f"Invalid SMILES ignored: {sm}")

    # Check if any valid molecules are found
    if not valid_smiles:
        raise ValueError("No valid optimized molecules found.")

    # Visualize molecules
    mols = [Chem.MolFromSmiles(sm) for sm in valid_smiles]
    visualization = Draw.MolsToGridImage(
        mols,
        molsPerRow=4,
        subImgSize=(200, 200),
        legends=[f"ΔpLogP: {round(imp, 2)} | Sim: {round(sim, 2)}"
                 for imp, sim in zip(valid_improvements, valid_similarities)]
    )

    return valid_smiles, valid_improvements, valid_similarities, visualization

input_smiles = "C1=CC2=C(C=C1)C=CC=C2"  # Replace with your SMILES input
valid_smiles, valid_improvements, valid_similarities, visualization = optimize_molecule(input_smiles)

# Save or display the visualization
visualization.save("optimized_molecules.png")


end_time = time.time()
execution_time = end_time - start_time
print(f"Execution time: {execution_time}")



Invalid SMILES ignored: [=N+1][C][\S@]
Invalid SMILES ignored: [B][C@H1][Br][\F]
Invalid SMILES ignored: [Si][=Branch2][#S][/S@@]
Invalid SMILES ignored: [N-1][=C][N]
Invalid SMILES ignored: [C][C][=Branch2]
[15:59:41] SMILES Parse Error: syntax error while parsing: [=N+1][C][\S@]
[15:59:41] SMILES Parse Error: Failed parsing SMILES '[=N+1][C][\S@]' for input: '[=N+1][C][\S@]'
[15:59:41] SMILES Parse Error: syntax error while parsing: [B][C@H1][Br][\F]
[15:59:41] SMILES Parse Error: Failed parsing SMILES '[B][C@H1][Br][\F]' for input: '[B][C@H1][Br][\F]'
[15:59:41] SMILES Parse Error: syntax error while parsing: [Si][=Branch2][#S][/S@@]
[15:59:41] SMILES Parse Error: Failed parsing SMILES '[Si][=Branch2][#S][/S@@]' for input: '[Si][=Branch2][#S][/S@@]'
[15:59:41] SMILES Parse Error: syntax error while parsing: [N-1][=C][N]
[15:59:41] SMILES Parse Error: Failed parsing SMILES '[N-1][=C][N]' for input: '[N-1][=C][N]'
[15:59:41] SMILES Parse Error: syntax error while parsing: [C][C][=Branch2]
[15:59:41] SMILES Parse Error: Failed parsing SMILES '[C][C][=Branch2]' for input: '[C][C][=Branch2]'
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[14], line 110
    107     return valid_smiles, valid_improvements, valid_similarities, visualization
    109 input_smiles = "C1=CC2=C(C=C1)C=CC=C2"  # Replace with your SMILES input
--> 110 valid_smiles, valid_improvements, valid_similarities, visualization = optimize_molecule(input_smiles)
    112 # Save or display the visualization
    113 visualization.save("optimized_molecules.png")

Cell In[14], line 95, in optimize_molecule(input_smiles)
     93 # Check if any valid molecules are found
     94 if not valid_smiles:
---> 95     raise ValueError("No valid optimized molecules found.")
     97 # Visualize molecules
     98 mols = [Chem.MolFromSmiles(sm) for sm in valid_smiles]

ValueError: No valid optimized molecules found.

its generting selfies not smiles