#!/usr/bin/env python
# coding: utf-8

# In[1]:


get_ipython().run_line_magic('load_ext', 'autotime')


# In[2]:


import numpy as np
import matplotlib.pyplot as plt
import control

M1 = 0
M2 = .01
R1 = 3
R2 = 5
I1 = 0
I2=0
g=9.8
a = M2*R2
c = I2 + M2*R2**2
b = c + M2*R2*R1

alp = a*g
bet = c
cap = M1*R1**2 + I1 + M2*R1**2
delt = M2**2*R2**2*R1**2
A = np.array([[0,1,0,0],
              [a*g/b,0,0,0],
              [0,0,0,1],
              [a*g/b,0,0,0]])
#no fric
# A = np.array([[0,1,0,0],
#               [alp*cap/(cap*bet+delt),0,0,0],
#               [0,0,0,1],
#               [-alp*M2*R1*R2/(cap*bet+delt),0,0,0]])
B = np.array([[0],
              [-a*R1/b],
              [0],
              [c/b]])
#print("A.shape=",A.shape," B.shape=",B.shape)
#print(A)
#print(B)
X = np.array([[.001],[0],[0.001],[0]])
Xvals = np.zeros((4,100))
#print(Xvals[:,0].shape,X.shape)
X_dot = np.zeros(X.shape)
K = control.place(A,B,[-1+.01j,-1-.01j,-2-.01j,-2+.01j])


# In[3]:


print(control.ctrb(A,B))
print(np.linalg.matrix_rank(control.ctrb(A,B)))


# In[4]:


#Goal of this block of code is to plot the set of generated state variables over time, generating a plot of the robot's bodies over time
#Robot consists of:
#Two disks
#A stick connecting the top disk to the center of the bottom one.
#The robot will exist upon a floor in Earth gravity
#Each disk has an angle and its derrivatives, a moment of inertia I, A mass M, and and a distance to its point of rotation, R
def get_coordinates_disk1(R1,theta1):
    return [R1*theta1,R1]
def get_coordinates_disk2(R1,theta1,R2,theta2):
    return [R1*theta1+R2*np.sin(theta2),R1+R2*np.cos(theta2)]
def plot_circle1(R1,theta1,fig):
    return plt.Circle(get_coordinates_disk1(R1,theta1), R1, color='r', animated=True, figure=fig)
def plot_circle2(R1,theta1,R2,theta2,fig):
    return plt.Circle(get_coordinates_disk2(R1,theta1,R2,theta2), R2/3, color='b',animated=True, figure=fig)


# In[5]:


import matplotlib.animation as animation
from IPython.display import HTML
fig, ax = plt.subplots(figsize=(10,10))
plt.hlines(0,-20,20)
ax.set_xlim(-20, 20)
ax.set_ylim(-R1-2, 40-R1-2)


X=np.matrix([[0],[2],[0],[0]])

ims = []
T=.003
dummy = []
#plt.show(block=True) 
for i in range(round(6*3/T)):
    dummy = X.copy()
    #print(dummy)

    dummy[0][0]=np.sin(dummy[0][0])
    #print (X.shape)
    X_dot = (A-B@K)@dummy + np.matrix([[0],[0],[0],[0]])
                                            #3*np.sin(10*i*2*np.pi/round(6*3/T))],[0],[0]]
    #print(X,i)
    circle1 = plot_circle1(R1,X[2][0],fig)
    circle2 = plot_circle2(R1,X[2][0],R2,X[0][0],fig)
    if(i%(int(round(1/(10*T)))) == 0 or False):
        ax.add_artist(circle1)
        ax.add_artist(circle2)
        ims.append([circle1,circle2])
    #print("res=",(X + T*X_dot)[0])
    X = X + T*X_dot
# for i in range(60):
#     x += np.pi / 15.
#     y += np.pi / 20.
#     im = plt.imshow(f(x, y), animated=True)
#     ims.append([im])
#print(ims)
#print(X[0,0]/np.pi)

plt.show()
print(len(ims))
ani = animation.ArtistAnimation(fig, ims, interval=30, repeat_delay=1000, blit=True)
HTML(ani.to_html5_video())


# In[ ]: