import numpy as np
import matplotlib.pyplot as plt



def add_gaussian_noise(image, stddev):
    # Add Gaussian noise to the image
    noisy_image = image + np.random.normal(0, stddev, image.shape)
    # Ensure data stays between 0 and 255
    noisy_image = np.clip(noisy_image, 0, 255)
    return noisy_image


def add_salt_pepper(image, amount, s_vs_p=0.5):
    # Figure out the amount of salt and pepper desired
    salt_amount = amount * s_vs_p
    pepper_amount = amount * (1-s_vs_p)

    # Calculate the thresholds required to snap to 0 or 255
    salt_threshold = 255*(1-salt_amount)
    pepper_threshold = 255*pepper_amount

    # Create random linear noise between 0 and 255
    spots = np.random.randint(0, 255+1, image.shape)
    
    # Snap to 255 if above the salt threshold
    salted_image = np.where(spots > salt_threshold, 255, image)
    # Snap to 0 if below the pepper threshold
    peppered_image = np.where(spots < pepper_threshold, 0, salted_image)

    return peppered_image