from ._ufuncs import _lambertw

import numpy as np


def lambertw(z, k=0, tol=1e-8):
    r"""
    lambertw(z, k=0, tol=1e-8)

    Lambert W function.

    The Lambert W function `W(z)` is defined as the inverse function
    of ``w * exp(w)``. In other words, the value of ``W(z)`` is
    such that ``z = W(z) * exp(W(z))`` for any complex number
    ``z``.

    The Lambert W function is a multivalued function with infinitely
    many branches. Each branch gives a separate solution of the
    equation ``z = w exp(w)``. Here, the branches are indexed by the
    integer `k`.

    Parameters
    ----------
    z : array_like
        Input argument.
    k : int, optional
        Branch index.
    tol : float, optional
        Evaluation tolerance.

    Returns
    -------
    w : array
        `w` will have the same shape as `z`.

    See Also
    --------
    wrightomega : the Wright Omega function

    Notes
    -----
    All branches are supported by `lambertw`:

    * ``lambertw(z)`` gives the principal solution (branch 0)
    * ``lambertw(z, k)`` gives the solution on branch `k`

    The Lambert W function has two partially real branches: the
    principal branch (`k = 0`) is real for real ``z > -1/e``, and the
    ``k = -1`` branch is real for ``-1/e < z < 0``. All branches except
    ``k = 0`` have a logarithmic singularity at ``z = 0``.

    **Possible issues**

    The evaluation can become inaccurate very close to the branch point
    at ``-1/e``. In some corner cases, `lambertw` might currently
    fail to converge, or can end up on the wrong branch.

    **Algorithm**

    Halley's iteration is used to invert ``w * exp(w)``, using a first-order
    asymptotic approximation (O(log(w)) or `O(w)`) as the initial estimate.

    The definition, implementation and choice of branches is based on [2]_.

    References
    ----------
    .. [1] https://en.wikipedia.org/wiki/Lambert_W_function
    .. [2] Corless et al, "On the Lambert W function", Adv. Comp. Math. 5
       (1996) 329-359.
       https://cs.uwaterloo.ca/research/tr/1993/03/W.pdf

    Examples
    --------
    The Lambert W function is the inverse of ``w exp(w)``:

    >>> import numpy as np
    >>> from scipy.special import lambertw
    >>> w = lambertw(1)
    >>> w
    (0.56714329040978384+0j)
    >>> w * np.exp(w)
    (1.0+0j)

    Any branch gives a valid inverse:

    >>> w = lambertw(1, k=3)
    >>> w
    (-2.8535817554090377+17.113535539412148j)
    >>> w*np.exp(w)
    (1.0000000000000002+1.609823385706477e-15j)

    **Applications to equation-solving**

    The Lambert W function may be used to solve various kinds of
    equations.  We give two examples here.

    First, the function can be used to solve implicit equations of the
    form

        :math:`x = a + b e^{c x}`

    for :math:`x`.  We assume :math:`c` is not zero.  After a little
    algebra, the equation may be written

        :math:`z e^z = -b c e^{a c}`

    where :math:`z = c (a - x)`.  :math:`z` may then be expressed using
    the Lambert W function

        :math:`z = W(-b c e^{a c})`

    giving

        :math:`x = a - W(-b c e^{a c})/c`

    For example,

    >>> a = 3
    >>> b = 2
    >>> c = -0.5

    The solution to :math:`x = a + b e^{c x}` is:

    >>> x = a - lambertw(-b*c*np.exp(a*c))/c
    >>> x
    (3.3707498368978794+0j)

    Verify that it solves the equation:

    >>> a + b*np.exp(c*x)
    (3.37074983689788+0j)

    The Lambert W function may also be used find the value of the infinite
    power tower :math:`z^{z^{z^{\ldots}}}`:

    >>> def tower(z, n):
    ...     if n == 0:
    ...         return z
    ...     return z ** tower(z, n-1)
    ...
    >>> tower(0.5, 100)
    0.641185744504986
    >>> -lambertw(-np.log(0.5)) / np.log(0.5)
    (0.64118574450498589+0j)
    """
    # TODO: special expert should inspect this
    # interception; better place to do it?
    k = np.asarray(k, dtype=np.dtype("long"))
    return _lambertw(z, k, tol)