######################################################
# Céline Berni, Christine Poulard, INRAE
#        v1.02     juillet 2021
# Démo pour le cours HySL2 Sciences de l'Eau ENTPE  1èreA
#
#  /ENG/  DEMO :  subcritical flow in a rectangular channel flowing into a reservoir
#
######################################################

import numpy as np  # Numpy for array computing

import matplotlib.pyplot as plt
from matplotlib.widgets import Slider  # widgets du module matplotlib !
from matplotlib.patches import Rectangle
from matplotlib.collections import LineCollection
import matplotlib.gridspec as gridspec

## Calcul d'une ligne d'eau en fluvial  pour un canal rectangulaire de grande largeur
# calcul d'aval en amont
# /ENG/   subcritical flow for a wide rectangular channel : computation from down- to up-stream

## Données d'entrée /ENG/  Input data
#  Q = débit     /ENG/  Q = discharge
Q = 50  # m3/s/ENG/
# Coefficient de Strickler   //  Strickler = 1 / Manning
K_ini = 60  # m**(1/3)/s
# pente  // longitudinal slope
I = 0.0012
# largeur  /ENG/  width
b = 20  # m
# gravité  /ENG/  gravity constant
g = 9.81  # m/s²

## point de départ   /ENG/  downstream point =  starting point for computations (abscissis, water depth)
x0 = 0  # m
h0 = 1.5  # m
## pas de calcul  /ENG/  computation step =  delta water depth
delta_h_ini = 0.025  # m

# FONCTIONS (pour être éventuellement réutilisées dans un autre code)  /ENG/  FUNCTIONS

## calcul de la hauteur normale   /ENG/ normal depth
def hauteur_normale(Q, K, b, I):
    hn = (Q / (K * b * np.sqrt(I))) ** (3 / 5)
    return hn


## calcul de la hauteur critique  /ENG/  critical depth
def hauteur_critique(Q, b):
    hc = (Q ** 2 / (g * b ** 2)) ** (1 / 3)
    return hc


def calcul_x_de_proche_en_proche(Q, K, b, I, h0, delta_h):
    """
    step-by-step estimation of h(x), by searching for x successively for water depth h from h0 by steps of delta_h
    """
    hn = hauteur_normale(Q, K, b, I)
    hc = hauteur_critique(Q, b)
    print(f"hauteur normale : {hn:.2f}")
    print(f"hauteur critique : {hc:.2f}")
    liste_h = None
    liste_x = None
    if hn < hc:
        print('le régime uniforme doit être fluvial : calcul impossible //  impossible : subcritical regime only.')
    elif hn > h0:
        print('la hauteur normale doit être inférieure à la hauteur imposée par la réservoir : calcul impossible.')
        print('//  impossible : downstream the channel, normal depth can not be higher than reservoir depth ')
    else:
        ## boucle de calcul pour calculer la hauteur d'eau par itération
        # initialisation
        hi = h0
        xi = x0

        liste_h = [h0]
        liste_x = [x0]

        while hi > hn:
            dhi = min(delta_h, hi - hn)  # pour ajuster le pas de calcul en dernière itération
            hip1 = hi - dhi  # hauteur d'eau au pas suivant i+1
            hm = hi - dhi / 2  # hauteur d'eau moyenne sur le segment i i+1 considéré
            dhdxi = I * (1 - (hn / hm) ** (10 / 3)) / (
                        1 - (hc / hm) ** 3)  # pente de la ligne d'eau (hyp grande largeur)
            xip1 = xi - dhi / dhdxi  # abscisse du point de hauteur d'eau hip1

            # itération iplus1
            xi = xip1
            hi = hip1

            print(f"itération : xi={xi:.3f} et hi={hi:.3f}")

            # stockage des valeurs de h et x dans un tableau
            liste_h.append(hi)
            liste_x.append(xi)

    return liste_h, liste_x


