from numpy import array, dot
import math
import matplotlib.pyplot as plt


def euler_step(f, h, x, y, threshold = 0):
    return x+h, y + h*f(x, y),h

def r_k_step(f, h, x, y,threshold = 0):
    k1 = h * f(x, y)
    k2 = h*f(x + h/2, y + k1/2)
    k3 = h*f(x + h/2, y + k2/2)
    k4 = h * f(x + h, y + k3)
    return x+h, y + k1/6 + k2/3 + k3/3 + k4/6, h

def r_k_adaptive_step(f, h, x, y, threshold):
    #do step, half and double
    y_step = r_k_step(f, h, x, y)[1]
    y_half_step = r_k_step(f, h/2, x, y)[1]
    y_two_half_step = r_k_step(f, h/2, x + h/2, y_half_step)[1]

    #if step is good enough, (step-half)<threshold
    if math.sqrt(dot((y_step - y_two_half_step),(y_step - y_two_half_step))) < h*threshold:
        # if bound is met, be optimistic -  double the step size 
        new_x, new_y, new_h = x + h, y_two_half_step, 2*h
    else:
        # halve step size until bound is met    
        met = False
        while not met:
            h = h /2
            y_step = y_half_step
            y_half_step = r_k_step(f, h/2, x, y)[1]
            y_two_half_step = r_k_step(f, h/2, x + h/2, y_half_step)[1]
            met = math.sqrt(dot((y_step - y_two_half_step),(y_step - y_two_half_step))) < h*threshold
        new_x, new_y, new_h = x + h, y_two_half_step, h
    return new_x, new_y, new_h
    






if __name__ == "__main__":
    def f(x, y):
        g = -9.8
        z_dd = 10 * math.sin(x)
        l = 4.0
        #return array([ y[1], -(g + z_dd)*math.sin(y[0])/l  ])
        return array([ y[1], -y[0]])
    
    xMax = 6*math.pi
    x1 = 0.0
    y1 = array([0.0, 2.0000001])
    h1 = .001
    threshold1 = 0.01
    pos_hist1=[]
    time_hist1=[]

    x2 = 0.0
    y2 = array([0.0, 2.0])
    h2 = 0.1
    threshold2 = 0.01
    pos_hist2=[]
    time_hist2=[]

    while x1<xMax:
        x1, y1, h1 = r_k_adaptive_step(f, h1, x1, y1, threshold1)
        pos_hist1.append(y1[0])
        time_hist1.append(x1)

    while x2<xMax:
        x2, y2, h2 = euler_step(f, h2, x2, y2)
        pos_hist2.append(y2[0])
        time_hist2.append(x2)
    
    pl1 = plt.plot(time_hist1, pos_hist1)
    pl2 = plt.plot(time_hist2, pos_hist2)
    plt.setp(pl1, linewidth=1.0, ls = "-", marker = "+")
    plt.setp(pl2, linewidth=1.0, ls = "-", marker = "+")
    
    plt.show()