import zipfile
import os
import sys
import pymongo
from optparse import OptionParser
from bson.objectid import ObjectId
import numpy as np
from bson.binary import Binary
import json
import pickle
import logging
import re
import Bio
import gridfs
import glob
import pandas as pd
from io import StringIO
from datetime import datetime, timedelta, timezone
from Parser_functions import *
from env_info import environment_var, UNZIP_FOLDER, UPLOAD_FOLDER, OUTPUT_FOLDER, STATIC_FOLDER, get_db
from Write_functions import clean_pdb, format_longer_pdb, single_user_write_sitecontext_result_file
logging.basicConfig(filename="siteContext.log", level=logging.DEBUG)
from Histogram_functions import compare_sites
from Features_functions import get_features_charges
from LRFPCA import LRFEstimatorPCA

def unzip_folder(input_file):
    zip_ref = zipfile.ZipFile(input_file, "r")
    extracted = zip_ref.namelist()
    zip_ref.extractall(UNZIP_FOLDER)
    zip_ref.close()
    path_name = extracted[0]
    if ".DS_Store" in path_name:
        path_name = path_name.split("/")[0]
    return path_name

def get_shape_feature_result(output_folder):
    shape, shape_feature = 0, 0
    shape_radius, shape_feature_radius = None, None
    sinkhorn_score_shape, sinkhorn_score_shape_feature = None, None
    final_RMSD_shape, final_RMSD_shape_feature = None, None
    scores_dict = {}
    for f in glob.glob(os.path.join(output_folder, '*.html')):
        basename = os.path.basename(f)
        if basename == '3Dmol_viewer_steric(shape).html':
            shape = 1
        elif basename == '3Dmol_viewer_steric(shape)_and_charges.html':
            shape_feature = 1
    for f in glob.glob(os.path.join(output_folder, '*.csv')):
        basename = os.path.basename(f)
        if basename == 'similarity_scores_using_steric(shape).csv':
            df = pd.read_csv(f)
            sinkhorn_score_shape = round(df.iloc[0, 1].astype(float), 2)
            shape_radius = round(df.iloc[0, 2].astype(int), 2)
            final_RMSD_shape = round(df.iloc[0, 4].astype(float), 2)
        elif basename == 'similarity_scores_using_steric(shape)_and_charges.csv':
            df = pd.read_csv(f)
            sinkhorn_score_shape_feature = round(df.iloc[0, 1].astype(float), 2)
            shape_feature_radius = round(df.iloc[0, 2].astype(int), 2)
            final_RMSD_shape_feature = round(df.iloc[0, 4].astype(float), 2)
    return shape, shape_radius, sinkhorn_score_shape, final_RMSD_shape, shape_feature, shape_feature_radius,sinkhorn_score_shape_feature, final_RMSD_shape_feature

def fileHandler(project_id):
    print("In fileHandler")
    db = get_db()
    project = db.SiteContext_projects.find_one({"_id": ObjectId(project_id)})
    if not project:
        print("Project is not found in DB.")
        return

    # update to your local environment_var
    if environment_var == "DEV":
        os.system(
            "python site_context_compute.py -p "
            + str(project_id)
        )
    else:
        os.system(
            "python site_context_compute.py -p "
            + str(project_id)
        )

    # update 
    db["SiteContext_projects"].find_one_and_update(
        {"_id": ObjectId(project_id)}, {"$set": {"job_status": "1", "expires_at": datetime.now(timezone.utc) + timedelta(hours=24),}}
    )

