#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@file               :optimize_sdf.py
@Description:       : 使用rdkit 将2d转3d，增加基本的坐标优化功能
@Date               :2024/09/07 09:34:41
@Author             :lyzeng
@Email              :pylyzeng@gmail.com
@version            :1.0
'''



import os
import click
from rdkit import Chem
from rdkit.Chem import AllChem

def generate_3d_coordinates(input_sdf: str, output_sdf: str):
    # 读取SDF文件
    suppl = Chem.SDMolSupplier(input_sdf)
    writer = Chem.SDWriter(output_sdf)

    for mol in suppl:
        if mol is None:
            continue

        # 生成3D坐标
        mol = Chem.AddHs(mol)  # 添加氢原子
        if AllChem.EmbedMolecule(mol) == 0:  # 生成3D坐标
            AllChem.UFFOptimizeMolecule(mol)  # 优化分子结构
            writer.write(mol)  # 将分子写入输出文件

    writer.close()

@click.command()
@click.argument('input_sdf', type=click.File('r'))
@click.option('--output_dir', type=click.Path(), default='.', help='输出目录，默认为当前目录')
def optimize_sdf(input_sdf, output_dir):
    """优化输入SDF文件中的3D坐标并输出到指定目录"""
    # 获取输入文件名并构造输出文件路径
    input_filename = os.path.basename(input_sdf.name)
    output_sdf = os.path.join(output_dir, f"optimized_{input_filename}")

    # 创建输出目录（如果不存在）
    os.makedirs(output_dir, exist_ok=True)

    # 生成3D结构并保存
    generate_3d_coordinates(input_sdf.name, output_sdf)

    click.echo(f"优化的SDF文件已保存到: {output_sdf}")

if __name__ == '__main__':
    optimize_sdf()

# python optimize_sdf.py /mnt/c/project/pymol/data/Structure2D_COMPOUND_CID_5280440.sdf --output_dir /mnt/c/project/pymol/data/output
# python optimize_sdf.py --help