from __future__ import division
from rdkit import RDConfig
import unittest
from rdkit.DataManip.Metric import rdMetricMatrixCalc as rdmmc
import numpy
import random
from rdkit import DataStructs


def feq(v1, v2, tol2=1e-4):
  return abs(v1 - v2) <= tol2


class TestCase(unittest.TestCase):

  def setUp(self):
    pass

  def test0DistsArray(self):
    exp = numpy.array([1., 1.414213, 1.0], 'd')

    # initialize a double array and check if get back the expected distances
    desc = numpy.zeros((3, 2), 'd')
    desc[1, 0] = 1.0
    desc[2, 0] = 1.0
    desc[2, 1] = 1.0

    dmat = rdmmc.GetEuclideanDistMat(desc)
    for i in range(numpy.shape(dmat)[0]):
      assert feq(dmat[i], exp[i])

    # repeat with an flaot array
    desc = numpy.zeros((3, 2), 'f')
    desc[1, 0] = 1.0
    desc[2, 0] = 1.0
    desc[2, 1] = 1.0

    dmat = rdmmc.GetEuclideanDistMat(desc)
    for i in range(numpy.shape(dmat)[0]):
      assert feq(dmat[i], exp[i])

    # finally with an interger array
    desc = numpy.zeros((3, 2), 'i')
    desc[1, 0] = 1
    desc[2, 0] = 1
    desc[2, 1] = 1

    dmat = rdmmc.GetEuclideanDistMat(desc)
    for i in range(numpy.shape(dmat)[0]):
      assert feq(dmat[i], exp[i])

  def ctest1DistsListArray(self):
    exp = numpy.array([1., 1.414213, 1.0], 'd')

    desc = [numpy.array([0.0, 0.0], 'd'), numpy.array([1.0, 0.0], 'd'),
            numpy.array([1.0, 1.0], 'd')]
    dmat = rdmmc.GetEuclideanDistMat(desc)

    for i in range(numpy.shape(dmat)[0]):
      assert feq(dmat[i], exp[i])

    # repeat the test with a list of numpy.arrays of floats
    desc = [numpy.array([0.0, 0.0], 'f'), numpy.array([1.0, 0.0], 'f'),
            numpy.array([1.0, 1.0], 'f')]
    dmat = rdmmc.GetEuclideanDistMat(desc)
    for i in range(numpy.shape(dmat)[0]):
      assert feq(dmat[i], exp[i])

    # repeat the test with a list of numpy.arrays of ints
    desc = [numpy.array([0, 0], 'i'), numpy.array([1, 0], 'i'), numpy.array([1, 1], 'i')]
    dmat = rdmmc.GetEuclideanDistMat(desc)
    for i in range(numpy.shape(dmat)[0]):
      assert feq(dmat[i], exp[i])

  def test2DistListList(self):
    exp = numpy.array([1., 1.414213, 1.0], 'd')

    desc = [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0]]
    dmat = rdmmc.GetEuclideanDistMat(desc)
    for i in range(numpy.shape(dmat)[0]):
      assert feq(dmat[i], exp[i])

    #test with ints
    desc = [[0, 0], [1, 0], [1, 1]]
    dmat = rdmmc.GetEuclideanDistMat(desc)
    for i in range(numpy.shape(dmat)[0]):
      assert feq(dmat[i], exp[i])

  def test3Compare(self):
    n = 30
    m = 5

    dscArr = numpy.zeros((n, m), 'd')
    for i in range(n):
      for j in range(m):
        dscArr[i, j] = random.random()
    dmatArr = rdmmc.GetEuclideanDistMat(dscArr)

    dscLL = []
    for i in range(n):
      row = []
      for j in range(m):
        row.append(dscArr[i, j])
      dscLL.append(row)
    dmatLL = rdmmc.GetEuclideanDistMat(dscLL)

    assert numpy.shape(dmatArr) == numpy.shape(dmatLL)

    for i in range(n * (n - 1) // 2):
      assert feq(dmatArr[i], dmatLL[i])

  def test4ebv(self):

    n = 30
    m = 2048
    dm = 800
    lst = []
    for i in range(n):
      v = DataStructs.ExplicitBitVect(m)
      for j in range(dm):
        v.SetBit(random.randrange(0, m))
      lst.append(v)

    dMat = rdmmc.GetTanimotoDistMat(lst)

    sMat = rdmmc.GetTanimotoSimMat(lst)

    for i in range(n * (n - 1) // 2):
      assert feq(sMat[i] + dMat[i], 1.0)

  def test5sbv(self):

    n = 30
    m = 2048
    dm = 800
    lst = []
    for i in range(n):
      v = DataStructs.SparseBitVect(m)
      for j in range(dm):
        v.SetBit(random.randrange(0, m))
      lst.append(v)

    dMat = rdmmc.GetTanimotoDistMat(lst)

    sMat = rdmmc.GetTanimotoSimMat(lst)

    for i in range(n * (n - 1) // 2):
      assert feq(sMat[i] + dMat[i], 1.0)


if __name__ == '__main__':
  unittest.main()