import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, AutoMinorLocator
from scipy.optimize import minimize

def msePolynomial(x, y, degree):
    a = np.vander(x, degree + 1)
    pa = np.linalg.pinv(a)
    w = pa @ y
    return w
    
def huberFunction(x):
    ax = abs(x)
    if ax < 1.:
        return ax*ax*0.5
    else:
        return ax - 0.5

def huberPolynomial(x, y, degree):
    m = len(x)
    assert m <= len(y)
    if m == 0:
        return None

    class errorFunction:
        def __init__(self, x, degree, y):
            self.x = x
            self.y = y
            self.degree = degree

        def __call__(self, pol):
            m = len(self.x)
            n = self.degree + 1
            res = 0.
            for i in range(m):
                v = np.polyval(pol, self.x[i])
                res += huberFunction(v - self.y[i])
            return res

    # Initial approximation: polynomial of degree zero
    meanValue = sum(y)/m
    n = degree + 1
    pol = np.zeros(n)
    pol[degree] = meanValue

    res = minimize(fun = errorFunction(x, degree, y), x0 = pol)
    print(
        "minimization: success =", res.success,
        "\n    message =", res.message
    );
    w = res.x
    return w

def main():
    deg = int(input("Enter a degree of polynomial: "))
    m = int(input("Enter the number of nodes: "))
    x = np.array([0.]*m)
    y = np.array([0.]*m)
    print("Enter array of nodes x, y:")
    for i in range(m):
        x[i] = float(input("x["+str(i)+"]: "))
        y[i] = float(input("y["+str(i)+"]: "))
    wMSE = msePolynomial(x, y, deg)
    wHuber = huberPolynomial(x, y, deg)
    print("Coefficients of MSE polynomial:")
    print(wMSE)
    print("Values of MSE polynomial in nodes:")
    print(np.polyval(wMSE, x))

    print("Coefficients of Huber polynomial:")
    print(wHuber)
    print("Values of Huber polynomial in nodes:")
    print(np.polyval(wHuber, x))
    
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    
    ax.set_xlim(-8., 8.)
    ax.set_ylim(-6., 6.)
    ax.set_aspect(1)
    
    ax.axhline(0, color="black")
    ax.axvline(0, color="black")
    
    ax.xaxis.set_major_locator(MultipleLocator(5))
    ax.xaxis.set_minor_locator(AutoMinorLocator(5))
    ax.yaxis.set_major_locator(MultipleLocator(5))
    ax.yaxis.set_minor_locator(AutoMinorLocator(5))

    ax.grid(which='major')
    ax.grid(which='minor', linestyle=":")
   
    ax.scatter(x, y, color="red")
    xx = np.arange(-9., 9., 0.05)
    yy = np.polyval(wMSE, xx)
    ax.plot(xx, yy, color="blue")
    yyy = np.polyval(wHuber, xx)
    ax.plot(xx, yyy, color="green")
    plt.show()

if __name__ == "__main__":
    main()