"""Created on Tue Mar 16 22:25:59 2021.

@author: mofarrag
"""
import numpy as np

CWEIR = 1.704
BWIDTH = 10
HLIN = 0.001
power = 2 / 3


def HBC(Sub, River, Hbnd, dt, dx, inih, storewl, MinQ):

    xs = np.zeros(shape=(Sub.XSno, 8))
    Diff_coeff = np.zeros(shape=(Sub.XSno, 3))
    XSarea = np.zeros(shape=(Sub.XSno))

    wl = np.zeros(shape=(Sub.XSno))
    hx = np.zeros(shape=(Sub.XSno))
    qx = np.zeros(shape=(Sub.XSno))

    h = np.zeros(shape=(Sub.XSno, River.TS))
    q = np.zeros(shape=(Sub.XSno, River.TS))

    storewl = np.zeros(shape=(River.XSno, 2))
    # OverTopFlow = np.zeros(shape=(River.XSno*2,24))
    # OverTopWL = np.zeros(shape=(River.XSno*2,24))
    Lateral_q = MinQ.values[:-1, :-2]

    q_outL = np.zeros(shape=(Sub.XSno, River.TS))
    q_outR = np.zeros(shape=(Sub.XSno, River.TS))
    q_out = np.zeros(shape=(Sub.XSno, River.TS))

    q_outLx = np.zeros(shape=(Sub.XSno))
    q_outRx = np.zeros(shape=(Sub.XSno))
    q_outx = np.zeros(shape=(Sub.XSno))

    for t in range(0, River.TS):
        dtrest = dt
        # IC
        if t == 0:
            # fill the water depth array with the initial condition at time step 1
            hx[:] = inih
        else:
            # water level in the previous time-step is now turned to be the current
            hx = h[:, t - 1]

        # boundary condition
        # water level at the first cross section (x=1]is the BC
        hx[0] = Hbnd[0]
        # hx is the water depth of the previous time step
        # except for the first cross section water depth is the boundary condition
        # for this time step
        # hx=[t,t-1,t-1,t-1,.............]
        adaptiveTS_counter = 0
        # for each time step use adaptive time step
        # adaptive time step -----------------------------------------------------
        while dtrest > 0.001:
            adaptiveTS_counter = adaptiveTS_counter + 1
            # calculate the area and perimeter based on the water depth of the previous time step
            # for the whole cross sections in the sub-basin
            for x in range(0, Sub.XSno):
                # xs[:,:] = 0
                # calculate the area & perimeter of the whole XS
                Coords = River.getVortices(
                    hx[x],
                    Sub.hl[x],
                    Sub.hr[x],
                    Sub.cl[x],
                    Sub.cr[x],
                    Sub.mw[x],
                    Sub.Dbf[x],
                )
                xs[x, 0:2] = River.polygonGeometry(Coords)

                # area & perimeter of the Upper part only
                if hx[x] <= Sub.Dbf[x]:
                    xs[x, 2] = 0
                    xs[x, 3] = 0
                else:
                    # area of upper part only
                    xs[x, 2] = xs[x, 0] - (Sub.Dbf[x] * Sub.mw[x])
                    # perimeter of upper part
                    xs[x, 3] = xs[x, 1] - (2 * Sub.Dbf[x] + Sub.mw[x])

                # area & perimeter of the lower part only
                xs[x, 4] = xs[x, 0] - xs[x, 2]
                xs[x, 5] = xs[x, 1] - xs[x, 3]

                # to calculate hydraulic radius perimeter is in the denominator so it has not to be zero
                if xs[x, 1] == 0 or xs[x, 5] == 0:
                    assert False, (
                        "at x= "
                        + str(Sub.xsid[x])
                        + " t= "
                        + str(t)
                        + " total perimiter is zero"
                    )

            # for each cross section -------------------------------
            for x in range(0, Sub.XSno):
                # print(*,'(a3,i4,a4,i6,a]'] 't= ',t,' x= ', xsid[x],'------------------'
                # forward difference
                # surface slope is calculated based on the previous time step water depth
                # except for the first cross section hx[1]is from current time step
                # while hx[2]is from previous time step
                if x < Sub.XSno - 1:
                    # friction slope = (so - dhx/dx] - diffusive wave
                    sf = (
                        hx[x] + Sub.bedlevel[x] - hx[x + 1] - Sub.bedlevel[x + 1]
                    ) / dx
                else:
                    # for LOWER BOUNDARY node
                    sf = (
                        hx[x - 1] + Sub.bedlevel[x - 1] - hx[x] - Sub.bedlevel[x]
                    ) / dx

                # if sf <= 0 :
                # print[800000,'(a]'] '----------'
                # print(*,'(a4,i4,2x,a6,G10.4]'] 'XS= ',xsid[x], ', sf= ', sf
                # print[800000,'(a7,i4,a5,i6,2x,a6,G10.4]'] 'min = ',t ,' XS= ',xsid[x], ', sf= ', sf

                if River.D1["ModelMode"] == 2:
                    # check for the dike overtopping
                    # to LEFT?
                    # overtoping dischargee is zero unless water level exceeds the dike level
                    q_outLx[x] = 0
                    # check first the water level in the 2d grid
                    # if water level in the 1D river is less than the water depth in
                    # nighbouring cell in the 2D grid don't calculate overtopping

                    if storewl[Sub.xsid[x]] - Sub.bedlevel[x] < hx[x]:

                        # if water depth is higher than the left dike height and the threshold
                        if hx[x] > max(
                            Sub.zl[x] - Sub.bedlevel[x], River.D1["MinDepth"]
                        ):
                            # weir formular - free flow
                            q_outLx[x] = (
                                CWEIR
                                * BWIDTH
                                * (
                                    hx[x]
                                    - max(
                                        Sub.zl[x] - Sub.bedlevel[x],
                                        River.D1["MinDepth"],
                                    )
                                )
                                ** 1.5
                            )

                    # to RIGHT?
                    q_outRx[x] = 0
                    # nxs is added to the ipf just because of the numbering system of the
                    # cross section
                    if storewl[Sub.xsid[x] + River.XSno] - Sub.bedlevel[x] < hx[x]:

                        if hx[x] > max(
                            Sub.zr[x] - Sub.bedlevel[x], River.D1["MinDepth"]
                        ):
                            q_outRx[x] = (
                                CWEIR
                                * BWIDTH
                                * (
                                    hx[x]
                                    - max(
                                        Sub.zr[x] - Sub.bedlevel[x],
                                        River.D1["MinDepth"],
                                    )
                                )
                                ** 1.5
                            )

                    # correction factor
                    q_outx[x] = q_outLx[x] + q_outRx[x]
                    if q_outx[x] == 0:
                        corrL = 0
                        corrR = 0
                    else:
                        corrL = q_outLx[x] / q_outx[x]
                        corrR = q_outRx[x] / q_outx[x]

                # compute mean area and hydraulic radius
                # all cross section except the last one
                if x < Sub.XSno - 1:
                    if hx[x] <= Sub.Dbf[x]:
                        # upper discharge
                        xs[x, 6] = 0
                        # diffusive wave coefficient for upper wave
                        Diff_coeff[x, 0] = 0
                        # calculate only lower part
                        # Lower mean area and hydraulic radius
                        Area_LM = (xs[x, 4] + xs[x + 1, 4]) / 2  # mean area
                        R_LM = (
                            xs[x, 4] / xs[x, 5] + xs[x + 1, 4] / xs[x + 1, 5]
                        ) / 2  # mean radius
                    else:
                        # Upper mean area and hydraulic radius
                        Area_UM = (xs[x, 2] + xs[x + 1, 2]) / 2  # mean area
                        # print(*,'(a7,f8.3]'] 'A_UM = ', Area_UM
                        if xs[x, 3] != 0 and xs[x + 1, 3] == 0:
                            R_UM = (xs[x, 2] / xs[x, 3]) / 2  # mean radius
                        elif xs[x, 3] == 0 and xs[x + 1, 3] != 0:
                            R_UM = (xs[x + 1, 2] / xs[x + 1, 3]) / 2  # mean radius
                        elif xs[x, 3] == 0 and xs[x + 1, 3] == 0:
                            R_UM = 0  # mean radius
                        else:
                            R_UM = (
                                xs[x, 2] / xs[x, 3] + xs[x + 1, 2] / xs[x + 1, 3]
                            ) / 2  # mean radius

                        # upper discharge
                        xs[x, 6] = (1.0 / Sub.mn[x]) * Area_UM * (R_UM**power)
                        # diffusive wave coefficient for upper wave
                        Diff_coeff[x, 0] = (1.0 / Sub.mn[x]) * (R_UM**power)

                        # Lower mean area and hydraulic radius
                        Area_LM = (xs[x, 4] + xs[x + 1, 4]) / 2  # mean area
                        R_LM = (
                            xs[x, 4] / xs[x, 5] + xs[x + 1, 4] / xs[x + 1, 5]
                        ) / 2  # mean radius
                else:  # the last cross section
                    if hx[x] <= Sub.Dbf[x]:
                        # upper discharge
                        xs[x, 6] = 0
                        # diffusive wave coefficient for upper wave
                        Diff_coeff[x, 0] = 0
                        # calculate only lower part
                        # Lower mean area and hydraulic radius
                        Area_LM = (
                            xs[x, 4] + xs[x - 1, 4]
                        ) / 2  # for LOWER BOUNDARY node
                        R_LM = (
                            xs[x, 4] / xs[x, 5] + xs[x - 1, 4] / xs[x - 1, 5]
                        ) / 2  # mean radius
                    else:  # for LOWER BOUNDARY node
                        # Upper mean area and hydraulic radius
                        Area_UM = (xs[x, 2] + xs[x - 1, 2]) / 2

                        if xs[x, 3] != 0 and xs[x - 1, 3] == 0:
                            R_UM = (xs[x, 2] / xs[x, 3]) / 2  # mean radius
                        elif xs[x, 3] == 0 and xs[x - 1, 3] != 0:
                            R_UM = (xs[x - 1, 2] / xs[x - 1, 3]) / 2  # mean radius
                        elif xs[x, 3] == 0 and xs[x - 1, 3] == 0:
                            R_UM = 0
                        else:
                            R_UM = (
                                xs[x, 2] / xs[x, 3] + xs[x - 1, 2] / xs[x - 1, 3]
                            ) / 2  # mean radius

                        # upper discharge
                        xs[x, 6] = (1.0 / Sub.mn[x]) * Area_UM * (R_UM**power)
                        # diffusive wave coefficient for upper wave
                        Diff_coeff[x, 0] = (1.0 / Sub.mn[x]) * (R_UM**power)

                        # Lower mean area and hydraulic radius
                        Area_LM = (
                            xs[x, 4] + xs[x - 1, 4]
                        ) / 2  # for LOWER BOUNDARY node
                        R_LM = (
                            xs[x, 4] / xs[x, 5] + xs[x - 1, 4] / xs[x - 1, 5]
                        ) / 2  # mean radius

                # calculation of fluxes
                if abs(sf) > HLIN / dx:
                    # diffusive wave coefficient for upper wave
                    Diff_coeff[x, 0] = Diff_coeff[x, 0] / np.sqrt(abs(sf))
                    # diffusive wave coefficient for Lower wave
                    Diff_coeff[x, 1] = (
                        (1.0 / Sub.mn[x]) * (R_LM**power) / np.sqrt(abs(sf))
                    )
                    # total diffusion coefficient
                    Diff_coeff[x, 2] = Diff_coeff[x, 0] + Diff_coeff[x, 1]
                    # discharge
                    # upper discharge
                    xs[x, 6] = xs[x, 6] * (np.sqrt(abs(sf)) * sf / abs(sf))
                    # Lower discharge
                    xs[x, 7] = (
                        (1.0 / Sub.mn[x])
                        * Area_LM
                        * (R_LM**power)
                        * (np.sqrt(abs(sf)) * sf / abs(sf))
                    )
                else:
                    # diffusive wave coefficient for upper wave
                    Diff_coeff[x, 0] = Diff_coeff[x, 0] / np.sqrt(HLIN / dx)

                    # diffusive wave coefficient for Lower wave
                    Diff_coeff[x, 1] = (
                        (1.0 / Sub.mn[x]) * (R_LM**power) / np.sqrt(HLIN / dx)
                    )
                    # total diffusion coefficient
                    Diff_coeff[x, 2] = Diff_coeff[x, 0] + Diff_coeff[x, 1]

                    # discharge
                    # upper discharge
                    xs[x, 6] = xs[x, 6] * (np.sqrt(dx / HLIN) * sf)
                    # xs[x,7] = xs[x,7] * (np.sqrt(HLINQ/dx]]
                    # Lower discharge
                    xs[x, 7] = (
                        (1.0 / Sub.mn[x])
                        * Area_LM
                        * (R_LM**power)
                        * (np.sqrt(dx / HLIN) * sf)
                    )
                    # xs[x,8] = (1.0/mn[x]] * Area_LM * (R_LM** (2.0/3.0]] * (np.sqrt(HLINQ/dx]]

                # total Discharge
                qx[x] = xs[x, 6] + xs[x, 7]

                # 			if qx[x] < 0:
                # print(*,'(a,i4]'] 'Q is -ve at XS= ',xsid[x]
                # print[800000,'(a7,i4,a17,i6,a4,f8.3]'] 'min = ',t ,' Q is -ve at XS= ',xsid[x],' Q= ',qx[x]
                # print[800000,'(a,f8.3]'] 'Q is -ve at Q= ',qx[x]
                # exit

                if River.D1["ModelMode"] == 2:
                    # if the overtopping is less than 0.5 the flow
                    if qx[x] > 2 * q_outx[x]:
                        # the whole overtopping will be considered
                        qx[x] = qx[x] - q_outx[x]
                    else:
                        # if the overtopping is greater than 0.5 the flow
                        # overtopping will be 0.5 the flow only
                        q_outx[x] = 0.5 * qx[x]
                        q_outLx[x] = corrL * q_outx[x]
                        q_outRx[x] = corrR * q_outx[x]
                        qx[x] = qx[x] - q_outx[x]

            # for each cross section -------------------------------
            # print(*,'(a17,f8.3]'] 'max diff coeff = ', maxval[Diff_coeff(:,3]]

            # adjusting the time step
            dto = (dx**2) / (2 * max(Diff_coeff[:, 2]))

            if dto > dt:
                dto = dt
            else:
                if dto == 0:
                    print("adaptive time step is almost zero dto= " + str(t))
                    # arr = Diff_coeff(:,3]
                    # call Find_FInArrayF[maxval[arr],arr,loc1[1]]
                    # loc =maxloc (Diff_coeff(:,3]]
                    # print(*,'(a37,i4]'] 'At XS= ', xsid[loc1[1]]
                    # print(*,'(a37,i4]'] 'Diff coeff= ', Diff_coeff[loc1[1],3]
                    # print(*,'(a37,i4]'] 'discharge= ', qx[loc1[1]]
                    assert False, "dto= " + str(dto)

            if dto > dtrest:
                dto = dtrest

            # print[800000,'(a3,f8.2,a7,f4.2,a7,f8.2]'] 'dto= ', dto,' ratio=', dto/dt, ' max Q= ',qx[Sub.XSno]

            # update the water level at each cross section except the upstream boundary node
            for x in range(1, (Sub.XSno) - 2):
                # check if the XS has laterals
                # FindInArrayF(real[xsid[x]],sub_XSLaterals,loc)
                if Sub.xsid[x] in Sub.LateralsTable:
                    loc = Sub.LateralsTable.index(Sub.xsid[x])
                    # have only the laterals
                    qlat = Lateral_q[t, loc]
                    # print(*,*] "Lat= ", qlat
                else:
                    qlat = 0

                # Forward time backward scpace
                XSarea[x] = (
                    xs[x, 0] - ((qx[x] - qx[x - 1]) * (dto / dx)) + qlat * (dto / dx)
                )  # dto # difference between two consecutive XS areas

                if XSarea[x] <= 0:
                    break
                    print(x)
                    print(XSarea[x])
                    assert False, "Error cross section area is less than zero"
                else:
                    # calculate the bankful area
                    bankful_area = Sub.Dbf[x] * Sub.mw[x]

                    # if the new area is less than the area of the min[hl,hr]
                    if XSarea[x] <= Sub.AreaPerLow[x, 0]:
                        # if the area is less than the bankful area
                        if XSarea[x] <= bankful_area:
                            hx[x] = XSarea[x] / Sub.mw[x]
                        else:
                            # Case 1
                            dummyarea = XSarea[x] - bankful_area
                            dummyh = QuadraticEqn(
                                Sub.cr[x] / Sub.hr[x] + Sub.cl[x] / Sub.hl[x],
                                2 * Sub.mw[x],
                                -2 * dummyarea,
                            )
                            hx[x] = dummyh + Sub.Dbf[x]

                    # if the new area is less than the area of the max[hl,hr]
                    elif XSarea[x] <= Sub.AreaPerHigh[x, 1]:
                        dummyarea = XSarea[x] - bankful_area
                        # check which dike is higher
                        if Sub.hl[x] < Sub.hr[x]:
                            # case 2
                            dummyh = QuadraticEqn(
                                Sub.cr[x] / Sub.hr[x] * 0.5,
                                Sub.mw[x] + Sub.cl[x],
                                -dummyarea - Sub.hl[x] * Sub.cl[x] * 0.5,
                            )
                        else:
                            # case 3
                            dummyh = QuadraticEqn(
                                Sub.cl[x] / Sub.hl[x] * 0.5,
                                Sub.mw[x] + Sub.cr[x],
                                -dummyarea - Sub.hr[x] * Sub.cr[x] * 0.5,
                            )

                        hx[x] = dummyh + Sub.Dbf[x]

                    else:
                        dummyarea = XSarea[x] - bankful_area
                        # case 4 (linear function]
                        dummyh = (
                            dummyarea
                            + Sub.cl[x] * Sub.hl[x] / 2.0
                            + Sub.cr[x] * Sub.hr[x] / 2.0
                        ) / (Sub.cr[x] + Sub.cl[x] + Sub.mw[x])
                        hx[x] = dummyh + Sub.Dbf[x]

                    # to prevent -ve Sf in the next time step calculation
                    if hx[x] + Sub.bedlevel[x] > hx[x - 1] + Sub.bedlevel[x - 1]:
                        # print[800000,'(a]'] '-----'
                        # print[800000,'(a,i5,a11,i2]'] 'time = ',t , ' round no= ',adaptiveTS_counter
                        # print[800000,'(a9,i5,a4,f10.3]'] 'Q at x = ', xsid[x],' Q= ' ,qx[x]
                        # print[800000,'(a9,i5,a4,f10.3]'] 'Q at x = ', xsid[x-1],' Q= ' ,qx[x-1]
                        # g = (hx[x]+Sub.bedlevel[x])- (hx[x-1]+Sub.bedlevel[x-1])
                        # print[800000,'(a,i5,a3,f10.4,a2]'] 'H is higher than previous XS= ',xsid[x], "by ", g,' m'
                        # print(*,'(a,i5]'] 'H is higher than previous XS= ',xsid[x]
                        hx[x] = hx[x - 1] + Sub.bedlevel[x - 1] - Sub.bedlevel[x] - 0.01

                    # this check to prevent the previous step from creating -ve values
                    # as after calculating the h[x]by solving the quadratic equation (QuadraticEqn]the step
                    # of hx[x] = hx[x-1] + bedlevel[x-1] - bedlevel[x] -0.01 can result in -ve hx
                    # so this check should be after not before
                    if hx[x] <= 0.01:
                        hx[x] = 0.01

            # water depth at the last point equals to wl at point before the last
            hx[Sub.XSno - 1] = hx[Sub.XSno - 2] - 0.001

            if min(hx) < 0:
                assert False, "calculated water depth is less than 0 "

            # update the discharge array & dt
            dtrest = dtrest - dto
            q[:, t] = q[:, t] + qx[:] * (dto / dt)

            if River.D1["ModelMode"] == 2:
                q_outL[:, t] = q_outL[:, t] + q_outLx * (dto / dt)
                q_outR[:, t] = q_outR[:, t] + q_outRx * (dto / dt)
                q_out[:, t] = q_outL[:, t] + q_outR[:, t]

        # adaptive time step ------------------------------------------------

        # store the calculated water level & water depth
        h[:, t] = hx
        wl = hx + Sub.bedlevel

    return q, h, wl


def QuadraticEqn(a, b, c):
    delta = (b**2) - 4 * a * c
    return (-b + np.sqrt(delta)) / (2 * a)
