# encoding:utf-8
from sklearn import model_selection
from sklearn import datasets
from sklearn import metrics
import tensorflow as tf
import numpy as np
from tensorflow.contrib.learn.python.learn.estimators.estimator import SKCompat
learn = tf.contrib.learn

# 1 定义softmax回归模型
def my_model(features, target):
    target = tf.one_hot(target, 3, 1, 0) # 每个目标标签one-hot维度是3(一共就3个类别标签), 当标签出现的位置用1填充，不出现的位置用0填充

    # 计算预测值及损失函数
    logits = tf.contrib.layers.fully_connected(features, 3, tf.nn.softmax)
    loss = tf.losses.softmax_cross_entropy(onehot_labels=target, logits=logits)

    # 优化

    train_op = tf.contrib.layers.optimize_loss(loss,
                                               tf.contrib.framework.get_global_step(),
                                               optimizer='Adam',
                                               learning_rate=0.01
                                               )
    return tf.arg_max(input=logits, dimension=1), loss, train_op

# 2 将数据变换成tf要求的tf.float32
iris = datasets.load_iris()
x_train, x_test, y_train, y_test = model_selection.train_test_split(
    iris.data, iris.target, test_size=0.2, random_state=0
)
x_train, x_test = map(np.float32, [x_train, x_test])

# 3训练
classifier = SKCompat(learn.Estimator(model_fn=my_model, model_dir='Models/model_logist'))
classifier.fit(x=x_train, y=y_train, steps=800)

y_predicted = [i for i in classifier.predict(x=x_test)]

score = metrics.accuracy_score(y_true=y_test, y_pred=y_predicted)
print(score)