import math
import numpy as np
import matplotlib.pyplot as plt

def sum_x_power(xs, power):
    """
    Compute the sum of the x values raised to a given power.
    @param xs the x values.
    @param power the given power.
    @return the sum.
    """
    sum = 0
    
    # Loop over the x values.
    for i in range(len(xs)):
        sum += xs[i]**power
    
    return sum

def sum_x_power_y(xs, ys, power):
    """
    Compute the sum of the x values raised to a given power
    and multiplied by the corresponding y value.
    @param xs the x values.
    @param yx the y values.
    @param power the given power.
    @return the sum.
    """
    sum = 0
    
    # Loop over the x and y values.
    for i in range(len(xs)):
        sum += (xs[i]**power)*ys[i]
    
    return sum

def polynomial_regression(degree, xs, ys):
    """
    Compute the coefficients of a polynomial regression.
    @param degree the degree of the regression equation.
    @param xs the x values.
    @param ys the y values.
    """
    size = degree + 1
    A = np.zeros([size, size])  # coefficient matrix
    b = np.zeros(size)          # right-hand side
    
    # Loop over each equation.
    for r in range(size):
        
        # The multiplier of the first term.
        sum = sum_x_power(xs, r)  
        
        # Copy that multiplier up the diagonal.
        j = 0
        for i in range(r, -1, -1):
            A[i][j] = sum
            j += 1
        
        # Right-hand side value.
        b[r] = sum_x_power_y(xs, ys, r)
        
    # Multipliers of the terms in the last equation.
    for c in range(1, size):
        
        # The multiplier of the second term.
        sum = sum_x_power(xs, degree + c)
        
        # Copy that multiplier up the diagonal.
        i = degree
        for j in range(c, size):
            A[i][j] = sum
            i -= 1
    
    # Solve the system to obtain the regression coefficients.
    a = np.linalg.solve(A, b)
    
    return [A, b, a]

def at(x, a, degree):
    """
    Compute the value of the regression function at a given x value.
    @param x the x value.
    @param a the coefficients of the regression function.
    @param degree the degree of the regression function.
    @return the function value.
    """
    n       = degree + 1
    x_power = 1
    y       = 0
    
    # Loop to compute the function value y.
    for i in range(n):
        y += a[i]*x_power
        x_power *= x
    
    return y

def plot_polynomial_regression(xs, ys, degree=0, y_axis_pos=0, size=(6, 6)):
    """
    Plot a polynomial regression line among a scatter plot of x, y values.
    @param xs the x values.
    @param ys the y values.
    @param degree the degree of the regression equation.
    @param y_pos the y position of the x axis.
    @param size the size of the graph.
    @return the vector of coefficients
    """
    fig = plt.figure(figsize=size)

    ax = fig.add_subplot(1, 1, 1)

    ax.spines['left'].set_position(('data', 0))
    ax.spines['right'].set_color('none')
    #ax.spines['bottom'].set_position(('data', y_axis_pos))
    ax.spines['top'].set_color('none')

    plt.scatter(xs, ys)
    
    if degree > 0:
        A, b, a  = polynomial_regression(degree, xs, ys)
        x_values = np.linspace(0, max(xs), 100)
        y_preds  = np.array([at(x, a, degree)for x in x_values])    
        plt.plot(x_values, y_preds, color='red')
    else:
        a = None
    
    plt.show()
    return a

def plot_polynomial(xs, ys, a, degree=0, x_min=0, x_max=math.nan, y_axis_pos=0, size=(6, 6)):
    """
    Plot a polynomial regression line among a scatter plot of x, y values.
    @param xs the x values.
    @param ys the y values.
    @param a  the vector of coefficients.
    @param degree the degree of the regression equation.
    @param y_pos the y position of the x axis.
    @param size the size of the graph.
    """
    fig = plt.figure(figsize=size)
    
    ax = fig.add_subplot(1, 1, 1)

    ax.spines['left'].set_position(('data', 0))
    ax.spines['right'].set_color('none')
    #ax.spines['bottom'].set_position(('data', y_axis_pos))
    ax.spines['top'].set_color('none')

    plt.scatter(xs, ys)
    
    if degree > 0:
        if math.isnan(x_max):
            x_max = max(xs)
        x_values = np.linspace(x_min, x_max, 100)
        y_preds  = np.array([at(x, a, degree)for x in x_values])    
        plt.plot(x_values, y_preds, color='red')
    else:
        a = None
    
    plt.show()