# -*- coding: utf-8 -*-
"""
Created on Sun Jul 23 08:12:43 2023

@author: Sen
"""
#%%
import os 
import argparse
import shutil
import logging
from openbabel import openbabel
parser = argparse.ArgumentParser()
parser.add_argument('-d', type=str, default=None, help='Input the dirpath')
args = parser.parse_args()

dirpath = args.d
if dirpath[-1] != '/' :
    dirpath = dirpath + '/'
#%%
def create_directory(dir_name):
    # 使用os.path.exists()检查目录是否存在
    if not os.path.exists(dir_name):
        # 如果目录不存在，使用os.makedirs()创建它
        os.makedirs(dir_name)

input_dirpath = dirpath  #文件夹以/结尾
dir_name = os.path.basename(os.path.dirname(input_dirpath))
dir_path = os.path.dirname(os.path.dirname(input_dirpath))
output_dirpath = os.path.join(dir_path,dir_name+'_min')
create_directory(output_dirpath)
logging.basicConfig(level=logging.CRITICAL)
openbabel.obErrorLog.StopLogging()
#%%
def sdf_min(input_sdf, output_sdf):
    # 创建一个分子对象
    mol = openbabel.OBMol()

    # 创建转换器，用于文件读写
    conv = openbabel.OBConversion()
    conv.SetInAndOutFormats("sdf", "sdf")

    # 从SDF文件中读取分子
    conv.ReadFile(mol, input_sdf)

    # 创建力场对象，使用MMFF94力场
    forcefield = openbabel.OBForceField.FindForceField("MMFF94")

    # 为分子设置力场
    success = forcefield.Setup(mol)
    if not success:
        raise Exception("Error setting up force field")

    # 进行能量最小化
    forcefield.SteepestDescent(10000)  # 原来是5000步最速下降法
    forcefield.GetCoordinates(mol)  # 将能量最小化后的坐标保存到分子对象

    # 将能量最小化后的分子写入到SDF文件
    conv.WriteFile(mol, output_sdf)
    
#%% 
from tqdm import tqdm
from concurrent.futures import ProcessPoolExecutor
import multiprocessing

def handle_file(filename):
    if '.sdf' == filename[-4:]:
        input_sdf_file = os.path.join(input_dirpath, filename)
        output_sdf_file = os.path.join(output_dirpath, filename)
        try:
            sdf_min(input_sdf_file, output_sdf_file)
        except Exception as e:
            print(f"An error occurred while processing the file '{input_sdf_file}': {e}")
            try:
                os.remove(output_sdf_file)
                print(output_sdf_file + ' was successfully removed')
            except Exception:
                if os.path.exists(output_sdf_file):
                    print('please remove '+output_sdf_file)
    else:
        src_file = os.path.join(input_dirpath, filename)
        dst_file = os.path.join(output_dirpath, filename)
        shutil.copy(src_file, dst_file)

def main():
    file_list = os.listdir(input_dirpath)
    # 获取系统的 CPU 核心数
    num_cores = multiprocessing.cpu_count()

    # 创建一个进程池
    with ProcessPoolExecutor(max_workers=num_cores) as executor:
        # 使用 tqdm 提供进度条功能
        list(tqdm(executor.map(handle_file, file_list), total=len(file_list)))



#%%
import pandas as pd
import os
import hashlib
class dir_check():
    def __init__(self,dirpath):
        self.dirpath = dirpath
        self.mapping_file = os.path.join(self.dirpath, 'hash_ligand_mapping.csv')
        self.mapping_data = pd.read_csv(self.mapping_file, header=None)
        self.smiles_list = self.mapping_data.iloc[:, 1].tolist()
        self.hash_list = self.mapping_data.iloc[:, 0].tolist()
        self.filename_list = os.listdir(self.dirpath)
        self.sdf_filename_list = [x for x in self.filename_list if x[-4:] == '.sdf']
        
    def mapping_file_check(self):
        for i in range(len(self.smiles_list)):
            if hashlib.sha1(self.smiles_list[i].encode()).hexdigest() != self.hash_list[i]:
                print('error in mapping file')
        print('mapping file check completed')
        
    def dir_file_check(self):
        filename_hash_list = [x[:-4] for x in self.sdf_filename_list]
        for filename_hash in filename_hash_list:
            if filename_hash not in self.hash_list:
                filename = filename_hash+'.sdf'
                print(filename + ' not in mapping file')
                os.remove(os.path.join(self.dirpath,filename))
                print('remove ' + os.path.join(self.dirpath,filename))
#%%      

if __name__ == '__main__':
    main()
  
    check = dir_check(output_dirpath)
    check.mapping_file_check()
    check.dir_file_check()