#!/usr/bin/python

import numpy as np
import matplotlib.pyplot as plt 




#  RHS
def RHS1(x, y1,y2):  # dot y1 = y2
  return (y2)

def RHS2(x,y1,y2):   # dot y2 = g - (beta/m) y2 - omega0^2 y1
  return(-1*y2-50.0*y1)


# Runge-Kutta
def RK4(x, y1, y2, h):

    k11 = h*RHS1(x,y1,y2)
    k12 = h*RHS2(x,y1,y2)

    k21 = h*RHS1(x+0.5*h,y1+0.5*k11,y2+0.5*k12)
    k22 = h*RHS2(x+0.5*h,y1+0.5*k11,y2+0.5*k12)

    k31 = h*RHS1(x+0.5*h,y1+0.5*k21,y2+0.5*k22)
    k32 = h*RHS2(x+0.5*h,y1+0.5*k21,y2+0.5*k22)

    k41 = h*RHS1(x+h,y1+k31,y2+k32)
    k42 = h*RHS2(x+h,y1+k31,y2+k32)

    out1 = y1+(1.0/6.0)*(k11+2*k21+2*k31+k41)
    out2 = y2+(1.0/6.0)*(k12+2*k22+2*k32+k42)

    return out1,out2



# initial values
x0 = 0.0
x1 = 10.0
h = 0.02
y01 = 0.0  # x0
y02 = -0.5  # u0 = dot x0

n = int((x1-x0)/h)

indep = np.linspace(x0,x1,num=n)
dep1 = np.zeros(n)
dep1[0] = y01
dep2 = np.zeros(n)
dep2[0] = y02



for i in range(1,len(indep)):
    dep1[i],dep2[i]=RK4(indep[i-1],dep1[i-1],dep2[i-1],h)

print(dep1)


plt.plot(indep, dep1)
plt.show()