import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
# Excel预测函数
# y = 1771.8 + 28.777x

# 数据集
x = [150, 200, 250, 300, 350, 400, 600]
y = [6450, 7450, 8450, 9450, 11450, 15450, 18450]
# 迭代过程直线起始y值
data = [[] for i in range(int(1E6))]


# 损失函数
def LossFunctionX(theta0, theta1):
    los = 0
    for i in range(len(x)):
        h = (theta0 + theta1 * x[i])  # 对theta0求偏导
        los = los + (h - y[i])
    return los


# 损失函数
def LossFunctionY(theta0, theta1):
    los = 0
    for i in range(len(x)):
        h = (theta0 + theta1 * x[i])  # 对theta1求偏导
        los = los + (h - y[i]) * x[i]
    return los


# GIF更新函数
def update(i):
    label = 'timestep {0}'.format(i)
    # print(label)
    line.set_xdata([x[0], x[len(x) - 1]])
    line.set_ydata(data[i])  # 重绘直线
    ax.set_xlabel(label)  # 重绘x轴标签
    return line, ax


# 梯度下降
def GradientDescent(theta0, theta1, iters, alphaX, alphaY):
    for i in range(iters):
        # 迭代更新theta值
        theta0 = theta0 - alphaX * LossFunctionX(theta0, theta1) / len(x)
        theta1 = theta1 - alphaY * LossFunctionY(theta0, theta1) / len(x)
        # print(theta0, theta1)  # 输出迭代过程中theta值

        # 画出迭代过程中拟合线性方程
        plt.plot(x, y, 'bo')
        plt.plot([x[0], x[len(x) - 1]], [theta0 + theta1 * x[0], theta0 + theta1 * x[len(x) - 1]], 'b')

        # 记录data值
        data[i].append(theta0 + theta1 * x[0])
        data[i].append(theta0 + theta1 * x[len(x) - 1])
    return theta0, theta1


if __name__ == '__main__':
    # 参数分别为(theta0初始值, theta1初始值, 迭代次数, theta0学习速率, theta1学习速率)
    a, b = GradientDescent(0, 0, int(4E3), 1E-1, 1E-5)

    # 绘制图形部分
    plt.plot(x, y, 'bo')  # 画出散点图
    plt.plot([x[0], x[len(x) - 1]], [a + b * x[0], a + b * x[len(x) - 1]], 'r')  # 画出拟合出的线性方程
    plt.show()  # 显示图形

    # 绘制GIF部分
    fig, ax = plt.subplots()
    fig.set_tight_layout(True)
    ax.scatter(x, y)  # 绘制散点图, 不刷新
    line, = ax.plot([x[0], 0], [x[len(x) - 1], 0], 'b')  # 绘制初始直线, update函数对其进行更新
    anim = FuncAnimation(fig, update, frames=10, interval=1000)  # frames为帧数, interval为每帧停留的的时间ms
    anim.save('line.gif', dpi=100, writer='pillow')  # dpi为分辨率

    # 输出部分
    print("theta0: {0}".format(a))
    print("theta1: {0}".format(b))