import numpy as np


def get_chezy_coeff(h, ks):
    # Dimensionless Chézy coefficient for rough turbulent flow (=1/sqrt(Cd))
    # based on the empirical White-Colebrook eq. for large estuaries (high Re)
    # In the hydro, the bulk drag Cd = 2.5-3 => min value of Chezy is 1/sqrt(2.5e-3)
    x = np.array(2.5 * np.log(11 * h / ks))
    x = np.append(x, 1.0 / np.sqrt(2.5e-3))
    return np.min(x)


def get_skin_roughness_stress(U, h, rho_w, D90):
    # Bottom shear stress due to skin roughness
    ks = 3 * D90
    Cs = get_chezy_coeff(h, ks)
    return rho_w * (U ** 2) / (Cs ** 2)


def get_form_roughness_stress(U, h, rho_w, eta):
    # Bottom shear stress due to bed form roughness
    # eta = bed form roughness height
    Cs = get_chezy_coeff(h, eta)
    return rho_w * (U ** 2) / (Cs ** 2)


def get_etab(M, h, D50):
    # Equilibrium bed form height eta_b according to Van Rijn (1984)
    # eta_b = 0 if D50 < 0.05mm, if M<1 or if M>24
    # if (D50<5e-5 or M<=1 or M>24):
    #     return 0.0
    # else:
    #     return 0.11 * h * ((D50/h)**0.3) * (1-np.exp(-0.5*(M-1))) * (24-M)
    res = np.zeros_like(M)
    if D50 < 5e-5:
        return res + 1e-2
    else:
        res = 0.11 * h * ((D50 / h) ** 0.3) * (1 - np.exp(-0.5 * (M - 1))) * (24 - M)
        res[np.logical_or(M <= 1, M > 24)] = 0
        return res + 1e-2


def get_ws(Dgr, D, nu):
    # Particle fall velocity ws (m/s) from Soulsby (1997)

    if (np.array(Dgr)).size == 1:
        if Dgr >= 0.672:
            tmp = np.sqrt(107.33 + 1.049 * (Dgr ** 3)) - 10.36
        else:
            tmp = 0.0077 * (Dgr ** 2)
    else:
        tmp = 0.0077 * (Dgr ** 2)
        index = Dgr >= 0.672
        tmp[index] = np.sqrt(107.33 + 1.049 * (Dgr[index] ** 3)) - 10.36
    return nu * tmp / D


def get_zs(h, ws, ust):
    # Height of the centroid of the suspended load
    # ws = particle fall get_bed_shear_velocity
    # ust = shear velocity from both shear and form drag
    kappa = 0.4
    p = 1.2 * np.log(ws / (kappa * ust)) - 0.4
    return 0.0398 * h * 10 ** (-1.08 * np.tanh(p))


def get_us(ust, zs, ks):
    # Advection speed of suspended load us (m/s)
    # ust = shear velocity associated with form drag only
    # ks = bed form roughness height
    return 2.5 * ust * np.log(30 * zs / ks)


########


def get_theta(tau, rho_w, rho_s, D):
    # Shields parameter
    # tau = skin shear stress
    g = 9.81
    s = rho_s / rho_w
    return tau / (rho_w * g * (s - 1) * D)


def get_Dgr(rho_w, rho_s, nu, D50):
    # Sediment dimensionless grain size
    # D50 = grain size at which 50% of the sediment is finer
    g = 9.81
    s = rho_s / rho_w
    return D50 * ((s - 1) * g / (nu ** 2)) ** (1 / 3)


def get_thetac(Dgr):
    # Critical Shields parameter from Soulsby and Whitehouse (1997)
    return 0.3 / (1 + 1.2 * Dgr) + 0.055 * (1 - np.exp(-0.02 * Dgr))


def get_bed_shear_velocity(U, Cs):
    return U / Cs


def get_ub(ust, M):
    # Advection speed of the bedload from Engelund and Fredsoe (1976)
    # ust = shear velocity associated with skin friction only
    # M = dimensionless mobility parameter
    ub = ust * (10 - 7 / np.sqrt(M))
    return np.maximum(ust * (10 - 7 / np.sqrt(M)), np.zeros_like(M))


def get_sediment_param(sed_str="medium_sand"):
    mu = 1e-6  # micrometer
    if sed_str == "coarse_sand":
        D90 = 1000 * mu
        D50 = 750 * mu
        D35 = 500 * mu
        rho_s = 1900.0
    elif sed_str == "medium_sand":
        D90 = 500 * mu
        D50 = 375 * mu
        D35 = 250 * mu
        rho_s = 1900.0
    elif sed_str == "fine_sand":
        D90 = 250 * mu
        D50 = 190 * mu
        D35 = 125 * mu
        rho_s = 1900.0
    elif sed_str == "coarse_silt":
        D90 = 31 * mu
        D50 = 23 * mu
        D35 = 16 * mu
        rho_s = 1700.0
    elif sed_str == "medium_silt":
        D90 = 16 * mu
        D50 = 12 * mu
        D35 = 8 * mu
        rho_s = 1700.0
    elif sed_str == "fine_silt":
        D90 = 8 * mu
        D50 = 6 * mu
        D35 = 4 * mu
        rho_s = 1700.0
    elif sed_str == "dissolved":
        D90 = 0
        D50 = 0
        D35 = 0
        rho_s = 1000.0
    else:
        print("Sediment type unknown")
        exit(0)
    return D90, D50, D35, rho_s


def get_sed_size_density(min_size, max_size, n):
    """generate random size and related density in given range for each particle"""
    mu = 1e-6
    D = np.random.uniform(min_size * mu, max_size * mu, n)  # in micrometers
    phi = -np.log2(D * 1e3)  # This equation expects D in mm
    rho_s = (
        2.374 - 0.175 * phi + 0.008 * phi * phi
    ) * 1e3  # multiplied by 1e3 to get kg/m3 instead of g/cm3
    return D, rho_s


def resuspensionProbability(theta, theta_crit):
    """Return a probability of resupsension for a settled particle for which M>1."""
    proba = np.random.uniform(0, (theta / theta_crit) ** 2)
    if proba > 0.8:
        return True
    else:
        return False