# -*- coding: utf-8 -*-
"""
Created on Mon Mar 20 22:56:32 2017

@author: Z.Y
"""

import numpy as np
import sklearn.preprocessing as prep
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data


def xavier_init(fan_in, fan_out, constant = 1):
    low  = -constant * np.sqrt(6.0 /(fan_in + fan_out))
    high = constant * np.sqrt(6.0/(fan_in + fan_out))
    return tf.random_uniform((fan_in, fan_out), minval = low, maxval= high , dtype = tf.float32)

class AdditiveGassionNoiseAntoencoder(object):
    def __init__(self, 
                 n_input, 
                 n_hidden, 
                 transfer_function = tf.nn.softplus, 
                 optimizer = tf.train.AdamOptimizer(),  #优化器，默认为 AdamOptimizer
                 scale = 0.1
                 ):
        self.n_input = n_input  #输入变量数
        self.n_hidden = n_hidden # 隐含节点数
        
        self.tranfer = transfer_function # 隐含层激活函数 ，默认为softplus
        
        self.scale = tf.placeholder(tf.float32) 
        self.traning_scale = scale  # 高斯噪声系数
        
        network_weights = self._initialiaze_weight()  # 参数初始化 定义初始化函数 
        self.weights = network_weights
        
        self.x = tf.placeholder(tf.float32,[None,self.n_input])
        self.hidden = self.tranfer(
                tf.add(tf.matmul(
                    self.x + scale*tf.random_normal((n_input,)),self.weights['w1']),
                    self.weights['b1']))
                        
        self.reconstruction = tf.add(tf.matmul(self.hidden,self.weights['w2']),self.weights['b2'])
        self.cost = 0.5 * tf.reduce_sum(tf.pow(tf.subtract(
                self.reconstruction,self.x),2.0)) 
        self.optimizer = optimizer.minimize(self.cost)
        
        init = tf.global_variables_initializer()
        self.sess = tf.Session()
        self.sess.run(init)
        
    def _initialiaze_weight(self):
        all_weight = dict()
        all_weight['w1'] = tf.Variable(xavier_init(self.n_input,self.n_hidden))
        all_weight['b1'] = tf.Variable(tf.zeros([self.n_hidden],dtype=tf.float32))
        all_weight['w2'] = tf.Variable(tf.zeros([self.n_hidden,self.n_input],dtype=tf.float32))
        all_weight['b2'] = tf.Variable(tf.zeros(self.n_input,dtype=tf.float32))
        return all_weight
    
    def partial_fit(self,X):
        cost, opt = self.sess.run((self.cost, self.optimizer),feed_dict = {self.x : X, self.scale : self.traning_scale})
        return cost
    
    def calc_total_cost(self, X):
        return self.sess.run(self.cost, feed_dict = {self.x :X ,self.scale : self.traning_scale})
    
    def transfrom(self, X):
        return self.sess.run(self.hidden, feed_dict = {self.x :X , self.scale : self.traning_scale})
    
    def generate(self, hidden = None):
        if hidden is None:
            hidden = np.random.normal(size = self.weights['b1'])
        return self.sess.run(self.reconstruction, feed_dict = {self.hidden:hidden})
    
    def reconstruction(self, X):
        return self.sess.run(self.reconstruction, feed_dict = {self.x:X,self.scale:self.training_scale})
    
    def getWeights(self):
        return self.sess.run(self.weights['w1'])
    
    def getBiases(self):
        return self.sess.run(self.weights['b1'])
    
mnist = input_data.read_data_sets('MNIST_data', one_hot = True)

def standard_scale(X_train, X_test):
    precprocessor = prep.StandardScaler().fit(X_train)
    X_train = precprocessor.transform(X_train)
    X_test = precprocessor.transform(X_test)
    return X_train, X_test

def get_random_block_from_data(data, batch_size):
    start_index = np.random.randint(0, len(data) - batch_size)
    return data[start_index:(start_index + batch_size)]

X_train, X_test = standard_scale(mnist.train.images,mnist.test.images)

n_samples = int(mnist.train.num_examples)
training_epcohs = 40
batch_size = 128
display = 1

antoencoder = AdditiveGassionNoiseAntoencoder(n_input=784,
                                              n_hidden = 400,
                                              transfer_function  = tf.nn.softplus,
                                              optimizer=tf.train.AdamOptimizer(learning_rate=0.001),
                                                                              scale=0.01)

for epoch in range(training_epcohs):
    avg_cost = 0
    total_batch = int(n_samples / batch_size)
    for i in range(total_batch):
        batch_xs = get_random_block_from_data(X_train, batch_size)
        cost  = antoencoder.partial_fit(batch_xs)
        avg_cost +=cost/n_samples*batch_size
        
    if epoch % display ==0:
        print((epoch+1),(avg_cost))
        
print(antoencoder.calc_total_cost(X_test))
    
        
    ```