#!/usr/bin/env python
#version    :       2023.09.07
#language   :       en
import time
from maix import camera, display, image
import os
import sys
sys.path.append('/root/')
from CocoPi import BUTTON

key_A = BUTTON(14)
key_B = BUTTON(8)
key_C = BUTTON(13)
key_D = BUTTON(7)

#image.load_freetype("/root/preset/fonts/CascadiaCodePL-Italic.ttf")
image.load_freetype("/root/preset/fonts/simhei.ttf")

import numpy as np

def sigmoid(x):
    return 1/(1+np.exp(-x))


def grad(x):
    return x*(1-x)


class NeuralNetwork:
    """
    三层全连接前馈神经网络
    """

    def __init__(self, inputnodes, hiddennodes, outputnodes, learningrate, active_function=sigmoid, gradient=grad, lambda_=0.1):
        """

        :param inputnodes: 输入层结点数
        :param hiddennodes: 隐藏层节点数
        :param outputnodes: 输出层节点数
        :param learningrate: 学习率
        :param active_function: 激活函数
        :param gradient: 激活函数的导数
        :param lambda_: L2正则化系数
        """
        self.inputnodes = inputnodes
        self.hiddennodes = hiddennodes
        self.outputnodes = outputnodes
        self.learningrate = learningrate
        self.active_function = active_function
        self.gradient = gradient
        self.lambda_ = lambda_

        # 权值矩阵
        self.weights_i_h = np.random.rand(
            self.hiddennodes, self.inputnodes) - 0.5
        self.weights_h_o = np.random.rand(
            self.outputnodes, self.hiddennodes) - 0.5

    def train_sgd(self, x, y):
        """梯度下降训练"""
        train_x = np.array(x).reshape(-1, 1)
        target = np.zeros((self.outputnodes, 1)) + 0.01
        target[y, 0] = 0.99

        hiddeninputs = np.dot(self.weights_i_h, train_x)
        hiddenoutputs = self.active_function(hiddeninputs)

        outputinputs = np.dot(self.weights_h_o, hiddenoutputs)
        final_outputs = self.active_function(outputinputs)

        error = target - final_outputs

        hidden_error = np.dot(self.weights_h_o.transpose(), error)

        self.weights_h_o += self.learningrate * error * \
            np.dot(self.gradient(final_outputs), hiddenoutputs.transpose())

        self.weights_i_h += self.learningrate * hidden_error * \
            np.dot(self.gradient(hiddenoutputs), train_x.transpose())

    def fit(self, train_x, targets):
        train_x = np.array(train_x)
        for i in range(train_x.shape[0]):
            self.train_sgd(train_x[i], targets[i])

    def query(self, inputs, debug=False):
        """单个值预测"""
        inputs = np.array(inputs).reshape(-1, 1)
        hidden_input = np.dot(self.weights_i_h, inputs)
        hidden_output = self.active_function(hidden_input)

        output_input = np.dot(self.weights_h_o, hidden_output)

        final_output = self.active_function(output_input)

        if debug:
            print('predict: ', final_output)

        return np.argmax(final_output)

    def predict(self, inputs):
        """批量预测"""
        res = []
        for x in inputs:
            res.append(self.query(x))
        return res

    def __str__(self):
        return "NeuralNetwork: \ninput_nodes = {0}, hidden_nodes = {1}, \noutputnodes = {2}, learningrate = {3}".format(
            self.inputnodes, self.hiddennodes, self.outputnodes, self.learningrate
        )


test_df = np.loadtxt("/root/preset/training/res/mnist_test_10.csv", delimiter=",", dtype=str)
test_df

train_df = np.loadtxt("/root/preset/training/res/mnist_train_100.csv", delimiter=",", dtype=str)
train_df

# 用测试数据测试
def accuracy(y_true, y_pred):
    """准确度"""
    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    return sum(y_true == y_pred)/y_true.shape[0]

# 用全部数据进行训练


def get_data():
    # train_df = np.loadtxt("mnist_train.csv", delimiter=",", dtype=str)
    # test_df = np.loadtxt("mnist_test.csv", delimiter=",", dtype=str)
    global train_df, test_df
    print(train_df.shape)
    print(test_df.shape)

    train_data = train_df.astype('int')
    train_x = train_data[:, 1:]
    train_y = train_data[:, 0]
    train_x = train_x / 255 * 0.99 + 0.01

    test_data = test_df.astype('int')
    test_x = test_data[:, 1:]
    test_y = test_data[:, 0]
    test_x = test_x / 255 * 0.99 + 0.01

    return train_x, train_y, test_x, test_y

