# faUtils.py
# Andrew Davison, ad@coe.psu.ac.th, June 2026
'''
  Utilities for function approximations using:
     * Taylor/Maclaurin series
     * Lagrange Polynomials
     * Chebyshev polynomials
     * Pade rationals
'''

import math, sys
from fractions import Fraction
import sympy as sp



def linspace(start, stop, num):
  # Generate evenly spaced values between start and stop
  if num == 1:
    return [start]
  step = (stop - start)/(num-1)
  return [start + i*step for i in range(num)]


def toFrac(x, maxDenom=10**4):
  return Fraction(x).limit_denominator(maxDenom)


# ------------- Taylor/Maclaurin ------------------


def seriesCoefs(symFn, x0=0, nTerms=10, verbose=False):
  """
  Return the coeffs of the degree-(nTerms-1) Taylor 
  polynomial of symFn around x=x0.
  """
  # x = sp.symbols('x')
  u = sp.Dummy('u')
  
  # Expand around x0 by shifting to u = x - x0
  series = sp.series(symFn(u+x0), u, 0, nTerms).removeO()
  poly = sp.Poly(series, u, domain='EX') 
                       # so x0 can be an expression
  if verbose:
    print("  ", poly)

  coefs = [poly.nth(k) for k in range(nTerms)]
  fCoefs = [float(coef) for coef in coefs]
  if verbose:
    print("\nCoefs in ascending power order:")
    print("  ", coefs)
  return fCoefs



def evalPowers(coefs, x, x0=0):
  # coefs are in ascending power order
  tot = 0
  u = x - x0
  for i, coef in enumerate(coefs):
    tot += coef * (u ** i)
  return tot


# --------------------- Lagrange ------------------


def lagrangeInterp(xs, ys, x):
  # compute the Lagrange basis polynomial at x
  n = len(xs)
  total = 0.0
  for k in range(n):
    term = ys[k]
    for j in range(n):  # 0 to n except for k
      if k != j:
        term *= (x - xs[j]) / (xs[k] - xs[j])
    total += term
  return total



# --------------------- Chebyshev ------------------


def chebFit(a, b, nTerms, f):
  '''  Given a function f, interval [a,b], 
    compute the nTerms coefficients c[0..nTerms-1] 
    of the chebyshev polynomial
    Based on chebft() in 'Numerical Recipes in C' Chapter 5.8 
  '''
  coefs = [0]*nTerms
  scale = (b-a)/2
  offset = (a+b)/2
  # sample f() at Chebyshev zeros
  fzs = [ f(offset + math.cos(math.pi*(k+0.5)/nTerms) * scale)
                         for k in range(nTerms) ]
          # 0 to nTerms-1 == index 1 to nTerms in the math
  # print("No. fzs:", len(fzs))
  fac = 2/nTerms
  # calculate nTerms Chebyshev coefs c_0 to c_n-1
  for j in range(nTerms):
    tot = 0
    for k in range(nTerms):  
      # 0 to nTerms-1 == index 1 to nTerms in the math
      tot += fzs[k] * math.cos(math.pi*j*(k+0.5)/nTerms)
    coefs[j] = fac * tot
  coefs[0] *= 0.5
  return coefs


def chebEval(a, b, coefs, nTerms, x):
  ''' Evaluate the nTerms terms Chebyshev approximation
         f(x) \approx sum c[k] * T_k(v)
  Note: NMR in C uses Clenshaw's recurrence.
  '''
  if (x-a)*(x-b) > 0:
    print("x not in range in chebEval()")
    sys.exit(1)

  xv = (2*x - (b+a))/(b-a)     # map x onto [-1,1]

  t0 = 1  # T_0(x) = 1
  approx = coefs[0] * t0
  if nTerms == 0:
    return approx

  t1 = xv   # T_1(x) = x
  approx += coefs[1] * t1

  for k in range(2, nTerms):
    # T_n(x) = 2*x*T_{n-1}(x) - T_{n-2}(x)
    t2 = 2*xv*t1 - t0
    approx += coefs[k] * t2
    t0 = t1      # shift vars
    t1 = t2
  return approx


def printCoefs(coefs):
  for c in coefs:
    # hide tiny floating-point noise
    # if abs(c) < 1.0e-12:
      # print(f"0 ", end="")
    # else:
    print(f"{c:.6f} ", end="")
  print("\n")



# --------------------- Pade ----------------------
 
