import os

import pytest

import oddt
from oddt.datasets import pdbbind, dude

test_data_dir = os.path.dirname(os.path.abspath(__file__))


def test_pdbbind():

    results = {
        'core': (['4yef', '10gs'],
                 [5.35, 6.4]),
        'refined': (['1nlp', '1imx', '4yef', '10gs'],
                    [4.96, 3.52, 5.35, 6.4]),
        'general_PL': (['1k9q', '1nlo', '1nlp', '1imx', '4yef', '10gs'],
                       [3.15, 5.47, 4.96, 3.52, 5.35, 6.4]),
    }

    with pytest.raises(ValueError):
        pdbbind(home=os.path.join(test_data_dir, 'data', 'pdbbind'))

    for year in [2007, 2013, 2016]:
        pdbbind_db = pdbbind(home=os.path.join(test_data_dir, 'data', 'pdbbind'),
                             version=year, default_set='core')

        for set_name, (ids, activities) in results.items():
            if set_name == 'general_PL' and year == 2007:
                set_name = 'general'
            pdbbind_db.default_set = set_name
            assert pdbbind_db.ids == ids
            assert pdbbind_db.activities == activities

            for pid in pdbbind_db:
                assert isinstance(pid.pocket, oddt.toolkit.Molecule)
                assert len(pid.pocket.atoms) > 0
                assert isinstance(pid.ligand, oddt.toolkit.Molecule)
                assert len(pid.ligand.atoms) > 0
                if pid.id == '10gs':
                    assert pid.protein is None
                else:
                    assert isinstance(pid.protein, oddt.toolkit.Molecule)
                    assert len(pid.protein.atoms) > 0

        # reset the pdbbind set
        pdbbind_db.default_set = 'refined'

        # getting by name
        assert pdbbind_db['1imx'].id == '1imx'

        # getting by id
        assert pdbbind_db[-3].id == '1imx'
        assert pdbbind_db[1].id == '1imx'

        with pytest.raises(KeyError):
            pdbbind_db['xxxx']
        with pytest.raises(KeyError):
            pdbbind_db[123456]
        with pytest.raises(KeyError):
            pdbbind_db[-123456]

        pid = pdbbind_db['1imx']
        # get ligand
        ligand = pid.ligand
        ligand.removeh()
        assert len(ligand.atoms) == 60

        # get pocket
        pocket = pid.pocket
        pocket.removeh()
        assert len(pocket.atoms) == 234

        # protein do exist
        protein = pid.protein
        protein.removeh()
        assert len(protein.atoms) == 478


def test_dude():
    results = {
        'fabp4': (1022, 36, 57, 2855),
        'inha': (1857, 22, 71, 2318),
    }

    dude_db = dude(home=os.path.join(test_data_dir, 'data', 'dude'))

    for target in dude_db:
        if target.dude_id == 'xiap':
            # different file names
            assert target.protein is None
            assert target.ligand is None
            assert target.actives is None
            assert target.decoys is None
            continue

        prot_atoms, lig_atoms, num_act, num_dec = results[target.dude_id]

        prot = target.protein
        prot.removeh()
        assert len(prot.atoms) == prot_atoms
        lig = target.ligand
        lig.removeh()
        assert len(lig.atoms) == lig_atoms

        assert len(list(target.actives)) == num_act
        for a in target.actives:
            assert len(a.atoms) > 0
        assert len(list(target.decoys)) == num_dec
        for d in target.decoys:
            assert len(d.atoms) > 0

        with pytest.raises(KeyError):
            dude_db['xxxx']