import numpy as np
import sys
from math import sqrt

def swap(v,i,j): #change rows
    if len(v.shape) == 1:
        v[i],v[j] = v[j],v[i]
    else:
        v[[i,j],:] = v[[j,i],:]
        
def swapRows(v,i,j): #change rows
    if len(v.shape) == 1:
        v[i],v[j] = v[j],v[i]
    else:
        v[[i,j],:] = v[[j,i],:]
        
def swapCols(v,i,j): #change coloumns
    v[:,[i,j]] = v[:,[j,i]] 
        
def error(string): #When an error occurs during execution of a program an exception is raised and the program stops
    print(string)
    input('Press return to exit')
    sys.exit()
    
"""
The function polyFit in this module sets up and solves the normal equations for the coefficients of a polynomial of degree m. 
It returns the coefficients of the polynomial.  
Coefficient matrix  are first stored in the vector s and then inserted into a. 
The normal equations are then solved by Gauss elimination with pivoting.
"""

def polyFit(xData,yData,m):
    a = np.zeros((m+1,m+1))
    b = np.zeros(m+1)
    s = np.zeros(2*m+1)
    for i in range(len(xData)):
        temp = yData[i]
        for j in range(m+1):
            b[j] = b[j] + temp
            temp = temp*xData[i]
        temp = 1.0
        for j in range(2*m+1):
            s[j] = s[j] + temp
            temp = temp*xData[i]
        for i in range(m+1):
            for j in range(m+1):
                a[i,j] = s[i+j]
                    
    def gaussPivot(a,b,tol=1.0e-12):
        n = len(b)        
        s = np.zeros(n)
        for i in range(n):
            s[i] = max(np.abs(a[i,:])) # max():Choose the biggest number in the series.
        for k in range(0,n-1):            
            p = np.argmax(np.abs(a[k:n,k])/s[k:n]) + k #argmax():Returns the indices of the maximum values along an axis.
            if abs(a[p,k]) < tol: error.err('Matrix is singular') #
            if p != k:
                swap.swapRows(b,k,p)
                swap.swapRows(s,k,p)
                swap.swapRows(a,k,p)

            for i in range(k+1,n):
                if a[i,k] != 0.0:
                    lam = a[i,k]/a[k,k]
                    a[i,k+1:n] = a[i,k+1:n] - lam*a[k,k+1:n]
                    b[i] = b[i] - lam*b[k]
        if abs(a[n-1,n-1]) < tol: error.err('Matrix is singular')
        
        b[n-1] = b[n-1]/a[n-1,n-1]
        for k in range(n-2,-1,-1):
            b[k] = (b[k] - np.dot(a[k,k+1:n],b[k+1:n]))/a[k,k]
        return b

    return gaussPivot(a,b)

"""
After the solution is found, the standard deviation σ can be computed with the function std-Dev. 
The polynomial evaluation in stdDev is carried out by the embedded function evalPoly.
(Analysis of the F (x) function.)
"""
def stdDev(c,xData,yData):
    def evalPoly(c,x):
        m = len(c) - 1
        p = c[m]
        for j in range(m):
            p = p*x + c[m-j-1] 
        return p
    
    n = len(xData) - 1
    m = len(c) - 1
    sigma = 0.0
    for i in range(n+1):
        p = evalPoly(c,xData[i])
        sigma = sigma + (yData[i] - p)**2
    sigma = sqrt(sigma/(n - m))
    return sigma

xData = np.asarray( [ 0 ,1, 2, 2.5, 3 ] )
yData = np.asarray( [ 2.9, 3.7, 4.1, 4.4, 5 ] )

degree = 1
coeff = polyFit( xData, yData, degree )
s = stdDev(coeff,xData,yData)
print ("degree= %5d; coeff[a  b]= %s; \
       standartdeviation= %10.6f;" % (degree,coeff,s))

"""
or solution:
"""
"""
while True:
    try:
        m = eval(input("\nDegree of polynomial ==> "))
        coeff = polyFit(xData,yData,m)
        print("Coefficients are:\n",coeff)
        print("Std. deviation =",stdDev(coeff,xData,yData))
    except SyntaxError: break
input("Finished. Press return to exit")

"""