import jax.numpy as jnp
import numpyro.distributions as ndist


def ensure_positive(data, indices=None):
    """
    Recursive function for ensuring that the scalar elements are positive

    :param data: tuple or dict
    :param indices: If specified, indicates that only these indices should be checked for positivity
    :return: True if all positive else False
    """
    if type(data) is dict:
        for v in data.values():
            if not ensure_positive(v, indices): return False
    elif type(data[0]) in (tuple, list):
        for v in data:
            if not ensure_positive(v, indices): return False
    else:
        if indices is None: indices = list(range(len(data)))
        for e, v in enumerate(data):
            if (e in indices) and (v < 0): return False
    return True


def Threshold(x, threshold, sharpness=100):
    """
    Returns a probabilistic soft threshold which can be used in place of a distribution

    :param data: The data vector to threshold
    :param threshold: The threshold
    :param sharpness: Governs how well this approximates a Step function
    :return: Distribution
    """
    p = 1/(1 + jnp.exp(-sharpness*(x - threshold)))
    return ndist.Bernoulli(probs=p)
