

import unittest
from operator import concat
import numpy as np
rng = np.random.RandomState()

import theano
import theano.tensor as T

from nccn import LogisticReg, HiddenLayer, ConvPoolLayer, CnnModel

import base_test_classes

#@unittest.skip("TEMP")
class TestLRLayer( base_test_classes.CNNLayerTester ):
    def setUp(self):
        super(TestLRLayer, self).setUp()

        self.nin = rng.randint(5, 10)
        self.nout = rng.randint(2, 4)
        self.bs = rng.randint(2, 8)

    def get_input(self):
        return np.asarray( self.rng.uniform(low=-1, high=1,
            size=(self.bs, self.nin)), dtype=np.float64 )

    def get_layer(self):
        l = LogisticReg(T.matrix("X"), self.nin, self.nout,
                self.rng, self.get_input().dtype)
        ## all is initialized to 0. revert that
        l._weights_[0].set_value( np.asarray( rng.rand(self.nin, self.nout),
            dtype=l.get_weights()[0].dtype ) )
        return l

    def test_rlinit(self):
        "lr weights are all init'ed to rnd"

        l = LogisticReg(T.matrix("X"), self.nin, self.nout,
                self.rng, self.get_input().dtype)
        self.assertGreaterEqual( np.max( np.abs(l.get_weights()[0]) ), 0 )
        self.assertLessEqual( np.max( np.abs(l.get_weights()[1]) ), 0 )

    def test_init(self):
        self._test_init_()

    def test_norms(self):
        self._test_norms_()

    def test_activation(self):
        layer = self.get_layer()
        out_f = theano.function(inputs=[layer.input], outputs=layer.activation() )
        i = self.get_input()

        self.assertFalse( np.all( out_f( i ) == 1. / self.nout) )

        ## 0.5
        layer.set_weights( self.zero_weights() )
        self.assertTrue( np.all( out_f( self.get_input() ) == 1. / self.nout ) )

        layer._weights_[0].set_value( np.ones( self.zero_weights()[0].shape, self.zero_weights()[0].dtype ) )
        print out_f(i) - 1.  /self.nout
        self.assertAlmostEqual( np.max( np.abs( out_f(i) - 1.  /self.nout ) ), 0 )

        wv = np.ones( self.zero_weights()[0].shape, self.zero_weights()[0].dtype )
        wv[:,0] = 0.0
        layer._weights_[0].set_value( wv )
        o = out_f( i )
        eo = np.exp (np.dot( i, wv ) )
        d = np.sum( eo, axis=1 )
        for idx in range(eo.shape[0]):
            eo[idx,:] /= d[idx]

        print i.shape, wv.shape
        print o
        print eo
        print eo - o

        self.assertAlmostEqual( np.max( np.abs( eo - o ) ), 0 )


#@unittest.skip("TEMP")
class TestHLayer( base_test_classes.CNNLayerTester ):
    def setUp(self):
        super(TestHLayer, self).setUp()

        self.nin = rng.randint(5, 10)
        self.nout = rng.randint(2, 4)
        self.bs = rng.randint(2, 10)

    def get_input(self):
        return np.asarray( self.rng.uniform(low=-1, high=1,
            size=(self.bs, self.nin)), dtype=np.float64 )

    def get_layer(self):
        return HiddenLayer(T.matrix("X"), self.nin, self.nout,
                self.rng, self.get_input().dtype)

    def test_hlinit(self):
        "hr weights are all init'ed to [-thr, thr]"

        self._weights_inrange( np.sqrt( 6./ (self.nin+self.nout) ), 0 )
        self._weights_inrange(0.0, 1)

    def test_init(self): self._test_init_()

    def test_norms(self): self._test_norms_()

    def test_activation(self):
        layer = self.get_layer()
        out_f = theano.function(inputs=[layer.input], outputs=layer.activation() )
        i = self.get_input()

        #W [ 0, 1, 1, ...]
        wv = np.ones( self.zero_weights()[0].shape, self.zero_weights()[0].dtype )
        wv[:,0] = 0.0
        layer._weights_[0].set_value( wv )

        o = out_f( i )
        eo = 1. / (1 + np.exp (-np.dot( i, wv ) ))

        print i.shape, wv.shape
        print o
        print eo
        print eo - o

        self.assertAlmostEqual( np.max( np.abs( eo - o ) ), 0 )

