# .:Simple autoencoder:.
# Untuk memberikan gambaran bagaimana membangan Auto-Encoder yang sederhana
#
# Alfan Farizki Wicaksono
# Fasilkom, Universitas Indonesia


import tensorflow as tf
import numpy as np
import sys

# Autoencoder:
# Tujuan: kita ingin membuat representasi yang lebih padat dari input x, yaitu z
# x nantinya bisa berupa image, kata, noun phrase, dsb
#
# encoder -> z = W1.x + b1
# 
# tentunya representasi z harus bisa dikembalikan ke bentuk awal x
#
# decoder -> x = W2.z + b2
# jadi z adalah representasi yang lebih padat dari x


# dense layer/fully-connected layer
# x: input berdimensi [batch, input unit]
# output_unit: integer, banyaknya output unit
def dense(x, output_unit, activation_fn, scope_name='dense', reuse=False):

    input_unit = x.get_shape()[-1]
    bias_shape = [output_unit]

    with tf.variable_scope(scope_name) as scope:

        # Jika kita ingin reuse parameters; atau pakai parameter dengan nama yang sama
        # kalau tidak di-reuse, maka kita tidak boleh membuat 2 atau lebih parameter
        # dengan nama scope yang sama!
        if reuse:
            scope.reuse_variables()

        # Membuat sebuah variable untuk "dense/weights"; hindari penggunaan tf.Variable(.)!!
        weights = tf.get_variable("weights", \
                                  [input_unit, output_unit], \
                                  initializer=tf.truncated_normal_initializer(stddev=0.02))

        # Membuat sebuah variable untuk "dense/biases"
        biases = tf.get_variable("biases", \
                                 bias_shape, \
                                 initializer=tf.constant_initializer(0.0))

        y = tf.matmul(x, weights) + biases # raw logits
        
        if activation_fn == None:
            return y
        return activation_fn(y)

# standard cross-entropy sebagai loss function
# yy_logits adalah nilai sebelum masuk activation function yang terakhir
def loss(yy, yy_out):
    #cross_entropy = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=yy, logits=yy_logits))
    #cross_entropy = tf.reduce_mean(-tf.reduce_sum(yy * tf.log(yy_out), reduction_indices=[1]))
    #cross_entropy biasanya hanya cocok untuk klasifikasi, dimana outputnya adalah probabilitas
    #
    #disini, cocoknya pakai mean-squared error
    mse = tf.reduce_mean(tf.square(yy - yy_out))
    return mse

# untuk setiap batch
# fungsi ini digunakan saat training untuk mengambil batch demi batch pada sample
def gen_batch(samples, batch_size=2):
    num_sample, vector_size = samples.shape

    i = 0
    while i <= num_sample - batch_size:
        yield samples[i:(i+batch_size),:]
        i = i + batch_size

##########################################################

# data dummy
xs = np.array([[1,0,1,0,1,0,1], \
               [0,1,1,1,0,1,1], \
               [1,1,1,1,1,1,0], \
               [0,0,1,0,0,0,1], \
               [1,1,0,1,1,1,1], \
               [1,1,0,0,0,0,1]])

batch_size = 1
sample_size = len(xs)
embed_size = 7
encoded_size = 4

# placeholder untuk satu batch dataset
# autoencoder dibangun dengan set input = output
X = tf.placeholder(shape=[batch_size, embed_size], dtype=tf.float32)
Y = tf.placeholder(shape=[batch_size, embed_size], dtype=tf.float32)

# encoder
Z = dense(X, encoded_size, activation_fn=tf.nn.sigmoid, scope_name='W1_dense')

# decoder
Y_out = dense(Z, embed_size, activation_fn=tf.nn.sigmoid, scope_name='W2_dense')

# bangun graph computation untuk menghitung loss function
out_loss = loss(Y, Y_out)

# optimizer
# pakai SGD, learning_rate = 0.01
optimizer = tf.train.GradientDescentOptimizer(0.05)
train_onestep = optimizer.minimize(out_loss) #mengembalikan None, ini untuk update parameter

# training and evaluation
def train(sess, epoch=30000, step=10):
    init = tf.global_variables_initializer()

    #inisialisasi variables/parameters
    sess.run(init)

    print 'training:'
    #mulai epoch
    for i in range(epoch):

        #untuk setiap batch pada dataset
        for x in gen_batch(xs, batch_size=batch_size):

            # INGAT! di autoencoder set input = output
            _,loss = sess.run([train_onestep, out_loss], feed_dict={X:x, Y:x})

            #cetak progress untuk setiap 'step' epoch
            if (i % step == 0):
                sys.stdout.write("\r Epoch-{}, Loss-value: {}".format(i+1, loss))
                sys.stdout.flush()

        if (i % step == 0):
            sys.stdout.write("\n")
            sys.stdout.flush()


##############################################

with tf.Session() as sess:

    ### training ###
    train(sess)

    ### pasca training ###
    # gunakan Encoder untuk lihat hasil encode data dummy nya
    print 'encode semua data:'
    
    # reuse learned parameter W1_dense/weights dan W1_dense/biases
    XX = tf.placeholder(shape=[sample_size, embed_size], dtype=tf.float32)
    ZZ = dense(XX, encoded_size, activation_fn=tf.nn.sigmoid, scope_name='W1_dense', reuse=True)

    encoded_data = sess.run([ZZ], feed_dict={XX:xs})

    print 'data asli:'
    print xs

    print 'data hasil encoding:'
    print encoded_data

    # coba kalau dibalikkan lagi...
    # coba latihan...bagaimana caranya mengembalikan encoded_data ke data awal ? :)