# test_lsfit.py3   uses   lsfit.py3  peval.py3  in this file, could import
from numpy import array

xt = array([0.0, 0.1,  0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9])
yt = array([(0.0) for j in range(20)]) # computed below

# lsfit.py3    p = lsfit(pwr, x, y)   p[0]..p[pwr+1] computed
# i=0, i<n     y[i] = p[0] + p[1]*x[i] + p[2]*x[i]^2 + ... p[pwr]*x[i]^pwr
# The problem is stated as follows : len(x) = len(y) >= pwr+1
# Given measured data for values of Y based on values of X. e.g.
#
#  i    y_actual         x
# ---   --------       -----
#  0     32.5           1.0
#  1      7.2           2.0
#  2      6.9           3.2
#  3      8.4           2.7
#  4      9.6           1.4
#
# Find p[] such that  
# y_approximate =  p[0] + p[1]*x[i] + p[2]*x[i]^2 ... +p[pwr]*x[i]^pwr
# and such that the sum of (y_actual - y_approximate) squared is minimized.
# To find the coefficients  p  solve the linear system of equations for pwr=3
# SUM k=0..pwr x is x(k)
#
# | SUM(1*1)   SUM(1*x)   SUM(1*x^2)   SUM(1*x^3)   | | p[0] | | SUM(1*y)   |
# | SUM(x*1)   SUM(x*x)   SUM(x*x^2)   SUM(x*x^3)   |x| p[1] |=| SUM(x*y)   |
# | SUM(x^2*1) SUM(x^2*x) SUM(x^2*x^2) SUM(x^2*x^2) | | p[2] | | SUM(x^2*y) |
# | SUM(x^3*1) SUM(x^3*x) SUM(x^3*x^2) SUM(x^3*X^2) | | p[3] | | SUM(x^3*y) |
#   solve  |A|*|p|=|Y|  for polynomial coefficients p
  
from numpy import array
from numpy.linalg import solve

def lsfit(pwr, x, y): # return polygon coefficients in p
  pn = pwr+1
  A = [[(0.0) for j in range(pn)] for i in range(pn)]
  Y = [(0.0) for i in range(pn)]
  n = len(x) # should be same as len(y)
  xpwr = [(0.0) for k in range(n)]
  for k in range(n):
    xpwr[0] = 1.0
    for i in range(1,pn): 
      xpwr[i] = xpwr[i-1]*x[k]

    # end i
    for i in range(pn):
      for j in range(pn):
        A[i][j] = A[i][j] + xpwr[i]*xpwr[j]

      # end j
      Y[i] = Y[i] + y[k]*xpwr[i]

    # end i

  # end k
  p = solve(A,Y)
  return p

# end lsfit.py3

# peval.py3  
import math

# from peval import peval
def peval(x, p):   # using Horners method
  # an n=len(p) order polynomial has n+1 coefficients 
  # stored in p[0] through p[n]                  
  # y = p[0] + p[1]*x + p[2]*x^2 +...+ p[n]*x^n
  n = len(p)-1  
  y = p[n]*x;
  if n<=0 :
    return p[0]
  for i in range(n-1,0,-1):
    y = (p[i]+y)*x;

  return y+p[0];

# end peval 

print("squire test_lsfit.py3 running")
for k in range(20):
  print("k=",k)
  yt[k] = xt[k]*xt[k]*xt[k] # build yt, simple pwr=3 polynomial
  print("xt=",xt[k]," yt=",yt[k]) 

# end k
print("  ")
print("test pwr=3 x^3")
poly = lsfit(3,xt, yt)
print("poly=",poly)
max_err = 0.0
avg_err = 0.0
rms_err = 0.0
for j in range(20):
  err = yt[j] - peval(xt[j],poly)
  print("xt=",xt[j]," yt=",yt[j]," err=",err)
  err = abs(err)
  max_err = max(max_err, err)
  avg_err = avg_err + err
  rms_err = rms_err + err*err

# end j
avg_err = avg_err/20
rms_err = math.sqrt(rms_err/20)
print("max_err=",max_err,"  avg_err=",avg_err,"  rms_err",rms_err)
print("end test_lsfit.py3")
      
# end