# cubicSplineMat.py
# Andrew Davison, Sept. 2025, ad@coe.psu.ac.th

'''
a globally smooth interpolant with continuous first and 
second derivatives, unlike cubics or polynomials of high degree.

S_i(x) = a_i + b_i (x - x_i) + c_i (x - x_i)^2 + d_i (x - x_i)^3.

The implementation uses natural spline conditions (second derivative = 0 at endpoints), so it only solves for the interior c_i coefficients.

Instead of using the specialized Thomas algorithm for tridiagonal systems, the code constructs the full coefficient matrix A and right-hand side vector b, then solves A * c = b using matrix operations.

The Thomas algorithm is O(n) for tridiagonal systems, while matrix inversion is O(n³). For large datasets, the Thomas algorithm would be more efficient.

The code builds the tridiagonal coefficient matrix explicitly:
* Diagonal elements: 2 * (h[i-1] + h[i])
* Sub/super-diagonal elements: h[i-1] and h[i]
* Right-hand side: The alpha values representing slope differences

'''

import math
import matplotlib.pyplot as plt
from Mat import Mat


def cubicSpline(xs, ys):
  # Compute cubic spline coefficients using Mat.py matrices
  n = len(xs)
  if n < 3:
    raise ValueError("Need at least 3 points for cubic spline")
  
  ais = ys.copy()
  bs = [0]*n
  cs = [0]*n
  ds = [0]*n

  # Compute spacing
  hs = [xs[i+1]-xs[i] for i in range(n-1)]
  
  # Create matrices and solve the system
  A, b = buildMats(ais, hs, n)
  try:
    cSoln = A.inverse() * b
    cs[0] = 0
    for i in range(cSoln.nRows):
      cs[i+1] = cSoln[i][0]
    cs[n-1] = 0
  except ValueError as e:
    raise ValueError(f"Failed to solve spline system: {e}")
  
  # Compute other coefs using the cs values
  for j in range(n-1):
    bs[j] = (ais[j+1] - ais[j]) / hs[j] - \
            hs[j] * (cs[j+1] + 2 * cs[j]) / 3
    ds[j] = (cs[j+1] - cs[j]) / (3 * hs[j])

  return ais, bs, cs, ds


def buildMats(ais, hs, n):
  # Set up A and b for the tridiagonal system A*c = b
  innerSz = n - 2
  aData = [[0.0] * innerSz for _ in range(innerSz)]
  bData = [[0.0] for _ in range(innerSz)]
  
  # Fill the tridiagonal matrix A and right-hand side b
  for i in range(innerSz):
    iP = i + 1   # offset added to access hs[] and ais[]
    
    # Diagonal element
    aData[i][i] = 2 * (hs[iP-1] + hs[iP])
    
    # Sub-diagonal element
    if i > 0:
      aData[i][i-1] = hs[iP-1]
    
    # Super-diagonal element
    if i < innerSz - 1:
      aData[i][i+1] = hs[iP]
    
    # Right-hand side
    alpha = (3 / hs[iP]) * (ais[iP+1] - ais[iP]) - \
           (3 / hs[iP-1]) * (ais[iP] - ais[iP-1])
    bData[i][0] = alpha
  
  A = Mat(aData)
  print("A ="); print(A)
  b = Mat(bData)
  print("b ="); print(b)

  return A, b


def splineEvaluate(x, xs, ais, bs, cs, ds):
  # Evaluate spline at x
  n = len(xs)
  i = n-2
  for j in range(n-1):
    if xs[j] <= x <= xs[j+1]:
      i = j
      break
  dx = x - xs[i]
  return ais[i] + bs[i]*dx + cs[i]*dx**2 + ds[i]*dx**3


# Test (select one of the ys lists)
xs = [0, 1, 2, 3, 4, 5]
ys = [0, 0.5, 2, 1.5, 0.5, 0]
# ys = [0, 1, 4, 9, 16, 25]
# ys = [1, math.e, math.e**2, math.e**3, math.e**4, math.e**5]

print("Input:")
for i, (x, y) in enumerate(zip(xs, ys)):
  print(f" ({x}, {y})", end='')
print()

ais, bs, cs, ds = cubicSpline(xs, ys)
print("\nSpline coefficients:")
for i in range(len(xs)-1):
  print(f"Interval [{xs[i]}, {xs[i+1]}]: a={ais[i]:.4f}, b={bs[i]:.4f}, c={cs[i]:.4f}, d={ds[i]:.4f}")

# Interpolation
xInter = [i * 0.05 for i in range(101)]
yInter = [splineEvaluate(x, xs, ais, bs, cs, ds) for x in xInter]

# Plot
plt.figure(figsize=(10, 6))
plt.plot(xs, ys, 'ro', markersize=8, label="Data")
plt.plot(xInter, yInter, 'b-', linewidth=2, label="Cubic Spline")
plt.legend()
plt.title("Cubic Spline Interpolation")
plt.xlabel("x")
plt.ylabel("y")
plt.grid(True, alpha=0.3)
plt.show()