import numpy as np
import pandas as pd
# import matplotlib.pyplot as plt


def get_flow(coefficient):
    index = 0
    # 读取 input.csv 文件
    input = pd.read_csv(f"./data/input/input{index+1}.csv", header=0)
    row = index
    info = pd.read_csv("./data/origin/source.csv", header=0)
    area = info.iloc[row, 5]
    true_peak = info.iloc[row, 8]
    Sr = info.iloc[row, 6]
    Ks = info.iloc[row, 7]
    print("Sr:", Sr)
    print("Ks:", Ks)
    print("area:", area)
    print("true peak:", true_peak)
    TIME = input.iloc[:, 0]  # 时间
    Rain = input.iloc[:, 1]  # 时段降雨量
    EM = input.iloc[:, 2]  # 时段蒸发量
    QS = input.iloc[:, 3]  # 时段预报流量
    QU = input.iloc[:, 4]  # 支流汇入，取QU=0

    # 输入预报流域面积
    Area = area  # 水库控制流域面积

    B, alpha0, JL, KKS, KKSS, CS = coefficient[:6]
    z = coefficient[6:9]
    ni = coefficient[9:12]
    H01 = coefficient[12:15]
    Hc1 = coefficient[15:18]
    # Hr1 = coefficient[18:21]
    delta = coefficient[21:24]
    LS, LSS, L = map(round, coefficient[24:27])
    Gam = coefficient[27:30]
    EC, KKG, LG = coefficient[30], coefficient[31], round(coefficient[32])

    Hs = ni * z
    H0 = ni * z * H01
    Hc = ni * z * Hc1
    # Hr = ni * z * Hr1

    H = H0.copy()

    h0 = H[0]
    nz = ni[0] * z[0]
    B0 = h0 / nz
    dt = TIME[1] - TIME[0]

    Nt = len(Rain)
    M = len(z)

    PE0 = np.zeros(Nt)
    # G = np.zeros(M)
    # Gr = np.zeros(M)
    E = np.zeros(M)
    RSS = np.zeros(Nt)
    RS = np.zeros(Nt)
    RG = np.zeros(Nt)
    RRR = np.zeros(Nt)
    balance = np.zeros(Nt)
    TRS = np.zeros(Nt)
    TRSS = np.zeros(Nt)
    TRG = np.zeros(Nt)

    Fm, T0, E0 = 0, 0, 0

    # 循环计算
    for i in range(Nt):
        P = Rain[i]
        # 计算净雨量（扣除蒸发和截留）
        if P >= JL:
            P -= JL
            JL = 0
        else:
            JL -= P
            P = 0
        if P <= EM[i]:
            DEM = EM[i] - P
            PE = 0
            PE0[i] = PE
        else:
            DEM = 0
            PE = P - EM[i]
            PE0[i] = PE

        # 计算地表径流
        B1 = 1 - B0
        Fm1 = Sr * B1 * (T0**0.5) + Ks * T0
        Fm2 = Sr * B1 * ((T0 + dt) ** 0.5) + Ks * (T0 + dt)
        Dfm = Fm2 - Fm1
        X = PE / Dfm

        B1 = B / (1 + B0)
        eta1 = alpha0 + (1 - alpha0) / (B1 + 1)
        if X <= alpha0:
            eta = X
        elif X < 1:
            eta = eta1 - (1 - alpha0) / (B1 + 1) * ((1 - X) / (1 - alpha0)) ** (B1 + 1)
        else:
            eta = eta1
        Df0 = eta * Dfm
        Drb = PE - Df0
        if Drb < 0:
            Drb = 0

        if X <= 1:
            Fm += PE
        else:
            Fm += Dfm

        S0 = Sr * (1 - B0)
        A0 = Ks
        SA = (S0**2 + 4 * A0 * Fm) ** 0.5
        T0 = ((SA - S0) / (2 * A0)) ** 2

        # 壤中流计算
        Drs = 0
        for k in range(M):
            if DEM > 0:
                if H[0] + Df0 >= DEM:
                    E[0], E[1], E[2] = DEM, 0, 0
                else:
                    E[0] = H[0] + Df0
                    if H[1] > EC * Hs[1]:
                        E[1] = (DEM - E[0]) * H[1] / Hs[1]
                        E[2] = 0
                    else:
                        if H[1] > EC * (DEM - E[0]):
                            E[1] = EC * (DEM - E[0])
                            E[2] = 0
                        else:
                            E[1] = H[1]
                            E[2] = EC * (DEM - E[0]) - E[1]
            else:
                E[0], E[1], E[2] = 0, 0, 0

            Hw = Df0 + H[k] - E[k]
            beta0 = Hc[k] / Hs[k] / (Gam[k] + 1)
            omega1 = beta0 + (1 - beta0) / (Gam[k] + 1)
            beta = Hw / Hs[k]
            omega = beta
            if beta <= beta0:
                omega = beta
            elif beta >= beta0 and beta < 1:
                omega = omega1 - ((1 - beta0) / (Gam[k] + 1)) * (
                    (1 - beta) / (1 - beta0)
                ) ** (Gam[k] + 1)
            elif beta >= 1:
                omega = omega1
            Ds = Hw - Hs[k] * omega
            Dr = Ds * delta[k]
            Drs += Dr
            DF = (1 - delta[k]) * Ds
            H[k] = H[k] + Df0 - Ds - E[k]
            Df0 = DF
            if H[k] < 0:
                H[k] = 0
            if H[k] > Hs[k]:
                H[k] = Hs[k]

        RSS[i] = Drs
        RS[i] = Drb
        RG[i] = Df0
        RRR[i] = Drs + Drb + Df0
        E0 += np.sum(E)
        balance[i] = np.sum(H - H0) + np.sum(RRR) - np.sum(PE0) + E0

    # 汇流计算
    U = Area / dt / 3.6
    TRS[:LS] = 0
    TRS[LS] = RS[0] * U

    for T in range(1 + LS, Nt):
        TRS[T] = TRS[T - 1] * KKS + RS[T - LS] * (1 - KKS) * U

    for T in range(1 + LSS, Nt):
        TRSS[T] = TRSS[T - 1] * KKSS + RSS[T - LSS] * (1 - KKSS) * U

    TRG[:LG] = QS[0] - QU[0]
    for T in range(1 + LG, Nt):
        TRG[T] = TRG[T - 1] * KKG + RG[T - LG] * (1 - KKG) * U

    TR = TRS + TRSS + TRG + QU
    QJ = TR.copy()
    if L < 0:
        L = 0

    for T in range(L + 1, Nt):
        QJ[T] = CS * QJ[T - 1] + (1 - CS) * TR[T - L]

    # # 绘图
    # plt.style.use("fivethirtyeight")
    # plt.plot(TIME, QJ, "b-", linewidth=2)
    # plt.xlabel("time")
    # plt.ylabel("flow")
    # plt.title("predicted flow")
    # plt.legend(["calculated flow"])
    # plt.show()
    difference = max(QJ) - true_peak
    print("peak flow:", max(QJ))
    print("difference:", difference)

    return QJ, difference