function AddRicianNoise(original_dti_final,SNR_percent)

% ----------------- ARGUMENTS -----------------
% original_dti_final: original (default SNR) NIFTI filename and path of processed, final diffusion-weighted image (EX: 'dti_final.nii.gz')
% SNR_percent: vector of percentage of original SNR values desired (EX: [90 80 70 60 50])

% ----------------- OUTPUT ----------------
% 'SNR_values_orig.txt': text file with original image SNR values
% NIFTI image files for each of SNR levels (Ex:'dti_final_80percentSNR.nii.gz')
% Text files with new SNR values for each of the SNR levels (Ex: 'SNR_values_80percentSNR.txt')

% ----------------- NOTES -----------------
%   - make sure the following are in the 'segmentations' folder: 't1_reg2dti_muscle_filled.nii.gz' and 'dti_bone_seg.nii.gz' (both in DTI space)
%   - Lines 58-60: replace with desired mu, sigma, and threshold values

%% Load Original Diffusion-Weighted Image

orig = load_untouch_nii(original_dti_final);    % Load the denoised NIFTI image as a structure
im_orig = double(orig.img);                     % Get the image from the structure
og_size = size(im_orig);                        % Find size of the image

%% Load Masks for SNR Calculations

muscle_mask = load_untouch_nii('segmentations/t1_reg2dti_muscle_filled.nii.gz');    % Load the segmented muscle region (DTI space) as a structure
muscle_img = double(muscle_mask.img);                                               % Get the image from the structure
bone_mask = load_untouch_nii('segmentations/dti_bone_seg.nii.gz');                  % Load the segmented bone region (DTI space) as a structure
bone_img = double(bone_mask.img);                                                   % Get the image from the structure

%% Calculate Original SNR Values, Save to a Text File

% Calculate signal and noise in non-diffusion-weighted (b=0) volume
signal_b0 = mean(nonzeros(muscle_img(:,:,:) .* im_orig(:,:,:,1)),'all');                        % signal mask*img --> signal only in masked region --> take mean in b0 volume
noise_b0 = mean(std(nonzeros(bone_img(:,:,:) .* im_orig(:,:,:,1)),0,'all'),'all');              % noise mask*img --> noise only in masked region --> take mean in b0 volume

% Calculate signal and noise across diffusion-weighted (b>0) volumes
for vol=2:og_size(4)
    signal_DW(vol-1) = mean(nonzeros(muscle_img(:,:,:) .* im_orig(:,:,:,vol)),'all');           % signal mask*img --> signal only in masked region --> take mean in b400 volume;
    noise_DW(vol-1) = mean(std(nonzeros(bone_img(:,:,:) .* im_orig(:,:,:,vol)),0,'all'),'all'); % noise mask*img --> noise only in masked region --> take mean in b400 volume
end
avg_sig_DW = mean(signal_DW);
avg_noise_DW = mean(noise_DW);

% Calculate SNR Values
snr_b0 = signal_b0 / noise_b0;          % SNR (b=0) volume only
snr_DW = avg_sig_DW / avg_noise_DW;     % SNR (DW) volumes only
snr_final = signal_b0 / avg_noise_DW;   % Final SNR: signal(b=0)/noise(DW)

SNR_values_orig = [signal_b0 noise_b0 snr_final; avg_sig_DW avg_noise_DW snr_DW];
writematrix(SNR_values_orig,'SNR_values_orig.txt','Delimiter','space')

%% Monte Carlo Loops to Add Rician Noise

% Read Original SNR Value
orig_SNR = SNR_values_orig(1,3);    % Row 1, Col 3 = Final SNR

% Set Constant Values
SNR_fraction = SNR_percent/100;     % Create vector of desired SNR fractions (ie. [0.9 0.8 0.7 0.6 0.5])
SNR_threshold = 0.015;              % Threshold of allowable error for SNR percent (1.5%)
mu=0;                               % Set mean for Gaussian noise distribution: 0
sigma_int = 0.25;                   % Set sigma for Gaussian noise distribution to be added each iteration: 0.25

