import numpy as np
import matplotlib.pyplot as plt
import math

#variables and constants
tmax=2
dt=0.01 #stepsize
g=9.81
l=1.45 # m
tau=math.sqrt(2*l/(3*g))
vPhi0=0 #zero initial phi-velocity

#method to solve the diff-eq.
def solver2(phi0,stepT):
    t=[] #generate lists
    phi=[]
    vPhi=[]

    # initial values
    phi.append(phi0)
    vPhi.append(vPhi0)
    i=0

    while True: #iterations according to Euler
        vPhi.append(vPhi[i] +stepT*math.sin(phi[i])/(tau**2))
        phi.append(phi[i]+stepT*vPhi[i])

        i+=1 #next step

        if phi[i] >= math.pi/2: # stick has reached ground at 90°
            break

    tmax=len(phi)*stepT #calculate elapsed time
    t=np.arange(0,tmax,stepT)
    return (tmax,t,phi) #return tupel with fall-time + time and angle arrays

#initial angles  
initPhi=np.arange(0.25,1.3,0.05)

def calcTs(step): #calc the fall time for different initial phi
    ts=[]

    for i in initPhi: #generate list of times
        ts.append(solver2(i,step)[0]) #first element in return tuple is tmax

    return ts
  
  
steps=np.array([0.1,0.05,0.01,0.001])

times=[]
for m in steps:
    times.append(calcTs(m))

    
# Erzeuge eine Figure.
fig = plt.figure(figsize=(7, 5))
fig.set_tight_layout(True)

# Plotte das Diagramm.
ax1 = fig.add_subplot(1, 1, 1)
ax1.set_xlabel("phi0 [rad]")
ax1.set_ylabel("T [s]")
ax1.grid()

ax1.plot(initPhi, times[0], ".b", label=f"dt={steps[0]}")
ax1.plot(initPhi, times[1], ".r", label=f"dt={steps[1]}")
ax1.plot(initPhi, times[2], ".g", label=f"dt={steps[2]}")
ax1.plot(initPhi, times[3], "xm", label=f"dt={steps[3]}")
ax1.legend()

#anzeigen
plt.show()