from numpy import pi, sin, cos, sqrt, array, arange, linspace
from pylab import plot, show, legend, xlabel, ylabel

m = 1.0
g = 1.0
l = 1.0



# solve df/dt = f(r, t)
def f(r, t):
    theta = r[0]
    omega = r[1] # angular velocity
    ftheta = omega # dtheta/dt = ftheta
    fomega = -g / l * sin(theta) # domega/dt = fomega

    return array([ftheta, fomega], float)

def findPeriod(theta0):

    omega0 = 0
    r = array([theta0, omega0], float)
    t0 = 0
    tf = 50.0
    N = 1000
    h = (tf - t0) / N # timestep
    tpoints = arange(t0, tf, h) # array of time values
    thetapoints = []
    omegapoints = []

    isFirstCrossRecorded = False
    isSecondCrossRecorded = False
    firstCross = 0 # in time steps
    secondCross = 0

    for t in tpoints:
        thetapoints.append(r[0])
        omegapoints.append(r[1])

        thetaold = thetapoints[len(thetapoints) - 2]
        if not isFirstCrossRecorded and (r[0] < 0 and thetaold > 0):
            isFirstCrossRecorded = True
            firstCross = t

        if not isSecondCrossRecorded and isFirstCrossRecorded and (r[0] > 0 and thetaold < 0):
            isSecondCrossRecorded = True
            secondCross = t

        # second order RK method
    #    rmid = r + h * f(r, t)
    #    r += h * f(rmid, t)

        # fourth order RK method
        k1 = h * f(r, t)
        k2 = h * f(r + 0.5 * k1, t + 0.5 * h)
        k3 = h * f(r + 0.5 * k2, t + 0.5 * h)
        k4 = h * f(r + k3, t + h)
        r += 1/6 * (k1 + 2 * k2 + 2 * k3 + k4)

    return secondCross - firstCross

theta0 = 0.1 * pi
print("Period is ", findPeriod(theta0), "for a theta0 of", theta0)