def insert_result_to_db(collection_molecule_score, results_dict, f, radius,
                        SASA_site_A_name, SASA_site_B_name, 
                        SASA_site_A_indexes, SASA_site_B_indexes,):
    q_idx = np.array(SASA_site_A_indexes)[results_dict['row_idx']]
    m_idx = np.array(SASA_site_B_indexes)[results_dict['col_idx']]
    un_m_idx = np.array(SASA_site_B_indexes)[results_dict['un_col_idx']]
    
    insert_fields = {
        "key": f"{SASA_site_A_name}#{SASA_site_B_name}#{f}", 
        "cost": results_dict['cost'],
        "cost_normalized": results_dict['cost_normalized'],
        "radius": radius,
        "length": results_dict['length'],
        "entireSite_row_idx": json.dumps(q_idx.tolist()),
        "entireSite_col_idx": json.dumps(m_idx.tolist()),
        "entireSite_un_col_idx": json.dumps(un_m_idx.tolist()),
        "row_idx": json.dumps(np.array(results_dict['row_idx']).tolist()),
        "col_idx": json.dumps(np.array(results_dict['col_idx']).tolist()),
        "un_col_idx": json.dumps(np.array(results_dict['un_col_idx']).tolist()),
        "corr_10": json.dumps(results_dict['corr_10'].tolist()),
        "corr_25": json.dumps(results_dict['corr_25'].tolist()),
        "corr_50": json.dumps(results_dict['corr_50'].tolist()), 
        "RMSD": results_dict['RMSD'], 
        "RMSD_normalized": results_dict['RMSD_normalized'],
        "R": json.dumps(results_dict['R'].tolist()),
        't': json.dumps(results_dict['t'].tolist()),
        'switch': results_dict['switch'],
        'rotate': results_dict['rotate'],
        'R_0': json.dumps(results_dict['R_0']),
        't_0': json.dumps(results_dict['t_0']),
    }
    collection_molecule_score.insert_one(insert_fields)


def get_SASA(pdb_name, pdb_path):
    # Load structures
    parser = Bio.PDB.PDBParser(QUIET=True)
    io = Bio.PDB.PDBIO()
    pdb_data = parser.get_structure(pdb_name, pdb_path)
    # get pdb_data_string
    io.set_structure(pdb_data)
    string_io = StringIO() # Create a string buffer
    io.save(string_io)  # Write the structure to the string buffer
    pdb_data_string = string_io.getvalue() # Get the string from the buffer
    string_io.close()
    # get SASA information
    binding_site_coords, SASA_atoms, SASA_site_coords, SASA_site_atoms_indexes_from_whole, SASA_site_atoms_indexes_from_site = get_SASA_biomolecule(pdb_name, pdb_data)
    SASA_features = get_features_charges(SASA_site_coords, SASA_atoms)
    
    return pdb_data_string, binding_site_coords, SASA_site_coords, SASA_site_atoms_indexes_from_whole, SASA_site_atoms_indexes_from_site, SASA_features


def apply_rotation_to_pdb_strings(pdb_A_string, pdb_B_string, switch, rotate, R, t):
    # Apply a pre-computed rotation/translation to the appropriate pdb string(s).
    # R and t may be lists or numpy arrays; ensure they're numpy arrays.
    try:
        R_arr = np.array(R)
        t_arr = np.array(t)
    except Exception:
        return pdb_A_string, pdb_B_string

    if switch:
        if rotate == 1:
            pdb_B_string = format_longer_pdb(pdb_B_string, R_arr, t_arr)
        elif rotate == 2:
            pdb_A_string = format_longer_pdb(pdb_A_string, R_arr, t_arr)
    else:
        if rotate == 1:
            pdb_A_string = format_longer_pdb(pdb_A_string, R_arr, t_arr)
        elif rotate == 2:
            pdb_B_string = format_longer_pdb(pdb_B_string, R_arr, t_arr)

    return pdb_A_string, pdb_B_string

def site_comparsion(SASA_site_A_indexes, SASA_site_B_indexes, 
                    SASA_coords_A, SASA_coords_B, SASA_lrf_A, SASA_lrf_B, downsampling, radius,
                    features_A, features_B):
    
    results_dict = molecules_score_s_f_n_single(SASA_coords_A, SASA_coords_B, SASA_lrf_A, SASA_lrf_B,
                                downsampling, radius, features_A,features_B)
    q_idx = np.array(SASA_site_A_indexes)[results_dict['row_idx']]
    m_idx = np.array(SASA_site_B_indexes)[results_dict['col_idx']]
    un_m_idx = np.array(SASA_site_B_indexes)[results_dict['un_col_idx']]
    results_dict["entireSite_row_idx"] = q_idx.tolist()
    results_dict["entireSite_col_idx"] = m_idx.tolist()
    results_dict["entireSite_un_col_idx"] = un_m_idx.tolist()
    return results_dict

