import func
from Numeric import *

BadInput = "Bad xa input to routine splint."

class Spline(func.FuncOps):
    def __init__(self, x_array, y_array, low_slope=None, high_slope=None):
	self.x_vals = x_array
	self.y_vals = y_array
	self.low_slope  = low_slope
	self.high_slope = high_slope
	# must be careful, so that a slope of 0 still works...
	if low_slope is not None:
	    self.use_low_slope  = 1
	else:
	    self.use_low_slope  = 0   # i.e. false
	if high_slope is not None:
	    self.use_high_slope = 1
	else:
	    self.use_high_slope = 0
	self.calc_ypp()
   
    def calc_ypp(self):
	x_vals = self.x_vals
	y_vals = self.y_vals
	n = len(x_vals)
	y2_vals  = zeros(n, Float)
	u        = zeros(n-1, Float)
	
	if self.use_low_slope:
	    u[0] = (3.0/(x_vals[1]-x_vals[0])) * \
		   ((y_vals[1]-y_vals[0])/
		    (x_vals[1]-x_vals[0])-self.low_slope)
	    y2_vals[0] = -0.5
	else:
	    u[0] = 0.0
	    y2_vals[0] = 0.0   # natural spline
	    
	for i in range(1, n-1):
	    sig = (x_vals[i]-x_vals[i-1]) / \
		  (x_vals[i+1]-x_vals[i-1])
	    p   = sig*y2_vals[i-1]+2.0
	    y2_vals[i] = (sig-1.0)/p
	    u[i] = (y_vals[i+1]-y_vals[i]) / \
		   (x_vals[i+1]-x_vals[i]) - \
		   (y_vals[i]-y_vals[i-1])/ \
		   (x_vals[i]-x_vals[i-1])
	    u[i] = (6.0*u[i]/(x_vals[i+1]-x_vals[i-1]) - 
		    sig*u[i-1]) / p
	    
	if self.use_high_slope:
	    qn = 0.5
	    un = (3.0/(x_vals[n-1]-x_vals[n-2])) * \
		 (self.high_slope - (y_vals[n-1]-y_vals[n-2]) /
		  (x_vals[n-1]-x_vals[n-2]))
	else:
	    qn = 0.0
	    un = 0.0    # natural spline
      
	y2_vals[n-1] = (un-qn*u[n-2])/(qn*y2_vals[n-1]+1.0)

	rng = range(n-1)
	rng.reverse()
	for k in rng:         # backsubstitution step
	    y2_vals[k] = y2_vals[k]*y2_vals[k+1]+u[k]
	self.y2_vals = y2_vals
      
      
    # compute approximation
    def __call__(self, arg):
	"Simulate a ufunc; handle being called on an array."
	if type(arg) == func.ArrayType:
	    return func.array_map(self.call, arg)
	else:
	    return self.call(arg)

    def call(self, x):
	"Evaluate the spline, assuming x is a scalar."

	# if out of range, return endpoint
	if x <= self.x_vals[0]:
	    return self.y_vals[0]
	if x >= self.x_vals[-1]:
	    return self.y_vals[-1]

	pos = searchsorted(self.x_vals, x)
      
	h = self.x_vals[pos]-self.x_vals[pos-1]
	if h == 0.0:
	    raise BadInput
      
	a = (self.x_vals[pos] - x) / h
	b = (x - self.x_vals[pos-1]) / h
	return (a*self.y_vals[pos-1] + b*self.y_vals[pos] + \
		((a*a*a - a)*self.y2_vals[pos-1] + \
		 (b*b*b - b)*self.y2_vals[pos]) * h*h/6.0)


class LinInt(func.FuncOps):
    def __init__(self, x_array, y_array):
	self.x_vals = x_array
	self.y_vals = y_array
      
    # compute approximation
    def __call__(self, arg):
	"Simulate a ufunc; handle being called on an array."
	if type(arg) == func.ArrayType:
	    return func.array_map(self.call, arg)
	else:
	    return self.call(arg)

    def call(self, x):
	"Evaluate the interpolant, assuming x is a scalar."

	# if out of range, return endpoint
	if x <= self.x_vals[0]:
	    return self.y_vals[0]
	if x >= self.x_vals[-1]:
	    return self.y_vals[-1]

	pos = searchsorted(self.x_vals, x)
      
	h = self.x_vals[pos]-self.x_vals[pos-1]
	if h == 0.0:
	    raise BadInput
      
	a = (self.x_vals[pos] - x) / h
	b = (x - self.x_vals[pos-1]) / h
	return a*self.y_vals[pos-1] + b*self.y_vals[pos]

def spline_interpolate(x1, y1, x2):
    """
    Given a function at a set of points (x1, y1), interpolate to
    evaluate it at points x2.
    """
    sp = Spline(x1, y1)
    return sp(x2)

def logspline_interpolate(x1, y1, x2):
    """
    Given a function at a set of points (x1, y1), interpolate to
    evaluate it at points x2.
    """
    sp = Spline(log(x1), log(y1))
    return exp(sp(log(x2)))


def linear_interpolate(x1, y1, x2):
    """
    Given a function at a set of points (x1, y1), interpolate to
    evaluate it at points x2.
    """
    li = LinInt(x1, y1)
    return li(x2)
