# we need some ipython magic commands to 
# 1. allow sympy to pretty print math equations inline
# 2. use pylab for inline graphing
%pylab inline
%load_ext sympy.interactive.ipythonprinting

def numerical_cubic_bezier_fn(control_pts, t):
    """ 
    Numerically calculate an x,y point for the bezier curve at position t. 

    Parameters: control_pts -- list of 4 (x,y) control points
                t -- 'time' value should be between 0 and 1

    Return: (xpt, ypt) -- tuple of x,y data on the curve
    """
    check = ((0 <= t) and (t <= 1))
    assert check
    # define the actual cubic bezier equation here
    fn = lambda c,t : c[0]*(1 - t)**3 + c[1]*3*t*(1 - t)**2 + c[2]*3*t**2*(1 - t) + c[3]*t**3
    xs = [x for x,y in control_pts]
    ys = [y for x,y in control_pts]
    # now calculate the x,y position from the bezier equation
    xpt = fn(xs, t)
    ypt = fn(ys, t)
    return xpt, ypt

# just to make sure that I understand, lets create a list of numbers for a specific bezier curve
# arbitrarily define my 4 control points here
p0 = (0.7,0.2); p1 = (2.3,1.0); p2 = (7.2,0.8); p3 = (10.0,0.0)
# create a list of 1000 x,y points that traces over the length of this curve
points = [numerical_cubic_bezier_fn(control_pts=[p0,p1,p2,p3], t=t) for t in arange(0.0,1.0,1./1000)]
# separate the x's and y's so matplotlib can plot them
xs = [x for x,y in points]
ys = [y for x,y in points]

# plot the resulting curve and the control points
plot(xs, ys, 'b-')
hold(True)
for pt in [p0,p1,p2,p3]:
    plot( pt[0], pt[1], 'ro')
hold(False)
t=title('example cubic bezier curve')
#legend(['curve','control pts'])

from sympy.printing.mathml import mathml
from sympy import Symbol, Matrix, summation, diff, expand, collect, preview, Sum, latex, Eq, factor

t = Symbol('t_i')
c0, c1, c2, c3 = [Symbol(name) for name in ['c0','c1','c2','c3']]

# the equation for a cubic bezier path b(t) over t in [0,1]
# Note: the c control points are (x,y) tuples, but we're going to compute
# the resulting x and y values independantly. x and y are independant of
# each other, but both are dependant on t which goes from 0 to 1
symbolic_cubic_bezier_fn = c0*(1 - t)**3 + c1*3*t*(1 - t)**2 + c2*3*t**2*(1 - t) + c3*t**3
print 'The cubic bezier function B(t) is:'
B = Symbol('B(t_i)')
Eq(B, symbolic_cubic_bezier_fn)

d_i = Symbol('d_i')
# the error between what the 'best fit' bezier curve produces and the hand drawn point is defined as
error = symbolic_cubic_bezier_fn - d_i

e_i = Symbol('e_i')
# where symbolic_cubic_bezier_fn is bezier(t, control_points) and data_i is the i'th hand drawn point
Eq(e_i, error)


# the least squares penalty function is the sum of all of the squared errors
penalty_fn = error**2

print 'The penalty function is :'
i = Symbol('i')
n = Symbol('n')
M = Symbol('M')
Eq(M, Sum( penalty_fn, (i,0,n) ))

df_dc0 = collect( expand(penalty_fn.diff(c0), deep=True), [c0,c1,c2,c3])
print 'Partial derivative wrt c0 is:'
Eq(0, Sum(df_dc0, (i,0,n)) )

df_dc1 = collect( expand(penalty_fn.diff(c1), deep=True), [c0,c1,c2,c3])
print 'Partial derivative wrt c1 is:'
Eq(0, Sum(df_dc1, (i,0,n)))

df_dc2 = collect( expand(penalty_fn.diff(c2), deep=True), [c0,c1,c2,c3])
print 'Partial derivative wrt c2 is:'
Eq(0, Sum(df_dc2, (i,0,n)) )

