# test_psimeq_dbg.py

import math
from numpy import array
from psimeq_dbg import psimeq
from pybarrier import *

print "test_psimeq_dbg.py"
idata = 0 # reset after each use of the data set
n_data = 20 # number of data points
max_err = 0.0
avg_err = 0.0
rms_err = 0.0
debug = True
y_data=[0.0 for i in range(n_data)]
x_data=[0.0 for i in range(n_data)]

def fit_pn(n, A, Y):  # compute C coefficients
  pwr = [0.0 for i in range(n)]
  for k in range(n_data):
    y = y_data[k]
    x = x_data[k]
    pwr[0] = 1.0
    for i in range(1,n):
      pwr[i] = pwr[i-1]*x
    for i in range(n):
      for j in range(n):
        A[i][j] = A[i][j] + pwr[i]*pwr[j]
      Y[i] = Y[i] + y*pwr[i]
  # end for k
  if debug:
    for i in range(n):
      for j in range(n):
        print "A[",
        print i,
        print "][",
        print j,
        print "]=",
        print A[i][j]
      print "Y[",
      print i,
      print "]=",
      print Y[i]

  C = psimeq(A, Y)
  if debug:
    for i in range(n):
      print "C[",
      print i,
      print "]=",
      print C[i]
  return C

def check_pn(n, C):  # compute error of fit
  maxe = 0.0 
  sum = 0.0
  sumsq = 0.0
  for k in range(n_data):
    y = y_data[k]
    x = x_data[k]
    if k==0:
      xmin=x
      xmax=x
      ymin=y
      ymax=y
      imax=0
      xbad=x
      ybad=y
    if x>xmax:
      xmax=x
    if x<xmin:
      xmin=x
    if y>ymax:
      ymax=y
    if y<ymin:
      ymin=y
    
    ya = C[n-1]*x
    for i in range(n-2,0,-1):
      ya = (C[i]+ya)*x
    ya = ya + C[0]
    diff = abs(y-ya)
    if diff>maxe: 
      maxe=diff
      imax=k
      xbad=x
      ybad=y
    sum = sum + diff
    sumsq = sumsq + diff*diff
  # end loop k
  if debug:
    print "xmin=",
    print xmin,
    print ", xmax=",
    print xmax,
    print ", ymin=",
    print ymin,
    print ", ymax=",
    print ymax
  max_err = maxe
  avg_err = sum/k
  rms_err = math.sqrt(sumsq/n_data)
  if debug:
    print "max=",
    print max_err,
    print " at i=",
    print imax,
    print ", xbad=",
    print xbad,
    print ", ybad=",
    print ybad
  return max_err,avg_err,rms_err

for i in range(n_data):
  x_data[i]=i/10.0
y_data[0]=0.0  
y_data[1]=6.0  
y_data[2]=14.1  
y_data[3]=5.01  
y_data[4]=4.326  
for i in range(5,n_data):
  y_data[i]=y_data[i-1]
y_data[n_data-1]=0.0
if debug:
  for i in range(0,n_data):
    print "x_data[",
    print i,
    print "]=",
    print x_data[i],
    print ", y_data[",
    print i,
    print "]=",
    print y_data[i]
# sample polynomial least square fit, nth power 
for n in range(4,6): # does 4 and 5, 6 was n_data
  A=array([[0.0 for i in range(n)] for i in range(n)])
  if debug:
    print "A="
    print A
  Y=array([0.0 for i in range(n)])
  C=array([0.0 for i in range(n)])
  if debug:
    print "C="
    print C
  print "fit data to ",
  print (n-1),
  print " degree polynomial"
  C = fit_pn(n, A, Y) # compute C coefficients
  max_err,avg_err,rms_err = check_pn(n, C) # compute error of fit
  for i in range(n):
    print "C[",
    print i,
    print "]=",
    print C[i]
  print "max_err=",
  print max_err
  print "avg_err=",
  print avg_err
  print "rms_err=",
  print rms_err

# end  test_psimeq_dbg.py