def padeApprox(synfn, n, m, x0=0):
  print(f"\nUsing the series:")
  coefs = seriesCoefs(synfn, x0=x0, nTerms=n+m+1, verbose=True)

  p, q = padeSolve(n, m, coefs)

  print(f"\nPade Numerator Coefs p[0] ... p[{n}]:")
  numers = []
  for v in p:
    f = toFrac(v)
    numers.append(f)
    print("  ", str(f), end="")
  print() 

  print(f"Pade Denominator Coefs q[0] ... q[{m}]:")
  denoms = []
  for v in q:
    f = toFrac(v)
    denoms.append(f)
    print("  ", str(f), end="")
  print() 

  latexNumer = frac2latexPoly(numers)
  latexDenom = frac2latexPoly(denoms)
  latexFnm = f"\\frac{{{latexNumer}}}{{{latexDenom}}}"
  # print(latexFnm)
  return p, q, latexFnm



def padeSolve(n, m, coefs):
  ''' Compute Padé rational approx coefs.
  n : degree of numerator poly
  m : degree of denominator poly
  coefs : Maclaurin coefs  c_0 c_1, ..., (at least n+m+1) 
  '''
  
  # Build linear system for q1..qm (m equations)
  # since we know q0 == 1
  # Sum(q[j] * coefs[n+k-j]) = -coefs[n+k] for k = 1..m
  arr = []
  for k in range(1, m+1):
    # k == 1 corresponds to x^{n+1}
    row = []
    for j in range(1, m+1):
      row.append(coefs[n+k-j])
    row.append(-coefs[n+k])
    arr.append(row)

  xs = gaussian(arr)  # Solve for q1..qm

  # extract q coefficients
  q = [0] * (m+1)
  q[0] = 1
  for j in range(1, m+1):
    q[j] = xs[j-1]

  # Compute p coefficients (p0 to pn) using the q's
  # p[k] = Sum(q[j] * coefs[k-j])
  p = [0] * (n+1)
  for k in range(n+1):
    tot = 0
    for j in range(min(k, m) + 1):
      tot += q[j] * coefs[k-j]
    p[k] = tot

  return p, q



def gaussian(arr):
  # gaussian elimination
  n = len(arr); xs = [0]*n
  for z in range(n-1):  # looking at row z
    # partial pivot
    maxVal = 0
    # look down col z for largest val, and save row val
    for r in range(z, n):
      if abs(arr[r][z]) > maxVal:
        maxRow = r 
        maxVal = abs(arr[r][z])
    if maxRow != z:
      # exchange z and maxRow rows
      arr[maxRow], arr[z] = arr[z], arr[maxRow]

    # reduce all later rows
    for r in range(z+1, n):
      pivot = arr[r][z]/arr[z][z]
      for s in range(z+1, n+1): 
        arr[r][s] -= pivot*arr[z][s]
        # subtract from all values on row
  
  # back substitution; work up the rows
  for r in range(n-1, -1, -1):
    pivot = arr[r][n]
    if r != n-1:
      for s in range(r+1, n):
        pivot -= arr[r][s]*xs[s]
    xs[r] = pivot/arr[r][r]
  return xs


def evalPade(x, p, q, x0=0):
  # Evaluate numerator and denominator
  num = 0
  u = x - x0
  for i in range(len(p)-1, -1, -1):
    num = num * u + p[i]
  den = 0
  for i in range(len(q)-1, -1, -1):
    den = den * u + q[i]
  return num / den




def frac2latexPoly(coefs):
  ''' Convert a list of fractions to a Latex
      string containing a sum of x power terms
      that use those fractions as coefficients. 
      The terms are in ascending power order.
  '''
  if not coefs or all(c == 0 for c in coefs):
    return "0"
  
  terms = []
  for power, coef in enumerate(coefs):
    if coef == 0:
      continue
        
    # Determine sign and abs value of fraction
    isNeg = coef < 0
    absCoef = abs(coef)
    
    # Format coef part
    if absCoef.denominator == 1:
      coefStr = str(absCoef.numerator)
    else:
      coefStr = f"\\frac{{{absCoef.numerator}}}{{{absCoef.denominator}}}"
        
    # Format x variable part
    if power == 0:
      xStr = ""
    elif power == 1:
      xStr = "x"
      # Omit '1' coefficient for clarity (e.g., just 'x' or '-x')
      if coefStr == "1":
        coefStr = ""
    else:
      xStr = f"x^{{{power}}}"
      if coefStr == "1":
        coefStr = ""
            
    # Handle signs and combining term
    termStr = f"{coefStr}{xStr}"
    
    if isNeg:
      # If it's the first term, prepend '-' without spaces
      if not terms:
        terms.append(f"-{termStr}")
      else:
        terms.append(f" - {termStr}")
    else:
      # For pos terms, add '+' unless it's the first term
      if not terms:
        terms.append(termStr)
      else:
        terms.append(f" + {termStr}")
            
  return "".join(terms)