import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
from numpy import sin


def calculate(omd):
    t=np.array([20],dtype=float)
    th=np.array([],dtype=float)
    om=np.array([],dtype=float)
    e=np.array([0],dtype=float)



    th0=0.05 #float(input("Give initial angle"))
    time=50 #float(input("Give total time"))
    om0=0
    dt=0.05 #float(input("Give time step"))
    l=1 #float(input("Give length of string"))
    n=int(time/dt)
    q=1
    omd=omd
    th=np.append(th,th0)
    om=np.append(om,om0)
    for i in range(n):
        #ANGULAR VELOCITY:
        temp1=om[i]-((9.8/l)*th[i]+q*om[i]-1*sin(omd*t[i]))*dt

        #ANGLE:
        temp2=th[i]+temp1*dt

        #TIME:
        t=np.append(t,50+(i+1)*dt)

        #ENERGY:
        temp3=e[i]+0.5*9.8*l*(om[i]**2+(9.8/l)*th[i]**2)*dt**2

        om=np.append(om,temp1)
        th=np.append(th,temp2)
        e=np.append(e,temp3)

    amp=(np.amax(th)-np.min(th))/2


    return amp


omda=np.linspace(0,X,1000) #X is the max value I want to look for, (20 or 0.5)

def om_iter(ampa,omda):
    for i in omda:
        temp1=calculate(i)
        ampa=np.append(ampa,temp1)
    return ampa,omda


ampa=np.array([],dtype=float)
ampa,omda=om_iter(ampa,omda)

plt.plot(omda,ampa)
plt.show()