def run_demo(output_folder, project):
    # load variables from project
    
    pdb_set1_single = json.loads(project['set1_molecules_name'])[0]
    pdb_set2_single = json.loads(project['set2_molecules_name'])[0]
    feature_type_list = json.loads(project['feature_type'])
    # Load structures
    pdb_path_1 = os.path.join(STATIC_FOLDER, "pdb_files", f'{pdb_set1_single}.pdb')
    pdb_path_2 = os.path.join(STATIC_FOLDER, "pdb_files", f'{pdb_set2_single}.pdb')
    pdb_1_string, pdb_1_binding_coords, pdb_set1_single_SASA_coord, pdb_set1_single_SASA_indexes_from_whole, pdb_set1_single_SASA_indexes_from_site, pdb_set1_single_SASA_features = get_SASA(pdb_set1_single, pdb_path_1)
    pdb_2_string, pdb_2_binding_coords, pdb_set2_single_SASA_coord, pdb_set2_single_SASA_indexes_from_whole, pdb_set2_single_SASA_indexes_from_site, pdb_set2_single_SASA_features = get_SASA(pdb_set2_single, pdb_path_2)
    if len(pdb_1_binding_coords) <= len(pdb_2_binding_coords):
        aligned_coords_2_site, RMSD, R, t = get_downsampled_attributes_parallel(pdb_1_binding_coords, pdb_2_binding_coords)
        pdb_2_binding_coords = aligned_coords_2_site
        aligned_coords = align_coordinates(R, t, pdb_set2_single_SASA_coord)
        pdb_set2_single_SASA_coord = aligned_coords
        pdb_2_string = format_longer_pdb(pdb_2_string, R, t)   
    else:
        aligned_coords_1_site, RMSD, R, t = get_downsampled_attributes_parallel(pdb_2_binding_coords, pdb_1_binding_coords)
        pdb_1_binding_coords = aligned_coords_1_site
        aligned_coords = align_coordinates(R, t, pdb_set1_single_SASA_coord)
        pdb_set1_single_SASA_coord = aligned_coords
        pdb_1_string = format_longer_pdb(pdb_1_string, R, t)   
    
    radius = get_dynamic_radius(pdb_set1_single_SASA_coord, pdb_set2_single_SASA_coord)
    max_radius = radius
    lrf_1, lrf_2 = get_lrf(pdb_1_binding_coords, pdb_set1_single_SASA_indexes_from_site, pdb_2_binding_coords, pdb_set2_single_SASA_indexes_from_site, 10, max_radius, True)
    
    SASA_coords_A, SASA_coords_B, downsampling, switch = get_coordinate_lengths(pdb_set1_single_SASA_coord, pdb_set2_single_SASA_coord)
    (SASA_site_A_name, SASA_site_B_name) = (pdb_set2_single, pdb_set1_single) if switch else (pdb_set1_single, pdb_set2_single)
    (SASA_site_A_indexes, SASA_site_B_indexes) =  (pdb_set2_single_SASA_indexes_from_whole, pdb_set1_single_SASA_indexes_from_whole) if switch else (pdb_set1_single_SASA_indexes_from_whole, pdb_set2_single_SASA_indexes_from_whole)
    (SASA_site_A_features, SASA_site_B_features) =  (pdb_set2_single_SASA_features, pdb_set1_single_SASA_features) if switch else (pdb_set1_single_SASA_features, pdb_set2_single_SASA_features)
    (pdb_data_string_A, pdb_data_string_B) = (pdb_2_string, pdb_1_string) if switch else (pdb_1_string, pdb_2_string)
    (SASA_site_A_lrf, SASA_site_B_lrf) = (lrf_2, lrf_1) if switch else (lrf_1, lrf_2)
    
    molecules_results = {}
    for f in feature_type_list:
        key = f"{SASA_site_A_name}#{SASA_site_B_name}#{f}"
        if f == 'shape':
            results_dict = site_comparsion(SASA_site_A_indexes, SASA_site_B_indexes, SASA_coords_A, SASA_coords_B, 
                                            SASA_site_A_lrf, SASA_site_B_lrf, downsampling, max_radius, features_A=None, features_B=None)
        else:
            results_dict = site_comparsion(SASA_site_A_indexes, SASA_site_B_indexes, SASA_coords_A, SASA_coords_B, SASA_site_A_lrf, SASA_site_B_lrf,
                                            downsampling, max_radius, features_A=SASA_site_A_features, features_B=SASA_site_B_features)
        molecules_results[key] = results_dict

    for k1, v1 in molecules_results.items():
        # get pdb string for both A and B
        SASA_site_A_name, SASA_site_B_name, f = k1.split('#')
        assert(f in ['shape', 'shape+feature'])
        k = SASA_site_A_name + '_and_' + SASA_site_B_name
        pdb_data_string_B_updated = format_longer_pdb(pdb_data_string_B, np.array(v1['R']), np.array(v1['t']))      
        if f == 'shape':
            f_text = 'steric(shape)'
        else:
            f_text = 'steric(shape)_and_charges'     
        distance = {k: {"cost": v1['cost_normalized'], "radius": max_radius+1,  "feature": f_text, 'RMSD': v1['RMSD']}}

        correspondence_list = []
        row_list, col_list = v1['entireSite_row_idx'], v1['entireSite_col_idx']
        corr_10, corr_25, corr_50 = v1['corr_10'], v1['corr_25'], v1['corr_50']
        for row_idx, col_idx in zip(row_list, col_list):
            correspondence_list.append([row_idx, col_idx])
        
        single_user_write_sitecontext_result_file(distance, SASA_site_A_name, SASA_site_B_name, f_text,
                                                pdb_data_string_A, pdb_data_string_B_updated, 
                                                correspondence_list, corr_10, corr_25, corr_50, v1['entireSite_un_col_idx'],
                                                output_folder)
            


