import math
    import matplotlib.pyplot as plt
    import numpy as np
    from particle import Particle
    
    xmax = 4400 # 全長 [m]
    So = 1/500 # 河床勾配
    Bs = (          # 川幅 [m]
        300, 300, 300, 300, 300,
        300, 300, 300, 300, 300,
        300, 220, 220, 220, 220,
        220, 220, 300, 300, 300,
        300, 300, 300)
    
    d = 2.0     # 粒径 [cm]
    n = 0.025   #  Manning の粗度係数
    dx = 200    # 計算距離刻み [m]
    dt = 600    # 計算時間刻み [m]
    Q = 1000    # 流量 [cu.m/s]
    Hd = 0      # 下流端水位
    void = 0.4  # 空隙率
    sz = 23     # 計算点数
    zs = np.linspace(-2.0, 6.8, sz) # 河床高
    z0 = np.copy(zs)
    
    # 汎用変数
    p10_3 = 10 / 3
    QQ = Q * Q
    QQ_2g = QQ / 19.6
    nnQQ = n * n * QQ
    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
    
    xs = np.linspace(0, xmax, sz)
    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)
    
    def calc_qbs():
    
        # 境界条件(下流端)の更新
        BB = Bs[0]**2
        h = Hd - zs[0]
        Sf = nnQQ / (BB * math.pow(h, p10_3)); invSf[0] = 1 / Sf
        E = h + QQ_2g / (BB * h * h) + Sf * dx_2
        ts = 98000 * h * Sf / sgd 
    
        qbs[0] = cnv * p.MPM(ts)
        Hs[0] = Hd
        tss[0] = ts
    
        cnst = zs[0] + E
        for i, z, B in zip(range(1, sz), zs[1:], Bs[1:]):
            BB = B * B
            qq_2g = QQ_2g / BB
            nnqq = nnQQ / 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()
    
    plt.plot(xs, Hs)
    plt.plot(xs, zs)
    plt.xlabel("$x$ (m)")
    plt.ylabel("$H, z$ (m)")
    plt.grid()
    plt.show()
    
    #print("time=0 hr.")
    #print(f"|$x$<br>[m]|$z$<br>[m]|$B$<br>[m]|$H$<br>[m]|$Hs$<br>[m]|$\\tau_*$|$q_b$<br>[m<sup>3</sup>/s/m]|$\Delta z$<br>[m]|")
    #print(f"|---:|---:|---:|---:|---:|---:|---:|")
    #for i, x, z, B, H, Hss, ts, qb in zip(range(sz), xs, zs, Bs, Hs, Hs-zs, tss, qbs):
        #dz = dz_dqb * (qbs[i+1] - qb) if x < xmax else 0
        #print(f"|{x:.0f}|{z:.1f}|{B:.1f}|{H:.3f}|{Hss:.3f}|{ts:.3f}|{qb:.3f}|{dz:.3f}|")
    
    
    # 設問 1 動的平衡
    print(f"ex21-1")
    plt.cla()
    plt.plot(xs, zs, label=" 0 hr.")
    
    t = 0
    tms = [3600 * hr for hr in [ 24, 48]]
    while t < tms[-1]:
    
        calc_qbs()
    
        # 河床高の更新
        for i in range(sz-2, -1 ,-1):
            dzs[i] = (qbs[i+1]*Bs[i+1] - qbs[i]*Bs[i]) / Bs[i] * dz_dqb
            zs[i] += dzs[i]
            
        t += dt
        if t in tms:    #出力
            plt.plot(xs, zs, label=f"{t // 3600:2d} hr.")
            if t in [24*3600]:
                zs1 = zs.copy()
            if t in [48*3600]:
                zs2 = zs.copy()
    
    print(f"|$x$<br>[m]|$z(t=0h)$<br>[m]|$z(t=24h)$<br>[m]|$z(t=48h)$<br>[m]|")
    print(f"|---:|---:|---:|---:|")
    for i, x, z, z1, z2 in zip(range(sz), xs, z0, zs1, zs2):
        print(f"|{x:6.0f}|{z:6.2f}|{z1:6.2f}|{z2:6.2f}|")
    
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z$ (m)")
    plt.legend()
    plt.grid()
    plt.show()
    
    
    # 設問 2 静的平衡
    print(f"ex21-2")
    zs = np.copy(z0)
    plt.cla()
    plt.plot(xs, zs, label=" 0 hr.")
    
    t = 0
    tms = [3600 * hr for hr in [ 24, 48]]
    while t < tms[-1]:
    
        calc_qbs()
    
        # 上流端は土砂供給なし(静的平衡)
        qbs[sz-1] = 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
            zs[i] += dzs[i]
            
        t += dt
        if t in tms:    #出力
            plt.plot(xs, zs, label=f"{t // 3600:2d} hr.")
            if t in [24*3600]:
                zs1 = zs.copy()
            if t in [48*3600]:
                zs2 = zs.copy()
    
    print(f"|$x$<br>[m]|$z(t=0h)$<br>[m]|$z(t=24h)$<br>[m]|$z(t=48h)$<br>[m]|")
    print(f"|---:|---:|---:|---:|")
    for i, x, z, z1, z2 in zip(range(sz), xs, z0, zs1, zs2):
        print(f"|{x:6.0f}|{z:6.2f}|{z1:6.2f}|{z2:6.2f}|")
    
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z$ (m)")
    plt.legend()
    plt.grid()
    plt.show()
    
    # 設問 3 KP2.6(i=13)に床止工
    print(f"ex21-3")
    zs = np.copy(z0)
    plt.cla()
    plt.plot(xs, zs, label=" 0 hr.")
    
    t = 0
    tms = [3600 * hr for hr in [ 24, 48]]
    while t < tms[-1]:
    
        calc_qbs()
    
        # 河床高の更新
        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 [13]:   # 床止上の流砂量補正
                zs00 = zs[i] + dzs[i]
                if zs00 < z0[i]:
                    dzs[i] = z0[i] - zs[i]
                    qbs[i] = ( -dzs[i] * Bs[i] / dz_dqb + qbs[i+1] * Bs[i+1] ) / Bs[i]
            zs[i] += dzs[i]
            
        t += dt
        if t in tms:    #出力
            plt.plot(xs, zs, label=f"{t // 3600:2d} hr.")
            if t in [24*3600]:
                zs1 = zs.copy()
            if t in [48*3600]:
                zs2 = zs.copy()
    
    print(f"|$x$<br>[m]|$z(t=0h)$<br>[m]|$z(t=24h)$<br>[m]|$z(t=48h)$<br>[m]|")
    print(f"|---:|---:|---:|---:|")
    for i, x, z, z1, z2 in zip(range(sz), xs, z0, zs1, zs2):
        print(f"|{x:6.0f}|{z:6.2f}|{z1:6.2f}|{z2:6.2f}|")
    
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z$ (m)")
    plt.legend()
    plt.grid()
    plt.show()
    
    
    # 設問 4 KP2.2,2.6,3.0(i=11,13,15)に床止工
    print(f"ex21-4")
    zs = np.copy(z0)
    plt.cla()
    plt.plot(xs, zs, label=" 0 hr.")
    
    t = 0
    tms = [3600 * hr for hr in [ 24, 48]]
    while t < tms[-1]:
    
        calc_qbs()
    
        # 河床高の更新
        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 [11,13,15]:   # 床止上の流砂量補正
                zs00 = zs[i] + dzs[i]
                if zs00 < z0[i]:
                    dzs[i] = z0[i] - zs[i]
                    qbs[i] = ( -dzs[i] * Bs[i] / dz_dqb + qbs[i+1] * Bs[i+1] ) / Bs[i]
            zs[i] += dzs[i]
            
        t += dt
        if t in tms:    #出力
            plt.plot(xs, zs, label=f"{t // 3600:2d} hr.")
            if t in [24*3600]:
                zs1 = zs.copy()
            if t in [48*3600]:
                zs2 = zs.copy()
    
    print(f"|$x$<br>[m]|$z(t=0h)$<br>[m]|$z(t=24h)$<br>[m]|$z(t=48h)$<br>[m]|")
    print(f"|---:|---:|---:|---:|")
    for i, x, z, z1, z2 in zip(range(sz), xs, z0, zs1, zs2):
        print(f"|{x:6.0f}|{z:6.2f}|{z1:6.2f}|{z2:6.2f}|")
    
    
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z$ (m)")
    plt.legend()
    plt.grid()
    plt.show()