#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
from rdkit import Chem

# 当前目录路径
current_dir = os.path.dirname(os.path.abspath(__file__))

# 输入SDF文件路径
sdf_file = os.path.join(current_dir, 'fgbar_vina_SP_1_pv.sdf')

# 输出SMILES文件路径
smiles_file = os.path.join(current_dir, 'molecules.txt')

def sdf_to_smiles(sdf_path, output_path):
    """读取SDF文件中的分子并将其转换为SMILES格式保存到文本文件"""
    # 读取SDF文件
    suppl = Chem.SDMolSupplier(sdf_path)
    
    # 计数有效分子数量
    valid_mol_count = 0
    
    # 打开输出文件
    with open(output_path, 'w') as f:
        # 遍历所有分子
        for i, mol in enumerate(suppl):
            if mol is not None:  # 确保分子有效
                # 获取SMILES
                smiles = Chem.MolToSmiles(mol)
                # 写入文件
                f.write(f"{smiles}\n")
                valid_mol_count += 1
    
    return valid_mol_count

def main():
    print(f"正在读取SDF文件: {sdf_file}")
    mol_count = sdf_to_smiles(sdf_file, smiles_file)
    print(f"成功处理了 {mol_count} 个分子")
    print(f"SMILES已保存到: {smiles_file}")

if __name__ == "__main__":
    main()