import unittest
import os
import tempfile
from deepchem.utils import rdkit_utils
from deepchem.utils import pdbqt_utils


class TestPDBQTUtils(unittest.TestCase):

    def setUp(self):
        current_dir = os.path.dirname(os.path.realpath(__file__))
        self.protein_file = os.path.join(current_dir,
                                         "../../dock/tests/1jld_protein.pdb")
        self.ligand_file = os.path.join(current_dir,
                                        "../../dock/tests/1jld_ligand.sdf")

    def test_pdbqt_to_pdb(self):
        """Test that a PDBQT molecule can be converted back in to PDB."""
        xyz, mol = rdkit_utils.load_molecule(self.protein_file,
                                             calc_charges=False,
                                             add_hydrogens=False)
        with tempfile.TemporaryDirectory() as tmp:
            out_pdb = os.path.join(tmp, "mol.pdb")
            out_pdbqt = os.path.join(tmp, "mol.pdbqt")

            rdkit_utils.write_molecule(mol, out_pdb, is_protein=True)
            rdkit_utils.write_molecule(mol, out_pdbqt, is_protein=True)

            pdb_block = pdbqt_utils.pdbqt_to_pdb(out_pdbqt)
            from rdkit import Chem
            pdb_mol = Chem.MolFromPDBBlock(pdb_block,
                                           sanitize=False,
                                           removeHs=False)

            xyz, pdbqt_mol = rdkit_utils.load_molecule(out_pdbqt,
                                                       add_hydrogens=False,
                                                       calc_charges=False)

        assert pdb_mol.GetNumAtoms() == pdbqt_mol.GetNumAtoms()
        for atom_idx in range(pdb_mol.GetNumAtoms()):
            atom1 = pdb_mol.GetAtoms()[atom_idx]
            atom2 = pdbqt_mol.GetAtoms()[atom_idx]
            assert atom1.GetAtomicNum() == atom2.GetAtomicNum()

    def test_convert_mol_to_pdbqt(self):
        """Test that a ligand molecule can be coverted to PDBQT."""
        from rdkit import Chem
        xyz, mol = rdkit_utils.load_molecule(self.ligand_file,
                                             calc_charges=False,
                                             add_hydrogens=False)
        with tempfile.TemporaryDirectory() as tmp:
            outfile = os.path.join(tmp, "mol.pdbqt")
            writer = Chem.PDBWriter(outfile)
            writer.write(mol)
            writer.close()
            pdbqt_utils.convert_mol_to_pdbqt(mol, outfile)
            pdbqt_xyz, pdbqt_mol = rdkit_utils.load_molecule(
                outfile, add_hydrogens=False, calc_charges=False)
        assert pdbqt_mol.GetNumAtoms() == pdbqt_mol.GetNumAtoms()
        for atom_idx in range(pdbqt_mol.GetNumAtoms()):
            atom1 = pdbqt_mol.GetAtoms()[atom_idx]
            atom2 = pdbqt_mol.GetAtoms()[atom_idx]
            assert atom1.GetAtomicNum() == atom2.GetAtomicNum()

    def test_convert_protein_to_pdbqt(self):
        """Test a protein in a PDB can be converted to PDBQT."""
        from rdkit import Chem
        xyz, mol = rdkit_utils.load_molecule(self.protein_file,
                                             calc_charges=False,
                                             add_hydrogens=False)
        with tempfile.TemporaryDirectory() as tmp:
            outfile = os.path.join(tmp, "mol.pdbqt")
            writer = Chem.PDBWriter(outfile)
            writer.write(mol)
            writer.close()
            pdbqt_utils.convert_protein_to_pdbqt(mol, outfile)
            pdbqt_xyz, pdbqt_mol = rdkit_utils.load_molecule(
                outfile, add_hydrogens=False, calc_charges=False)
        assert pdbqt_mol.GetNumAtoms() == pdbqt_mol.GetNumAtoms()
        for atom_idx in range(pdbqt_mol.GetNumAtoms()):
            atom1 = pdbqt_mol.GetAtoms()[atom_idx]
            atom2 = pdbqt_mol.GetAtoms()[atom_idx]
            assert atom1.GetAtomicNum() == atom2.GetAtomicNum()
