#!/usr/bin/python
import numpy as np
import matplotlib.pyplot as plt
import scipy.integrate as ode

k=5.
m=0.1
omega2=k/m
alpha=0.1
beta=alpha/(2.*m)

def func(y,t):
    dy=np.zeros_like(y)
    dy[0]=y[1]
    dy[1]=-omega2*y[0]-2.*beta*y[1]
    return dy

y0=[0.1,0.]
time=np.arange(0.,10.,0.001)
res=ode.odeint(func,y0,time)
ekin=0.5*m*res[:,1]**2
epot=0.5*k*res[:,0]**2
etot=ekin+epot

n=np.arange(0,len(time)-1)
too=np.zeros_like(time)
tsum=0.
for i in n:
    tsum=tsum-alpha*res[i,1]*(res[i+1,0]-res[i,0])
    too[i]=tsum
#******** visual ************************************
fig=plt.figure(figsize=(10,10))
p1=plt.subplot(221)
p1.set_xlabel("t")
p1.set_ylabel("x")
p1.plot(time,res[:,0],color="red",label="x")
p1.grid(color="black")
p1.legend(loc='upper right')

p2=plt.subplot(222)
p2.set_xlabel("t")
p2.set_ylabel("v")
p2.plot(time,res[:,1],color="black",label="kiirus")
p2.grid(color="black")
p2.legend(loc='upper right')

p3=plt.subplot(223)
p3.set_xlabel("t")
p3.set_ylabel("enegy")
p3.plot(time,ekin,color="blue",label="ekin")
p3.plot(time,epot,color="red",label="epot")
p3.plot(time,etot,color="green",label="etot")
p3.plot(time,etot-too,color="black",label="ekorr")
p3.grid(color="black")
p3.legend(loc='upper right')

plt.show()