df_dc3 = collect( expand(penalty_fn.diff(c3), deep=True), [c0,c1,c2,c3])
print 'Partial derivative wrt c3 is:'
Eq(0, Sum(df_dc3, (i,0,n)) )

def pull_arguments_by_var(eqn, variables):
    """
    For a equation such as : a*2 + b*3*x + 5*x + 5
    this function will return a dictionary of the arguments that are multiplied by each variable.
    So if variables=[a,b], the dict returned would be:
        a:[2],
        b:[3*x],
        'other':[5*x, 5]

    Parameters: eqn -- a sympy equation that has already been term collected
                variables -- a list of variables to search for arguments that are multiplied by each variable

    Returns: vardict -- a dictionary containing the arguments for each variable and 'other'
    """
    # try to find what is multiplied by the variables in variables and other
    cdict = {}
    for var in variables:
        for arg in eqn.args:
            if var in arg.atoms(Symbol):
                cdict.setdefault(var, []).append(arg/var)
                #print '%s : %s' %(c, arg)
    for arg in eqn.args:
        # test for non-cx arguments
        if not any([(cc in arg) for cc in variables]):
            cdict.setdefault('other', []).append(arg)
            #print 'other : %s' %(arg)
    return cdict   

# from the 1st partial derivative, reshuffle the equations so the c's are on the left and the others are on the right
# also redistribute the sums
cvars = [c0,c1,c2,c3]
arg_dict = pull_arguments_by_var(eqn=df_dc0, variables=cvars)
leftside = c0*Sum(factor(arg_dict[c0][0]), (i,0,n)) + \
           c1*Sum(factor(arg_dict[c1][0]), (i,0,n)) + \
           c2*Sum(factor(arg_dict[c2][0]), (i,0,n)) + \
           c3*Sum(factor(arg_dict[c3][0]), (i,0,n))
rightside = arg_dict['other'][0]
for arg in arg_dict['other'][1:]:
    rightside = rightside + arg
print '1st Equation:'
Eq(leftside, -1*Sum(factor(rightside), (i,0,n)))

# from the 2nd partial derivative, reshuffle the equations so the c's are on the left and the others are on the right
# also redistribute the sums
cvars = [c0,c1,c2,c3]
arg_dict = pull_arguments_by_var(eqn=df_dc1, variables=cvars)
leftside = c0*Sum(factor(arg_dict[c0][0]), (i,0,n)) + \
           c1*Sum(factor(arg_dict[c1][0]), (i,0,n)) + \
           c2*Sum(factor(arg_dict[c2][0]), (i,0,n)) + \
           c3*Sum(factor(arg_dict[c3][0]), (i,0,n))
rightside = arg_dict['other'][0]
for arg in arg_dict['other'][1:]:
    rightside = rightside + arg
print '2nd Equation:'
Eq(leftside, -1*Sum(factor(rightside), (i,0,n)))

# from the 3rd partial derivative, reshuffle the equations so the c's are on the left and the others are on the right
# also redistribute the sums
cvars = [c0,c1,c2,c3]
arg_dict = pull_arguments_by_var(eqn=df_dc2, variables=cvars)
leftside = c0*Sum(factor(arg_dict[c0][0]), (i,0,n)) + \
           c1*Sum(factor(arg_dict[c1][0]), (i,0,n)) + \
           c2*Sum(factor(arg_dict[c2][0]), (i,0,n)) + \
           c3*Sum(factor(arg_dict[c3][0]), (i,0,n))
rightside = arg_dict['other'][0]
for arg in arg_dict['other'][1:]:
    rightside = rightside + arg
print '3rd Equation:'
Eq(leftside, -1*Sum(factor(rightside), (i,0,n)))

