"""Return the indices of the elements that are non-zero."""
import numpy
import numpoly

from ..dispatch import implements


@implements(numpy.nonzero)
def nonzero(x, **kwargs):
    """
    Return the indices of the elements that are non-zero.

    Args:
        x (numpoly.ndpoly):
            Input array.

    Returns:
        indices (Union[bool, numpoly.ndarray]):
            Indices of elements that are non-zero.

    Examples:
        >>> q0, q1 = numpoly.variable(2)
        >>> poly = numpoly.polynomial([[3*q0, 0, 0],
        ...                            [0, 4*q1, 0],
        ...                            [5*q0+q1, 6*q0, 0]])
        >>> poly
        polynomial([[3*q0, 0, 0],
                    [0, 4*q1, 0],
                    [q1+5*q0, 6*q0, 0]])
        >>> numpoly.nonzero(poly)
        (array([0, 1, 2, 2]), array([0, 1, 0, 1]))
        >>> poly[numpoly.nonzero(poly)]
        polynomial([3*q0, 4*q1, q1+5*q0, 6*q0])

    """
    x = numpoly.aspolynomial(x)
    return numpy.nonzero(numpy.any(x.coefficients, axis=0))
