import numpy as np
import tensorflow as tf
from tensorflow.contrib.factorization import KMeans

# 本程式碼演示K均值的用法， tensorflow版本必須大於等於V1.1.0
# 程式碼專案：Project: https://github.com/aymericdamien/TensorFlow-Examples/

# 由於tensorflow實現的K均值演算法無法從GPU中獲得額外好處，所以我們忽略GPU裝置
import os

os.environ['CUDA_VISIBLE_DEVICES'] = ''

# 匯入MNIST資料集
from tensorflow.examples.tutorials.mnist import input_data

mnist = input_data.read_data_sets('../data', one_hot=True)
full_data_x = mnist.train.images # shape:(55000, 784)，注意記住這個55000，理解後面會用到

# 模型超引數
num_steps = 50  # 訓練的總步數
batch_size = 1024  # 每個batch的樣本數
k = 25  # K的大小
num_classes = 10  # 十個數字，這也是模型最終分類的個數
num_features = 784  # 每個圖片都是28X28，共784個畫素

# 輸入圖片
X = tf.placeholder(tf.float32, shape=[None, num_features])
# 標註
Y = tf.placeholder(tf.float32, shape=[None, num_classes])

# K-Means的引數，其實是從庫裡使用提前封裝好的圖
kmeans = KMeans(inputs=X, num_clusters=k, distance_metric='cosine', use_mini_batch=True)

# 構建K-Means的計算圖
training_graph = kmeans.training_graph()

if len(training_graph) > 6:  # tensorflow 1.4及以上版本
    (all_scores, cluster_idx, scores, cluster_cnters_initialized,
     cluster_cnters_var, init_op, train_op) = training_graph
else:
    (all_scores, cluster_idx, scores, cluster_cnters_initialized,
     init_op, train_op) = training_graph

cluster_idx = cluster_idx[0] # 存放所有資料的圖心序號
avg_distance = tf.reduce_mean(scores) # 存放平均距離

# 初始化變數
init_vars = tf.global_variables_initializer()

# 建立一個tensorflow會話
sess = tf.Session()

# 執行初始化操作
sess.run(init_vars, feed_dict={X: full_data_x})
sess.run(init_op, feed_dict={X: full_data_x})

# 訓練
for i in range(1, num_steps + 1):
    _, d, idx = sess.run([train_op, avg_distance, cluster_idx], feed_dict={X: full_data_x})

    if i % 10 == 0 or i == 1:
        print('步驟 %i， 平均距離是：%f' % (i, d))

# 給每個圖心分配一個標籤
# 計算每個圖心的樣本個數，把樣本歸入離它最近的圖心（使用idx）
counts = np.zeros(shape=(k, num_classes))  # counts的shape是(25, 10),用於存放25個圖心分類的頻率計數
for i in range(len(idx)):
    # idx的shape是(55000,),每個成員都是0~24之間的值，對應所屬圖心的編號
    counts[idx[i]] += mnist.train.labels[i]
    # mnist.train.labels的shape是(55000, 10), 每個成員都是獨熱編碼，用來標註屬於哪個數字

# 將最高頻的標註分配給圖心。 len(labels_map)是25,也就是每個圖心一個成員，記錄每個圖心所屬的數字分類
labels_map = [np.argmax(c) for c in counts]
# 轉換前，labels_map的shape為(25,)
labels_map = tf.convert_to_tensor(labels_map)
# 此時labels_map變成了一個const op，輸出就是上面(25,)包含的值

# 評估模型。下面開始構建評估計算圖
# 注意：centroid_id就是對應label
cluster_label = tf.nn.embedding_lookup(labels_map, cluster_idx)
# cluster_idx輸出的tensor，每個成員都對映到labels_map的一個值。
# cluster_label的輸出就是對映的label值，後面用來跟標註比較計算準確度

# 計算準確率
correct_prediction = tf.equal(cluster_label, tf.cast(tf.argmax(Y, 1), tf.int32))
accuracy_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

# 測試模型
test_x, test_y = mnist.test.images, mnist.test.labels
print("測試準確率：", sess.run(accuracy_op, feed_dict={X: test_x, Y: test_y}))