'''
Parameters
----------
image : ndarray
Input image data. Will be converted to float.
mode : str
One of the following strings, selecting the type of noise to add:

    'gauss'     Gaussian-distributed additive noise.
    'poisson'   Poisson-distributed noise generated from the data.
    's&p'       Replaces random pixels with 0 or 1.
    'speckle'   Multiplicative noise using out = image + n*image,where
'''

import numpy as np
import os
#import cv2

def noisy(noise_typ, image):
    if noise_typ == "gauss":
        row, col, ch = image.shape
        mean = 0
        var = 0.1
        sigma = var ** 0.5
        gauss = np.random.normal(mean, sigma, (row, col, ch))
        gauss = gauss.reshape(row, col, ch)
        noisy = image + gauss
        return noisy
    elif noise_typ == "s&p":
        row, col, ch = image.shape
        s_vs_p = 0.5
        amount = 0.01
        out = np.copy(image)
        #out = np.asarray(image)
        blackTreshold = 3
        whiteTreshold = 253
        # Salt mode
        #num_salt = np.ceil(amount * image.size * s_vs_p)
        rmatrix = np.random.randint(0,255,(row,col))
        for x, y in np.ndenumerate(out):
            if(rmatrix[x[0]][x[1]] < blackTreshold):
                out[x[0]][x[1]] = (0,0,0,255)
            if(rmatrix[x[0]][x[1]] > whiteTreshold):
                out[x[0]][x[1]] = (255, 255, 255,255)

        return out
    elif noise_typ == "poisson":
        vals = len(np.unique(image))
        vals = 2 ** np.ceil(np.log2(vals))
        noisy = np.random.poisson(image * vals) / float(vals)
        return noisy
    elif noise_typ == "speckle":
        row, col, ch = image.shape
        gauss = np.random.randn(row, col, ch)
        gauss = gauss.reshape(row, col, ch)
        noisy = image + image * gauss
        return noisy