#! /usr/bin/env python
# $Id: grid_splines.py

# Chichi Lalescu
# 2010 March 22

""" grid_splines.py
    Compute spline polynomials
"""

#to find the centered differences, use a taylor expansion
#determine the coefficients of the taylor expansion by imposing that it is equal to the function on the grid
#thus the number of grid nodes directly determines the number of centered differences that can be computed:
#if you have n grid points, you have access to n-1 approximations for the derivatives
#after you find the coefficients, they need to be replaced in the original formula for the splines

from sympy import *

def get_beta_polynomials(x, spline_order, grid_points):
    #check that spline_order is compatible with grid_points
    # order of the highest continuous derivative is m
    m = (spline_order - 1)/2
    # the highest order centered difference that can be computed is given by
    gm = grid_points - 2
    if m > gm:
        print 'number of grid points is too small'
        return None
    #Generate "actual values" of function on grid points
    # these are only used locally, to obtain the polynomials
    F = [Symbol("F_0"), Symbol("Fp1")]
    for i in range(1, grid_points/2):
        F.append(Symbol("Fp{0}".format(i+1)))
    for i in range(1, grid_points/2):
        F.append(Symbol("Fm{0}".format(grid_points/2-i)))
    #Generate arrays with the derivatives of the function
    # values of f and its derivatives in 0
    fderivs_0 = [Symbol("f0_0")]
    # values of f and its derivatives in 1
    fderivs_1 = [Symbol("f0_1")]
    # first terms of taylor expansions
    tef_0 = fderivs_0[0]
    tef_1 = fderivs_1[0]
    for i in range(1,gm+1):
        fderivs_0.append(Symbol("f{0}_0".format(i)))
        fderivs_1.append(Symbol("f{0}_1".format(i)))
        tef_0 += fderivs_0[i]*x**i/factorial(i)
        tef_1 += fderivs_1[i]*x**i/factorial(i)
    #Generate the spline approximation
    scoeff = [Symbol("s0")]
    spline_formula = scoeff[0]
    for i in range(1,spline_order+1):
        scoeff.append(Symbol("s{0}".format(i)))
        spline_formula += scoeff[i]*x**i
    #Generate the equations to obtain the centered differences
    equations_centered_differences = [tef_0.subs(x,0) - F[0], tef_1.subs(x,0) - F[1]]
    for i in range(1, grid_points/2):
        equations_centered_differences += [tef_0.subs(x,-i) - F[ -i], tef_0.subs(x,i) - F[  i],
                                           tef_1.subs(x,-i) - F[1-i], tef_1.subs(x,i) - F[1+i]]
    #generate the equations to obtain the coefficients in the spline approximation
    equations_spline = [spline_formula.subs(x,0) - fderivs_0[0], spline_formula.subs(x,1) - fderivs_1[0]]
    for i in range(1, m+1):
        equations_spline += [spline_formula.diff(x,i).subs(x,0) - fderivs_0[i], spline_formula.diff(x,i).subs(x,1) - fderivs_1[i]]
    all_unknowns = fderivs_0 + fderivs_1 + scoeff
    solution = solve(equations_centered_differences+equations_spline, all_unknowns)
    explicit_spline_formula = spline_formula.subs(solution)
    #get the beta polynomials by obtaining the coefficients of the "actual values"
    beta = [explicit_spline_formula.diff(F[0])]
    for i in range(1,len(F)):
        beta.append(explicit_spline_formula.diff(F[i]))
    return beta

def belongs(x, minx, maxx):
    return Heaviside(x-minx)*Heaviside(maxx-x)

class omega:
    def __init__(self, spline_order, grid_points):
        self.x = Symbol("x")
        self.spline_order = spline_order
        self.grid_points = grid_points
        self.beta = get_beta_polynomials(self.x, self.spline_order, self.grid_points)
        return None
    def __call__(self, xval):
        if abs(xval) >= self.grid_points/2:
            return 0.0
        else:
            i = int(floor(xval)) + 1
            return self.beta[i].subs(self.x, i - xval)

def main():
    x = Symbol("x")
    resolution = 10
    nintervals = 4
    totalpoints = resolution*(nintervals+1)
    spline_order = 3
    om = omega(spline_order,nintervals)
    out = open("omega.txt",'w')
    for i in range(totalpoints):
        xval =  (i - totalpoints*0.5)/resolution
        out.write('{0}\t{1}\n'.format(xval, om(xval)))
    out.close()
    return None

import cProfile
import pstats

if __name__ == '__main__':
    cProfile.run('main()', 'profile')
    p = pstats.Stats('profile')
    p.sort_stats('cumulative').print_stats(10)