'''
Description: Training Multiple Linear Regression Models Using Gradient Descent Methods
'''
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 定义损失函数
def loss_function(theta0, theta1, theta2, x1, x2, y):
    total_error = 0
    for i in range(len(x1)):
        total_error += (y[i] - (theta1 * x1[i] + theta2 * x2[i] + theta0)) ** 2
    error = total_error / (2 * len(x1))
    return error

# 数据标准化
def Scaler(feature):
    gap, min_val = max_min(feature)
    feature = (feature - min_val) / gap
    return feature

# 计算特征最值
def max_min(x):
    max_val = x.max()
    min_val = x.min()
    gap = max_val - min_val
    return gap, min_val

# 梯度下降
def gradientDescent(theta0, theta1, theta2, x1, x2, y, learning_rate, n_iterables, loss_history):
    m = len(x1)
    for i in range(n_iterables):
        h_theta = theta1 * x1 + theta2 * x2 + theta0  # 计算模型的预测值
        grad1 = np.sum((h_theta - y) * x1) / m  # 计算theta1的梯度
        grad2 = np.sum((h_theta - y) * x2) / m  # 计算theta2的梯度
        grad0 = np.sum(h_theta - y) / m  # 计算theta0的梯度
        theta1 = theta1 - learning_rate * grad1  # 更新theta1
        theta2 = theta2 - learning_rate * grad2  # 更新theta2
        theta0 = theta0 - learning_rate * grad0  # 更新theta0
        loss_history.append(loss_function(theta0, theta1, theta2, x1, x2, y))  # 记录损失函数值
    return theta0, theta1, theta2

def main():
    data = np.loadtxt('data2.txt', delimiter=',')  # 加载数据
    x1 = data[:, 0].reshape(-1, 1)
    x2 = data[:, 1].reshape(-1, 1)
    y = data[:, 2].reshape(-1, 1)
    theta1 = 0
    theta2 = 0
    theta0 = 0
    learning_rate = 0.1  # 设置学习率
    epoch = 1000  # 设置迭代次数

    x1 = Scaler(x1)
    x2 = Scaler(x2)
    y = Scaler(y)

    loss_history = []  # 初始化损失函数历史记录
    theta0, theta1, theta2 = gradientDescent(theta0, theta1, theta2, x1, x2, y, learning_rate, epoch, loss_history)

    print(f'theta0={theta0}, theta1={theta1}, theta2={theta2}')

    # 绘制损失函数历史
    plt.figure(figsize=(10, 5))
    plt.plot(loss_history, label='Loss History')
    plt.xlabel('Iterations')
    plt.ylabel('Loss')
    plt.title('Loss History')
    plt.legend()
    plt.show()

    # 绘制3D数据和拟合平面
    fig = plt.figure(figsize=(10, 8))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(x1, x2, y, color='blue', label='Data')

    x1_plot = np.linspace(x1.min(), x1.max(), 10)
    x2_plot = np.linspace(x2.min(), x2.max(), 10)
    X1, X2 = np.meshgrid(x1_plot, x2_plot)
    Y_plot = theta1 * X1 + theta2 * X2 + theta0
    ax.plot_surface(X1, X2, Y_plot, color='red', alpha=0.5, label='Fitted Plane')

    plt.title('3D Data and Fitted Plane')
    plt.xlabel('x1')
    plt.ylabel('x2')
    plt.show()

if __name__ == "__main__":
    main()