def update_from_sliders(val):
    global plot_fb, ligne_normale

    # on prend comme valeur de K la valeur donnée par le slider, mais en convertissant en entier
    K_slider = int(slider_K.val)
    dh_slider = slider_dh.val
    liste_h, liste_x = calcul_x_de_proche_en_proche(Q, K_slider, b, I, h0, dh_slider)

    try:
        plot_fb.remove()
        ligne_normale.remove()
    except:
        pass
    hn = hauteur_normale(Q, K_slider, b, I)
    xy_normale = (min_de_x, - I * min_de_x + hn), (max_de_x, - I * max_de_x + hn)
    lcol_normale = LineCollection([xy_normale], color=["red"], lw=2, ls="--", label=f"Zf+ hn (pour K={K_slider})")
    ligne_normale = ax_profile.add_collection(lcol_normale)

    if liste_h is not None and liste_x is not None:
        # on a pu calculer
        x = np.asarray(liste_x)
        h = np.asarray(liste_h)
        # altitude du fond (on a besoin de tous les points pour le tracé)
        zf = - I * x
        # altitude de la surface libre
        zsl = zf + h
        # charge
        H = zf + h + Q ** 2 / (2 * g * b ** 2 * h ** 2)

        plot_zsl.set_data(x, zsl)
        plot_H.set_data(x, H)
        plot_fb = ax_profile.fill_between(x, y1=zf, y2=zsl, color="lightblue", alpha=0.75)

        fig.suptitle(f'profil en long // longitudinal profile, K={K_slider:.2f}')
    else:
        # on n'a pas pu calculer
        plot_zsl.set_data(x0, h0)
        plot_H.set_data(x0, h0)
        fig.suptitle(f'profil en long, K={K_slider:.2f} ; calcul impossible')

        if hn < h_critique:
            # print('le régime uniforme doit être fluvial : calcul impossible.')
            plot_fb = ax_profile.fill_between([-1000, 0], y1=0, y2=h0, color="pink", alpha=0.75)

        elif hn > h0:
            # print('la hauteur normale doit être inférieure à la hauteur imposée par la réservoir : calcul impossible.')

            plot_fb = ax_profile.fill_between([-1400, 0], y1=[1400 * I, 0], y2=[1400 * I + h0, h0], color="pink", alpha=0.75)

    # infos sur les sliders

    text_K.set_text(f'K={K_slider}')
    text_dh.set_text(f'delta_h={dh_slider:.3f}')

    fig.canvas.draw_idle()


# CORPS DU PROGRAMME

# initialisations

liste_h, liste_x = calcul_x_de_proche_en_proche(Q, K_ini, b, I, h0, delta_h_ini)

## pour la visualisation : on transforme les listes en vecteur numpy, pour ensuite travailler directement en VECTEURS
x = np.asarray(liste_x)
h = np.asarray(liste_h)
# vecteur des altitudes du fond
zf = - I * x  # définition avec un point par point de x alors que 2 points suffisent pour définir la droite
min_de_x = min(x)
max_de_x = max(x)
xy_zfond = (min_de_x, - I * min_de_x) , (max_de_x, - I *max_de_x)
lcol_zfond = LineCollection([xy_zfond], color=["sienna"], lw=2, label='fond du canal \n Z_channel bottom')
# vecteur des altitudes de la surface libre
zsl = zf + h

# vecteur des charges
H = zf + h + Q ** 2 / (2 * g * b ** 2 * h ** 2)

h_critique = hauteur_critique(Q, b)
# Takes list of lines, where each line is a sequence of coordinates
xy_critique =(min_de_x, - I * min_de_x + h_critique), (max_de_x, - I * max_de_x+ h_critique)
lcol_critique = LineCollection([xy_critique], color=["orange"], lw=2, ls=":", label="Z_fond + hc (indépendant de K) \n Z_bottom + hc")

