import time

from orionclient.types import Dataset
from drconvert import MolFileConverter
from orionclient.links import RecordLink
from orionclient.session import APISession
from datarecord import Types, OEField, OEPrimaryMolField, create_link

# Note: oechem must be imported before OpenEye toolkits
from openeye.oechem import OEMol, OEGraphMol

from openeye.oedocking import (
    OEDock,
    OEHybrid,
    OEIsReceptor,
    OEReceptorHasBoundLigand,
    OEDockingReturnCode_Success,
)

# Upload the receptor dataset, as linked records must exist in
# Orion prior to linking
receptor_ds = Dataset.upload(APISession, "receptors", "rec1a28.oeb.gz")

# Wait until the receptor dataset is finalized
while receptor_ds.dirty:
    APISession.refresh_resource(receptor_ds)
    time.sleep(1)


# Retrieve the records for the ligands, in this case just one
records = []
for rec in MolFileConverter("1a28_flynn_refined_lig.sdf"):
    records.append(rec)

mol_field = OEPrimaryMolField()
score_field = OEField("docking score", Types.Float)
# Field that will link to the receptor record
link_field = OEField("receptor link", Types.Link)

# Iterate through the receptors and ligands and
# set the score and link field.
for receptor_rec in receptor_ds.records():
    receptor = receptor_rec.get_value(mol_field)
    if receptor is None or not OEIsReceptor(receptor):
        continue
    receptor = OEGraphMol(receptor)
    if OEReceptorHasBoundLigand(receptor):
        dock = OEHybrid()
    else:
        dock = OEDock()
    dock.Initialize(receptor)
    for lig_rec in records:
        ligand = lig_rec.get_value(mol_field)
        if ligand is None:
            continue
        out_mol = OEMol()
        if (
            dock.DockMultiConformerMolecule(out_mol, ligand)
            != OEDockingReturnCode_Success
        ):
            continue
        score = out_mol.GetEnergy()
        prev_score = lig_rec.get_value(score_field)
        if prev_score is None or prev_score < score:
            lig_rec.set_value(score_field, score)
            # Set the link_field to point to the receptor field
            lig_rec.set_value(
                link_field, create_link(record=receptor_rec, field=mol_field)
            )

for lig_rec in records:
    link = lig_rec.get_value(link_field)
    if link is None:
        continue
    # Use OrionClient to retrieve the link object
    link_obj = APISession.get_link(link)
    if not isinstance(link_obj, RecordLink):
        continue
    # Get the associated record
    receptor_rec = link_obj.get_record()
    # Retrieve the receptor from the linked record
    receptor = receptor_rec.get_value(link_obj.get_field())