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

# Initialize the GEKKO model
m = GEKKO(remote=False)
m.time = np.linspace(0, 20, 21)  # Time grid
mass = 500
b = m.Param(value=50)
K = m.Param(value=0.8)
p = m.MV(value=0, lb=0, ub=100)  # Manipulated Variable (gas pedal)
v = m.CV(value=0)  # Controlled Variable (velocity)

# Define the dynamic equation
m.Equation(mass * v.dt() == -v * b + K * b * p)

# Set optimization mode for control
m.options.IMODE = 6  # Control mode

# Tuning parameters for MV (gas pedal)
p.STATUS = 1  # Allow optimizer to change
p.FSTATUS = 0
p.DCOST = 0.1  # Smooth out gas pedal movement
p.DMAX = 20  # Limit rate of change

# Tuning parameters for CV (velocity)
v.STATUS = 1  # Include set point in the objective
m.options.CV_TYPE = 2  # L2 norm
v.SP = 40  # Set point for velocity
v.TR_INIT = 1  # Enable set point trajectory
v.TAU = 5  # Time constant of trajectory

# List to store frames for GIF
frames = []

# Run the optimization in a loop and save plots
for i in range(25):
    if i==15:
        v.SP = 20

    print(f'MPC cycle: {i} of 25')
    m.solve(disp=False)  # Solve the optimization problem

    # Create a plot for each iteration
    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(6, 8))

    ax1.plot(m.time+i, p.value, 'b-', lw=2)
    ax1.set_ylabel('Gas Pedal (%)')
    ax1.set_title(f'MPC Cycle {i+1}')
    ax1.set_ylim([0,100])

    ax2.plot(m.time+i, v.value, 'r--', lw=2)
    ax2.plot(m.time+i,np.ones_like(m.time)*v.SP,'k:',lw=2)
    ax2.set_ylim([0,50])
    ax2.set_ylabel('Velocity (m/s)')
    ax2.set_xlabel('Time (s)')
    fig.tight_layout()

    # Save the plot as an image frame
    fig.canvas.draw()
    frame = np.frombuffer(fig.canvas.tostring_rgb(), dtype='uint8')
    frame = frame.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    frames.append(Image.fromarray(frame))

    plt.close(fig)  # Close the figure to save memory

# Save the frames as a GIF
frames[0].save('control_simulation.gif', format='GIF',
               append_images=frames[1:], save_all=True, duration=500, loop=0)

print("GIF saved as 'control_simulation.gif'")