# import every package needed throughout the tutorial once at the beginning
import scipy.integrate as integrate
import scipy.optimize as root
import numpy as np
import matplotlib.pylab as plt


#1a Numerics: Inegrating-----------------------------------------------------------------

def f(x, a, b):
    return (a*x**3-x**2)/(x**7+b*x**(-1))

def IntGeneric(function, lower, upper, parameters):
    result, error =  integrate.quad(function, lower, upper, args=parameters)
    return result

a = 2.
b = 4.5
lower = 0
upper = 100.
parameters = (a, b)
result = IntGeneric(f, lower, upper, parameters)
print 'The integral yields: %g' %result

#1b Numerics: Root finding-----------------------------------------------------------------

def f(x):
    return x**2-7

def g(x):
    return 1./4.*np.log10(x)

def findRoot(function, initialguess):
    return root.newton(lambda t: function(t), initialguess)

initialguess=1
result = findRoot(lambda x: f(x)-g(x),initialguess)
print 'The nummerical solver gives %g' %result

#1c Numerics: Solving an equation including an integral numerically--------------------

def f(x, a, b):
    return (a*x**3-x**2)/(x**7+b*x**(-1))

def IntGeneric(function, lower, upper, parameters):
    result, error =  integrate.quad(function, lower, upper, args=parameters)
    return result

def g(x, a):
    return a/4.*np.log10(x)

def findRootComplx(function, upper,initialguess):
    return root.newton(lambda z: function(z), initialguess)

parameters = (3.,2.)
upper = 4.
initialguess = 1.
result = findRootComplx(lambda z: IntGeneric(f, z, upper, parameters)-g(z, parameters[0]), upper, initialguess)
print 'The nummerical solver gives %g' %result

#2a Plotting: Simple example--------------------------------------------------------

def f(x, a, b):
    return (a*x**3-x**2)/(x**7+b*x**(-1))

def IntGeneric(function, lower, upper, parameters):
    result, error =  integrate.quad(function, lower, upper, args=parameters)
    return result

parameters = (1.,1.)
X = np.linspace(-4.,10.,1000)
Y = []
for x in X:
    y = IntGeneric(f,-4.,x,parameters)
    Y.append(y)


fig = plt.figure()
ax = plt.axes()

color = '#008141'

ax.plot(X,Y,c=color,label='Integral')

ax.set_xlim([-4,10])
ax.set_ylim([0,1])
ax.set_xlabel(r'$x$')
ax.set_ylabel(r'$f(x)$')
ax.legend()
fig.savefig('simpleExample.pdf',dpi=450)
print 'saved'

#2b Plotting: Advanced example--------------------------------------------------------

def f(x, a, b):
    return (a(x)*x**3-x**2)/(x**7+b*x**(-1))

def IntGeneric(function, lower, upper, parameters):
    result, error =  integrate.quad(function, lower, upper, args=parameters)
    return result

A = [
    lambda x: x**2+1.,
    lambda x: 4.,
    lambda x: x**2-x**(-1),
    lambda x: (x+1.)**(2),
    ]    

Labels = [
        r'$a=y^2+1$',
        r'$a=4$',
        r'$a=y^2-y^{-1}$',
        r'$a=(y+1)^{2}$',
        ]

Colors = ['#008141','#fdc513','#8b0a50','#4f7687']

X = np.linspace(-4.,10.,1000)

YY = []
for a in A:
    Y = []
    for x in X:
        parameters = (a,1.)
        y = IntGeneric(f,-4.,x,parameters)
        Y.append(y)
    YY.append(Y)

fig = plt.figure()
ax = plt.axes()

for Y,c,l in zip(YY,Colors,Labels):
    ax.plot(X,Y,c=c,label=l)

ax.set_xlim([-4,10])
ax.set_ylim([0,3.5])
ax.set_xlabel(r'$x$')
ax.set_ylabel(r'$f(x)$')
ax.legend()
fig.savefig('advancedExample.pdf',dpi=450)
print 'saved'