import control
import numpy as np
import math

def tanks2i_rhs(t, x, u, params):
    F1, F2, k11, k22 = 0.5, 0.6, 0.8, 0.5
    h1, h2 = x
    q01, q02 = u
    h2 = 0 if h2 < 0 else h2   # numerical issues
    h1 = h2 if h1 < h2 else h1
    dh1 = (q01 - k11 * math.sqrt(h1 - h2)) / F1
    dh2 = (q02 + k11 * math.sqrt(h1 - h2) - k22 * math.sqrt(h2)) / F2
    print(x, u, dh1, dh2)
    return [dh1, dh2]


tanks2i_io = control.NonlinearIOSystem(
    tanks2i_rhs, None, inputs=('q01','q02'), outputs=('h1', 'h2'),
    states=('h1', 'h2'), name='tanks2i')

ts = 1.5
Tsample = np.linspace(0, ts, 10)
q01 = np.empty_like(Tsample); q01.fill(1.007)
q02 = np.empty_like(Tsample); q02.fill(0)
xinit2 = [5.38, 3.84]
uopen = np.vstack((q01, q02))
ti, yi = control.input_output_response(tanks2i_io, Tsample, uopen, xinit2)