import os

from tempfile import mkdtemp, NamedTemporaryFile

import pytest

from numpy.testing import assert_array_equal, assert_array_almost_equal

import pandas as pd

import oddt

from oddt.utils import method_caller

from oddt.spatial import rmsd

from oddt.scoring import scorer

from oddt.scoring.functions import rfscore, nnscore

from oddt.virtualscreening import virtualscreening

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

# common file names

dude_data_dir = os.path.join(test_data_dir, 'data', 'dude', 'xiap')

xiap_crystal_ligand = os.path.join(dude_data_dir, 'crystal_ligand.sdf')

xiap_protein = os.path.join(dude_data_dir, 'receptor_rdkit.pdb')

xiap_actives_docked = os.path.join(dude_data_dir, 'actives_docked.sdf')

def test_vs_scoring_vina():

    """VS scoring (Vina) tests"""

    vs = virtualscreening(n_cpu=1)

    vs.load_ligands('sdf', xiap_crystal_ligand)

    vs.score(function='autodock_vina', protein=xiap_protein)

    mols = list(vs.fetch())

    assert len(mols) == 1

    mol_data = mols[0].data

    assert 'vina_affinity' in mol_data

    assert 'vina_gauss1' in mol_data

    assert 'vina_gauss2' in mol_data

    assert 'vina_hydrogen' in mol_data

    assert 'vina_hydrophobic' in mol_data

    assert 'vina_repulsion' in mol_data

    assert mol_data['vina_affinity'] == '-3.57594'

    assert mol_data['vina_gauss1'] == '63.01213'

    assert mol_data['vina_gauss2'] == '999.07625'

    assert mol_data['vina_hydrogen'] == '0.0'

    assert mol_data['vina_hydrophobic'] == '26.12648'

    assert mol_data['vina_repulsion'] == '3.63178'

def test_vs_docking():

    """VS docking (Vina) tests"""

    vs = virtualscreening(n_cpu=1)

    vs.load_ligands('sdf', xiap_crystal_ligand)

    # bad docking engine

    with pytest.raises(ValueError):

        vs.dock('srina', 'prot.pdb')

    vs.dock(engine='autodock_vina',

            protein=xiap_protein,

            auto_ligand=xiap_crystal_ligand,

            exhaustiveness=1,

            energy_range=6,

            num_modes=7,

            size=(20, 20, 20),

            seed=0)

    mols = list(vs.fetch())

    assert len(mols) == 7

    mol_data = mols[0].data

    assert 'vina_affinity' in mol_data

    assert 'vina_rmsd_lb' in mol_data

    assert 'vina_rmsd_ub' in mol_data

    vina_scores = [-6.3, -6.0, -5.8, -5.8, -3.9, -3.0, -1.1]

    assert_array_equal([float(m.data['vina_affinity']) for m in mols], vina_scores)

    # verify the SMILES of molecules

    ref_mol = next(oddt.toolkit.readfile('sdf', xiap_crystal_ligand))

    if oddt.toolkit.backend == 'ob':

        # OB 2.3.2 will fail the following, since Hs are removed, etc.

        # OB 2.4 recognizes the smiles chirality wrong

        pass

    else:

        vina_rmsd = [8.153314, 5.32554, 8.514586, 8.510169, 9.060128, 8.995098,

                     8.626776]

        assert_array_equal([mol.smiles for mol in mols],

                           [ref_mol.smiles] * len(mols))

        assert_array_almost_equal([rmsd(ref_mol, mol, method='min_symmetry')

                                   for mol in mols], vina_rmsd)

def test_vs_empty():

    vs = virtualscreening(n_cpu=1)

    with pytest.raises(StopIteration, match='no molecules loaded'):

        vs.fetch()

def test_vs_docking_empty():

    vs = virtualscreening(n_cpu=1)

    vs.load_ligands('smi', os.path.join(dude_data_dir, 'actives_rdkit.smi'))

    vs.dock(engine='autodock_vina',

            protein=xiap_protein,

            auto_ligand=xiap_crystal_ligand,

            exhaustiveness=1,

            energy_range=5,

            num_modes=9,

            size=(20, 20, 20),

            seed=0)

    with pytest.raises(ValueError, match='has no 3D coordinates'):

        next(vs.fetch())

def test_vs_multithreading_fallback():

    vs = virtualscreening(n_cpu=8)

    vs.load_ligands('sdf', xiap_crystal_ligand)

    vs.score(function='autodock_vina', protein=xiap_protein)

    with pytest.warns(UserWarning, match='Falling back to sub-methods multithreading'):

        method_caller(vs, 'fetch')

if oddt.toolkit.backend == 'ob':  # RDKit rewrite needed

    def test_vs_filtering():

        """VS preset filtering tests"""

        vs = virtualscreening(n_cpu=1)

        vs.load_ligands('sdf', xiap_actives_docked)

        vs.apply_filter('ro5', soft_fail=1)

        assert len(list(vs.fetch())) == 49

        vs.load_ligands('sdf', xiap_actives_docked)

        vs.apply_filter('ro3', soft_fail=2)

        assert len(list(vs.fetch())) == 9

def test_vs_pains():

    """VS PAINS filter tests"""

    vs = virtualscreening(n_cpu=1)

    # TODO: add some failing molecules

    vs.load_ligands('sdf', xiap_actives_docked)

    vs.apply_filter('pains', soft_fail=0)

    assert len(list(vs.fetch())) == 100

