import numpy as np

from foxes.tools.abl.stable import logz 

def psi(height, mol):
    """
    The Psi function

    Parameters
    ----------
    height : float
        The height value
    mol : float
        The Monin-Obukhov height
    
    Returns
    -------
    psi : float
        The Psi function value

    """
    x = ( 1.- 16. * height / mol )**0.25
    return    2. * np.log( ( 1. + x ) / 2. ) + np.log( ( 1. + x**2 ) / 2. ) \
            - 2. * np.arctan(x) + np.pi / 2.

def ustar(ws_ref, h_ref, z0, mol, kappa=0.41):
    """
    Calculates the friction velocity,
    based on reference data.

    Parameters
    ----------
    ws_ref : float
        The reference wind speed
    h_ref : float
        The reference height
    z0 : float
        The roughness length
    mol : float
        The Monin-Obukhov height
    kappa : float
        The van-Karman constant
    
    Returns
    -------
    ustar : float
        The friction velocity

    """
    return ws_ref * kappa / ( logz(h_ref, z0) - psi(h_ref, mol) )

def calc_ws(height, z0, ustar, psi, kappa=0.41):
    """
    Calculate wind speeds at given height

    Parameters
    ----------
    height : float
        The evaluation height
    z0 : float
        The roughness length
    ustar : float
        The friction velocity
    psi : float
        The Psi function values
    kappa : float
        The van-Karman constant

    Returns
    -------
    ws : float
        The wind speed

    """
    return ustar / kappa * ( logz(height, z0) - psi ); 