# -*- coding: utf-8 -*-
"""
Created on Wed Oct 16 09:26:54 2019

@author: student
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
from matplotlib.animation import FuncAnimation

m1 = 100
m2 = 100
k01 = 50
k12 = 50
k23 = 50
def F1(t):
    return 10*np.sin(t)
def F2(t):
    return 10*np.cos(t)
def D(X, t):
    D1 = X [2]
    D2 = X[3]
    D3 = (1/m1)* (F1(t) - k01*X[0] - k12*(X[0] - X[1]))
    D4 = (1/m2)* (F2(t) - k12*(X[1] - X[0]) - k23*X[1])
    return np.array ([D1, D2, D3, D4])
X0 = [0, 0,0,0]
tstop = 20
nt = 500
T = np.linspace (0, tstop, nt)
R = odeint (D, X0, T)

X1 = R[:,0]
X2 = R[:,1]
V1 = R[:,2]
V2 = R[:,3]

plt.plot (T, X1, label ='pomak m1')
plt.plot (T, X2, label ='pomak m2')
plt.legend()
plt.xlabel ('t[s]')
plt.ylabel ('x[m]')
plt.title ('rješenje')

plt.figure ()
plt.plot (X1, V1, label = 'm1')
plt.plot (X2, V2, label ='m2')
plt.legend()
plt.xlabel ('x[m]')
plt.ylabel ('v[m/s]')

sl = plt.figure ()
l01 = 100/k01
l12 = 100/k12
l23 = 100/k23
def scena (i):
    sl.clf()
    x0 = 0
    x1 = x0 + l01 + X1 [i]
    x2 = x1 + l12 + X2[i]
    x3 = l01 + l12 + l23
    plt.plot ([x0,x1,x2,x3],
              np.zeros(4), 'o-',ms=10)
    plt.arrow (x1, 0, 0.1*F1(T[i]), 0,
              width = 0.02)

    plt.arrow (x2, 0, 0.1*F2(T[i]), 0,
              width = 0.02, overhang = 10)
    plt.xlim([x0,x3])
    plt.title(f't = {np.round(T[i], 2)}s')
a = FuncAnimation (sl, scena, frames=np.size(T))
a.save ('sustavmk.avi', fps =nt/tstop)