def insert_data_mongodb(folder, project):
    # insert data to mongodb database
    print("IN insert data mongodb")
    db = get_db()
    collection = db["molecules_SASA"]
    collection_name = "molecules_PDB"
    fs = gridfs.GridFS(db, collection=collection_name)

    exception_list = []
    for sub_folder in os.listdir(folder):
        if sub_folder[0] == ".":
            continue
        if sub_folder == "set1" or sub_folder == "set2":
            pdb_files = [f for f in glob.glob(folder + "/" + sub_folder + "/*.pdb")]
            for pdb_path in pdb_files:
                pdb_name = os.path.splitext(os.path.basename(pdb_path))[0]
                if (
                    db["molecules_PDB.files"].count_documents({"filename": pdb_name})
                    == 0
                ):
                    parser = Bio.PDB.PDBParser(QUIET=True)
                    io = Bio.PDB.PDBIO()
                    # Load structures
                    pdb_data = parser.get_structure(pdb_name, pdb_path)
                    io.set_structure(pdb_data)
                    string_io = StringIO() # Create a string buffer
                    io.save(string_io)  # Write the structure to the string buffer
                    pdb_data_string = string_io.getvalue() # Get the string from the buffer
                    string_io.close()
                    fs.put(pdb_data_string.encode('utf-8'), filename=pdb_name)
                    print("Insert one %s molecules_PDB" % pdb_name)
                else:
                    pdb_data_string = fs.find_one({"filename": pdb_name}).read().decode('utf-8')
                    parser = Bio.PDB.PDBParser(QUIET=True)
                    string_io = StringIO(pdb_data_string)
                    pdb_data = parser.get_structure(pdb_name, string_io)
                    string_io.close()
                    print("Load one %s molecules_PDB" % pdb_name)

                if collection.count_documents({"PDB_id": pdb_name}) == 0:
                    
                    try:
                        # extract sites and features
                        binding_site_coords, SASA_atoms, SASA_site_coords, SASA_site_atoms_indexes_from_whole, SASA_site_atoms_indexes_from_site = get_SASA_biomolecule(pdb_name, pdb_data)
                        SASA_features = get_features_charges(SASA_site_coords, SASA_atoms)
                        # convert to pickle
                        binding_pickle = pickle.dumps(np.array(binding_site_coords))
                        features_pickle = pickle.dumps(np.array(SASA_features))
                        coords_pickle = pickle.dumps(SASA_site_coords)
                        indexes_from_whole_pickle = pickle.dumps(SASA_site_atoms_indexes_from_whole)
                        indexes_from_site_pickle = pickle.dumps(SASA_site_atoms_indexes_from_site)
                        
                        collection.insert_one(
                            {
                                "PDB_id": pdb_name,
                                "Binding_coords_0": Binary(binding_pickle, subtype=128),
                                "SASA_coords_0": Binary(coords_pickle, subtype=128),
                                "SASA_indexes_from_whole_0": Binary(indexes_from_whole_pickle, subtype=128), 
                                "SASA_indexes_from_site_0": Binary(indexes_from_site_pickle, subtype=128),
                                "SASA_features_0": Binary(features_pickle, subtype=128),
                            }
                        )
                        print("Insert one %s molecules_SASA" % pdb_name)
                    except Exception as e:
                        exception_list.append(f'The error is {pdb_name} in {sub_folder}, {str(e)}')

    return exception_list


