import numpy as np
import SimpleITK as sitk

#paths
# path_b800 = '/data/pca-rad/sfransen/UMCG_009.nii.gz'
# path_seg_noise = '/data/pca-rad/sfransen/seg_umcg_0003.nii.gz'

path_b800 = '/data/pca-rad/sfransen/b800_003.nii.gz'
path_seg_noise = '/data/pca-rad/sfransen/segment_003.nii.gz'
#load image/seg_signal/seg_noise
b800_img = sitk.ReadImage(path_b800, sitk.sitkFloat32)
b800_arr = sitk.GetArrayFromImage(b800_img)

noise_img = sitk.ReadImage(path_seg_noise, sitk.sitkFloat32)
noise_arr = sitk.GetArrayFromImage(noise_img)

noise_list = np.matrix.flatten(np.multiply(noise_arr,b800_arr))
#calculate standard deviation
values = [i for i in np.matrix.flatten(np.multiply(noise_arr,b800_arr)) if i > 0]
std_old = np.std(values)
print('std old',std_old)
A1 = 12
A2 = 4
#calculate new std 
std_new = np.multiply(std_old,np.divide(np.sqrt(A1),np.sqrt(A2)))
print('std new',std_new)
# std_2 = np.divide(std_1,np.sqrt(2))
std_noise = np.sqrt(np.abs(np.power(std_old,2)-np.power(std_new,2)))
# std_noise = np.multiply(std_1,std_2)
print('std_noise',std_noise)
noise = np.random.normal(0,std_noise,np.shape(b800_arr))
b800_arr_noise = np.add(b800_arr,noise)

values_check = [i for i in np.matrix.flatten(np.multiply(noise_arr,b800_arr_noise)) if i > 0]
std_new_check = np.std(values_check)
print('std new check',std_new_check)

b800_img_2 = sitk.GetImageFromArray(b800_arr_noise)
b800_img_2.CopyInformation(b800_img)
sitk.WriteImage(b800_img_2, f"../test_4.nii.gz")
b800_img = sitk.ReadImage(path_b800, sitk.sitkFloat32)
sitk.WriteImage(b800_img, f"../test_true.nii.gz")