#!/usr/bin/env python

import sys
import paddle.v2 as paddle

input_dim=3 # 0/1/2
class_dim=3 # 0/1/2
emb_dim=64
hid_dim=64

#"""
layer_attr = paddle.attr.Extra(drop_rate=0.5)
bias_attr = paddle.attr.Param(initial_std=0., l2_rate=0.)
relu = paddle.activation.Relu()
linear = paddle.activation.Linear()
data = paddle.layer.data("status",paddle.data_type.integer_value_sequence(input_dim))
emb = paddle.layer.embedding(input=data, size=emb_dim)
fc1 = paddle.layer.fc(
                input=emb, size=hid_dim, act=linear, bias_attr=bias_attr)
lstm = paddle.layer.lstmemory(
                input=fc1, act=relu, bias_attr=bias_attr, layer_attr=layer_attr)

output = paddle.layer.fc(
                input=lstm,
                        size=class_dim,
                        act=paddle.activation.Softmax()
                        )

label = paddle.layer.data(
            name='label', type=paddle.data_type.integer_value(class_dim))
cost = paddle.layer.classification_cost(input=output, label=label)
parameters = paddle.parameters.create(cost)
#"""

def train_reader():
    def load_data():
        fname = "train.txt"
        for line in file(fname):
            line = line[:-1]
            labelstr,feastr = line.split(" ")
            label = int (labelstr)
            fea = map(int,feastr.split(","))
            #print "-"
            #print fea,[label]
            for i in xrange(0,len(fea)):
                fea[i] += i*3
    
            #print fea,[label]
            fea2 = map(lambda x:[x],fea)
            #print fea2,[[label]]
            yield fea,[label]
            #yield fea2,[label]

    return load_data

paddle.init(use_gpu=False, trainer_count=1)
adam_optimizer = paddle.optimizer.Adam(
                learning_rate=2e-3,
                        regularization=paddle.optimizer.L2Regularization(rate=8e-4),
                                model_average=paddle.optimizer.ModelAverage(average_window=0.5))

def event_handler(event):
        if isinstance(event, paddle.event.EndIteration):
            if event.batch_id % 100 == 0:
                print "\nPass %d, Batch %d, Cost %f, %s" % (
                    event.pass_id, event.batch_id, event.cost, event.metrics)
            else:
                sys.stdout.write('.')
                sys.stdout.flush()
        if isinstance(event, paddle.event.EndPass):
            result = trainer.test(reader=test_reader, feeding=feeding)
            print "\nTest with Pass %d, %s" % (event.pass_id, result.metrics)
feeding = {'status': 0, 'label': 0}

trainer = paddle.trainer.SGD(
        cost=cost, parameters=parameters, update_equation=adam_optimizer)
trainer.train(
        reader=paddle.batch(paddle.reader.shuffle(train_reader(), buf_size=100),batch_size=2),
        event_handler=event_handler,
        feeding=feeding,
        num_passes=2
        )