import math

# 列主元高斯消元法
def SovleRowMain(A,b):
    ks = 0.00000001
    n = len(A)
    if len(A[0]) != n:
        print("A要为方阵")
        return None,None,None,None
    if len(b) != n:
        print("b与A的行数不匹配")
        return None,None,None,None
    p = list(range(n))
    for i in range(n):
        row_max = abs(A[i][i])
        row_max_index = i
        for j in range(i + 1, n):
            if abs(A[j][i]) > row_max:
                row_max = abs(A[j][i])
                row_max_index = j
        A[i], A[row_max_index] = A[row_max_index], A[i]
        b[i], b[row_max_index] = b[row_max_index], b[i]
        p[i], p[row_max_index] = p[row_max_index], p[i]

        if abs(A[i][i]) < ks:
            print("A矩阵奇异，无法进行高斯消元")
            return None,None,None,None
        for j in range(i + 1, n):
            m = A[j][i] / A[i][i]
            A[j][i] = m
            for k in range(i + 1, n):
                A[j][k] -= m * A[i][k]
            b[j] -= m * b[i]

    if abs(A[n - 1][n - 1]) < ks:
        print("A矩阵奇异，无法进行高斯消元")
        return None,None,None,None
    
    # 回代求解
    b[n - 1] /= A[n - 1][n - 1]
    for i in range(n - 2, -1, -1):
        for j in range(i + 1, n):
            b[i] -= A[i][j] * b[j]
        b[i] /= A[i][i]
    return b


# 最小二乘法拟合
def LeastSquares(list_x,list_y,n):
    m = len(list_x)
    for i in range(2*n):
        tl = [list_x[j]**(i+1) for j in range(m)]
        print(f"x^{i+1} :",tl)
        print("sum:", sum(tl))
    for i in range(n+1):
        tl = [list_y[j]*list_x[j]**i for j in range(m)]
        print(f"y*x^{i}: ",tl)
        print("sum:", sum(tl))
    x_n = []
    b = []
    for i in range(2*n+1):
        x_n.append(0)
        b.append(0)
        for j in range(m):
            x_n[i]+=(list_x[j]**i)
            b[i]+=(list_y[j]*list_x[j]**i)
    b = b[:n+1]
    A = []
    for i in range(n+1):
        tmp = []
        for j in range(n+1):
            tmp.append(x_n[i+j])
        A.append(tmp)
    print("A:", A)
    print("b:", b)
    result = SovleRowMain(A, b)
    print("result:", result)
    return result

# 计算多项式在给定x值上的值
def CalculateY(list_x, coeff):
    re = []
    for i in range(len(list_x)):
        re.append(0)
        for j in range(len(coeff)):
            re[i] += coeff[j]*list_x[i]**j
    print("y预测:", re)
    return re

# 计算均方根误差
def MeanSquareErr(list_y,list_y_approx):
    m = len(list_y)
    err = 0
    print("y差值: ")
    for i in range(m):
        print(f"y-y测={list_y[i] - list_y_approx[i]}, (y-y测)^2={(list_y[i] - list_y_approx[i])**2}")
        err += (list_y[i] - list_y_approx[i])**2
    return err**0.5

# 计算最大误差
def MaxErr(list_y,list_y_approx):
    m = len(list_y)
    err = 0
    for i in range(m):
        if abs(list_y[i] - list_y_approx[i]) > err:
            err = abs(list_y[i] - list_y_approx[i])
    return err

# 打印拟合方程
def PrintEquation(coeff):
    n = len(coeff)
    str_ = str(coeff[0]) + "+"
    for i in range(0,n-1):
        str_ += str(coeff[i+1])
        str_ += "*x^"+str(i+1)+"+"
    str_ = str_[0:len(str_)-1]
    print(str_)


if __name__ == "__main__":
    ##########################################################################
    # 下面的数值根据题意改成适合的，比如要取对数就先取对数，下面的x和y是直接用于拟合的
    list_x = [2,4,6,8]        # 已给出的x数值，与y数值对应
    list_y = [2,11,28,40]     # 已给出的y数值，与x数值对应

    # 记得改下面最后一个参数，为拟合阶数
    print("一次拟合")
    coeff = LeastSquares(list_x,list_y,1)
    PrintEquation(coeff)
    y_approx = CalculateY(list_x, coeff)
    print("均方根误差:",MeanSquareErr(list_y,y_approx))
    print("最大误差:",MaxErr(list_y,y_approx))

    print()
    
    print("二次拟合")
    coeff = LeastSquares(list_x,list_y,2)
    PrintEquation(coeff)
    y_approx = CalculateY(list_x, coeff)
    print("均方根误差:",MeanSquareErr(list_y,y_approx))
    print("最大误差:",MaxErr(list_y,y_approx))