# preparation of the numerical environment
%matplotlib inline
from numpy import sqrt, cos, pi, arctan, linspace
from matplotlib import pyplot as plt

def radix(q_0,E_0,g):
    """
    Compute analitically the roots (Y_sb, Y_sp) of the equation:
    E_0 - Y - q_0^2/(2 g Y^2) = 0

    Parameters
    ----------
    Input:                             |  Output:
        q_0 : float                    |     Y_sb : float
          [m^2/2] specific discarge    |       [m] subcritical depth
        E_0 : float                    |     Y_sp : float
          [m] specific energy          |       [m] supercritical depth
        g : float                      |
          [m/s^2] gravity              |
    """
    
    Y_c = (q_0**2/g)**(1.0/3.0)   # critical depth
    E_c = 3.0/2.0*Y_c             # crtitical energy

    # check the esistence of physical solutions!
    if E_0 > E_c:
        Gamma_0 = E_0/E_c
    else:
        print('no physical solutions!')
        return 0.0, 0.0

    alpha = arctan(sqrt(Gamma_0**3 - 1.0))

    # non-dimensional subcritical depth
    eta_sb = Gamma_0 / 2.0 * (1.0 + 2.0*cos(( pi - 2.0*alpha)/3.0))
    # non-dimensional supercritical depth
    eta_sp = Gamma_0 / 2.0 * (1.0 + 2.0*cos(( pi + 2.0*alpha)/3.0))

    Y_sb = eta_sb * Y_c # subcritical depth
    Y_sp = eta_sp * Y_c # supercritical depth

    return Y_sb, Y_sp 

gg = 9.81    # [m/s^2] gravity
qq = 2.00    # [m^2/2] specific discarge
EE = 2.50    # [m]     specific energy

Y_ss = radix(qq,EE,gg)
    
Y = linspace(0.2,3,100)
E = Y + qq**2 / (2.0*gg*Y**2)
    
plt.figure(2,figsize=(10, 7))
plt.title('Speficic Energy vs. Depth')
plt.plot(E,Y,label='E = E(Y)')
plt.plot((EE,EE),Y_ss,'or',label='solutions')
plt.xlabel('E [m]')
plt.ylabel('Y [m]')
plt.axis((1.0,4.0,0.0,3.0))
plt.grid('on')
plt.legend()
plt.show()