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


# In[24]:


'''
统计多个化合物在docking方法中出现次数超过targets_num的共同靶标，输出txt文件。
运行方法：
文件夹下需要放以下内容：该脚本（windows运行需要修改folder_path参数为.\，同时指定targets_num的值），以化合物名字命名的子文件夹，子文件夹内包含learning，docking，similarity（可选）的xlsx文件，文件名最好使用网站
自动生成的，文件名中必须包含方法的名字，否则无法识别。
'''

import os
import glob
import pandas as pd

#共同靶标数量
targets_num = 4
folder_path ='./'

# 存储蛋白质名字的字典
protein_counts = {}

# 遍历文件夹及其子文件夹，查找包含“docking”的xlsx文件
for foldername, subfolders, filenames in os.walk(folder_path):
    for filename in filenames:
        if 'docking' in filename or 'Docking' in filename and filename.endswith('.xlsx'):
            # 读取xlsx文件中的数据
            file_path = os.path.join(foldername, filename)
            df = pd.read_excel(file_path, nrows=100)
            
            # 统计不同蛋白质的出现次数
            for protein_name in df['uniprot id']:
                if protein_name in protein_counts:
                    protein_counts[protein_name]['count'] += 1
                    protein_counts[protein_name]['files'].append((foldername, filename))
                else:
                    protein_counts[protein_name] = {'count': 1, 'files': [(foldername, filename)]}
                    
# 输出出现次数超过targets_num的蛋白质名字、行和文件信息到txt文件
with open(f'common_{targets_num}targets.txt', 'w') as f:
    for protein_name, info in protein_counts.items():
        if info['count'] > targets_num:
            f.write(f"Protein Name: {protein_name}\n")
            for foldername, filename in info['files']:
                f.write(f"    Chemical Name: {foldername}\n")
            f.write('\n')