# %% Import of packages:

import matplotlib.pyplot as plt
import numpy as np

# %% Model parameters:

A = 2000  # [m2]
h_min = 0  # [m]
h_max = 4  # [m]

# %% Simulation time settings:

dt = 1  # [s]
t_start = 0  # [s]
t_stop = 4000  # [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)

# %% Initialization:

h_k = h_init = 2.0  # [m]

# %% Simulation loop:

for k in range(0, N_sim):

    t_k = k*dt  # Time

    # State limitiation:
    h_k = np.clip(h_k, h_min, h_max)

    # Selecting input:
    if (t_k >= t_start and t_k < 1000):
        F_i_k = 2
        F_u_k = 2
    else:
        F_i_k = 4
        F_u_k = 2

    # Time-derivative
    dh_dt_k = (1/A)*(F_i_k - F_u_k)

    # State updates using the Euler method:
    h_kp1 = h_k + dh_dt_k*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

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

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

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

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

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