%%%%% Assignment 2, part 3.

% set noise variance
sigma = 0.05;
% set histogram range
x = [-0.3:0.001:0.3];

% generate noise image
noise_im2 = randn(256,256)*sigma;
noise_im = imread('/tmp/test.jpg');

% load in real images
orig_im = imread('einstein.pgm'); % image we will denoise
%prior_im = imread('feynman.pgm'); % image used to build prior
prior_im = imread('einstein.pgm'); % image used to build prior

% make images 0...1
orig_im = double(orig_im)/256;
prior_im = double(prior_im)/256;
noise_im = double(noise_im)/256;

% create signal image = noise +  real image
%signal_im = orig_im + noise_im2;
signal_im = noise_im;

% form steerable pyramids of noise image, signal image, prior image 
[n_pyr,n_ind]=buildSFpyr(noise_im,3,3);
[p_pyr,p_ind]=buildSFpyr(prior_im,3,3);
[s_pyr,s_ind]=buildSFpyr(signal_im,3,3);

figure;
for a=1:5 %%% only use high freq. bands
  
  % set sub-bands of prior and noise image
  p = pyrBand(p_pyr,p_ind,a);
  n = pyrBand(n_pyr,n_ind,a);
  
  % histogram coefficients
  % be very careful with regularization. if we have 0 in histogram bins
  % we get into big trouble. note 0.1 for prior histogram, but only eps
  % for noise. Have to ensure that noise regularization is always a lot
  % less than prior image regularization, otherwise estimator deviates
  % from diagonal.
  
  hp = hist(p(:),x)+0.3;
  hn = hist(n(:),x)+eps;
 
  % compute numerator of eqn. 1
  y = conv2(x.*hp,hn,'same');
  % compute denominator of eqn. 1
  y2  = conv2(hp,hn,'same');
  % combine
  yy(:,a) = y ./ y2;
  keyboard
  % plot estimator out
  plot(x,yy(:,a));
  hold on;
  axis('equal');

  %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
  %  Now use estimator on the signal image

  % get sub-band
  s = pyrBand(s_pyr,s_ind,a);
  
  % remap coefficient values using estimator
  s_new = interp1(x,yy(:,a),s(:));
  s_new = reshape(s_new,size(s));
  
  % insert new band into pyramid
  s_pyr = setPyrBand(s_pyr,s_ind,s_new,a);
  
end


% reconstruct steerable pyramid
output_im = reconSFpyr(s_pyr,s_ind);

% compute baseline using Wierner filtering.
baseline_im = wiener2(signal_im,[3 3],sigma^2);

% compute PSNRs
signal_psnr=PSNR(orig_im,signal_im);
baseline_psnr=PSNR(orig_im,baseline_im);
output_psnr=PSNR(orig_im,output_im);

% plot images
figure; imagesc([orig_im,signal_im,baseline_im,output_im]); colormap(gray);
title(sprintf('Original,Noisy: %f, Wiener: %f, Cored: %f\n',signal_psnr,baseline_psnr,output_psnr));

