% Contact ythomas@csail.mit.edu or msabuncu@csail.mit.edu for bugs or questions 
%
%=========================================================================
%
%  Copyright (c) 2008 Thomas Yeo and Mert Sabuncu
%  All rights reserved.
%
%Redistribution and use in source and binary forms, with or without
%modification, are permitted provided that the following conditions are met:
%
%    * Redistributions of source code must retain the above copyright notice,
%      this list of conditions and the following disclaimer.
%
%    * Redistributions in binary form must reproduce the above copyright notice,
%      this list of conditions and the following disclaimer in the documentation
%      and/or other materials provided with the distribution.
%
%    * Neither the names of the copyright holders nor the names of future
%      contributors may be used to endorse or promote products derived from this
%      software without specific prior written permission.
%
%THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
%ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
%WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
%DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
%ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
%(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
%LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
%ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
%(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
%SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.    
%
%=========================================================================
%
% This function performs forward and inverse wavelet transform of the
% coordinates of MARS_sbjMesh.
%
% The function
%           (1) Interpolates coordinates of MARS_sbjMesh onto lat-lon grid
%           (2) Performs fast spherical harmonic transform using the yawtb and S2kit toolkits.
%           (3) In the spherical harmonic domain, perform convolution of coordinate functions with
%               analysis filters to get wc_x,wc_y,wc_z.
%           (4) In the spherical harmonic domain, perform convolution of wc_x, wc_y, wc_z with
%               synthesis filters to get rc_x,rc_y,rc_z.
%
function [wc_x, wc_y, wc_z, rc_x, rc_y, rc_z] = WaveletTransformMesh(MARS_sbjMesh, analysis_filters, synthesis_filters, checkBool)

% [analysis_outputs, synthesis_outputs] = WaveletTransformMesh(MARS_sbjMesh, analysis_filters, synthesis_filters, checkBool)

if(nargin < 4)
   checkBool = 0; 
end

if(size(analysis_filters, 2) ~= size(synthesis_filters, 2))
   error('We expect analysis and synthesis filters to be of same bandwidth');
end

if(size(analysis_filters, 1) ~= size(synthesis_filters, 1))
   error('We expect the number of analysis and synthesis filters to be the same');
end


% Compute initials
num_filters = size(analysis_filters, 1);
BW = 1/2 *(-1 + sqrt(1+8*size(analysis_filters, 2)));

if( BW - round(BW) ~= 0)
   error('BW not correct!'); 
end

% Fast spherical harmonic transform on mesh
disp('Interpolate mesh coordinates onto lat-lon grid');
mesh_vertices_s2kit = ResampleMeshOntoS2kitGrid(MARS_sbjMesh, BW);
mesh_harm_coefficients = zeros(3, size(analysis_filters, 2));
for i = 1:3
    disp('Performing fast spherical harmonic transform of mesh coordinates');
    mesh_yawtb_compatible = transpose(reshape(mesh_vertices_s2kit(i, :), BW*2, BW*2));
    mesh_s2kit_coeffs = s2fst(double(mesh_yawtb_compatible));

    if(checkBool)
        disp('================ Check that inverse spherical harmonic transform is invertible ================');
        mesh_yawtb_compatible2 = s2ifst(mesh_s2kit_coeffs);
        if(i == 1)
            disp(['Imag Part of recon x (should be 0): ' num2str(max(abs(imag(mesh_yawtb_compatible2(:))))) ', diff in real part of recon x (should be non-zero): '  ...
                num2str(max(abs(mesh_yawtb_compatible2(:) - mesh_yawtb_compatible(:))))]);
        elseif(i == 2)
            disp(['Imag Part of recon y (should be 0): ' num2str(max(abs(imag(mesh_yawtb_compatible2(:))))) ', diff in real part of recon y (should be non-zero): '  ...
                num2str(max(abs(mesh_yawtb_compatible2(:) - mesh_yawtb_compatible(:))))]);
        else
            disp(['Imag Part of recon z (should be 0): ' num2str(max(abs(imag(mesh_yawtb_compatible2(:))))) ', diff in real part of recon z (should be non-zero): '  ...
                num2str(max(abs(mesh_yawtb_compatible2(:) - mesh_yawtb_compatible(:))))]);
        end
    end
    mesh_harm_coefficients(i, :) = ConvertS2kitFormat2MyFormat(mesh_s2kit_coeffs);
end

% Convolve Mesh with analysis filters
disp('Wavelet Transform');
wc_x = zeros(size(analysis_filters));
wc_y = zeros(size(analysis_filters));
wc_z = zeros(size(analysis_filters));

for i = 1:num_filters
   wc_x(i, :) =  convolveSphere(transpose(mesh_harm_coefficients(1, :)), transpose(analysis_filters(i, :)), 1);
   wc_y(i, :) =  convolveSphere(transpose(mesh_harm_coefficients(2, :)), transpose(analysis_filters(i, :)), 1);
   wc_z(i, :) =  convolveSphere(transpose(mesh_harm_coefficients(3, :)), transpose(analysis_filters(i, :)), 1);
end

% Convolve wavelet coeffs with synthesis filters
disp('Inverse Wavelet Transform');
rc_x = zeros(size(analysis_filters));
rc_y = zeros(size(analysis_filters));
rc_z = zeros(size(analysis_filters));
for i = 1:num_filters
    rc_x(i, :) =  convolveSphere(transpose(wc_x(i, :)), transpose(synthesis_filters(i, :)), 0);
    rc_y(i, :) =  convolveSphere(transpose(wc_y(i, :)), transpose(synthesis_filters(i, :)), 0);
    rc_z(i, :) =  convolveSphere(transpose(wc_z(i, :)), transpose(synthesis_filters(i, :)), 0);
end

% check invertibility
if(checkBool)
    
    disp('================ Check that Inverse Wavelet Transform is indeed invertible ================');
    disp(['Following should be close to 0 for invertibility: ' num2str(max(abs(sum(rc_x, 1) - mesh_harm_coefficients(1, :)))) ',' ...
        num2str(max(abs(sum(rc_y, 1) - mesh_harm_coefficients(2, :)))) ',' num2str( max(abs(sum(rc_z, 1) - mesh_harm_coefficients(3, :))))]);
   
    
    mat = ConvertMyFormat2S2kitFormat(transpose(sum(rc_x, 1)));
    mesh_s2kit_recon = s2ifst(mat);   
    mesh_yawtb_compatible = double(transpose(reshape(mesh_vertices_s2kit(1, :), BW*2, BW*2)));
    
    disp(['Imag Part of recon x: ' num2str(max(abs(imag(mesh_s2kit_recon(:))))) ', diff in real part of recon x (should be non-zero, but same as x above):'  ...
            num2str(max(abs(real(mesh_s2kit_recon(:)) - mesh_yawtb_compatible(:))))]);
    
    mat = ConvertMyFormat2S2kitFormat(transpose(sum(rc_y, 1)));
    mesh_s2kit_recon = s2ifst(mat);   
    mesh_yawtb_compatible = double(transpose(reshape(mesh_vertices_s2kit(2, :), BW*2, BW*2)));
    disp(['Imag Part of recon y: ' num2str(max(abs(imag(mesh_s2kit_recon(:))))) ', diff in real part of recon y (should be non-zero, but same as y above):'  ...
            num2str(max(abs(real(mesh_s2kit_recon(:)) - mesh_yawtb_compatible(:))))]);
    
    
    mat = ConvertMyFormat2S2kitFormat(transpose(sum(rc_z, 1)));
    mesh_s2kit_recon = s2ifst(mat);   
    mesh_yawtb_compatible = double(transpose(reshape(mesh_vertices_s2kit(3, :), BW*2, BW*2)));
    disp(['Imag Part of recon z: ' num2str(max(abs(imag(mesh_s2kit_recon(:))))) ', diff in real part of recon z (should be non-zero, but same as z above):'  ...
            num2str(max(abs(real(mesh_s2kit_recon(:)) - mesh_yawtb_compatible(:))))]);
    
end