def test_vs_similarity():

    """VS similarity filter (USRs, IFPs) tests"""

    ref_mol = next(oddt.toolkit.readfile('sdf', xiap_crystal_ligand))

    receptor = next(oddt.toolkit.readfile('pdb', xiap_protein))

    # following toolkit differences is due to different Hs treatment

    vs = virtualscreening(n_cpu=1, chunksize=10)

    vs.load_ligands('sdf', xiap_actives_docked)

    vs.similarity('usr', cutoff=0.4, query=ref_mol)

    if oddt.toolkit.backend == 'ob':

        assert len(list(vs.fetch())) == 11

    else:

        assert len(list(vs.fetch())) == 6

    vs = virtualscreening(n_cpu=1)

    vs.load_ligands('sdf', xiap_actives_docked)

    vs.similarity('usr_cat', cutoff=0.3, query=ref_mol)

    if oddt.toolkit.backend == 'ob':

        assert len(list(vs.fetch())) == 16

    else:

        assert len(list(vs.fetch())) == 11

    vs = virtualscreening(n_cpu=1)

    vs.load_ligands('sdf', xiap_actives_docked)

    vs.similarity('electroshape', cutoff=0.45, query=ref_mol)

    if oddt.toolkit.backend == 'ob':

        assert len(list(vs.fetch())) == 55

    else:

        assert len(list(vs.fetch())) == 95

    vs = virtualscreening(n_cpu=1)

    vs.load_ligands('sdf', xiap_actives_docked)

    vs.similarity('ifp', cutoff=0.8, query=ref_mol, protein=receptor)

    assert len(list(vs.fetch())) == 33

    vs = virtualscreening(n_cpu=1)

    vs.load_ligands('sdf', xiap_actives_docked)

    vs.similarity('sifp', cutoff=0.8, query=ref_mol, protein=receptor)

    assert len(list(vs.fetch())) == 33

    # test wrong method error

    with pytest.raises(ValueError):

        vs.similarity('sift', query=ref_mol)

def test_vs_scoring():

    protein = next(oddt.toolkit.readfile('pdb', xiap_protein))

    protein.protein = True

    data_dir = os.path.join(test_data_dir, 'data')

    home_dir = mkdtemp()

    pdbbind_versions = (2007, 2013, 2016)

    pdbbind_dir = os.path.join(data_dir, 'pdbbind')

    for pdbbind_v in pdbbind_versions:

        version_dir = os.path.join(data_dir, 'v%s' % pdbbind_v)

        if not os.path.isdir(version_dir):

            os.symlink(pdbbind_dir, version_dir)

    filenames = []

    # train mocked SFs

    for model in [nnscore(n_jobs=1)] + [rfscore(version=v, n_jobs=1)

                                        for v in [1, 2, 3]]:

            model.gen_training_data(data_dir, pdbbind_versions=pdbbind_versions,

                                    home_dir=home_dir)

            filenames.append(model.train(home_dir=home_dir))

    vs = virtualscreening(n_cpu=-1, chunksize=10)

    vs.load_ligands('sdf', xiap_actives_docked)

    # error if no protein is fed

    with pytest.raises(ValueError):

        vs.score('nnscore')

    # bad sf name

    with pytest.raises(ValueError):

        vs.score('bad_sf', protein=protein)

    vs.score('nnscore', protein=xiap_protein)

    vs.score('nnscore_pdbbind2016', protein=protein)

    vs.score('rfscore_v1', protein=protein)

    vs.score('rfscore_v1_pdbbind2016', protein=protein)

    vs.score('rfscore_v2', protein=protein)

    vs.score('rfscore_v3', protein=protein)

    vs.score('pleclinear', protein=protein)

    vs.score('pleclinear_p5_l1_s65536_pdbbind2016', protein=protein)

    # use pickle directly

    vs.score(filenames[0], protein=protein)

    # pass SF object directly

    vs.score(scorer.load(filenames[0]), protein=protein)

    # pass wrong object (sum is not an instance of scorer)

    with pytest.raises(ValueError):

        vs.score(sum, protein=protein)

    mols = list(vs.fetch())

    assert len(mols) == 100

    mol_data = mols[0].data

    assert 'nnscore' in mol_data

    assert 'rfscore_v1' in mol_data

    assert 'rfscore_v2' in mol_data

    assert 'rfscore_v3' in mol_data

    assert 'PLEClinear_p5_l1_s65536' in mol_data

    vs = virtualscreening(n_cpu=-1, chunksize=10)

    vs.load_ligands('sdf', xiap_actives_docked)

    vs.score('nnscore', protein=protein)

    vs.score('rfscore_v1', protein=protein)

    vs.score('rfscore_v2', protein=protein)

    vs.score('rfscore_v3', protein=protein)

    with NamedTemporaryFile('w', suffix='.sdf') as molfile:

        with NamedTemporaryFile('w', suffix='.csv') as csvfile:

            vs.write('sdf', molfile.name, csv_filename=csvfile.name)

            data = pd.read_csv(csvfile.name)

            assert 'nnscore' in data.columns

            assert 'rfscore_v1' in data.columns

            assert 'rfscore_v2' in data.columns

            assert 'rfscore_v3' in data.columns

            mols = list(oddt.toolkit.readfile('sdf', molfile.name))

            assert len(mols) == 100

            vs.write_csv(csvfile.name, fields=['nnscore', 'rfscore_v1',

                                               'rfscore_v2', 'rfscore_v3'])

            data = pd.read_csv(csvfile.name)

            assert len(data.columns) == 4

            assert len(data) == len(mols)

            assert 'nnscore' in data.columns

            assert 'rfscore_v1' in data.columns

            assert 'rfscore_v2' in data.columns

            assert 'rfscore_v3' in data.columns

    # remove files

    for f in filenames:

        os.unlink(f)

    # remove symlinks

    for pdbbind_v in pdbbind_versions:

        version_dir = os.path.join(data_dir, 'v%s' % pdbbind_v)

        if os.path.islink(version_dir):

            os.unlink(version_dir)