% Pre-Allocate Containers
im_SNR_fraction = cell(size(SNR_percent));          
SNR_values_decreased = cell(size(SNR_percent));
num_iter_needed = zeros(2,size(SNR_percent,2));

% Set Starting Values for Iteration Number 1
sigma = 0;              % SD starts at 0
rounds_tab=1;           % Number of loop rounds needed
iter_num = 1;           % Iteration starts at 1
im = im_orig;           % Start with original image
new_SNR = orig_SNR;     % Start with original SNR

disp(strcat('Round ',num2str(rounds_tab)));

for n = 1:length(SNR_percent)
    while (new_SNR > (SNR_fraction(n)+SNR_threshold)*orig_SNR) || (new_SNR < (SNR_fraction(n)-0.1)*orig_SNR)

        for i = 1:og_size(4)                                                        % For each diffusion volume
            for j = 1:og_size(3)                                                    % For each z slice
                im_2d = im(:,:,j,i);                                                % One slice at one diffusion condition --> 2-D array
                realchannel = normrnd(mu,sigma,size(im_2d)) + im_2d;                % Create "real" Gaussian noise
                imaginarychannel = normrnd(mu,sigma,size(im_2d));                   % Create "imaginary" Gaussian noise
                im_noisy(:,:,j,i) = sqrt(realchannel.^2 + imaginarychannel.^2);     % Add real and imaginary to create Rician noise distribution
            end
        end

        % Calculate New SNR
        [new_SNR,SNR_values_noisy] = Calculate_SNR_Values_Noisy(muscle_img,bone_img,im_noisy);      % See end of script for function
        
        % If new SNR falls within the threshold limit of the desired SNR percentage, save the image, SNR values, and number of iterations/rounds
        if (new_SNR < (SNR_fraction(n)+SNR_threshold)*orig_SNR) && (new_SNR > (SNR_fraction(n)-SNR_threshold)*orig_SNR)     
            im_SNR_fraction{n} = im_noisy;
            SNR_values_decreased{n} = SNR_values_noisy;
            num_iter_needed(1,n) = iter_num;
            num_iter_needed(2,n) = rounds_tab;
        end

        % Print Iteration # and New SNR
        disp(strcat('At Iteration #',num2str(iter_num),', the SNR is at: ',num2str((new_SNR/orig_SNR)*100),'% of original SNR'));

        % Add to starting values for the next iteration
        im = im_noisy;
        sigma = sigma + sigma_int;
        iter_num = iter_num + 1;

    end
end

% End Monte Carlo Loop (if all images were found)
idx_no_result = find(num_iter_needed(1,:)==0);  % Find SNR percents that did not have a result in the first round
if isempty(idx_no_result)                       % If all SNR percents had a result, loop is done
    disp('SNR Decreasing Done!');
end

% Start Loop Again (if some images were not found)
while ~isempty(idx_no_result)                   
    disp(strcat('Round ',num2str(rounds_tab)));

    % Set Starting Values
    sigma = 0;
    iter_num = 1;

    for n = 1:length(idx_no_result)
        im = im_SNR_fraction{idx_no_result(n)-1};
        new_SNR = SNR_values_decreased{idx_no_result(n)-1}(1,3);

        while (new_SNR > (SNR_fraction(idx_no_result(n))+SNR_threshold)*orig_SNR) || (new_SNR < (SNR_fraction(idx_no_result(n))-0.1)*orig_SNR)
    
            for i = 1:og_size(4)                                                        % For each diffusion volume
                for j = 1:og_size(3)                                                    % For each z slice
                    im_2d = im(:,:,j,i);                                                % One slice at one diffusion condition --> 2-D array
                    realchannel = normrnd(mu,sigma,size(im_2d)) + im_2d;                % Create "real" Gaussian noise
                    imaginarychannel = normrnd(mu,sigma,size(im_2d));                   % Create "imaginary" Gaussian noise
                    im_noisy(:,:,j,i) = sqrt(realchannel.^2 + imaginarychannel.^2);     % Add real and imaginary to create Rician noise distribution
                end
            end
    
            % Calculate New SNR
            [new_SNR,SNR_values_noisy] = Calculate_SNR_Values_Noisy(muscle_img,bone_img,im_noisy);  % See end of script for function

            % If new SNR falls within the threshold limit of the desired SNR percentage, save the image, SNR values, and number of iterations/rounds
            if (new_SNR < (SNR_fraction(idx_no_result(n))+SNR_threshold)*orig_SNR) && (new_SNR > (SNR_fraction(idx_no_result(n))-SNR_threshold)*orig_SNR)
                im_SNR_fraction{idx_no_result(n)} = im_noisy;
                SNR_values_decreased{idx_no_result(n)} = SNR_values_noisy;
                num_iter_needed(1,idx_no_result(n)) = iter_num;
                num_iter_needed(2,idx_no_result(n)) = rounds_tab;
            end
    
            % Print Iteration # and New SNR
            disp(strcat('At Iteration #',num2str(iter_num),', the SNR is at: ',num2str((new_SNR/orig_SNR)*100),'% of original SNR'));
    
            % Add to starting values for next iteration
            im = im_noisy;
            sigma = sigma + sigma_int;
            iter_num = iter_num + 1;
        end
    end

    % End Monte Carlo Loop (if all images were found)
    idx_no_result = find(num_iter_needed(1,:)==0);  % Find SNR percents that did not have a result in the first round
    if isempty(idx_no_result)                       % If all SNR percents had a result, loop is done
        disp('SNR Decreasing Done!');
    end
    rounds_tab=rounds_tab+1;