# from the 4th partial derivative, reshuffle the equations so the c's are on the left and the others are on the right
# also redistribute the sums
cvars = [c0,c1,c2,c3]
arg_dict = pull_arguments_by_var(eqn=df_dc3, variables=cvars)
leftside = c0*Sum(factor(arg_dict[c0][0]), (i,0,n)) + \
           c1*Sum(factor(arg_dict[c1][0]), (i,0,n)) + \
           c2*Sum(factor(arg_dict[c2][0]), (i,0,n)) + \
           c3*Sum(factor(arg_dict[c3][0]), (i,0,n))
rightside = arg_dict['other'][0]
for arg in arg_dict['other'][1:]:
    rightside = rightside + arg
print '4th Equation:'
Eq(leftside, -1*Sum(factor(rightside), (i,0,n)))

# build the actual matrices this time in a loop
cvars = [c0,c1,c2,c3]
a_vals = []
x_vals = cvars
b_vals = []
for eqn in [df_dc0, df_dc1, df_dc2, df_dc3]:
    arg_dict = pull_arguments_by_var(eqn=eqn, variables=cvars)
    for var in cvars:
        a_vals.append( Sum(factor(arg_dict[var][0]), (i,0,n)) )
    rightside = arg_dict['other'][0]
    for arg in arg_dict['other'][1:]:
        rightside = rightside + arg
    b_vals.append( -1*Sum(factor(rightside), (i,0,n)) )
    
    
A = Matrix(4,4,a_vals)
x = Matrix(4,1,x_vals)
b = Matrix(4,1,b_vals)

print 'A = '
A

print 'x = '
x

print 'b = '
b

from pylab import randn

# define the control points of a bezier curve we wish to generate
p0 = (0.7,0.2); p1 = (2.3,1.0); p2 = (7.2,0.8); p3 = (10.0,0.0)
# create a list of 1000 x,y points that traces over the length of this curve
points = [numerical_cubic_bezier_fn(control_pts=[p0,p1,p2,p3], t=t) for t in arange(0.0,1.0,1./1000)]
# separate the x's and y's so matplotlib can plot them; also add some noise
xs = [x+0.15*randn() for x,y in points]
ys = [y+0.15*randn() for x,y in points]

# plot the resulting 'hand drawn' curve
plot(xs, ys, 'b-')
t=title('noisy (hand drawn) curve')

# we already have out d_i's in xs and ys from above
# make a t_i list
ts = arange(0.0, 1.0, 1.0/len(points))

# make the summation functions for A (16 of them)
A_fns = [None for i in range(16)]
A_fns[0] = lambda time_list, data_list : sum([2*(t_i - 1)**6 for t_i, d_i in zip(time_list, data_list)])
A_fns[1] = lambda time_list, data_list : sum([-6*t_i*(t_i - 1)**5 for t_i, d_i in zip(time_list, data_list)])
A_fns[2] = lambda time_list, data_list : sum([6*t_i**2*(t_i - 1)**4 for t_i, d_i in zip(time_list, data_list)])
A_fns[3] = lambda time_list, data_list : sum([-2*t_i**3*(t_i - 1)**3 for t_i, d_i in zip(time_list, data_list)])
A_fns[4] = lambda time_list, data_list : sum([-6*t_i*(t_i - 1)**5 for t_i, d_i in zip(time_list, data_list)])
A_fns[5] = lambda time_list, data_list : sum([18*t_i**2*(t_i - 1)**4 for t_i, d_i in zip(time_list, data_list)])
A_fns[6] = lambda time_list, data_list : sum([-18*t_i**3*(t_i - 1)**3 for t_i, d_i in zip(time_list, data_list)])
A_fns[7] = lambda time_list, data_list : sum([6*t_i**4*(t_i - 1)**2 for t_i, d_i in zip(time_list, data_list)])
A_fns[8] = lambda time_list, data_list : sum([6*t_i**2*(t_i - 1)**4 for t_i, d_i in zip(time_list, data_list)])
A_fns[9] = lambda time_list, data_list : sum([-18*t_i**3*(t_i - 1)**3 for t_i, d_i in zip(time_list, data_list)])
A_fns[10] = lambda time_list, data_list : sum([18*t_i**4*(t_i - 1)**2 for t_i, d_i in zip(time_list, data_list)])
A_fns[11] = lambda time_list, data_list : sum([-6*t_i**5*(t_i - 1) for t_i, d_i in zip(time_list, data_list)])
A_fns[12] = lambda time_list, data_list : sum([-2*t_i**3*(t_i - 1)**3 for t_i, d_i in zip(time_list, data_list)])
A_fns[13] = lambda time_list, data_list : sum([6*t_i**4*(t_i - 1)**2 for t_i, d_i in zip(time_list, data_list)])
A_fns[14] = lambda time_list, data_list : sum([-6*t_i**5*(t_i - 1) for t_i, d_i in zip(time_list, data_list)])
A_fns[15] = lambda time_list, data_list : sum([2*t_i**6 for t_i, d_i in zip(time_list, data_list)])
    