h_normale = hauteur_normale(Q, K_ini, b, I)
xy_normale = (min_de_x, - I * min_de_x + h_normale), (max_de_x, - I * max_de_x+ h_normale)
lcol_normale = LineCollection([xy_normale], color=["red"], lw=2, ls="--", label = f"Zf+ hn (pour K={K_ini}) \n Z_bottom + hn(K={K_ini})")


## tracé du profil en long  //  PLOT

# Grille //  Layout (subplots)
gs = gridspec.GridSpec(5, 2, width_ratios = [1, 2], height_ratios = [1,10, 2, 1, 1],
    left = 0, right = 1, bottom = 0.1, top = 1,  hspace = 0, wspace = 0)

fig = plt.figure()
fig.canvas.manager.set_window_title("Backwater, demo")

# //  profile plot     ;  items = free surface, head, bottom line, normal and critical depths
# on ajoute les courbes si possible dans l'ordre haut vers bas (cas de l'exercice)
ax_profile = plt.subplot(gs[1, 1], facecolor = 'white')  # main plot : free surface profile
ax_profile.plot(x0, h0, 'v', label='cote aval imposée par le réservoir \n water level imposed by reservoir')
plot_H, = ax_profile.plot(x, H, '--', marker='.', label='ligne de charge \n head', color='green')

plot_zsl, = ax_profile.plot(x, zsl, marker='.', label='surface libre \n free surface', color='blue')
ligne_normale = ax_profile.add_collection(lcol_normale)
# lignes fixes   //  fixed lines (bottom ad critical depth)
ax_profile.add_collection(lcol_critique)
ax_profile.add_collection(lcol_zfond)
plot_fb = ax_profile.fill_between(x, y1=zf, y2=zsl, color="lightblue", alpha=0.75)
rect = Rectangle((x0, -I * x0), 100, h0, color='darkblue')
ax_profile.add_patch(rect)
ax_profile.annotate("réservoir", (x0 + 10, h0 * 1.1), color="darkblue")
ax_profile.set_xlabel('x (m)')
ax_profile.set_ylabel('z (m)')
legend_profile = ax_profile.legend()
legend_profile.set_visible(False)
handles, labels = ax_profile.get_legend_handles_labels()

ax_legend = plt.subplot(gs[1, 0], facecolor='lightcyan')  # separate space devoted to legend
ax_legend.legend(handles, labels, ncol=1,  prop={"size":8}, loc='lower left', bbox_to_anchor=(0.1, 0.1, 0.8, 0.8))

# K
ax_slider_K = plt.subplot(gs[3, 1], facecolor='lemonchiffon')  # space for K slider (Strickler = 1 / Manning)
slider_K = Slider(
    ax_slider_K, "Rugosité de Strickler \n roughness coefficient (1/Manning)", valmin=0.01, valmax=100, valfmt='%0.0f', valinit=K_ini, color="sienna")
slider_K.on_changed(update_from_sliders)

# dh
ax_slider_dh = plt.subplot(gs[4, 1], facecolor = 'lemonchiffon')  # space for delta_h slider
slider_dh = Slider(
    ax_slider_dh, "Pas d'espace //  discretisation step", valmin=0.001, valmax=0.5, valfmt='%0.3f', valinit=delta_h_ini, color="silver")
slider_dh.on_changed(update_from_sliders)

text_K = ax_slider_K.text(0.5, 0.5, f'K={K_ini}', transform=ax_slider_K.transAxes,
                 fontsize=10, fontweight='bold', va='center', ha='center')
text_dh = ax_slider_dh.text(0.5, 0.5, f'delta_h={delta_h_ini:.3f}', transform=ax_slider_dh.transAxes,
                  fontsize=10, fontweight='bold', va='center', ha='center')

for ax_s in [ax_legend, ax_slider_K, ax_slider_dh]:
    ax_s.xaxis.set_visible(False)
    ax_s.yaxis.set_visible(False)

for pos in ['right', 'top', 'bottom', 'left']:
    ax_legend.spines[pos].set_visible(False)


fig.suptitle(f'profil en long // longitudinal profile')



plt.show()