import numpy as np
from scipy import *
from scipy.special import *
from scipy.interpolate import *
import matplotlib.pylab as plt
import cannonspeed

def mittelpunkt(f, a, b, N=100):
    integral = 0.0
    dt = abs( (b-a)/N )
    time = 0.5*dt
    for i in range(N):
        integral += f(time)*dt
        time +=dt
    
    return integral

def romberg(f, a, b, N=(2,5,10,20)):
    # aktuelle Schrittweite
    h_list = []
    tf_list = []
    for i in N:
        tf_list.append(mittelpunkt(f, a, b, i))
        h_list.append(1.0/i)
    tf = np.array(tf_list)
    h = np.array(h_list)
    return lagrange(h**2, tf)(0)

def solve_mittelpunkt(N=100):
    endpos = np.array([0.0,0.0])
    time = 0.0
    while (endpos[1]>=0.0):
        endpos = mittelpunkt(cannonspeed.speed, 0.0, time, N)
        time += 0.02
    return time, mittelpunkt(cannonspeed.speed, 0.0, time, N)

def solve_romberg(N=(2,5,10,20)):
    endpos=0.0
    time = 0.0
    while (endpos>=0.0):
        endpos = romberg(cannonspeed.speed, 0.0, time, N)
        time += 0.02
    return mittelpunkt(cannonspeed.speed, 0.0, time)

mittel = []

N_mittel = (4,8,16,32,64)
N_romberg = (4,8,16)

for steps in N_mittel:
    print "Integrating for", steps, "steps."
    time, position = solve_mittelpunkt(steps)
    mittel.append(position[0])
    
print "Executing Romberg method."
rom = np.ones(len(N_mittel))*solve_romberg(N_romberg)[0]
print "Done."

plt.figure()
plt.plot(N_mittel, mittel, 'o-')
plt.plot(N_mittel, rom, '-')
plt.legend(('Mittelpunkt', 'Romberg'))
plt.show()