import numpy as np

import matplotlib.pyplot as plt

from math import sqrt, asin, tan

L= 5000
b = 7.5
k= 40
S0 = 10**(-4)
m = 1
g= 9.81

#weir A
C_A = 1
p_A = 0.8
b_A = 3.5
w_A = 0.6

#weir B
C_B = 1
p_B = 1.1
b_B = 2
w_B = 0

def a(b,m,y):
a = (b*m*y)*y
return a

def P(b,m,y):
P = b + (2*y*(1+m**2)**(0.5))
return P

c_A = w_A + p_A

n_A = c_A*20

c_B = w_B + p_B

n_B = c_B*20

y1_A = np.linspace(0,c_A,n_A)
y2_A = np.linspace(c_A,(3),(20))

y1_B = np.linspace(0,3,20)
y2_B = np.linspace(c_B,(3),(20))

Q1_A = C_A*b_A*(2/3)*np.sqrt((2/3)*g)*((y1_A-p_A)**(3/2))
Q2_A = (2/3)*C_A*b_A*w_A*((2*g)**(0.5))*((y2_A-p_A-(2/3)*w_A)**(0.5))
Q1_B = C_B*b_B*(2/3)*np.sqrt((2/3)*g)*((y1_B-p_B)**(3/2))
Q2_B = (2/3)*C_B*b_B*w_B*((2*g)**(0.5))*((y2_B-p_B-(2/3)*w_B)**(0.5))

#Q1_B[np.isnan(Q1_B)] = 0

plt.figure(figsize=(10,5))
plt.plot(y1_A,Q1_A)
plt.plot(y2_A,Q2_A)
plt.plot(y1_B,Q1_B)
#plt.plot(y2_B,Q2_B)