import os
import pandas as pd
import csv
import oddt
from oddt.scoring.functions import RFScore
oddt.toolkit = oddt.toolkits.rdk

# create RFScore model from trained model
# rf = RFScore.load('/path/to/your/model.pkl')

rf = RFScore.rfscore(version=1, n_jobs=-1)

with open('result_csv.csv', 'w', newline='') as csvfile:
    fieldnames = ['pdb_id', 'rf_score_v1']
    writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
    writer.writeheader()

    for root, dirs, _ in os.walk('./coreset'):
        for dir in dirs:
            pdb_id = dir 
            ligands_file = os.path.join(root, dir, f'{pdb_id}_ligand.sdf')
            proteins_file = os.path.join(root, dir, f'{pdb_id}_pocket.pdb')
            lig_mol = next(oddt.toolkit.readfile('sdf', ligands_file))
            prot_mol = next(oddt.toolkit.readfile('pdb', proteins_file))
            score = rf.predict(lig_mol, prot_mol)
            print(f'{pdb_id}_ligand.sdf', score)