# quadReg.py
# https://stackoverflow.com/questions/21158765/how-to-calculate-a-quadratic-equation-that-best-fits-a-set-of-given-data

import math, random
import matplotlib.pylab as plt


def estimateCoefs(xs, ys):
  # estimating coefficients in y = ax^2 + bx + c
  n = len(ys)
  sumX = sum(xs)
  sumY = sum(ys)
  sumXY = sum([x*y for (x,y) in zip(xs,ys)] )
  sumX2 = sum([x**2 for x in xs])
  sumX2Y = sum([x**2*y for (x,y) in zip(xs,ys)])
  sumX3 = sum([x**3 for x in xs])
  sumX4 = sum([x**4 for x in xs])
  
  denom = sumX*sumX - n*sumX2
  
  u = ((sumX*sumX2 - n*sumX3)*
       (sumX3 - (sumX*sumX2)/n))/denom - \
      sumX4 + (sumX2**2)/n
  
  v = ((sumX*sumY - n*sumXY)*
       (sumX*sumX2/n - sumX3))/denom + \
      sumX2Y - (sumX2*sumY)/n
  
  A = -v/u 
  B = (sumX*sumY - n*sumXY - 
       A*(sumX*sumX2 - n*sumX3))/denom 
  C = (sumY - A*sumX2 - B*sumX)/n
  return A,B,C


def calcCorrCoef(ys, yPreds):
  # calculate correlation coef
  n = len(ys)
  my = sum(ys)/n
  rss = sum([(y - yP)**2 for (y,yP) in zip(ys,yPreds)])
    # residual sum of squares
  mspread = sum( [(y-my)**2 for y in ys])
    # spread of model around the mean
  return math.sqrt((mspread-rss)/mspread)


def plotReg(xs, ys, yPreds, equStr):
  plt.scatter(xs, ys, color = "m", 
               marker = "o", s = 30, label="data")
  plt.plot(xs, yPreds, color = "g", label=equStr)
  plt.xlabel('x')
  plt.ylabel('y')
  plt.legend()
  plt.show()


def quadFn(a, b, c, x):
  return a*x**2 + b*x + c


# data
# xs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
# ys = [1, 3, 2, 5, 7, 8, 8, 9, 10, 12]

xs = [1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5]
ys = [0,  5,  8, 11,  14, 15, 16, 17, 17]

A, B, C = estimateCoefs(xs, ys)
print(f"Est. quad coefs: {A:.3f} {B:.3f} {C:.3f}")

yPreds = [ quadFn(A,B,C,x) for x in xs]

r = calcCorrCoef(ys, yPreds)
print(f"Corr. coef:  {r:.3f}")
equStr = f"y={A:.3f}x^2 + {B:.3f}x + {C:.3f} (r={r:.3f})"
plotReg(xs, ys, yPreds, equStr)