from random_transformation import row_random_transformation

import unittest
from theano import compile
from theano import gradient

from theano.sparse import _is_dense, _is_sparse, _is_dense_result, _is_sparse_result
from theano.sparse import _mtypes, _mtype_to_str
from theano.sparse import as_sparse

from theano.tensor import as_tensor
from theano.scalar import as_scalar

import random
import numpy.random

class T_RowRandomTransformation(unittest.TestCase):
    def setUp(self):
        random.seed(44)
        numpy.random.seed(44)

    def test_basic(self):
        rows = 4
        cols = 20
        fakeseed = 0
        length = 3 
        md = numpy.random.rand(rows, cols)
        for mtype in _mtypes:
            m = as_sparse(mtype(md))
            o = row_random_transformation(m, length, initial_seed=fakeseed)
            y = compile.eval_outputs([o])
            expected = "[[ 0.88239119  1.03244463 -1.29297503]\n [ 0.02644961  1.50119695 -0.025081  ]\n [-0.60741013  1.25424625  0.30119422]\n [-1.08659967 -0.35531544 -1.38915467]]"
            self.failUnless(str(y) == expected)

    def test_length(self):
        """ Test that if length is increased, we obtain the same results
        (except longer). """

        for i in range(10):
            mtype = random.choice(_mtypes)
            rows = random.randint(1, 20)
            cols = random.randint(1, 20)
            fakeseed = random.randint(0, 100)
            length = random.randint(1, 10)
            extralength = random.randint(1, 10)

            m = as_sparse(mtype(numpy.random.rand(rows, cols)))
            o1 = row_random_transformation(m, length, initial_seed=fakeseed)
            o2 = row_random_transformation(m, length + extralength, initial_seed=fakeseed)

            y1 = compile.eval_outputs([o1])
            y2 = compile.eval_outputs([o2])

            self.failUnless((y1 == y2[:,:length]).all())

    def test_permute(self):
        """ Test that if the order of the rows is permuted, we obtain the same results. """
        for i in range(10):
            mtype = random.choice(_mtypes)
            rows = random.randint(2, 20)
            cols = random.randint(1, 20)
            fakeseed = random.randint(0, 100)
            length = random.randint(1, 10)

            permute = numpy.random.permutation(rows)


            m1 = numpy.random.rand(rows, cols)
            m2 = m1[permute]
            for r in range(rows):
                self.failUnless((m2[r] == m1[permute[r]]).all())
            s1 = as_sparse(mtype(m1))
            s2 = as_sparse(mtype(m2))
            o1 = row_random_transformation(s1, length, initial_seed=fakeseed)
            o2 = row_random_transformation(s2, length, initial_seed=fakeseed)
            y1 = compile.eval_outputs([o1])
            y2 = compile.eval_outputs([o2])

            self.failUnless(y1.shape == y2.shape)
            for r in range(rows):
                self.failUnless((y2[r] == y1[permute[r]]).all())

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