import unittest
import numpy
import os
import io

from rdkit.six.moves import cPickle

from rdkit import RDConfig, RDRandom
from rdkit.ML.InfoTheory import rdInfoTheory as rdit
from rdkit import DataStructs


def feq(a, b, tol=1e-4):
  return abs(a - b) < tol


class TestCase(unittest.TestCase):

  def setUp(self):
    pass

  def test0GainFuns(self):
    arr = numpy.array([9, 5])
    self.assertTrue(feq(rdit.InfoEntropy(arr), 0.9403))
    arr = numpy.array([9, 9])
    self.assertTrue(feq(rdit.InfoEntropy(arr), 1.0000))
    arr = numpy.array([5, 5])
    self.assertTrue(feq(rdit.InfoEntropy(arr), 1.0000))
    arr = numpy.array([5, 0])
    self.assertTrue(feq(rdit.InfoEntropy(arr), 0.0000))
    arr = numpy.array([5, 5, 5])
    self.assertTrue(feq(rdit.InfoEntropy(arr), 1.5850))
    arr = numpy.array([2, 5, 5])
    self.assertTrue(feq(rdit.InfoEntropy(arr), 1.4834))

    mat2 = numpy.array([[6, 2], [3, 3]])
    self.assertTrue(feq(rdit.InfoGain(mat2), 0.0481))
    self.assertTrue(feq(rdit.ChiSquare(mat2), 0.9333))

    mat3 = numpy.array([[1, 1], [2, 1]])
    self.assertTrue(feq(rdit.InfoGain(mat3), 0.0200))

    mat4 = numpy.array([[2, 0], [1, 2]])
    self.assertTrue(feq(rdit.InfoGain(mat4), 0.4200))

    mat5 = numpy.array([[0, 0], [0, 0]])
    self.assertTrue(feq(rdit.InfoGain(mat5), 0.0000))

    mat6 = numpy.array([[1, 0], [1, 0]])
    self.assertTrue(feq(rdit.InfoGain(mat6), 0.0000))

  def test1ranker(self):
    nbits = 100
    ninst = 100
    dm = 50
    nact = 10
    nc = 2
    rn = rdit.InfoBitRanker(nbits, nc, rdit.InfoType.ENTROPY)
    fps = []
    na = 0
    ni = 0
    for i in range(ninst):
      v = DataStructs.SparseBitVect(nbits)
      for j in range(dm):
        v.SetBit(RDRandom.randrange(0, nbits))

      if (RDRandom.randrange(0, ninst) < nact):
        na += 1
        rn.AccumulateVotes(v, 1)
        fps.append((v, 1))
      else:
        ni += 1
        rn.AccumulateVotes(v, 0)
        fps.append((v, 0))

    res = rn.GetTopN(50)

    rn2 = rdit.InfoBitRanker(nbits, nc)
    for fp in fps:
      rn2.AccumulateVotes(fp[0], fp[1])

    res2 = rn2.GetTopN(50)
    self.assertTrue((res == res2).all())

    rn3 = rdit.InfoBitRanker(nbits, nc, rdit.InfoType.BIASENTROPY)
    #rn3.SetBiasList([0])
    for fp in fps:
      rn3.AccumulateVotes(fp[0], fp[1])

    res3 = rn3.GetTopN(50)
    for i in range(50):
      fan = res3[i, 2] / na
      fin = res3[i, 3] / ni
      self.assertTrue(fan > fin)

  def test2ranker(self):
    nbits = 100
    ninst = 100
    dm = 50
    nact = 10
    nc = 2
    RDRandom.seed(23)
    rn = rdit.InfoBitRanker(nbits, nc, rdit.InfoType.ENTROPY)
    rn.SetMaskBits([63, 70, 15, 25, 10])
    fps = []
    na = 0
    ni = 0
    for i in range(ninst):
      v = DataStructs.SparseBitVect(nbits)
      for j in range(dm):
        v.SetBit(RDRandom.randrange(0, nbits))
      if (RDRandom.randrange(0, ninst) < nact):
        na += 1
        rn.AccumulateVotes(v, 1)
        fps.append((v, 1))
      else:
        ni += 1
        rn.AccumulateVotes(v, 0)
        fps.append((v, 0))
    res = rn.GetTopN(5)
    ids = [int(x[0]) for x in res]
    ids.sort()
    self.assertTrue(ids == [10, 15, 25, 63, 70])
    with self.assertRaisesRegexp(Exception, ""):
      res = rn.GetTopN(10)

  def test3Issue140(self):
    nbits = 2
    examples = [[0, 0, 0], [1, 1, 0], [0, 0, 1], [1, 1, 1]]
    rn = rdit.InfoBitRanker(2, 2, rdit.InfoType.ENTROPY)
    for example in examples:
      act = example.pop(-1)
      bv = DataStructs.ExplicitBitVect(2)
      for i in range(2):
        bv[i] = example[i]
      rn.AccumulateVotes(bv, act)
    try:
      res = rn.GetTopN(1)
    except Exception:
      res = None
    self.assertTrue(res is not None)

  def test4Issue237(self):
    with open(
        os.path.join(RDConfig.RDBaseDir, 'Code', 'ML', 'InfoTheory', 'Wrap', 'testData',
                     'Issue237.pkl'), 'r') as inTF:
      buf = inTF.read().replace('\r\n', '\n').encode('utf-8')
      inTF.close()
    with io.BytesIO(buf) as inF:
      examples, avail, bias, nB, nPoss = cPickle.load(inF, encoding='bytes')
    ranker = rdit.InfoBitRanker(nB, nPoss, rdit.InfoType.BIASENTROPY)
    ranker.SetMaskBits(avail)
    for ex in examples:
      ranker.AccumulateVotes(ex[1], ex[-1])
    # this dumps core on linux if the bug isn't fixed:
    v = ranker.GetTopN(1)
    self.assertTrue(int(v[0][0]) == 12)


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