import qcportal as ptl
from openff.toolkit.topology import Molecule

client= ptl.FractalClient()
# list the collections available
client.list_collections()

"""
Explanation: QCArchive Interface
Here we show how to create OpenFF molecules safely from data in the QCArchive using the cmiles entries, specifically we want to use the canonical_isomeric_explicit_hydrogen_mapped_smiles data which is metadata stored at the entry-level of a collection.
First load up the client you wish to connect to, in this case, we use the public instance.
End of explanation
"""


ds = client.get_collection('OptimizationDataset', 'Kinase Inhibitors: WBO Distributions')

"""
Explanation: Now let us grab a molecule from an optimization dataset
End of explanation
"""


entry = ds.get_entry(ds.df.index[0])

"""
Explanation: Take the first entry from the collection.
End of explanation
"""


entry.dict()

"""
Explanation: We can view the entry in detail by looking at the dictionary representation.
End of explanation
"""


# first make a molecule using this record object
mol_record = Molecule.from_qcschema(entry)

# we could have also used the dictionary representation of the object
mol_dict = Molecule.from_qcschema(entry.dict(encoding='json'))

# we check that the molecule has been ordered to match the ordering used in the data base
# by printing out the atomic numbers of both objects in order

# first lets get the initial molecule from the database
initial_mol = client.query_molecules(id=entry.initial_molecule)[0]

for atoms in zip(mol_record.atoms, initial_mol.atomic_numbers):
    print(atoms[0].atomic_number, atoms[1])

# we can also check that the molecules are the same regardless of how they are made 
assert mol_dict == mol_record

# we can also compare the graph representations of the molecules to make sure they are in the same order
import networkx as nx

# make a graph of the initial molecule using newtorkx and the data in the record
initial_network = nx.Graph()
for i, atom_num in enumerate(initial_mol.atomic_numbers):
    initial_network.add_node(i, atomic_number=atom_num)
    
for bond in initial_mol.connectivity:
    initial_network.add_edge(*bond[:2])
# now we can use the new isomorphic check to get the atom mapping
isomorphic, atom_map = Molecule.are_isomorphic(mol_record, 
                                               initial_network, 
                                               return_atom_map=True,
                                               aromatic_matching=False,
                                               formal_charge_matching=False,
                                               bond_order_matching=False,
                                               bond_stereochemistry_matching=False,
                                               atom_stereochemistry_matching=False)

# we can print if the graph was found to be isomorphic and then the atom mapping
# the atoms are in the same order here as the idexes are the same in the mapping
print(isomorphic)
print(atom_map)

"""
Explanation: Now we can make a molecule using a few different input options.
End of explanation
"""


# check there is no geometry for the molecule
assert mol_record.n_conformers == 0

# if we also want the input geometry for the molecule, we just need to pass the relavent client instance
mol_dict = Molecule.from_qcschema(entry.dict(encoding='json'), client=client)

# check that there is a conformer
mol_dict.n_conformers

# Thanks to the qcschema method we also get visualisation for free, along with being able to compute
# properties like energy, gradient and hessian with qcengine using QM, rdkit, openmm, or ANI1
mol_dict.to_qcschema()

"""
Explanation: Now that we have seen how to make the molecule, lets look at also getting the geometry as currently we have none.
End of explanation
"""


# for example this molecules energy can be computed using qcengine and RDKit
import qcengine
from openff.toolkit.typing.engines.smirnoff import ForceField

# set up the RDKit task
rdkit_task = {"schema_name": "qcschema_input",
            "schema_version": 2,
            "molecule": mol_dict.to_qcschema(),
            "driver": "energy",
            "model": {"method": 'uff', "basis": None},
            "keywords": {"scf_type": "df"}}

# now lets compute the energy using qcengine and RDKit and print the result
result = qcengine.compute(rdkit_task, 'rdkit')

# note the result is in QC units of hartrees
print(result.return_result)

"""
Explanation: Here we will try and compute the energy using RDKit (only run this cell if qcengine is installed.)
End of explanation
"""


