% (c) Julie Digne 2018
% This program is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% This program is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with this program.  If not, see <https://www.gnu.org/licenses/>.


function [Phi, Phicorr, an, indices] = compute_wavejets(p, neighbors, order, s)
%[Phi, Phicorr, an, indices] = compute_wavejets(p, neighbors, order, s)
%
%DISCLAIMER: this is a simplified octave/matlab reimplementation of the code from the Wavejets paper
%
%Wavejets: A Local Frequency Framework for Shape Details Amplification,
%Yohann Béarzi, Julie Digne, Raphaëlle Chaine, Proc. Eurographics 2018
%(Computer Graphics Forum).
%
%It simply computes the wavejets decomposition and performs normal
%correction before computing the integral invariants.
%
% compute the wavejets coefficients around a point given its neighbors
%args:
% p: center point 1x3
% neighbors: set of neighbors nx3
% order: polynomial order of the wavejets (frequency order will be deduced)
% s: neighborhood radius
%returns:
% Phi: wavejets coefficients estimated from the PCA parametrization
% Phicorr: wavejets coefficients after tangent plane correction
% an set of integral invariants
% indices: used indices

nneigh = size(neighbors,1);

% %compute PCA normal
m = mean(neighbors);
C = 1/nneigh*(neighbors'*neighbors) - (m'*m);

[~,~,V]= svd(C);

t1 = V(:,1);
t2 = V(:,2);
normal = V(:,3);


%express neighbors in local coord system
locneighbors = (neighbors - repmat(p,nneigh,1))*[t1 t2 normal];

%switch to polar coordinates
r = sqrt(sum(locneighbors(:,1:2).^2,2));
theta = atan2(locneighbors(:,2), locneighbors(:,1));
z = locneighbors(:,3);

%compute Phi
ncolM = order^2/2+3*order/2+1;

M=zeros(nneigh,ncolM);
indices=zeros(ncolM,2);
idx = 1;
normalized_r=r/s;
z=z/s;
for k=0:order
  rk = normalized_r.^k;
  w = exp(-normalized_r.^2/18);
  for n = -k:2:k
          M(:,idx) = rk.*exp(1i*n*theta).*w;
          indices(idx,:) = [k n];
          idx=idx+1;
  end
end


Phi = linsolve(M,z.*w);


%correct normal
u = -t1*imag(Phi(2)-Phi(3))+t2*real(Phi(3)+Phi(2));%axis
u = u/norm(u);
gamma = -atan(2*abs(Phi(3)));%angle

nc = cos(gamma)*normal+(1-cos(gamma))*(normal'*u)*u+sin(gamma)*cross(u,normal);
nc = nc/norm(nc);
t1c = t1-(t1'*nc)*nc;
t1c = t1c/norm(t1c);
t2c = cross(nc,t1c); 
t2c = t2c/norm(t2c);

%redo everything (instead of correcting the coefficients - quick and dirty
%method, to be improved later)

% %express neighbors in new local coord system
locneighborsc = (neighbors - repmat(p,nneigh,1))*[t1c t2c nc];

%switch to polar coordinates
r = sqrt(sum(locneighborsc(:,1:2).^2,2));
theta = atan2(locneighborsc(:,2), locneighborsc(:,1));
z = locneighborsc(:,3);

%compute Phicorr: TODO: implement the real iterative correction algorithm
idx = 1;
M=zeros(nneigh,ncolM);
normalized_r=r/s;
z=z/s;
for k=0:order
  rk = normalized_r.^k;
  w = exp(-normalized_r.^2/18);
  for n = -k:2:k
        M(:,idx) = rk.*exp(1i*n*theta).*w;
        idx=idx+1;
  end
end

Phicorr = linsolve(M,z.*w);


%compute real(a0) and |an| n>=1
idx = 1;
an=zeros(order+1,1);
for k=0:order
    for n = -k:2:k
      if n>=0  
       an(n+1) = an(n+1) + Phicorr(idx)/(k+2);
      end
      idx=idx+1;
    end
end

an(1)=real(an(1));%imaginary part is 0
an(2:end)=abs(an(2:end));

