# %% Import of packages:

import matplotlib.pyplot as plt
import numpy as np

# %% Def of PI controller function:

def fun_pi(y_sp_k, y_k, u_i_km1, pi_params, dt):
    
    (Kp, Ti, u_man, u_min, u_max) = controller_params
    e_k = y_sp_k - y_k  # Control error
    u_p_k = Kp*(y_sp_k - y_k)  # P term
    u_i_k = u_i_km1 + (Kp*dt/Ti)*e_k  # PI term
    u_i_k = np.clip(u_i_k, -u_max, u_max)  # Limit ui
    u_k = u_man + u_p_k + u_i_k  # Total control signal
    u_k = np.clip(u_k, u_min, u_max)  # Limit of control
    
    return (u_i_k, u_k)

# %% Def of process simulator

def process_sim(h_k, F_u_k, process_params, dt):

    (A, h_min, h_max) = process_params
    h_k = np.clip(h_k, h_min, h_max) 
    dh_dt_k = (1/A)*(F_i_k - F_u_k)  # Time deriv.
    h_kp1 = h_k + dt*dh_dt_k  # Euler integration
    
    return h_kp1

# %% Model parameters:

A = 2000  # [m2]
h_min = 0  # [m] State limit
h_max = 4  # [m] State limit
process_params = (A, h_min, h_max)

# %% Controller parameters:

Kp = -2.0  # [(m3/s)/m]
Ti = 2000  # [s]

u_man = 3.0  # [m3/s]
u_min = 0  # [m3/s]
u_max = 8  # [m3/s]

controller_params = (Kp, Ti, u_man, u_min, u_max)

# %% Simulation time settings:

dt = 1.0  # [s]
t_start = 0  # [s]
t_stop = 12000  # [s]
N_sim = int((t_stop - t_start)/dt) + 1  # Num time-steps

# %% Preallocation of arrays for plotting:

t_array = np.zeros(N_sim)
h_array = np.zeros(N_sim)
F_i_array = np.zeros(N_sim)
F_u_array = np.zeros(N_sim)
h_sp_array = np.zeros(N_sim)

# %% Params of input signals:

h_sp = 2.0  # [m]

# %% Initialization:

h_k = h_init = 2.0  # [m]
u_i_km1 = 0.0  # [m3/s]

# %% Simulation loop:

for k in range(0, N_sim):

    t_k = k*dt  # Time

    # Selecting inputs:
    h_sp_k = h_sp
    if (0 <= t_k < 2000): F_i_k = 3
    else: F_i_k = 4
        
    # PI controller:
    (u_i_k, u_k) = fun_pi(h_sp_k, h_k, u_i_km1,
                          controller_params, dt)
    
    # Process sim (Euler integration):
    F_u_k = u_k
    h_kp1 = process_sim(h_k, F_u_k, process_params, dt)

    # Arrays for plotting:
    t_array[k] = t_k
    h_array[k] = h_k
    F_i_array[k] = F_i_k
    F_u_array[k] = F_u_k
    h_sp_array[k] = h_sp_k

    # Time shift:
    h_k = h_kp1
    u_i_km1 = u_i_k
    
# %% Plotting:

plt.close('all')
plt.figure(1, figsize=(12, 9))

plt.subplot(2, 1, 1)
plt.plot(t_array, h_sp_array, 'r', label='h_sp')
plt.plot(t_array, h_array, 'b', label='h')
plt.legend()
plt.xlabel('t [s]')
plt.ylabel('[m]')
plt.grid()
plt.ylim(1.5, 2.5)

plt.subplot(2, 1, 2)
plt.plot(t_array, F_i_array, 'm', label='F_i')
plt.plot(t_array, F_u_array,'g', label='F_u = u')
plt.legend()
plt.grid()
plt.ylim(2, 5)
plt.xlabel('t [s]')
plt.ylabel('[m3/s]')

# plt.savefig('plot_sim_pi_level_control_tank.pdf')
plt.show()