import math
    import matplotlib.pyplot as plt
    import numpy as np
    from particle import Particle
    
    xmax = 3000 # 全長 [m]
    So = 1/1000 # 河床勾配
    B = 100     # 川幅 [m]
    d = 0.5     # 粒径 [cm]
    n = 0.02    #  Manning の粗度係数
    dx = 100    # 計算距離刻み [m]
    dt = 10     # 計算時間刻み [m]
    Q = 1000    # 流量 [cu.m/s]
    
    void = 0.4  # 空隙率
    sz = 31     # 計算点数
    zs = np.linspace(0, 3, sz) # 初期河床高
    z0 = np.copy(zs); zc = zs[15]
    zs[10:16] += np.linspace(0, 0.5, 6) # マウンド (下流)
    zs[15] = zc
    zs[15:21] += np.linspace(0.5, 0, 6) # マウンド (上流)
    
    # 汎用変数
    p10_3 = 10 / 3
    q = Q / B
    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)
    qbs = np.empty(sz)
    
    # 境界条件
    ho = math.pow(n * q / math.sqrt(So), 0.6)
    Eo = ho + qq_2g / (ho * ho) + So * dx_2
    ts = 98000 * ho * So / sgd 
    qbs[0] = cnv * p.MPM(ts)
    
    def calc_qbs(isFirst=False):
    
        h = ho
        cnst = zs[0] + Eo 
        for i, z in zip(range(1, sz), zs[1:]):
            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 # 残差 
                if abs(er) < 1e-4: break # 収束
                h -= er / (1 + (p10_3 * Sfdx_2 - 2 * vv_2g) / h)
            H  = z + h
            ts = 98000 * h * Sf / sgd
            qbs[i] = cnv * p.MPM(ts) # 単位幅掃流砂量 [cu.m/s/m]
            if isFirst:
                Hs[i]  = H
                tss[i] = ts 
            cnst = H + vv_2g + Sfdx_2
    
    # 設問 1-4
    Hs  = np.empty(sz)
    tss = np.empty(sz)
    Hs[0]  = ho
    tss[0] = ts
    
    calc_qbs(True)
    
    plt.plot(xs, Hs)
    plt.plot(xs, zs)
    plt.xlabel("$x$ (m)")
    plt.ylabel("$H, z$ (m)")
    plt.grid()
    plt.show()
    
    cnv2 = 1e6 * dz_dqb # 1μ = 1e6 m
    print(f"|$x$<br>[m]|$z$<br>[m]|$H$<br>[m]|$\\tau_*$|$q_b$<br>[cm<sup>3</sup>/s/cm]|$\Delta z$<br>[μ]|")
    print(f"|---:|---:|---:|---:|---:|---:|")
    for i, x, z, H, ts, qb in zip(range(sz), xs, zs, Hs, tss, qbs):
        dz = cnv2 * (qbs[i+1] - qb) if x < xmax else 0
        print(f"|{x:.0f}|{z:.1f}|{H:.3f}|{ts:.3f}|{qb / cnv:.3f}|{dz:.3f}|")
    
    del Hs, tss
    
    # 設問 5
    plt.cla()
    plt.plot(xs, zs, label=" 0 hr.")
    
    t = 0
    tms = [3600 * hr for hr in [1, 4, 12, 24]]
    while t < tms[-1]:
    
        calc_qbs()
    
        # 河床高の更新
        zs[:-1] -= (qbs[:-1] - qbs[1:sz]) * dz_dqb
    
        t += dt
        if t in tms:
            plt.plot(xs, zs, label=f"{t // 3600:2d} hr.")
    
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z$ (m)")
    plt.legend()
    plt.grid()
    plt.show()

download