import os
from distutils.version import LooseVersion

import numpy as np
import pytest
from numpy.testing import assert_array_equal
from skimage import __version__ as skimage_version

import oddt
from oddt.surface import (generate_surface_marching_cubes,
                          find_surface_residues)

test_data_dir = os.path.dirname(os.path.abspath(__file__))
protein = next(oddt.toolkit.readfile('pdb', os.path.join(
    test_data_dir, 'data/dude/xiap/receptor_rdkit.pdb')))
protein.protein = True
protein.addh(only_polar=True)


def test_generate_surface_marching_cubes():
    """Tests generating surfaces"""
    verts1, faces1 = generate_surface_marching_cubes(protein, scaling=1., probe_radius=1.4, remove_hoh=False)
    verts2, faces2 = generate_surface_marching_cubes(protein, scaling=2., probe_radius=1.4, remove_hoh=False)
    verts3, faces3 = generate_surface_marching_cubes(protein, scaling=1., probe_radius=1.4, remove_hoh=True)
    verts4, faces4 = generate_surface_marching_cubes(protein, scaling=1., probe_radius=0, remove_hoh=True)

    # Higher scaling should result in a higher number of vertices
    assert len(verts2) > len(verts1), ('Higher scaling should result in '
                                       'a higher number of vertices')

    # versions of skimage older than 0.12 use a slightly different version of the marching cubes algorithm
    # producing slightly different results
    if LooseVersion(skimage_version) >= LooseVersion('0.13'):
        if (oddt.toolkit.backend == 'ob' or
                oddt.toolkit.backend == 'rdk' and oddt.toolkits.rdk.__version__ >= '2019.09'):
            ref_vert_shape_1 = (9040, 3)
            ref_face_shape_1 = (18094, 3)
            ref_vert_shape_2 = (35950, 3)
            ref_face_shape_2 = (71926, 3)
            ref_vert_shape_3 = (9040, 3)
            ref_face_shape_3 = (18094, 3)
            ref_vert_shape_4 = (14881, 3)
            ref_face_shape_4 = (30468, 3)
        else:
            ref_vert_shape_1 = (9044, 3)
            ref_face_shape_1 = (18102, 3)
            ref_vert_shape_2 = (35788, 3)
            ref_face_shape_2 = (71578, 3)
            ref_vert_shape_3 = (9044, 3)
            ref_face_shape_3 = (18102, 3)
            ref_vert_shape_4 = (15035, 3)
            ref_face_shape_4 = (30848, 3)
    else:
        if oddt.toolkit.backend == 'ob':
            ref_vert_shape_1 = (5923, 3)
            ref_face_shape_1 = (11862, 3)
            ref_vert_shape_2 = (20819, 3)
            ref_face_shape_2 = (41634, 3)
            ref_vert_shape_3 = (5923, 3)
            ref_face_shape_3 = (11862, 3)
            ref_vert_shape_4 = (10263, 3)
            ref_face_shape_4 = (21658, 3)
        else:
            ref_vert_shape_1 = (5916, 3)
            ref_face_shape_1 = (11848, 3)
            ref_vert_shape_2 = (20845, 3)
            ref_face_shape_2 = (41686, 3)
            ref_vert_shape_3 = (5916, 3)
            ref_face_shape_3 = (11848, 3)
            ref_vert_shape_4 = (10243, 3)
            ref_face_shape_4 = (21686, 3)

    assert ref_vert_shape_1 == verts1.shape
    assert ref_face_shape_1 == faces1.shape

    assert ref_vert_shape_2 == verts2.shape
    assert ref_face_shape_2 == faces2.shape

    assert ref_vert_shape_3 == verts3.shape
    assert ref_face_shape_3 == faces3.shape

    assert ref_vert_shape_4 == verts4.shape
    assert ref_face_shape_4 == faces4.shape

    with pytest.raises(TypeError):
        generate_surface_marching_cubes(molecule=1)
    with pytest.raises(ValueError):
        generate_surface_marching_cubes(molecule=protein, probe_radius=-1)
    with pytest.raises(ValueError):
        generate_surface_marching_cubes(molecule=protein, scaling=0.1)


def test_find_surface_residues():
    """Tests finding residues on the surface"""
    atom_dict_0 = find_surface_residues(protein, max_dist=0, scaling=1)
    atom_dict_1 = find_surface_residues(protein, max_dist=2, scaling=1)
    atom_dict_2 = find_surface_residues(protein, max_dist=3, scaling=1)
    atom_dict_3 = find_surface_residues(protein, max_dist=None, scaling=1)
    atom_dict_4 = find_surface_residues(protein, max_dist=None, scaling=2)

    assert atom_dict_0.size == 0
    assert len(atom_dict_1) > len(atom_dict_0), ('Increasing max_dist should '
                                                 'result in more/equal number '
                                                 'of atoms found')

    assert len(atom_dict_2) >= len(atom_dict_1), ('Increasing max_dist should '
                                                  'result in more/equal number '
                                                  'of atoms found')
    assert_array_equal(np.intersect1d(atom_dict_1['id'], atom_dict_2['id']), atom_dict_1['id'])

    if oddt.toolkit.backend == 'ob':
        ref_len_1 = 762
        ref_len_2 = 968
        ref_len_3 = 654
        ref_len_4 = 379
    elif oddt.toolkit.backend == 'rdk' and oddt.toolkits.rdk.__version__ >= '2019.09':
        ref_len_1 = 762
        ref_len_2 = 968
        ref_len_3 = 664
        ref_len_4 = 393
    else:
        ref_len_1 = 759
        ref_len_2 = 966
        ref_len_3 = 735
        ref_len_4 = 489

    assert len(atom_dict_1) == ref_len_1
    assert len(atom_dict_2) == ref_len_2
    assert len(atom_dict_3) == ref_len_3
    assert len(atom_dict_4) == ref_len_4

    # Adding hydrogen atoms should have no effect on the result
    protein.addh()
    atom_dict_addh = find_surface_residues(protein, max_dist=2, scaling=1)
    assert_array_equal(atom_dict_addh['id'], atom_dict_1['id'])

    with pytest.raises(TypeError):
        find_surface_residues(molecule=1)
    with pytest.raises(ValueError):
        find_surface_residues(molecule=protein, max_dist='a')
    with pytest.raises(ValueError):
        find_surface_residues(molecule=protein, max_dist=[1, 1, 1])