end

%% Save New Images

for n = 1:length(SNR_percent)
    NIFTI_im = make_nii(im_SNR_fraction{n},[1.25 1.25 6.5]);                        % Make NIFTI structure with correct voxel dimensions
    NIFTI_im.hdr = orig.hdr;                                                        % Save with original image NIFTI header
    filename = strcat('dti_final_',num2str(SNR_percent(n)),'percentSNR.nii.gz');    
    save_nii(NIFTI_im,filename);                                                    % Save new SNR Values Image
    fprintf('File saved as %s\n',filename);
    writematrix(SNR_values_decreased{n},strcat('SNR_values_',num2str(SNR_percent(n)),'percentSNR.txt'),'Delimiter','space');    % Save new SNR Values
end
writematrix(num_iter_needed,strcat('NumIterationsNeeded.txt'),'Delimiter','space'); % Save number of iterations needed

exit;

end

%% Calculate SNR of New Noise Added Images

function [new_SNR,SNR_values_noisy] = Calculate_SNR_Values_Noisy(muscle_img,bone_img,im_noisy)

    % Calculate signal and noise in non-diffusion-weighted (b=0) volume
    signal_b0 = mean(nonzeros(muscle_img(:,:,:) .* im_noisy(:,:,:,1)),'all');                           % signal mask*img --> signal only in masked region --> take mean in b0 volume
    noise_b0 = mean(std(nonzeros(bone_img(:,:,:) .* im_noisy(:,:,:,1)),0,'all'),'all');                 % noise mask*img --> noise only in masked region --> take mean in b0 volume

    % Calculate signal and noise across diffusion-weighted (b>0) volumes
    for vol=2:size(im_noisy,4)
        signal_DW(vol-1) = mean(nonzeros(muscle_img(:,:,:) .* im_noisy(:,:,:,vol)),'all');              % signal mask*img --> signal only in masked region --> take mean in b400 volume;
        noise_DW(vol-1) = mean(std(nonzeros(bone_img(:,:,:) .* im_noisy(:,:,:,vol)),0,'all'),'all');    % noise mask*img --> noise only in masked region --> take mean in b400 volume
    end
    avg_sig_DW = mean(signal_DW);
    avg_noise_DW = mean(noise_DW);

    % Calculate SNR Values
    snr_b0 = signal_b0 / noise_b0;          % SNR (b=0) volume only
    snr_DW = avg_sig_DW / avg_noise_DW;     % SNR (DW) volumes only
    snr_final = signal_b0 / avg_noise_DW;   % Final SNR: signal(b=0)/noise(DW)

    % Save SNR Values
    SNR_values_noisy = [signal_b0 noise_b0 snr_final; avg_sig_DW avg_noise_DW snr_DW];
    new_SNR = snr_final;

end