% An azimuth-resolved image wavelet transform.
function [Waz, W, theta0] = wazlet(image, Nscales, Naz)
% Returns the wavelet transform of image, divided into Nscales
% layers, plus a final remainder that has all the larger scales
% up to DC; and Naz azimuths. The layers are just what you'd
% get from Algorithme Á Trous. They correspond to annuli in the
% Fourier plane. The signal from each layer is then divided up
% by the azimuth of the wavenumber, using a cos^2 shaped window
% in theta.
%
% The image is treated with periodic boundary conditions, with
% no attempt to ameliorate edge effects. If you are concerned
% about edge effects, try unfolding/mirroring the image or
% padding it in some way.
% 
% The inverse transform is the sum of all layers and azimuths:
%     image = sum(sum(Waz, 4), 3) = sum(W, 3)
%
% CCK 2018-Mar-17


isize = size(image);
Nx = isize(1); % number of rows
Ny = isize(2); % number of columns


% Create Gaussian smoothing kernel, k1.
ksigma = 1.0/sqrt(2); % kernel standard deviation
% Coordinate system for the kernel. Note that the arrangement of positive and 
% negative coordinates is just like that in k-space (wavenumber.m).
x = [0:floor(Nx/2), -floor((Nx-1)/2):-1]' * ones(1,Ny);
y = ones(Nx,1) * [0:floor(Ny/2), -floor((Ny-1)/2):-1];
r2 = x.^2 + y.^2;
k1 = exp(-r2/(2*ksigma^2)); % Gaussian smoothing kernel
k1 = k1/sum(sum(k1)); % normalization

% storage for a trous wavelet transform (and its layer-wise FFT):
W = zeros(Nx,Ny,Nscales+1);
F_W = zeros(Nx,Ny,Nscales+1); 

% Necessary Fourier transforms. Note the F_ prefix notation.
F_image = fft2(double(image));
   % Since image data types are often uint8, double() makes sure
   % we are working with floating point numbers.
F_k1 = fft2(k1);    %%%%%%%%%% commented out broadening!   .^4;  %%%%%% BROADEN THE KERNEL. %%%%%%

F_image_blur = F_image;
for j=1:Nscales
   F_image_blur_next = F_image .* ( F_k1.^(2.^(j-1)) );
   F_W(:,:,j) = F_image_blur - F_image_blur_next;
   W(:,:,j) = real(ifft2(F_W(:,:,j)));
   F_image_blur = F_image_blur_next;
end

% The final layer:
F_W(:,:,Nscales+1) = F_image_blur;
W(:,:,Nscales+1) = real(ifft2(F_image_blur));

% Calculate azimuth angles in Fourier space.
kx = wavenumber(Nx)' * ones(1,Ny);
ky = ones(Nx,1) * wavenumber(Ny);
theta = atan2(ky, kx);

Waz = zeros(Nx,Ny,Nscales,Naz); % storage for wazlet transform

% Apply azimuth filters
theta0 = (0:(Naz-1))*pi/Naz; % central angle for each azimuth filter
for az=1:Naz
   relative_theta = mod(theta-theta0(az), 2*pi) - pi;
      % This coordinate has its origin at theta = theta0,
      % and it runs from -pi to pi.
   azfilt = (cos(0.5*Naz*relative_theta)) .^ 2; % azimuth filter
   % now exclude all but the central lobes of cos^2, which
   % occur in the neighborhood of zero and ±π :
   outside = find(                          ...
      ( abs(relative_theta) > pi/Naz )   &  ...
      ( abs(relative_theta) < pi*(1-1/Naz) )...
      );
   azfilt(outside) = 0;
   
   % Plot the azimuthal filters (for debugging)
%{
   figure(10+az)
   hold off
   imagesc(azfilt)
   colormap gray
%}
   
   % Apply the azimuth filter to each layer
   for j=1:Nscales+1
      Waz(:,:,j,az) = real(ifft2(azfilt .* F_W(:,:,j)));
   end
end

