import streamlit as st
from rdkit import Chem
from rdkit.Chem import Draw
import random

# 假設的生成 smiles 的 function


def generate_smiles(input_smiles):
    # 這裡應該調用 Hugging Face 模型生成 20 個 SMILES 字串
    # 為了示範，隨機生成 20 個相似的 SMILES
    return [input_smiles + str(random.randint(0, 9)) for _ in range(20)]

# 計算相似度的假設 function


def calculate_similarity(smiles1, smiles2):
    # 簡單的示範相似度計算
    return round(random.uniform(0.5, 1.0), 2)

# 畫圖的 function


def draw_molecules(smiles_list):
    mols = [Chem.MolFromSmiles(smi) for smi in smiles_list]
    img = Draw.MolsToGridImage(mols, molsPerRow=5, subImgSize=(200, 200))
    return img


st.sidebar.title("SMILES Input")
input_smiles = st.sidebar.text_input("Enter SMILES string:")
submit_button = st.sidebar.button("Submit")

st.title("(WIP 🚧) Drug-like Molecule Generation")
st.write("""
This app uses the `dpt-moses-ver2` model from Hugging Face to generate drug-like molecules.
You can input a SMILES string in the sidebar, and upon submission, the app will generate 20 similar SMILES strings,
display them, and provide their similarity to the original input. You can also download the generated SMILES strings.
""")

if submit_button and input_smiles:
    # 生成 20 個 smiles
    generated_smiles = generate_smiles(input_smiles)

    # 顯示 smiles 和相似度
    st.write("Generated SMILES:")
    for i, smi in enumerate(generated_smiles):
        similarity = calculate_similarity(input_smiles, smi)
        st.write(f"{i + 1}. {smi} (Similarity: {similarity})")

    # 畫分子圖
    st.image(draw_molecules(generated_smiles),
             caption="Generated Molecules", use_column_width=True)

    # 下載 smiles 按鈕
    smiles_str = "\n".join(generated_smiles)
    st.download_button(
        label="Download SMILES",
        data=smiles_str,
        file_name='generated_smiles.txt',
        mime='text/plain',
    )