import sys, time, unittest, copy

import numpy,time
import numpy as N

from theano.tests import unittest_tools as utt

from theano import function, Mode
import theano.tensor as T
from pylearn.sandbox.scan_inputs_groups import FillMissing
import theano.compile.mode as mode_module

class TestFillMissing(unittest.TestCase):
    def setUp(self):
        utt.seed_rng()

        #we need to desactivate the check for NaN value as we have them in input
        #TODO: Make an option to don't check NaN value in input only, bug check in output.
        m=mode_module.default_mode
        if m=="DEBUG_MODE":
            m=copy.copy(mode_module.predefined_modes[m])
            m.check_isfinite=False
        self.mode = m

    def test_vector(self):
        print "test_vector"
        n=100000
        v=T.dvector()
        def t(prob,val,fill):
            op=FillMissing(fill)(v)
            f=function([v],op,mode=self.mode)
            nb_missing=0
            for i in range(n):
                if prob[i]<0.1:
                    nb_missing+=1
                    val[i]=N.nan
            t=time.time()
            out=f(val)
            print "time %.3fs"%(time.time()-t)
            for i in range(n):
                if N.isnan(val[i]):
                    if isinstance(fill,N.ndarray):
                        assert abs(out[0][i]-fill[i])<1e-6
                    else:
                        assert out[0][i]==fill
                else:
                    assert out[0][i]==val[i]
                    assert out[1][i]==1                

        prob=N.random.random(n)
        val=N.random.random(n)

        fill=0
        t(prob,val,fill)#test with fill a constant

        fill=N.random.random(n)
        t(prob,val,fill)#test with fill a vector

#TODO: test fill_with_array!
    def test_matrix(self):
        print "test_matrix"
        shp=(100,10)
        v=T.dmatrix()
        def t(prob,val,fill):
            op=FillMissing(fill)(v)
            f=function([v],op,mode=self.mode)
        
            nb_missing=0
            for i in range(shp[0]):
                for j in range(shp[1]):
                    if prob[i,j]<0.1:
                        nb_missing+=1
                        val[i,j]=N.nan
            t=time.time()
            out=f(val)
            print "time %.3fs"%(time.time()-t)
            for i in range(shp[0]):
                for j in range(shp[1]):
                    if N.isnan(val[i,j]):
                        if isinstance(fill,N.ndarray):
                            assert abs(out[0][i,j]-fill[j])<1e-6
                        else:
                            assert out[0][i,j]==fill
                    else:
                        assert out[0][i,j]==val[i,j]
                        assert out[1][i,j]==1


        prob=N.random.random(N.prod(shp)).reshape(shp)
        val=N.random.random(shp)

        fill=0
        t(prob,val,fill)#test with fill a constant
        
        fill=N.random.random(shp[1])
        t(prob,val,fill)#test with fill a vector

#TODO: test fill_with_array!
    def test_matrix3d(self):
        print "test_matrix3d"
        shp=(10,100,100)
        v= T.TensorType('float64', (False, False, False))()
        op=FillMissing()(v)
        fct=function([v],op,mode=self.mode)
        
        prob=N.random.random(N.prod(shp)).reshape(shp)
        val=N.random.random(prob.shape)
        nb_missing=0
        for i in range(shp[0]):
            for j in range(shp[1]):
                for k in range(shp[2]):
                    if prob[i,j,k]<0.1:
                        nb_missing+=1
                        val[i,j,k]=N.nan
        t=time.time()
        out=fct(val)
        print "time %.3fs"%(time.time()-t)
        for i in range(shp[0]):
            for j in range(shp[1]):
                for k in range(shp[2]):
                    if N.isnan(val[i,j,k]):
                        assert out[0][i,j,k]==0
                        assert out[1][i,j,k]==0
                    else:
                        assert out[0][i,j,k]==val[i,j,k]
                        assert out[1][i,j,k]==1

if __name__ == '__main__':
    t = TestFillMissing("test_vector")
    t.test_vector()
#    from theano.tests import main
#    main("test_sp")