import numpy as np

# Fit the data points (0, 0), (x0, u0), (x1, u1) using a polynomial of
# the form u(x) = x**n (a + b*x).  Function pfit returns the
# coefficients a and b.  Functions peval and pevalp evaluate the
# function u and its derivative r = u' at the point x (could be a
# numpy array).

def pfit(x0, x1, u0, u1, n):
    b = (u0/x0**n - u1/x1**n)/(x0 - x1)
    a = u0/x0**n - b*x0
    return a,b

def peval(x, a, b, n):
    return x**n*(a+b*x)
    
def pevalp(x, a, b, n):
    return x**(n-1)*(n*a+(n+1)*b*x)

def smooth1(x, u, r, xlim, delx, n, which):

    # Replace points near one edge with a polynomial.
        
    ix = np.where(np.abs(x-xlim) <= 1.001*delx)[0]
    ix0,ix1 = ix[0],ix[-1]
    if which == 0:
        xf = x[ix0]
        a, b = pfit(x[ix1]-xf, x[ix1-1]-xf, u[ix1], u[ix1-1], n)
    else:
        xf = x[ix1]
        a, b = pfit(x[ix0]-xf, x[ix0+1]-xf, u[ix0], u[ix0+1], n)
    u[ix0:ix1+1] = peval(x[ix0:ix1+1]-xf, a, b, n)
    r[ix0:ix1+1] = pevalp(x[ix0:ix1+1]-xf, a, b, n)
        
def mysmooth(x, u, r, xmin, xmax, delx, n):
    smooth1(x, u, r, xmin, delx, n, 0)
    smooth1(x, u, r, xmax, delx, n, 1)