# make the summation functions for b (4 of them)
b_fns = [None for i in range(4)]
b_fns[0] = lambda time_list, data_list : -1.0 * sum([2*d_i*(t_i - 1)**3 for t_i, d_i in zip(time_list, data_list)])
b_fns[1] = lambda time_list, data_list : -1.0 * sum([-6*d_i*t_i*(t_i - 1)**2 for t_i, d_i in zip(time_list, data_list)])
b_fns[2] = lambda time_list, data_list : -1.0 * sum([6*d_i*t_i**2*(t_i - 1) for t_i, d_i in zip(time_list, data_list)])
b_fns[3] = lambda time_list, data_list : -1.0 * sum([-2*d_i*t_i**3 for t_i, d_i in zip(time_list, data_list)])

def solve_for_cs(time_series, data_series):
    """
    Take an input series of t_i values and the corresponding d_i values,
    compute the summation values that should go into the matrices and
    solve for the 4 unknown variables.

    Parameters: time_series -- t_i in increasing values
                data_series -- d_i corresponding to each t_i

    Returns: solution -- matrix containing the 4 solutions from solving the linear equations
    """
    # compute the data we will put into matrix A
    A_values = []
    for fn in A_fns:
        A_values.append(fn(time_series, data_series))
    # fill the A matrix with data
    A_numerical = Matrix(4,4, A_values)

    # compute the data we will put into the b vector
    b_values = []
    for fn in b_fns:
        b_values.append(fn(time_series, data_series))
    # fill the b vector with data
    b_numerical = Matrix(4,1, b_values)

    # solve for the unknowns in vector x
    x_numerical = A_numerical.inv() * b_numerical
    #print 'A = \n', A_numerical
    #print 'b = \n', b_numerical
    #print 'x = \n', x_numerical
    return x_numerical

# solve for the best fit in the x dimension
x_solutions = solve_for_cs(time_series=ts, data_series=xs)
# solve for the best fit in the y dimension
y_solutions = solve_for_cs(time_series=ts, data_series=ys)

# place the solutions into control points
best_fit_control_pts = [(x,y) for x,y in zip(x_solutions, y_solutions)]

# create a list of 1000 x,y points that traces over the length of our best fit curve
best_fit_points = [numerical_cubic_bezier_fn(control_pts=best_fit_control_pts, t=t) for t in arange(0.0,1.0,1./1000)]
# separate the x's and y's so matplotlib can plot them
best_fit_xs = [x for x,y in best_fit_points]
best_fit_ys = [y for x,y in best_fit_points]

# plot the original noisy curve
plot(xs, ys, 'b-')
hold(True)
# plot the best fit curve over top
plot(best_fit_xs, best_fit_ys, 'r-')
hold(False)
t=title('best fit cubic bezier curve')

# just for comparisons sake, print the 'clean' control points and then print the best 
# fit control points taken from the noisy data
print 'Clean Control Points = ',
print [p0, p1, p2, p3]
print 'Best Fit Control Points = '
print best_fit_control_pts