#@unittest.skip("TEMP")
class TestCPLayer( base_test_classes.CNNLayerTester ):
    def setUp(self):
        super(TestCPLayer, self).setUp()

        self.imsh = (4, 1, 11, 11)
        self.fsh = (2, 1, 4, 4)
        self.psh = (2,2)

    def get_layer(self):
        return ConvPoolLayer(T.dtensor4("X"), self.fsh, self.imsh,
                self.rng, (2,2), self.get_input().dtype)

    def get_input(self):
        return np.asarray( self.rng.uniform(low=-1, high=1, size=self.imsh),
                dtype=np.float64 )

    def test_init(self): self._test_init_()

    def test_norms(self): self._test_norms_()

    def test_activation(self):
        layer = self.get_layer()
        out_f = theano.function(inputs=[layer.input], outputs=layer.activation() )
        i = self.get_input()


        # set a 0 and a 1 RF
        wv = np.ones( self.zero_weights()[0].shape, self.zero_weights()[0].dtype )
        wv[0,0] = 0.0
        wv[1,0] = 1.0
        layer._weights_[0].set_value( wv )

        o = out_f( i )
        eo = np.ones( (4, 2, 4, 4), dtype=self.get_input().dtype) / 2

        print i.shape, wv.shape, o.shape
        print o[:,0]
        print eo[:,0]
        print eo[:,0] - o[:,0]

        self.assertAlmostEqual( np.max( np.abs( eo[:,0] - o[:,0] ) ), 0 )

