#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy

try:
  import matplotlib.pyplot as plt
  canplot = True
except:
  canplot = False

# interpolation function

def ntrp(x,xa,xb,ya,yb): return (x-xa) * (yb-ya) / (xb-xa) + ya

# perform regression

def polyregress(xdata,ydata,degree):
  return numpy.polynomial.polynomial.polyfit(xdata,ydata,degree)
  
# acquire polynomial result for x argument

def evaluate(x,coeffs):
  y = 0
  m = 1
  for c in coeffs:
    y += c * m
    m *= x
  return y

# declare data and degree
  
xdata = [-1,0,1,2,3,5,7,9]
ydata = [-1,3,2.5,5,4,2,5,4]

degree = 5

# acquire coefficients

coeffs = polyregress(xdata,ydata,degree)

print(f'Polynomial coefficients: {coeffs}')

# plot result if library matplotlib has been installed

if canplot:
 
  # generate regression curve data
  xplot = []
  yplot = []
  plotpoints = 100
  xmin = min(xdata)
  xmax = max(xdata)

  for n in range(plotpoints+1):
    x = ntrp(n,0,plotpoints,xmin,xmax)
    y = evaluate(x,coeffs)
    xplot += [x]
    yplot += [y]
  
  # set graph size
  plt.figure(figsize=(10,6))
  
  # plot data points
  plt.scatter(xdata,ydata, color = 'red',s=10)
  
  # plot regression curve
  plt.plot(xplot,yplot)
  plt.grid()
  plt.xlabel('X data')
  plt.ylabel('Y data')
  plt.title(f'Polynomial Regression Degree {degree}')
  
  # save the result
  plt.savefig(f'regression_degree_{degree}.png')
  
  # show the result
  plt.show()