import math
    import matplotlib.pyplot as plt
    import numpy as np
    from particle import Particle
    
    xmax = 2000 # 全長 [m]
    So = 1/1000 # 河床勾配(下流端)
    Bs = (          # 川幅 [m]
        500, 500, 500, 500, 500, 500,
        400, 400, 400, 400, 400, 400,
        100, 100, 100, 100, 100, 100)
    Qs = (          # 流量 [cu.m/s]
        4000, 4000, 4000, 4000, 4000, 4000,
        3500, 3500, 3500, 3500, 3500, 3500,
         500,  500,  500,  500,  500,  500)
    zs = (          # 河床高 [m]
        0.00, 0.20, 0.40, 0.60, 0.80, 1.00,
        1.00, 1.25, 1.50, 1.75, 2.00, 2.25,
        1.00, 1.40, 1.80, 2.20, 2.60, 3.00)
    xs = (          # 追加距離 [m]
          0, 200, 400, 600, 800,1000,
       1000,1200,1400,1600,1800,2000,
       1000,1200,1400,1600,1800,2000)
    
    d = 2.0     # 粒径 [cm]
    n = 0.025   #  Manning の粗度係数
    dx = 200    # 計算距離刻み [m]
    dt = 600    # 計算時間刻み [m]
    
    void = 0.4  # 空隙率
    sz = 18     # 計算点数
    z0 = np.copy(zs)
    
    # 汎用変数
    p10_3 = 10 / 3
    nn = n * n
    dx_2 = dx / 2
    dz_dqb = dt / ((1 - void) * dx)
    cnv = 1e-4 # 1 sq.cm = 1e-4 sq.m
    
    p = Particle(d)
    sgd = p.sgd
    
    invSf = np.empty(sz)
    qbs = np.empty(sz)
    Hs  = np.empty(sz)
    tss = np.empty(sz)
    dzs = np.empty(sz)
    zs1 = np.empty(sz)
    zs2 = np.empty(sz)
    
    j1 = 0          # A点(A-B)
    j2 = j1-1+6     # B点(A-B)
    j3 = j2+1       # B点(B-C)
    j4 = j3-1+6     # C点(B-C)
    j5 = j4+1       # B点(B-D)
    j6 = j5-1+6     # D点(B-D)
    
    def calc_qbs(s1=1,s2=sz,h=0):
    
        # 境界条件(下流端)の更新
        BB = Bs[s1-1]**2
        QQ = Qs[s1-1]**2
        if h == 0:
            h = math.pow(n * Qs[s1-1] /Bs[s1-1] / math.sqrt(So), 0.6)
        Sf = nn * QQ / (BB * math.pow(h, p10_3)); invSf[s1-1] = 1 / Sf
        E = h + QQ /19.6 / (BB * h * h) + Sf * dx_2
        ts = 98000 * h * Sf / sgd 
    
        qbs[s1-1] = cnv * p.MPM(ts)
        Hs[s1-1] = h + zs[s1-1]
        tss[s1-1] = ts
        cnst = zs[s1-1] + E
        for i, z, B, Q in zip(range(s1, s2), zs[s1:s2], Bs[s1:s2], Qs[s1:s2]):
            BB = B * B
            QQ = Q * Q
            qq_2g = QQ /19.6 / BB
            nnqq = nn * QQ / BB
            cnst -= z
            while True:
                vv_2g = qq_2g / (h * h)        # 速度水頭
                Sf = nnqq / math.pow(h, p10_3) # 摩擦勾配
                Sfdx_2 = Sf * dx_2             # 損失水頭
                er = h + vv_2g - Sfdx_2 - cnst # 残差 f(h1)
                if abs(er) < 1e-3: break # 収束
                # ニュートン法による補正
                h -= er / (1 + (p10_3 * Sfdx_2 - 2 * vv_2g) / h) # h1'=h1-f(h1)/fd(h1)
            H  = z + h
            ts = 98000 * h * Sf / sgd
            qbs[i] = cnv * p.MPM(ts) # 単位幅掃流砂量 [cu.m/s/m]
            Hs[i]  = H
            tss[i] = ts 
            cnst = H + vv_2g + Sfdx_2
    
    # 初期水理量の出力(dzは600s後)
    
    calc_qbs( j1+1, j2+1, 0)    # A-B区間
    h2 = Hs[j2] - zs[j2]        # 合流点の水位
    calc_qbs( j3+1, j4+1, h2)   # B-C区間
    calc_qbs( j5+1, j6+1, h2)   # B-D区間
    
    plt.plot(xs[j1:j4+1], Hs[j1:j4+1], "-", color="blue", label=" H-Main")
    plt.plot(xs[j1:j4+1], zs[j1:j4+1], "-", color="black", label=" z-Main")
    plt.plot(xs[j5:j6+1], Hs[j5:j6+1], "--", color="blue", label=" H-Tributary")
    plt.plot(xs[j5:j6+1], zs[j5:j6+1], "--", color="black", label=" z-Tributary")
    plt.xlabel("$x$ (m)")
    plt.ylabel("$H, z$ (m)")
    plt.legend()
    plt.grid()
    plt.show()
    
    print("time=0 hr.")
    print(f"|$i$|$Hs$<br>[m]|$H$<br>[m]|$z$<br>[m]|$\Delta z$<br>[m]|$\\tau_*$|$q_b$<br>[m<sup>3</sup>/s/m]|")
    print(f"|---:|---:|---:|---:|---:|---:|---:|")
    for i, x, z, H, Hss, ts, qb in zip(range(sz), xs, zs, Hs, Hs-zs, tss, qbs):
        dz = 0
        if i in [j3,j5]:
            print(f"||||||||")
        print(f"|{i:.0f}|{Hss:.3f}|{H:.3f}|{z:.3f}|{dz:.3f}|{ts:.4f}|{qb:.4e}|")
    
    zs = np.copy(z0)
    plt.cla()
    plt.plot(xs[j1:j4+1], zs[j1:j4+1], "-", color="black", label=" 0 hr. Main")
    plt.plot(xs[j5:j6+1], zs[j5:j6+1], "--", color="black", label=" 0 hr. Tributary")
    
    t = 0
    tms = [3600 * hr for hr in [ 5*24 ]]
    while t < tms[-1]:
    
        calc_qbs( j1+1, j2+1, 0)    # A-B区間
        h2 = Hs[j2] - zs[j2]        # 合流点の水深
        calc_qbs( j3+1, j4+1, h2)   # B-C区間
        calc_qbs( j5+1, j6+1, h2)   # B-D区間
    
        # 支川上流端は土砂供給なし(静的平衡)
        qbs[j6] = 0
        
        # 河床高の更新
        for i in range(sz-2, -1 ,-1):
            dzs[i] = (qbs[i+1]*Bs[i+1] - qbs[i]*Bs[i]) / Bs[i] * dz_dqb
            if i in [j4,j6]:   # 上流端の河床高
                dzs[i] =0
            if i in [j2]:   # 合流点の河床高
                dzs[i] = (qbs[j3+1]*Bs[j3+1] + qbs[j5+1]*Bs[j5+1] - qbs[i]*Bs[i]) / Bs[i] * dz_dqb
            zs[i] += dzs[i]
    
        zs[j3] = zs[j2]     # 合流点の河床高(B-C側)
        zs[j5] = zs[j2]     # 合流点の河床高(B-D側)
    
        t += dt
        if t in tms:    #出力
            plt.plot(xs[j1:j4+1], zs[j1:j4+1], "-", color="red", label=f"{t // 3600:2d} hr. Main")
            plt.plot(xs[j5:j6+1], zs[j5:j6+1], "--", color="red", label=f"{t // 3600:2d} hr. Tributary")
    
    print(f"time={t/3600:.0f} hr.")
    print(f"|$i$|$Hs$<br>[m]|$H$<br>[m]|$z$<br>[m]|$\Delta z$<br>[m]|$\\tau_*$|$q_b$<br>[m<sup>3</sup>/s/m]|")
    print(f"|---:|---:|---:|---:|---:|---:|---:|")
    for i, x, z, H, Hss, ts, qb in zip(range(sz), xs, zs, Hs, Hs-zs, tss, qbs):
        dz = z - z0[i]
        if i in [j3,j5]:
            print(f"||||||||")
        print(f"|{i:.0f}|{Hss:.3f}|{H:.3f}|{z:.3f}|{dz:.3f}|{ts:.4f}|{qb:.4e}|")
    
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z$ (m)")
    plt.legend()
    plt.grid()
    plt.show()