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

# In[20]:


'''
自动将三种方法的预测结果进行初步整理，生成一个包含多个sheet的summary.xlsx文件，sheet中包括：AI，docking方法根据打分的前100个结果,这些结果中两方法的共同靶标（sheet名带common的），
相似性的结果。
运行方法：
文件夹下需要放以下内容：该脚本（windows运行需要修改folder_path参数为.\），以化合物名字命名的子文件夹，子文件夹内包含learning，docking，similarity（可选）的xlsx文件，文件名最好使用网站
自动生成的，文件名中必须包含方法的名字，否则无法识别。
'''

import os
import pandas as pd
pd.io.formats.excel.header_style = None
# 路径符
folder_path ='./'
num_rows = 100


# 遍历文件夹下的所有子文件夹
for subdir, dirs, files in os.walk(folder_path):
    uniprot_ids_docking, uniprot_ids_deeplearning = [], []
    df_summary1 = pd.DataFrame()
    df_summary2 = pd.DataFrame()
    df_similarity = []
    for file in files:
        file_path = os.path.join(subdir, file)
        # 如果文件名包含'docking'字符
        if 'docking' in file or 'Docking' in file:
            # 读取前记录
            if 'csv' in file_path:
                df_docking = pd.read_csv(file_path,nrows=num_rows)
            else:
                df_docking = pd.read_excel(file_path,nrows=num_rows)
            # 获取uniprot id列的值
            uniprot_ids_docking = df_docking['uniprot id'].tolist()
        # 如果文件名包含'deeplearning'字符
        elif 'learning' in file or 'Learning' in file:
            # 读取记录
            if 'csv' in file_path:
                df_deeplearning = pd.read_csv(file_path,nrows=num_rows)
            else:
                df_deeplearning = pd.read_excel(file_path,nrows=num_rows)
            # 获取uniprot ID列的值
            uniprot_ids_deeplearning = df_deeplearning['Uniprot ID'].tolist()
        elif 'similarity' in file or 'Similarity' in file:
            if 'csv' in file_path:
                df_similarity = pd.read_csv(file_path,nrows=num_rows)
            else:
                df_similarity = pd.read_excel(file_path,nrows=num_rows)
        
    if uniprot_ids_docking and uniprot_ids_deeplearning:
        # 找到相同的uniprot ID并输出
        common_ids = set(uniprot_ids_docking).intersection(set(uniprot_ids_deeplearning))
        #print(subdir,':\n',common_ids)
        output_path = os.path.join(subdir, 'summary.xlsx')
        for id_ in common_ids:
            df_summary1 = pd.concat([df_summary1, df_deeplearning[df_deeplearning['Uniprot ID'] == id_]])
            df_summary2 = pd.concat([df_summary2, df_docking[df_docking['uniprot id'] == id_]])
            
        
            
        writer = pd.ExcelWriter(output_path)
        df_summary1.to_excel(writer, sheet_name='common_AI', index=False)
        df_summary2.to_excel(writer, sheet_name='common_docking', index=False)
        df_deeplearning.to_excel(writer, sheet_name='AI_100', index=False)
        df_docking.to_excel(writer, sheet_name='docking_100', index=False)
        if df_similarity:
            df_similarity.to_excel(writer, sheet_name='ligandsimilarity', index=False)
        writer.save()