import cPickle as pickle
import numpy as np
import argparse

# sys.path.append(os.path.join(os.path.dirname(__file__), '../../build/python'))
from singa import layer
from singa import loss
from singa import device
from singa import tensor
from singa import optimizer
from singa import initializer
from singa.proto import model_pb2
from tqdm import tnrange
from word2tensor import load_data,numpy2tensors,convert,labelconvert

import time

def get_lr(epoch):
        return 0.001 / float(1 << (epoch / 50))

if __name__ == "__main__":
        # SGD with L2 gradient normalization 
        vocab_size=7000
        opt = optimizer.RMSProp(constraint=optimizer.L2Constraint(5))
        cuda = device.create_cuda_gpu_on(1)
        encoder = layer.LSTM(name='lstm1', hidden_size=64, num_stacks=5, dropout=0.5, input_sample_shape=(vocab_size,))
        decoder = layer.LSTM(name='lstm2', hidden_size=64, num_stacks=5, dropout=0.5, input_sample_shape=(vocab_size,))
        encoder.to_device(cuda)
        decoder.to_device(cuda)
        encoder_w = encoder.param_values()[0]
        encoder_w.uniform(-0.08, 0.08)
        decoder_w = decoder.param_values()[0]
        decoder_w.uniform(-0.08, 0.08)

        dense = layer.Dense('dense', vocab_size, input_sample_shape=(64,))
        dense.to_device(cuda)
        dense_w = dense.param_values()[0]
        dense_b = dense.param_values()[1]
        initializer.uniform(dense_w, dense_w.shape[0], 0)
        dense_b.set_value(0)

        #g_encoder_w = tensor.Tensor(encoder_w.shape, cuda)
        #g_encoder_w.set_value(0.0)
        g_dense_w = tensor.Tensor(dense_w.shape, cuda)
        g_dense_b = tensor.Tensor(dense_b.shape, cuda)
        g_dense_w.set_value(0.0)
        g_dense_b.set_value(0.0)

        lossfun = loss.SoftmaxCrossEntropy()
        batch_size=50
        maxlength=22
        num_train_batch=5000
        num_epoch=5
        metadata, idx_q, idx_a=load_data()
        print idx_q.shape, idx_a.shape
        for epoch in range(num_epoch):
                train_loss = 0
                bar =range(num_train_batch)
                for b in range (num_train_batch):
                        batcha=idx_a[b * batch_size: (b + 1) * batch_size]
                        batchq=idx_q[b * batch_size: (b + 1) * batch_size]
                        inputs=convert(batchq,batch_size,20,vocab_size,cuda)
                        #print 'origin input:', len(inputs), inputs[0].shape
                        inputs.append(tensor.Tensor())
                        inputs.append(tensor.Tensor())
                        outputs = encoder.forward(model_pb2.kTrain, inputs)[-2:]
                        #print 'output for origin input:', len(outputs), outputs[0].shape, outputs[1].shape
                        inputs2=convert(batcha,batch_size,22,vocab_size,cuda)[:-1]
                        inputs2.extend(outputs)
                        #print len(inputs2), inputs2[-2].shape, inputs2[-1].shape
                        labels=labelconvert(batcha,cuda)[1:]
                        #print 'gt:', len(labels), labels[0].shape
                        grads=[]
                        batch_loss=0
                        outputs2 = decoder.forward(model_pb2.kTrain, inputs2)[0:-2]
                        #print 'output from labeled data:', len(outputs2), outputs2[0].shape
                        for output,label in zip(outputs2,labels):
                                act=dense.forward(model_pb2.kTrain, output)
                                lvalue=lossfun.forward(model_pb2.kTrain,act,label)
                                batch_loss += lvalue.l1()

                                grad=lossfun.backward()
                                grad/=batch_size
                                grad,gwb=dense.backward(model_pb2.kTrain,grad)
                                grads.append(grad)
                                g_dense_w += gwb[0]
                                g_dense_b += gwb[1]

                        train_loss += batch_loss
                        print '\nbatch %d loss is %f' % (b,batch_loss / maxlength)

                        grads.append(tensor.Tensor())
                        grads.append(tensor.Tensor())
                        g_decoder_w = decoder.backward(model_pb2.kTrain, grads)[1][0]
                        g_encoder_w = encoder.backward(model_pb2.kTrain, grads[:-1])[1][0]
                        dense_w, dense_b = dense.param_values()
                        opt.apply_with_lr(epoch, get_lr(epoch), g_decoder_w, decoder_w,'decoderw')
                        #opt.apply_with_lr(epoch, 0.0, g_encoder_w, encoder_w,'encoderw')
                        opt.apply_with_lr(epoch, get_lr(epoch), g_dense_w, dense_w, 'dense_w')
                        opt.apply_with_lr(epoch, get_lr(epoch), g_dense_b, dense_b, 'dense_b')

                        #time.sleep(0.2)
                        g_dense_w.set_value(0.0)
                        g_dense_b.set_value(0.0)
                print '\nEpoch %d, train loss is %f' % (epoch, train_loss /num_train_batch / maxlength)

@@ -1672,74 +1672,8 @@
 uda)
-%0A        g_dense_w.set_value(0.0)%0A        g_dense_b.set_value(0.0)
 %0A%0A  
@@ -3067,16 +3067,114 @@
 _loss=0%0A
+                        g_dense_w.set_value(0.0)%0A                        g_dense_b.set_value(0.0)%0A
         
@@ -4959,32 +4959,32 @@
 .set_value(0.0)%0A
-
                 
@@ -5059,20 +5059,717 @@
 _batch / maxlength)%0A
+                with open('%25d_model64single.bin'%25(epoch),'wb')as fd:%0A                  if epoch%2510==0:%0A                    print 'saving model'%0A                    d=%7B%7D%0A                    for name, w in zip(%5B'encoder_w','decoder_w', 'dense_w', 'dense_b'%5D,%5Bencoder_w,decoder_w, dense_w, dense_b%5D):%0A                        d%5Bname%5D = tensor.to_numpy(w)%0A                    d%5B'hidden_size'%5D = 64%0A                    d%5B'num_stacks'%5D = 1%0A                    d%5B'dropout'%5D = 0.5%0A                    pickle.dump(d, fd)%0A                  else:%0A                    print 'it is just a test model'%0A        np.savetxt(%22batchresult3.txt%22, batchlosslist);%0A        np.savetxt(%22trainresult3.txt%22, trainlosslist);%0A

Update model.py