import glob

from sample_for_pocket import call, OUT_DIR
from rdkit import Chem
import json


def load(value: str):
    full_pdb_path = '../MolCRAFT/data/test_set/' + value
    sdf_path = glob.glob(full_pdb_path.rstrip('.pdb') + '*.sdf')
    assert len(sdf_path) == 1
    sdf_path = sdf_path[0]
    return [full_pdb_path, sdf_path]


pdb_file = 'ABL2_HUMAN_274_551_0/4xli_B_rec.pdb'

protein_fn, sdf_fn = load(pdb_file)

print(protein_fn, sdf_fn)

ref_mol = Chem.SDMolSupplier(sdf_fn, removeHs=False)[0]
ref_lig_pos = ref_mol.GetConformer().GetPositions()
print('[DEBUG] read raw file')
print(ref_lig_pos.shape, ref_lig_pos.mean(axis=0))


def generate(value: str):
    protein_path, ligand_path = load(value)

    call(protein_path, ligand_path)
    
    out_fns = sorted(glob.glob(f'{OUT_DIR}/*.sdf'))
    print(out_fns)
    out_fn = out_fns[0]
    sdf_mol = Chem.SDMolSupplier(out_fn, removeHs=False)[0]
    # get all properties from sdf_mol
    props = sdf_mol.GetPropsAsDict()

    return [protein_path, out_fn], json.dumps(props, indent=4)


generate(pdb_file)