# $Id$
#
# Copyright (C) 2007,2008 Greg Landrum
#
#  @@ All Rights Reserved @@
#
import os,sys
import io
import unittest
from rdkit.six.moves import cPickle
from rdkit import RDConfig
from rdkit import DataStructs as ds

def feq(v1,v2,tol=1e-4):
  return abs(v1-v2)<tol
class TestCase(unittest.TestCase):
  def setUp(self) :
    pass

  def test1Int(self):
    """

    """
    v1 = ds.IntSparseIntVect(5)
    self.assertRaises(IndexError,lambda:v1[5])
    v1[0]=1
    v1[2]=2
    v1[3]=3
    self.assertTrue(v1==v1)
    self.assertTrue(v1.GetLength()==5)

    v2= ds.IntSparseIntVect(5)
    self.assertTrue(v1!=v2)
    v2|=v1
    self.assertTrue(v2==v1)

    v3=v2|v1
    self.assertTrue(v3==v1)

    onVs = v1.GetNonzeroElements()
    self.assertTrue(onVs=={0:1,2:2,3:3})


  def test2Long(self):
    """

    """
    l=1<<42
    v1 = ds.LongSparseIntVect(l)
    self.assertRaises(IndexError,lambda:v1[l])
    v1[0]=1
    v1[2]=2
    v1[1<<35]=3
    self.assertTrue(v1==v1)
    self.assertTrue(v1.GetLength()==l)

    v2= ds.LongSparseIntVect(l)
    self.assertTrue(v1!=v2)
    v2|=v1
    self.assertTrue(v2==v1)

    v3=v2|v1
    self.assertTrue(v3==v1)

    onVs = v1.GetNonzeroElements()
    self.assertTrue(onVs=={0:1,2:2,1<<35:3})

  def test3Pickle1(self):
    """

    """
    l=1<<42
    v1 = ds.LongSparseIntVect(l)
    self.assertRaises(IndexError,lambda:v1[l+1])
    v1[0]=1
    v1[2]=2
    v1[1<<35]=3
    self.assertTrue(v1==v1)

    v2=  cPickle.loads(cPickle.dumps(v1))
    self.assertTrue(v2==v1)
    
    v3=  ds.LongSparseIntVect(v2.ToBinary())
    self.assertTrue(v2==v3)
    self.assertTrue(v1==v3)

    #cPickle.dump(v1,file('lsiv.pkl','wb+'))
    with open(
      os.path.join(RDConfig.RDBaseDir,
                   'Code/DataStructs/Wrap/testData/lsiv.pkl'), 
      'r'
      ) as tf:
      buf = tf.read().replace('\r\n', '\n').encode('utf-8')
      tf.close()
    with io.BytesIO(buf) as f:
      v3 = cPickle.load(f)
      self.assertTrue(v3==v1)
    
  def test3Pickle2(self):
    """

    """
    l=1<<21
    v1 = ds.IntSparseIntVect(l)
    self.assertRaises(IndexError,lambda:v1[l+1])
    v1[0]=1
    v1[2]=2
    v1[1<<12]=3
    self.assertTrue(v1==v1)

    v2=  cPickle.loads(cPickle.dumps(v1))
    self.assertTrue(v2==v1)
    
    v3=  ds.IntSparseIntVect(v2.ToBinary())
    self.assertTrue(v2==v3)
    self.assertTrue(v1==v3)

    #cPickle.dump(v1,file('isiv.pkl','wb+'))
    with open(
      os.path.join(RDConfig.RDBaseDir, 
                   'Code/DataStructs/Wrap/testData/isiv.pkl'),
      'r'
      ) as tf:
      buf = tf.read().replace('\r\n', '\n').encode('utf-8')
      tf.close()
    with io.BytesIO(buf) as f:
      v3 = cPickle.load(f)
      self.assertTrue(v3==v1)

  def test4Update(self):
    """

    """
    v1 = ds.IntSparseIntVect(5)
    self.assertRaises(IndexError,lambda:v1[6])
    v1[0]=1
    v1[2]=2
    v1[3]=3
    self.assertTrue(v1==v1)

    v2 = ds.IntSparseIntVect(5)
    v2.UpdateFromSequence((0,2,3,3,2,3))
    self.assertTrue(v1==v2)
    
  def test5Dice(self):
    """

    """
    v1 = ds.IntSparseIntVect(5)
    v1[4]=4;
    v1[0]=2;
    v1[3]=1;
    self.assertTrue(feq(ds.DiceSimilarity(v1,v1),1.0))

    v1 = ds.IntSparseIntVect(5)
    v1[0]=2;
    v1[2]=1;
    v1[3]=4;
    v1[4]=6;
    v2 = ds.IntSparseIntVect(5)
    v2[1]=2;
    v2[2]=3;
    v2[3]=4;
    v2[4]=4;
    self.assertTrue(feq(ds.DiceSimilarity(v1,v2),18.0/26.))
    self.assertTrue(feq(ds.DiceSimilarity(v2,v1),18.0/26.))

  def test6BulkDice(self):
    """

    """
    sz=10
    nToSet=5
    nVs=6
    import random
    vs = []
    for i in range(nVs):
      v = ds.IntSparseIntVect(sz)
      for j in range(nToSet):
        v[random.randint(0,sz-1)]=random.randint(1,10)
      vs.append(v)

    baseDs = [ds.DiceSimilarity(vs[0],vs[x]) for x in range(1,nVs)]
    bulkDs = ds.BulkDiceSimilarity(vs[0],vs[1:])
    for i in range(len(baseDs)):
      self.assertTrue(feq(baseDs[i],bulkDs[i]))

  def test6BulkTversky(self):
    """

    """
    sz=10
    nToSet=5
    nVs=6
    import random
    vs = []
    for i in range(nVs):
      v = ds.IntSparseIntVect(sz)
      for j in range(nToSet):
        v[random.randint(0,sz-1)]=random.randint(1,10)
      vs.append(v)

    baseDs = [ds.TverskySimilarity(vs[0],vs[x],.5,.5) for x in range(1,nVs)]
    bulkDs = ds.BulkTverskySimilarity(vs[0],vs[1:],0.5,0.5)
    diceDs = [ds.DiceSimilarity(vs[0],vs[x]) for x in range(1,nVs)]
    for i in range(len(baseDs)):
      self.assertTrue(feq(baseDs[i],bulkDs[i]))
      self.assertTrue(feq(baseDs[i],diceDs[i]))
    
    bulkDs = ds.BulkTverskySimilarity(vs[0],vs[1:],1.0,1.0)
    taniDs = [ds.TanimotoSimilarity(vs[0],vs[x]) for x in range(1,nVs)]
    for i in range(len(bulkDs)):
      self.assertTrue(feq(bulkDs[i],taniDs[i]))
    taniDs = ds.BulkTanimotoSimilarity(vs[0],vs[1:])
    for i in range(len(bulkDs)):
      self.assertTrue(feq(bulkDs[i],taniDs[i]))
    
    
if __name__ == '__main__':
    unittest.main()