import numpy as np
import matplotlib.pyplot as plt

#define the constants
l = 1.45 #length of the broomstick
g = 9.81 
tau = np.sqrt((2 * l) / (3 * g))

#define the second derivation of phi
def dd_phi(phi):
    a = np.sin(phi) / (tau ** 2)
    return a

#define the algorythmen
def iteration(phi0, t0, d_phi0, dd_phi, steps_count, steps_wide):
    values = [[phi0, t0, d_phi0]]
  
    i = 0
    t = t0
    while i <= steps_count:
        
        next_d_phi = values[i][2] + dd_phi(values[i][0]) * steps_wide
        
        next_phi = values[i][0] + values[i][2] * steps_wide

        #break, when stick reaches the ground
        if next_phi >= (np.pi / 2.0):
            break
        
        t = values[i][1] + steps_wide
        
        new_values = [next_phi, t, next_d_phi]
        values.append(new_values)
        
        i += 1 
    return(values) 

#define function to reproduce graphic 3
def time_ankle():
    c = 0.2 #starting ankel
    ankel = [] #empty list for ankels
    time = [] #empty list for times
    while c <= 1.4:
        iterate_values = iteration(c, 0, 0, dd_phi, 10 ** 7, 10 ** -2)
        ankel.append(c)
        
        new_time = [i[1] for i in iterate_values]
        last_new_time = new_time[-1] #last time in list is time needed to fall
        time.append(last_new_time)
        
        c = c + 0.05
    return(ankel, time)

x, y = time_ankle()


fig = plt.figure()
ax = fig.add_subplot(1,1,1)
ax.set_title('Fallzeit')#title of the plot
ax.set_xlabel('Startwinkel in rad')#label of x axes
ax.set_ylabel('Zeit in s')#label of y axes
ax.set_xlim(0, 1.4) #interval of x value
ax.set_ylim(0, 1) #inerval of y value
ax.grid()
ax.plot(x, y)
plt.show()

#gives the needed time to fall for starting ankle
def diagram_points(phi0):
    values = iteration(phi0, 0, 0, dd_phi, 10 ** 7, 10 ** -2)
    
    new_time = [i[1] for i in values]
    last_new_time = new_time[-1] 
    diagramm = [phi0, last_new_time]

    return (diagramm)

#same ankels as in the experiment
d = diagram_points(0.1308996938995747)
e = diagram_points(0.2617993877991494)
f = diagram_points(0.4363323129985824)
g = diagram_points(0.6981317007977318)
h = diagram_points(0.8726646259971648)
i = diagram_points(1.0471975511965976)
j = diagram_points(1.2217304763960306)


print(d, e, f, g, h, i, j)