import numpy as np
import random
from utils.myutils import dim_prod
import matplotlib.pyplot as plt
def add_gaus_noise(input_tensor,sigma):
    tensor_degrd = input_tensor.copy()
    [n1, n2, n3] = input_tensor.shape
    noise_tensor = np.float32(np.random.normal(0, sigma / 255.,[n1,n2,n3]))
    tensor_degrd+=noise_tensor
    return tensor_degrd,noise_tensor
def add_gaus_noise_general(input_tensor,sigma):
    tensor_degrd = input_tensor.copy()
    noise_tensor = np.float32(np.random.normal(0, sigma / 255.,input_tensor.shape))
    tensor_degrd+=noise_tensor
    return tensor_degrd,noise_tensor

def random_sample_general(data_gt,sample_rate,sigma):
    size_list = data_gt.shape
    tensor_degrd = data_gt.copy()
    tensor_degrd,noise_tensor = add_gaus_noise_general(tensor_degrd,sigma)
    a = (np.random.rand(dim_prod(data_gt)) < sample_rate)
    data_ob = np.zeros_like(data_gt)
    data_ob[a.reshape(size_list)] = tensor_degrd[a.reshape(size_list)]
    return data_ob, a,noise_tensor
def random_sample(data_gt,sample_rate,sigma):
    [n1, n2, n3] = data_gt.shape
    tensor_degrd = data_gt.copy()
    tensor_degrd,noise_tensor = add_gaus_noise(tensor_degrd,sigma)
    a = (np.random.rand(dim_prod(data_gt)) < sample_rate)
    data_ob = np.zeros_like(data_gt)
    data_ob[a.reshape([n1, n2, n3])] = tensor_degrd[a.reshape([n1, n2, n3])]
    return data_ob, a,noise_tensor

