"""
Noise functions used to train Denoising Auto-Associators.

Functions in this module often include a `noise_lvl` argument that controls the amount of noise
that the function applies.
The noise contract is simple: noise_lvl is a symbolic variable going from 0 to 1.
0: no change.
1: maximum noise.
"""
import theano
import tags
s="""
* A latex mathematical description of the formulas(for picture representation in generated documentation)
* Tags(for searching):
   * a list of lower lovel fct used
   * category(name of the submodule itself)
* Tell if we did some work to make it more numerical stable. Do theano do the optimization needed?
* Tell if the grad is numericaly stable? Do theano do the optimization needed?
* Tell if work on gpu/not/unknow
* Tell alternate name
* Tell the domaine, range of the input/output(range should use the english notation of including or excluding)
"""

@tags.tags('noise','binomial','salt')
def binomial_noise(theano_rng, input, noise_lvl, noise_value=0):
    """
    Return `inp` with randomly-chosen elements set to zero.

    TODO: MATH DEFINITION

    :type input: Theano tensor variable
    :param input: input
    :type noise_lvl: float
    :param noise_lvl: The probability of setting each element to zero.
    :type noise_value: Theano scalar variable
    :param noise_value: The value that we want when their is noise.
    """
    mask = theano_rng.binomial(
            size = input.shape,
            n = 1,
            p =  1 - noise_lvl,
            dtype=input.dtype)
    value = theano.tensor.as_tensor_variable(noise_value)
    if value.type.ndim!=0:
        raise Exception('binomial_noise only support scalar noise_value')
    if noise_value==0:
        return mask * input
    else:
        return mask * input + noise_value*(not mask)


@tags.tags('noise','binomial NLP','pepper','salt')
def pepper_and_salt_noise(theano_rng,inp,noise_lvl):
    """ This add pepper and salt noise to inp

    :type inp: Theano variable
    :param inp: The input that we want to add noise
    :type noise_lvl: tuple(float,float)
    :param noise_lvl: The probability of changing each element to zero or one. 
                      (prob of salt, prob of pepper)

    :note: The sum of the prob of salt and prob of pepper should be less then 1.
    """
    assert inp.dtype in ['float32','float64']
    return theano_rng.binomial( size = inp.shape, n = 1, p =  1 - noise_lvl[0], dtype=inp.dtype) * inp \
                        + (theano.tensor.eq(inp,0)) * theano_rng.binomial( size = inp.shape, n = 1, p =  noise_lvl[1], dtype=inp.dtype)

@tags.tags('noise','gauss','gaussian')
def gaussian_noise(theano_rng,inp,noise_lvl):
    """ This add gaussian NLP noise to inp

    :type inp: Theano variable
    :param inp: The input that we want to add noise
    :type noise_lvl: float
    :param noise_lvl: The standard deviation of the gaussian.
    """
    assert inp.dtype in ['float32','float64']
    return theano_rng.normal( size = inp.shape, std = noise_lvl, dtype=inp.dtype) + inp