from rdkit import Chem
from rdkit.Chem import AllChem, Draw
from rdkit.Chem import rdFMCS

# 샘플 화합물 리스트 (SMILES 형식)
smiles_list = [
    "CCO",  # 에탄올
    "CCN",  # 에틸아민
    "CC(O)C",  # 이소프로판올
    "CCCC",  # 부탄
    "CC(C)C"  # 이소부탄
]

# SMILES 문자열을 RDKit Mol 객체로 변환
molecules = [Chem.MolFromSmiles(smiles) for smiles in smiles_list]

# 검색하고자 하는 서브 구조 (SMILES 형식)
query_smiles = "CC"
query_mol = Chem.MolFromSmiles(query_smiles)

# 서브 구조 검색
matches = []
for mol in molecules:
    if mol.HasSubstructMatch(query_mol):
        matches.append(mol)

# 검색 결과 출력
print(f"서브 구조 '{query_smiles}'를 포함하는 화합물 수: {len(matches)}")
for i, match in enumerate(matches):
    smiles = Chem.MolToSmiles(match)
    print(f"{i+1}: {smiles}")

# 검색된 화합물 그림 출력
img = Draw.MolsToGridImage(matches, molsPerRow=3, subImgSize=(200, 200))
img.show()