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

# Create Gekko model
m = GEKKO(remote=False)

# Time discretization
nt = 1000
T = 10.0
m.time = np.linspace(0, T, nt)

# Parameters
R_b, V_b, R_m, K_m = 0.05, 150.0, 0.03, 0.27
L_m, r, K_r, M = 0.05, 0.33, 10.0, 250.0
g, K_f, rho, S, C_x = 9.81, 0.03, 1.293, 2.0, 0.4

# Manipulated variable (p) as a step input
p = m.MV(value=0)
p.STATUS = 0

# Define the step input for current control
p_values = np.zeros(nt)
p_values[(m.time >= 1) & (m.time < 2)] = 100 / V_b
p.VALUE = p_values

# Differential states
i = m.Var(value=0)
omega = m.Var(value=0)
x = m.Var(value=0)
energy = m.Var(value=0)

# Differential equations
m.Equations([
    i.dt() == (p * V_b - R_m * i - K_m * omega) / L_m,
    omega.dt() == (K_r**2 / (M * r**2)) * (
        K_m * i - (r / K_r) * (M * g * K_f + 
        0.5 * rho * S * C_x * omega**2 * r**2 / (K_r**2))),
    x.dt() == omega * r / K_r,
    energy.dt() == (p * i * V_b + R_b * i**2)/1000
])

# Solver options
m.options.IMODE = 4  # Simulation mode
m.options.SOLVER = 3  # IPOPT solver

# Solve
m.solve(disp=True)

# Plotting the starting response
plt.figure(figsize=(7,5))
plt.subplot(4, 1, 1)
plt.plot(m.time, p.value, 'b-', label='MV (p)')
plt.grid(); plt.legend()

plt.subplot(4, 1, 2)
plt.plot(m.time, i.value, 'g--', label='Current (i)')
plt.plot(m.time, omega.value, 'r-', label='Angular Velocity (omega)')
plt.grid(); plt.legend()

plt.subplot(4, 1, 3)
plt.plot(m.time, x.value, color='orange', label='Position (x)')
plt.grid(); plt.legend()

plt.subplot(4, 1, 4)
plt.plot(m.time, energy.value, 'k:', label='Energy Consumed (kJ)')
plt.xlabel('Time (s)')
plt.grid(); plt.legend()

plt.tight_layout()
plt.savefig('sim.png',dpi=300)
plt.show()