#@unittest.skip("TEMP")
class TestCnnModel( base_test_classes.ModelTester, base_test_classes.NpyTester ):
    def setUp(self):
        super(TestCnnModel, self).setUp()

        self.bs = rng.randint(5, 10)
        self.arch=[(2,), ((4,4),), ((2,2),)]
        self.imsh = (self.bs, 1, 11, 11)
        self.lrs = 12
        self.G = 2

        self.CW_sh = (self.arch[0][0], 1, self.arch[1][0][0], self.arch[1][0][1])
        self.Cb_sh = (self.arch[0][0], )
        self.HW_sh = (32, 12)
        self.Hb_sh = (12, )
        self.RW_sh = (12, 2)
        self.Rb_sh = (2, )

    def get_model(self):
        return CnnModel(self.arch, self.lrs, self.imsh,
                self.G, rng, self.get_input().dtype, "int32")

    def get_input(self):
        return np.asarray( rng.rand(*self.imsh) , dtype=np.float64 )

    def get_output(self):
        return np.asarray(np.zeros((self.imsh[0],)), dtype=np.int32)

    def test_topcpidx(self):
        model = self.get_model()
        self.assertEqual( model.top_cp_idx, len(self.arch[0])-1 )
        self.assertEqual( model.top_cp_idx, len(model)-3)

    #@unittest.skip("TEMP")
    def test_cost(self):
        M = self.get_model()
        l1 = rng.randint(0, 10) / 10.
        l2 = rng.randint(0, 10) / 10.
        y = self.get_output()
        x = self.get_input()


        ## check dependence on weight norms
        cf = theano.function(inputs=[M.X, M.Y], outputs=M.cost(l1, l2))
        cf_00 = theano.function(inputs=[M.X, M.Y], outputs=M.cost(0, 0))
        w_summ = l1 * M.weight_norm("l1")
        w_summ_sq = l2 * M.weight_norm("l2")
        self.assertNotEqual( 0, w_summ + w_summ_sq )
        self.assertAlmostEqual( cf(x,y) - cf_00(x,y), w_summ + w_summ_sq )

        ## check that cost is zero on right answers
        # this will make output = 0
        nw = np.copy ( M[-1]._weights_[0].get_value() )
        nw[:,0] = 1
        M[-1]._weights_[0].set_value( nw )
        out_f = theano.function(inputs=[M.X], outputs=M[-1].p_y_given_x)
        y = np.asarray(np.argmax(out_f(x), axis=1), dtype=y.dtype)
        print cf(x, y)
        print cf_00(x, y)
        print map(lambda l: l.weight_norm("l1"), M)
        print map(lambda l: l.weight_norm("l2"), M)
        print M[-1].get_weights()
        self.assertTrue( np.all(cf_00(x,y) < 0.1) )

    #@unittest.skip("TEMP")
    def test_activation_vs_i(self):
        "test activation function, sensitive to input"

        model = self.get_model()
        out_f = theano.function( inputs=[model[0].input],
                outputs=model[0].activation() )
        x = self.get_input()
        self.assertDifferentArray( out_f(x), out_f(x-x) )


    #@unittest.skip("TEMP")
    def test_activation_vs_w(self):
        "test activation function, sensitive to weights"

        model = self.get_model()
        out_f = theano.function( inputs=[model[0].input],
                outputs=model[0].activation() )
        i = self.get_input()

        print i
        o1 = out_f(i)
        for l in range(len(model)):
            model[l].set_weights( self.zero_model()[l].get_weights() )
            self.assertFalse( np.all( o1 == out_f( i ) ) )

    #@unittest.skip("TEMP")
    def test_init(self):
        model = self.get_model()
        self.assertNotEqual( 0, model.weight_norm("l1"))
        self.assertNotEqual( 0, model.weight_norm("l2"))

        layer = model[0]
        self.assertFalse( all( map(lambda s: np.all(s==0), layer.get_weights()) ) )
        self.assertNotEqual( 0, layer.weight_norm("l1"))
        self.assertNotEqual( 0, layer.weight_norm("l2"))
        self.assertTrue( all( map(lambda s: np.all(s==0), layer.get_speeds()) ) )

        layer = model[1]
        self.assertFalse( all( map(lambda s: np.all(s==0), layer.get_weights()) ) )
        self.assertNotEqual( 0, layer.weight_norm("l1"))
        self.assertNotEqual( 0, layer.weight_norm("l2"))
        self.assertTrue( all( map(lambda s: np.all(s==0), layer.get_speeds()) ) )

        layer = model[2]
        self.assertFalse( all( map(lambda s: np.all(s==0), layer.get_weights()) ) )
        self.assertTrue( all( map(lambda s: np.all(s==0), layer.get_speeds()) ) )


    #@unittest.skip("TEMP")
    def test_update(self):
        model = self.get_model()

        x = theano.shared( value=self.get_input(), name="X" )
        y = theano.shared( value=self.get_output(), name="Y" )
        idx = T.iscalar("batch_idx")
        params = reduce(concat, map(lambda l: l.get_params(), model))
        grad_f = theano.function(inputs=[idx],
                outputs=T.grad(model.cost(0, 0), wrt=params),
                givens={model.X : x[idx:self.bs * (idx+1)], model.Y : y[idx:self.bs * (idx+1)]})

        rho = rng.randint(0, 10) / 10.
        mu = rng.randint(0, 10) / 10.

        grad_v = grad_f(0)
        all_weights = reduce(concat, model.get_weights())
        all_speeds = reduce(concat, model.get_speeds())
        model.update_params([0], grad_f, mu, rho)
        for layer in model:
            grads = map(lambda i: grad_v.pop(0), range(len(layer.get_params())))
            weights = map(lambda i: all_weights.pop(0), range(len(layer.get_params())))
            speeds = map(lambda i: all_speeds.pop(0), range(len(layer.get_params())))
            print grads[0].shape
            print weights[0].shape
            print speeds[0].shape
            for i in range(len(grads)):
                print i, rho, mu
                print np.max( np.abs((layer.get_weights()[i] - layer.get_speeds()[i]) - weights[i]))
                print np.max( np.abs( layer.get_speeds()[i] + rho * grads[i] ) )
                self.assertTrue( np.all(layer.get_speeds()[i] == -rho * grads[i]) )
                self.assertAlmostEqual( np.max( np.abs((layer.get_weights()[i] - layer.get_speeds()[i]) - weights[i])), 0 )


        print
        print

        grad_v = grad_f(0)
        all_weights = reduce(concat, model.get_weights())
        all_speeds = reduce(concat, model.get_speeds())
        model.update_params([0], grad_f, mu, rho)
        for layer in model:
            grads = map(lambda i: grad_v.pop(0), range(len(layer.get_params())))
            weights = map(lambda i: all_weights.pop(0), range(len(layer.get_params())))
            speeds = map(lambda i: all_speeds.pop(0), range(len(layer.get_params())))
            print grads[0].shape
            print weights[0].shape
            print speeds[0].shape
            print
            for i in range(len(grads)):
                print i, rho, mu
                print np.max( np.abs((layer.get_weights()[i] - layer.get_speeds()[i]) - weights[i]))
                print np.max( np.abs( layer.get_speeds()[i] + rho * grads[i] ) )
                self.assertTrue( np.all(layer.get_speeds()[i] == -rho * grads[i] + mu*speeds[i]) )
                self.assertAlmostEqual( np.max( np.abs((layer.get_weights()[i] - layer.get_speeds()[i]) - weights[i])), 0 )

    def test_io(self):
        self._test_io()


