import numpy as np
from gekko import GEKKO
import matplotlib.pyplot as plt

mass = [2,5,10]
nm = len(mass)
plt.figure()

for j in range(len(mass)):
    m1 = 10
    m2 = mass[j]
    eps = m2/(m1+m2)
    A = np.array([[0,1,0,0],
                 [0,0,eps,0],
                 [0,0,0,1],
                 [0,0,-1,0]])
    B = np.array([[0],
                 [1],
                 [0],
                 [-1]])
    C = np.identity(4)

    m = GEKKO()
    x,y,u = m.state_space(A,B,C,D=None)
    m.time = np.linspace(0,8,101)
    fn = [0 if m.time[i]<6.2 else 1 for i in range(101)]      
    final = m.Param(value=fn)
    u[0].status = 1
    m.Obj(final * ((y[0]-1)**2 + y[1]**2 + y[2]**2 + y[3]**2))    
    m.options.IMODE = 6
    m.solve()

    plt.subplot(len(mass),1,j+1)
    plt.ylabel(f'Ball={mass[j]} kg')
    plt.plot(m.time,y[0], label = r"Position ($y$)")
    plt.plot(m.time,y[1], label = r"Velocity ($v=\dot y$)")
    plt.plot(m.time,y[2], label = r"Angle ($\theta$)")
    plt.plot(m.time,y[3], label = r"Ang Vel ($q=\dot \theta$)")
    plt.plot(m.time,u[0], label = r"Force Input ($u$)")
    plt.legend(loc=1)
plt.show()