def run_batch_molecules(project_id):

    print("IN BATCH MOLECULE")
    db = get_db()
    project = db.SiteContext_projects.find_one({"_id": ObjectId(project_id)})
    assert(project['job_type'] == 'BATCH')
    
    # input_folder
    input_folder = os.path.join(UNZIP_FOLDER, project['file']).replace(".zip", "")
    # output_folder
    output_folder = os.path.join(OUTPUT_FOLDER, str(project["_id"]))
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    # save data into mongodb
    response = insert_data_mongodb(input_folder, project)
    need_return = 0
    need_remove = 0
    molecules_need_to_remove = []
    if response != []:
        with open(os.path.join(output_folder, 'Error_messages.txt'), 'w') as f:
            f.write("Certain pairs could not be compared due to the following:\n")
            for r in response:
                pdb_name_with_sub_folder, error = r.split(',')
                match = re.match(r"The error is (.+?) in (.+)", pdb_name_with_sub_folder)
                if match:
                    pdb_name = match.group(1)
                    sub_folder = match.group(2)
                else:
                    pdb_name, sub_folder = None, None
                if sub_folder == 'set1':
                    need_return = 1
                if sub_folder == 'set2':
                    molecules_need_to_remove.append(pdb_name.strip())
                    need_remove = 1
                f.write(f"{r}\n")
    if need_return:
        return
    if need_remove:
        PDB_set2_names = json.loads(project['set2_molecules_name'])
        #assert(type(PDB_molecules_name), list)
        PDB_set2_names_new = [f for f in PDB_set2_names if f not in molecules_need_to_remove]
        assert(len(PDB_set2_names_new) == (len(PDB_set2_names) - len(molecules_need_to_remove)))
        project["set2_molecules_name"] = json.dumps(PDB_set2_names_new)
    
    molecules_results = molecules_score(project)
    # Get PDB
    collection_PDB = "molecules_PDB"
    fs = gridfs.GridFS(db, collection=collection_PDB)
    parser = Bio.PDB.PDBParser(QUIET=True)
    df_list = []
    for k, v in molecules_results.items():
        output_folder_update = os.path.join(output_folder, k)
        # get pdb string for both A and B
        pdb_data_string_A = None
        pdb_data_string_B = None
        for k1, v1 in v.items():
            
            SASA_site_A_name, SASA_site_B_name, f = k1.split('#')
            switch = v1.get('switch')
            rotate = v1.get('rotate')
            R_0 = json.loads(v1.get('R_0'))
            t_0 = json.loads(v1.get('t_0'))
            radius = v1.get('radius')
            assert(f in ['shape', 'shape+feature'])

            if pdb_data_string_A is None:
                pdb_data_string_A = fs.find_one({"filename": SASA_site_A_name}).read().decode('utf-8')
            if pdb_data_string_B is None:
                pdb_data_string_B = fs.find_one({"filename": SASA_site_B_name}).read().decode('utf-8')

            # Apply any pre-computed rotation/translation (R_0, t_0) to the stored pdb strings.
            pdb_data_string_A_updated, pdb_data_string_B_updated = apply_rotation_to_pdb_strings(
                pdb_data_string_A, pdb_data_string_B, switch, rotate, R_0, t_0
            )

            pdb_data_string_B_updated = format_longer_pdb(pdb_data_string_B_updated, json.loads(v1['R']), json.loads(v1['t']))
            
            if f == 'shape':
                f_text = 'steric(shape)'
            else:
                f_text = 'steric(shape)_and_charges'     
            distance = {k: {"cost": v1['cost_normalized'], "radius": radius+1, "feature": f_text, 'RMSD': v1['RMSD']}}
            df_list.append(pd.DataFrame.from_dict(distance, orient="index"))
            correspondence_list = []
            row_list, col_list = json.loads(v1['entireSite_row_idx']), json.loads(v1['entireSite_col_idx'])
            corr_10, corr_25, corr_50 = json.loads(v1['corr_10']), json.loads(v1['corr_25']), json.loads(v1['corr_50'])
            for row_idx, col_idx in zip(row_list, col_list):
                correspondence_list.append([row_idx, col_idx])
            
            single_user_write_sitecontext_result_file(distance, SASA_site_A_name, SASA_site_B_name, f_text,
                                                    pdb_data_string_A_updated, pdb_data_string_B_updated, 
                                                    correspondence_list, corr_10, corr_25, corr_50, json.loads(v1['entireSite_un_col_idx']),
                                                    output_folder_update)
            
    df_all = pd.concat(df_list)
    df_all.to_csv(os.path.join(output_folder, 'similarity_scores_all.csv'))
    return


def run_single_molecules(output_folder, project):

    print("IN SINGLE MOLECULE")
    molecules_results = molecules_score(project)
    
    db = get_db()
    collection_PDB = "molecules_PDB"
    fs = gridfs.GridFS(db, collection=collection_PDB)
    parser = Bio.PDB.PDBParser(QUIET=True)
    
    for k, v in molecules_results.items():
        # get pdb string for both A and B
        pdb_data_string_A = None
        pdb_data_string_B = None
        for k1, v1 in v.items():
            SASA_site_A_name, SASA_site_B_name, f = k1.split('#')
            switch = v1.get('switch')
            rotate = v1.get('rotate')
            R_0 = json.loads(v1.get('R_0'))
            t_0 = json.loads(v1.get('t_0'))
            radius = v1.get('radius')
            assert(f in ['shape', 'shape+feature'])

            if pdb_data_string_A is None:
                pdb_data_string_A = fs.find_one({"filename": SASA_site_A_name}).read().decode('utf-8')
            if pdb_data_string_B is None:
                pdb_data_string_B = fs.find_one({"filename": SASA_site_B_name}).read().decode('utf-8')
            pdb_data_string_A_updated, pdb_data_string_B_updated = apply_rotation_to_pdb_strings(
                pdb_data_string_A, pdb_data_string_B, switch, rotate, R_0, t_0)

            pdb_data_string_B_updated = format_longer_pdb(pdb_data_string_B_updated, json.loads(v1['R']), json.loads(v1['t']))
            
            if f == 'shape':
                f_text = 'steric(shape)'
            else:
                f_text = 'steric(shape)_and_charges'     
            distance = {k: {"cost": v1['cost_normalized'], "radius": radius+1, "feature": f_text, 'RMSD': v1['RMSD']}}
            correspondence_list = []
            row_list, col_list = json.loads(v1['entireSite_row_idx']), json.loads(v1['entireSite_col_idx'])
            corr_10, corr_25, corr_50 = json.loads(v1['corr_10']), json.loads(v1['corr_25']), json.loads(v1['corr_50'])
            for row_idx, col_idx in zip(row_list, col_list):
                correspondence_list.append([row_idx, col_idx])
            
            single_user_write_sitecontext_result_file(distance, SASA_site_A_name, SASA_site_B_name, f_text,
                                                    pdb_data_string_A_updated, pdb_data_string_B_updated, 
                                                    correspondence_list, corr_10, corr_25, corr_50, json.loads(v1['entireSite_un_col_idx']),
                                                    output_folder)