canvas = image.new(size=(320, 240),color = (15,21,46),mode = "RGB")
#检查numpy.random模块是否已加载...
canvas.draw_string(6, 8 , "Checking numpy module state...", scale = 1, color = (255,255,255), thickness = 2)
display.show(canvas)
if os.path.exists('/usr/lib/python3.8/site-packages/numpy/random/__init__.py'):
    canvas.draw_string(6, 28 , "Numpy.random module have't been loaded ...", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    canvas.draw_string(6, 48 , "Loading numpy.random module...", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    os.system('rm -rf /usr/lib/python3.8/site-packages/numpy/random/__init__.py')
    canvas.draw_string(6, 68 , "Loaded numpy.random module successfully!", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    canvas.draw_string(6, 88 , "The program will exit to complete configuration.", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    offTime=3
    while offTime:
        canvas.draw_rectangle(12, 108, 230, 127, color=(15,21,46),thickness=-1)
        canvas.draw_string(6, 108 , "The program will exit in {} seconds.".format(offTime), scale = 1, color = (255,255,255), thickness = 2)
        display.show(canvas)
        time.sleep(1)
        offTime=offTime-1
else:
    canvas.draw_string(6, 28 , "Numpy.random module have been loaded ...", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    #如果未加载，选择提示...
    #如果已加载，显示Numpy已加载
    canvas.draw_string(6, 48 , "Building neural network training services...", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    #构建神经网络训练服务
    train_x, train_y, test_x, test_y = get_data()

    NN = NeuralNetwork(784, 100, 10, 0.3)
    NN.fit(train_x, train_y)
    y_pred = NN.predict(test_x)
    print("accuracy%.2f%%" % (100*accuracy(test_y, y_pred)))

    hiddennodes = [512, 256, 128]
    lrs = [0.1, 0.2, 0.3]
    for node in hiddennodes:
        for lr in lrs:
            NN = NeuralNetwork(784, node, 10, lr)
            NN.fit(train_x, train_y)
            y_pred = NN.predict(test_x)
            print("The number of hidden layer nodes%d,Learning rate%f,accuracy%.2f%%" %
                  (node, lr, 100*accuracy(test_y, y_pred)))

    import pickle
    # 最佳参数
    # 隐藏层节点数128,学习率0.100000,准确度70.00%
    NN = NeuralNetwork(784, 128, 10, 0.1)
    #开启神经网络训练，记录时长
    # 训练10次，每3次训练下降一次学习率
    canvas.draw_string(6, 68 , "Start NN training for 10 training iterations...", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    for e in range(1, 11):
        if e % 3 == 0:
            NN.learningrate /= 2
        NN.fit(train_x, train_y)
        y_pred = NN.predict(test_x)
        #"第%d次训练,准确度%.2f%%"
        #文件已保存至
        canvas.draw_rectangle(6, 108, 319, 147, color=(15,21,46),thickness=-1)
        canvas.draw_string(6, 108 , "The %dth training,accuracy:%.2f%%" % (e, 100*accuracy(test_y, y_pred)), scale = 1, color = (209,72,54), thickness = 2)
        canvas.draw_string(6, 128 , 'Model file path:/root/user/model/NN{}.pkl'.format(e), scale = 1, color = (209,72,54), thickness = 2)
        display.show(canvas)
        print("The %d th training,accuracy%.2f%%" % (e, 100*accuracy(test_y, y_pred)))
        with open('/root/user/model/NN{}.pkl'.format(e), 'wb') as f:  # 保存模型
            pickle.dump(pickle.dumps(NN), f)
    canvas.draw_string(6, 168 , "Model training completed!", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    os.system('touch /usr/lib/python3.8/site-packages/numpy/random/__init__.py')
    canvas.draw_string(6, 188 , "Numpy.random module have been unloaded ...", scale = 1, color = (255,255,255), thickness = 2)
    display.show(canvas)
    #训练已完成，是否要关闭numpy.random模块？
    #是，numpy.random模块已关闭，开机速度已优化
    #否，numpy.ramdom模块将持续占用系统内存，开机时将需要更多时间
    offTime=5
    while offTime:
        canvas.draw_rectangle(6, 208, 230, 227, color=(15,21,46),thickness=-1)
        canvas.draw_string(6, 208 , "The program will exit in {}seconds.".format(offTime), scale = 1, color = (255,255,255), thickness = 2)
        display.show(canvas)
        time.sleep(1)
        offTime=offTime-1