from sklearn import datasets
import numpy as np
import tensorflow as tf
from matplotlib import pyplot as plt

if __name__=="__main__":
    # 1: 准备数据
    x_data = datasets.load_iris().data
    y_data = datasets.load_iris().target

    # 数据集乱序
    np.random.seed(116)
    np.random.shuffle(x_data)
    np.random.seed(116)
    np.random.shuffle(y_data)
    tf.random.set_seed(116)

    # 数据集分出永不相见的训练集合测试集
    x_train = x_data[:-30]
    y_train = y_data[:-30]
    x_test = x_data[-30:]
    y_test = y_data[-30:]

    # 转换x的数据类型,否则后面矩阵相乘时会因数据类型不一致报错
    x_train = tf.cast(x_train, tf.float32)
    x_test = tf.cast(x_test, tf.float32)

    # 配成对
    train_db = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(32)
    test_db = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(32)

    # 定义神经网络中所有可调参数, 4个输入特征故，输入层为4个节点，因分3类，故输出层为3个神经元
    # 用tf.Variable定义可训练参数
    # 使用seed使每次s生成的随机数相同（方便教学，使大家结果一致，现实生活中不用写seed）stddev标准偏差
    w1 = tf.Variable(tf.random.truncated_normal([4, 3], stddev=0.1, seed=1))
    b1 = tf.Variable(tf.random.truncated_normal([3], stddev=0.1, seed=1))

    # acc 准确率  loss损失值
    lr = 0.1  # 学习率
    train_loss_results = []
    test_acc = []
    epoch = 500  # 循环500轮
    loss_all = 0  # 每轮分4个step, loss_all记录四个step生成的4个loss的和

    # 调用部分
    for epoch in range(epoch):
        for step, (x_train, y_train) in enumerate(train_db):
            with tf.GradientTape() as tape:  # with结构记录梯度信息
                y = tf.matmul(x_train, w1) + b1  # 神经网络乘加运算
                y = tf.nn.softmax(y)  # 使输出y符合概率分布
                y_ = tf.one_hot(y_train, depth=3)  # 将标签值转化为独热码格式
                loss = tf.reduce_mean(tf.square(y_ - y))  # 采用均方误差损失函数
                print(loss)
                print(loss.numpy())
                loss_all += loss.numpy()  # 将每个step计算的loss累加
            # 计算loss对每个参数的梯度,变化率
            grads = tape.gradient(loss, [w1, b1])

            # 实现梯度更新 w1=w1-lr*wl_grad  b=b-lr*b_grad
            w1.assign_sub(lr * grads[0])  # 参数w1自更新 权重
            b1.assign_sub(lr * grads[1])  # 参数b1自更新 偏置项

        # 每个epoch，打印loss信息
        print("epoch:{},loss:{}".format(epoch, loss_all/4))
        train_loss_results.append(loss_all/4)  # 将4个step的loss求平均记录在此变量中
        loss_all = 0  # loss_all归0， 为计算下一个epoch的loss准备

        # 测试部分
        # total_corrent为预测对的样本数量，total_number为总样本数
        total_corrent, total_number = 0, 0
        for x_test, y_test in test_db:
            # 使用更新后的预测参数
            y = tf.matmul(x_test, w1) + b1
            y = tf.nn.softmax(y)
            pred = tf.argmax(y, axis=1)  # 返回y中最大值的索引，即预测的分类
            # 将pred转化为y_test的数据类型
            pred = tf.cast(pred, dtype=y_test.dtype)
            # 若分类正确，则corrent=1 否则为0，将bool型的结果转化为int型
            corrent = tf.cast(tf.equal(pred, y_test), dtype=tf.int32)
            # 将每个batch的corrent数相加
            corrent = tf.reduce_sum(corrent)
            # 将所有batch的corrent数相加
            total_corrent += int(corrent)
            # 将total_number为测试样本总数，也就是x_test的行数，shape[0]返回变量的行数
            total_number += x_test.shape[0]
        # 总的准确率为total_corrent/total_number
        acc = total_corrent/total_number
        test_acc.append(acc)
        print("test_acc:", acc)
        print("-----------------------")


    # 绘制loss曲线
    plt.title("loss")
    plt.xlabel("epoch")
    plt.ylabel("loss")
    plt.plot(train_loss_results, label="loss")
    plt.legend()
    plt.show()

    # 绘制acc曲线
    plt.title("acc")
    plt.xlabel("epoch")
    plt.ylabel("acc")
    plt.plot(test_acc, label="acc")
    plt.legend()  # 画出曲线图标
    plt.show()