def molecules_score(project):
    print("In molecule_score function")
    db = get_db()
    collection_molecule_score = db['molecules_score']
    collection_SASA = db["molecules_SASA"]
    
    feature_type_list = json.loads(project['feature_type'])
    molecule_scores = {}
    # get set1 and set2 PDB from database
    PDB_set1_names = json.loads(project['set1_molecules_name'])
    PDB_set2_names = json.loads(project['set2_molecules_name'])
    print(f'PDB_set1_names: {PDB_set1_names}')
    print(f'PDB_set2_names: {PDB_set2_names}')
    for pdb_set1_single in PDB_set1_names:
        # load the SASA of one pdb from set1
        pdb_set1_single_dict = collection_SASA.find_one({"PDB_id": pdb_set1_single})
        pdb_set1_single_binding = pickle.loads(pdb_set1_single_dict["Binding_coords_0"])
        pdb_set1_single_SASA = pickle.loads(pdb_set1_single_dict["SASA_coords_0"])
        pdb_set1_single_SASA_indexes_from_whole = pickle.loads(pdb_set1_single_dict["SASA_indexes_from_whole_0"])
        pdb_set1_single_SASA_indexes_from_site = pickle.loads(pdb_set1_single_dict["SASA_indexes_from_site_0"])
        pdb_set1_single_SASA_features = pickle.loads(pdb_set1_single_dict["SASA_features_0"])

        
        for pdb_set2_single in PDB_set2_names:
            # load the SASA of one pdb from set2
            print(f'Comparing {pdb_set1_single} and {pdb_set2_single}')
            
            pdb_set2_single_dict = collection_SASA.find_one({"PDB_id": pdb_set2_single})
            pdb_set2_single_binding = pickle.loads(pdb_set2_single_dict["Binding_coords_0"])
            pdb_set2_single_SASA = pickle.loads(pdb_set2_single_dict["SASA_coords_0"])
            pdb_set2_single_SASA_indexes_from_whole = pickle.loads(pdb_set2_single_dict["SASA_indexes_from_whole_0"])
            pdb_set2_single_SASA_indexes_from_site = pickle.loads(pdb_set2_single_dict["SASA_indexes_from_site_0"])
            pdb_set2_single_SASA_features = pickle.loads(pdb_set2_single_dict["SASA_features_0"])

            pdb1_binding_local = np.array(pdb_set1_single_binding)
            pdb1_SASA_local = np.array(pdb_set1_single_SASA)
            pdb2_binding_local = np.array(pdb_set2_single_binding)
            pdb2_SASA_local = np.array(pdb_set2_single_SASA)
            pre_result = {}
            molecule_scores[pdb_set1_single + "_and_" + pdb_set2_single] = {}
            if len(pdb1_binding_local) <= len(pdb2_binding_local):
                aligned_coords_2_site, RMSD, R, t = get_downsampled_attributes_parallel(pdb1_binding_local, pdb2_binding_local)
                pdb2_binding_local = aligned_coords_2_site
                aligned_coords = align_coordinates(R, t, pdb2_SASA_local)
                pdb2_SASA_local = aligned_coords
                rotate = 2
                pre_result = {"rotate": rotate, "R_0": R, "t_0": t,}
            else:
                aligned_coords_1_site, RMSD, R, t = get_downsampled_attributes_parallel(pdb2_binding_local, pdb1_binding_local)
                pdb1_binding_local = aligned_coords_1_site
                aligned_coords = align_coordinates(R, t, pdb1_SASA_local)
                pdb1_SASA_local = aligned_coords
                rotate = 1
                pre_result = {"rotate": rotate, "R_0": R, "t_0": t,}
            radius = get_dynamic_radius(pdb1_SASA_local, pdb2_SASA_local)
            max_radius = radius
            lrf_1, lrf_2 = get_lrf(pdb1_binding_local, pdb_set1_single_SASA_indexes_from_site, pdb2_binding_local, pdb_set2_single_SASA_indexes_from_site, 10, max_radius, True)

            SASA_coords_A, SASA_coords_B, downsampling, switch = get_coordinate_lengths(pdb1_SASA_local, pdb2_SASA_local)
            (SASA_site_A_name, SASA_site_B_name) = (pdb_set2_single, pdb_set1_single) if switch else (pdb_set1_single, pdb_set2_single)
            (SASA_site_A_indexes, SASA_site_B_indexes) =  (pdb_set2_single_SASA_indexes_from_whole, pdb_set1_single_SASA_indexes_from_whole) if switch else (pdb_set1_single_SASA_indexes_from_whole, pdb_set2_single_SASA_indexes_from_whole)
            (SASA_site_A_features, SASA_site_B_features) =  (pdb_set2_single_SASA_features, pdb_set1_single_SASA_features) if switch else (pdb_set1_single_SASA_features, pdb_set2_single_SASA_features)
            (SASA_site_A_lrf, SASA_site_B_lrf) =  (lrf_2, lrf_1) if switch else (lrf_1, lrf_2)
            
            pre_result["switch"] = switch
            for f in feature_type_list:

                key =  f"{SASA_site_A_name}#{SASA_site_B_name}#{f}"
                if collection_molecule_score.count_documents({"key": key}) == 0:
                    molecules_score_using_shape_features(SASA_site_A_name, SASA_site_B_name, f, max_radius, 
                                                    SASA_coords_A, SASA_coords_B, SASA_site_A_lrf, SASA_site_B_lrf,
                                                    downsampling, SASA_site_A_features, SASA_site_B_features,
                                                    SASA_site_A_indexes, SASA_site_B_indexes,
                                                    collection_molecule_score, pre_result=pre_result)
                score_dict = collection_molecule_score.find_one({"key": key})
                molecule_scores[pdb_set1_single + "_and_" + pdb_set2_single][key] = score_dict
    return molecule_scores  

