import math
    import numpy as np
    import matplotlib.pyplot as plt
    from particle import Particle
    
    n  = 0.020      # Manning の祖度係数
    Q  = 500        # 流量 [cu.m/s]
    Hd = 12         # 下流端の水位 [m]
    zd = 0          # 下流端の河床高 [m]
    dx = 200        # 距離刻み [m]
    dt = 10         # 時間刻み [s]
    So = 1 / 700    # 河床勾配
    B  = 100        # 川幅 [m]
    X  = 10000      # 水路長 [m]
    d  = 0.01       # 粒径 [cm]
    void = 0.4      # 空隙率
    
    sz = int(X / dx) + 1
    xs = np.linspace(0, X, sz)
    zs = xs * So + zd
    Hs = np.empty(sz)
    qss = np.empty(sz)
    
    # 汎用変数
    p10_3 = 10 / 3
    q = Q / B
    qq = q * q
    qq_2g = qq / 19.6
    nnqq = n * n * qq
    dx_2 = dx / 2
    K = nnqq * dx_2
    q_dx = 100 * q / dx # cu.cm/s/sq.cm
    dt_ve = dt / (1 - void)
    
    p = Particle(d)
    sgd = p.sgd
    wf = p.wf
    fac = 1 / (wf + q_dx)
    
    def calc_qss(isFirst=False):
    
        # 境界条件
        h = Hd - zs[0]
        E = Hd + qq_2g / h**2
        Sf = nnqq / math.pow(h, p10_3)
        ts = 98000 * h * Sf / sgd
    
        Hs[0] = Hd
        qss[0] = p.Itakura(ts)
    
        for i, H, z in zip(range(1, sz), Hs[1:], zs[1:]):
            cnst = E + Sf * dx_2 - z
            if not isFirst: h = H - z 
            while True:
                vv_2g = qq_2g / (h * h)        # 速度水頭
                Sfdx_2 = K / math.pow(h, p10_3)    # 損失水頭
                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 = h + z
            E = H + vv_2g
            Sf = Sfdx_2 / dx_2
            ts = 98000 * h * Sf / sgd
    
            Hs[i]  = H
            qss[i] = p.Itakura(ts) # cu.cm/s/sq.cm
    
    # 設問 1-3
    
    calc_qss(True)
    
    cs  = np.empty(sz)
    dzs = np.empty(sz)
    
    c_up = qss[-1] / wf
    cs[-1]  = c_up
    dzs[-1] = 0
    for i, qs in zip(range(sz - 2, -1, -1), qss[-2::-1]):
        c = fac * (qs + c_up * q_dx)
        cs[i]  = c
        dzs[i] = (wf * c - qs) * dt_ve
        c_up = c
    
    cnv = 1e4 # 1cm = 0.01m = 1e4 μ
    print("|$x$<br>[m]|$z_b$<br>[m]|$H$<br>[m]|$c$<br>[%]|$\Delta z$<br>[μ]|")
    print("|---|---|---|---|---|")
    for x, z, H, c, dz in zip(xs, zs, Hs, cs, dzs):
        print(f"|{x:.0f}|{z:.3f}|{H:.3f}|{100 * c:.4f}|{cnv * dz:.3f}|")
    
    del cs, dzs
    
    plt.plot(xs, zs, label="$z_b$")
    plt.plot(xs, Hs, label="$H$")
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z_b, H$ (m)")
    plt.legend(loc="lower right")
    plt.grid()
    plt.show()
    
    # 設問 4
    tms_plt = [t * 3600 for t in [1, 4, 12, 24]]
    fac = 1 / (wf + q_dx)
    cnv = dt_ve / 100
    
    plt.cla()
    plt.plot(xs, zs, label="initial")
    
    t = dt; isFirst = True
    while t <= tms_plt[-1]:
    
        calc_qss(isFirst)
    
        c_up = qss[-1] / wf
        for i, qs in zip(range(sz - 2, -1, -1), qss[-2::-1]):
            c = fac * (qs + c_up * q_dx)
            zs[i] -= cnv * (qs - wf * c)
            c_up = c
        if t in tms_plt:
            lbl = f"{t // 3600} hr"
            plt.plot(xs, zs, label=lbl)
        t += dt
        isFirst = False
    
    plt.xlabel("$x$ (m)")
    plt.ylabel("$z_b$ (m)")
    plt.legend(loc="lower right")
    plt.grid()
    plt.show()