######################################################
# Céline Berni, Christine Poulard, INRAE
#             juin 2021
# Démo pour le cours Sciences de l'Eau ENTPE  1èreA
######################################################

import numpy as np   # Numpy for array computing

import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.collections import LineCollection

## Calcul d'une ligne d'eau en fluvial 
## pour un canal rectangulaire de grande largeur :
#calcul d'aval en amont

## Données d'entrée
# débit
Q = 50 # m3/s
# Coefficient de Strickler
K = 60 # m**(1/3)/s
# pente
I = 0.0012
# largeur
b = 20 # m
# gravité
g = 9.81 # m/s²

## point de départ
x0 = 0 # m
h0 = 1.5 # m
## pas de calcul
dh = 0.025 # m

# Fonctions (pour être éventuellement réutilisées dans un autre code)
## calcul de la hauteur normale
def hauteur_normale(Q, K, b, I):
    hn = (Q/(K*b*np.sqrt(I)))**(3/5)
    return hn

## calcul de la hauteur critique
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):
    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.')
    elif hn>h0:
        print('la hauteur normale doit être inférieure à la hauteur imposée par la réservoir : calcul impossible.')
    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(dh,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

# CORPS DU PROGRAMME

# initialisations

liste_h, liste_x = calcul_x_de_proche_en_proche(Q, K, b, I, h0)



h_critique = hauteur_critique(Q, b)

## tracé du profil en long
# on crée une figure composée de 3 "vignettes" (subplots)
fig, ax = plt.subplots()

# éléments toujours traçables
ax.plot(x0, h0, 'v', label='niveau imposé par le réservoir')

rect = Rectangle((x0,-I*x0), 100, h0, color='darkblue')
ax.add_patch(rect)
ax.annotate("réservoir", (x0+ 10, h0*1.1), color="darkblue")
ax.set_xlabel('x (m)')
ax.set_ylabel('z (m)')

# la ligne d'eau peut être calculable ou pas selon la valeur de K
if liste_x is not None:
    # tout est calculable est traçable avec comme abscisse le vecteur x
    ## 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
    # vecteur des altitudes de la surface libre
    zsl = zf + h
    # vecteur des charges
    H = zf + h + Q ** 2 / (2 * g * b ** 2 * h ** 2)
    ax.plot(x, H, '--', marker='.', label='ligne de charge', color='green')
    ax.plot(x,zsl,marker='.', label='surface libre', color='blue')
    ax.plot(x, zf + hauteur_normale(Q, K, b, I), marker='None', ls='--', color='orange',
            label="Zf+ hn (dépendant de K)")
    ax.plot(x, zf + h_critique, marker='None', ls=':', color='red', label="Zf+ hc (indépendant de K)")
    ax.fill_between(x, y1=zf,y2=zsl,color="lightblue", alpha=0.75)
    ax.plot(x, zf, label='fond du canal', color='sienna')
else:
    # la ligne d'eau et donc la ligne d'énergie ne sont pas calculables
    #   on peut tracer quelques infos en se passant du vecteur x
    x_gauche = -1000
    x_droite = 0

    h_normale = hauteur_normale(Q, K, b, I)
    xy_normale = (x_gauche, - I * x_gauche + h_normale), (x_droite, - I * x_droite + h_normale)
    lcol_normale = LineCollection([xy_normale], color=["red"], lw=2, ls="--", label="Zf+ hn (dépendant de K)")
    ax.add_collection(lcol_normale)

    xy_critique = (x_gauche, - I * x_gauche + h_critique), (x_droite, - I * x_droite + h_critique)
    lcol_critique = LineCollection([xy_critique], color=["orange"], lw=2, ls=":",
                                   label="z_fond + hc (indépendant de K)")
    ax.add_collection(lcol_critique)

    xy_zfond = (x_gauche, - I * x_gauche), (x_droite, - I * x_droite)
    lcol_zfond = LineCollection([xy_zfond], color=["sienna"], lw=2, label='fond du canal')
    ax.add_collection(lcol_zfond)
    ax.set_xlim(x_gauche, x_droite + 50)

ax.legend()

plt.show()