def molecules_score_using_shape_features(SASA_site_A_name, SASA_site_B_name, f, radius, 
                                        SASA_coords_A, SASA_coords_B, SASA_lrf_A, SASA_lrf_B,
                                        downsampling, features_A, features_B,
                                        SASA_site_A_indexes, SASA_site_B_indexes,
                                        collection_molecule_score, pre_result=None):
    if f == 'shape':
        results_dict = molecules_score_s_f_n_single(SASA_coords_A, SASA_coords_B, SASA_lrf_A, SASA_lrf_B,
                                                    downsampling, radius, features_A=None,features_B=None,)
    else:
        results_dict = molecules_score_s_f_n_single(SASA_coords_A, SASA_coords_B, SASA_lrf_A, SASA_lrf_B,
                                                    downsampling, radius, features_A=features_A, features_B=features_B)
    if pre_result is not None:
        results_dict["R_0"] = pre_result["R_0"].tolist()
        results_dict["t_0"] = pre_result["t_0"].tolist()
        results_dict["rotate"] = pre_result["rotate"]
        results_dict["switch"] = pre_result["switch"]
    insert_result_to_db(collection_molecule_score, results_dict, f, radius, 
                        SASA_site_A_name, SASA_site_B_name,
                        SASA_site_A_indexes, SASA_site_B_indexes)

def molecules_score_s_f_n_single(SASA_coords_A, SASA_coords_B, lrf_A, lrf_B, downsampling, radius, features_A, features_B):
    results_dict = compare_sites(SASA_coords_A, SASA_coords_B, lrf_A, lrf_B, radius, features_A, features_B)      
    return results_dict


def main():
    parser = OptionParser()
    parser.add_option("-p", "--project", dest="project", help="Project id", type=str)
    (options, args) = parser.parse_args()
    print("RUNNING BATCH")
    run_batch_molecules(options.project)


if __name__ == "__main__":
    main()
