# polyLS.py
# Andrew Davison, Sept. 2025, ad@coe.psu.ac.th
'''
Least-Squares Polynomial Fitting:

Find a polynomial of (usually lower) degree m that minimizes 
the sum of squared errors:
\min \sum_{i=0}^{N} (y_i - P_m(x_i))^2
Does not necessarily pass through all points. 
Often used when data is noisy or overdetermined (N>m+1).
'''

import math
import sys
import matplotlib.pyplot as plt
from Mat import Mat  # Import the Mat class


def createMats(xs, ys, m):
  """
  Initializes matrices for polynomial fitting.
  xs, ys: data points; m: order of the polynomial + 1
  """
  # matrix A and vector B
  A_data = [[0 for _ in range(m)] for _ in range(m)]
  B_data = [[0] for _ in range(m)]

  # higher-order equations
  for i in range(m):
    for j in range(m):
      sumIJ = sum(x ** (i + j) for x in xs)
      A_data[i][j] = sumIJ
    
    sumY = sum(y * (x ** i) for x, y in zip(xs, ys))
    B_data[i][0] = sumY

  A = Mat(A_data)
  B = Mat(B_data)
  return A, B


def fitCurve(xs, ys, coefs):
  # Generate x values
  xMin = min(xs)
  xMax = max(xs)
  nPts = 100
  step = (xMax - xMin) / (nPts - 1)
  xsFit = [xMin + i * step for i in range(nPts)]
  
  # Calculate y values
  ysFit = []
  m = len(coefs)
  for xi in xsFit:
    yi = 0.0
    for i in range(m):
      yi += coefs[i] * (xi ** i)
    ysFit.append(yi)

  return xsFit, ysFit



def goodnessOfFit(xs, ys, coefs):
  """
  Calculates the goodness of fit ($R^2$) for the curve.
  xs, ys: original data points; coefs: solved coefs
  """
  mean_y = sum(ys) / len(ys)

  # Total Sum of Squares
  ssTot = sum([(yi - mean_y) ** 2 for yi in ys])

  # Residual Sum of Squares
  ssRes = 0.0
  m = len(coefs)
  for i in range(len(xs)):
    # the fitted y value
    yFit = 0.0
    for k in range(m):
      yFit += coefs[k] * (xs[i] ** k)
    
    # Add squared residual
    ssRes += (ys[i] - yFit) ** 2

  rSqr = 1 - (ssRes / ssTot)
  return rSqr


def formatPolynomial(coefs):
  # Create a readable polynomial string
  terms = []
  for i, c in enumerate(coefs):
    if abs(c) < 1e-6:
      continue
    if i == 0:
      terms.append(f"{c:.2f}")
    elif i == 1:
      terms.append(f"{c:.2f}x")
    else:
      terms.append(f"{c:.2f}x^{i}")
  return " + ".join(terms) if terms else "0"


# ----------------------------------

# Data
xs = [0.1, 0.4, 0.5, 0.7, 0.7, 0.9]
ys = [0.61, 0.92, 0.99, 1.52, 1.47, 2.03]
nPts = len(xs)
order = int(input(f"Order of polynomial (n < {nPts-1}): "))
m = order + 1

A, B = createMats(xs, ys, m)
print("Matrix A:")
print(A)
print("Vector B:")
print(B)

# Solve system of linear equations A*X = B
try:
  AInv = A.inverse()
  X = AInv * B
  # print(X)
  coefs = X.transpose().toList()[0]
  # print(coefs)
except ValueError as e:
  print(f"Error: {e}")
  sys.exit(1)

xsFit, ysFit = fitCurve(xs, ys, coefs)
rSqr = goodnessOfFit(xs, ys, coefs)

equStr = formatPolynomial(coefs)
title = f"y = {equStr}\nR^2 = {rSqr:.4f}"
print(title)

plt.plot(xs, ys, 'o', label='Original Data')
plt.plot(xsFit, ysFit, '-', label='Fitted: '+ title)

plt.title(f'Polynomial Curve Fitting (order {order})')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.grid(True)
plt.show()