clear all; close all;

% ---- swiss roll ----
if 0
    % ---- data generation ----
    N = 1000;
    r = linspace(0,1,N);
    tc = (3*pi/2)*(1+2*r);
    x = tc.*cos(tc);
    y = tc.*sin(tc);
    z = 1*rand(1,N);

    X = [x', y', z'];
    D = compute_distance_matrix(X');
    norm((D+D')/2 - D) %== 0
else
    % use pre-generated data
    load data.mat;
end
% ---- swiss roll ----


% ---- diffusion map ----
data_dissim_sqr=D.*D';
d = sort(data_dissim_sqr,2);
d = sum(d);
d = d./length(d);
t = 8;
vareps = 0.5;
data_dissim_sqr = data_dissim_sqr./vareps;

K=exp(-data_dissim_sqr);

[i,j]=find(K > 1e-8);
ii = sub2ind(size(K),i,j);
sparse_K = sparse(i,j,K(ii));
assert(norm(full(sparse_K-K)) < 1e-6);
clear i j ii;

%-------------------
nbrEigVecs = 10;
%-------------------
sparse_q = sum(sparse_K,2);
sparse_K = sparse_K./(sparse_q*sparse_q');

sparse_D = sum(sparse_K,2);
sparse_pi = sparse_D./sum(sparse_D);
sparse_D_sqrt = sqrt(sparse_D);
sparse_A = sparse_K./(sparse_D_sqrt*sparse_D_sqrt');

if nbrEigVecs < (size(sparse_A,1)-1)
    opts.disp=0;
    [sparse_V_A,sparse_E_A]=eigs(sparse_A,nbrEigVecs,'LM',opts);
else
    [sparse_V_A,sparse_E_A]=eig(full(sparse_A));
end

% .............. need sorting...........................
[sparse_V_A,sparse_E_A]=SortEigen(sparse_V_A,sparse_E_A);
% .............. need sorting............................

sparse_V_P_left=[];
for i=1:size(sparse_V_A,2)
    sparse_V_P_left(:,i) = sparse_V_A(:,i)./sqrt(sparse_pi);
end
sparse_E_P_cent = sparse_E_A;

sparse_embedding = [];
t_power = 2^t;
sparse_embedding = sparse_V_P_left*(sparse_E_P_cent^(t_power));
% ---- diffusion map ----




%------------- KMEAN --------------------
if 0
    % compute k-mean data
    % ---------------
    % the algorithm is obviously sensitive to the kmean clusters
    % need multiple runs if need be (see paper)
    % ---------------
    NbrOfClusters = 150;
    repeat = 50; done = false;
    while (repeat > 0 && ~done)
        try
            p = sparse_embedding;
            [point2cluster,clustCent] = kmeans(p, NbrOfClusters);
            done = true;
        catch e
            repeat = repeat - 1;
        end
    end
    if (~done)
        fprintf('KMean failed!\n');
        assert(false);
    end

    indCluster = zeros(size(clustCent,1),1);
    indGroupSize = zeros(size(clustCent,1),1);
    embed_tree = kdtree_build(p);
    for i=1:size(clustCent,1)
        % find an actual point, rather than a mean as center
        indCluster(i) = kdtree_k_nearest_neighbors(embed_tree, clustCent(i,:),1 );
        indGroupSize(i) = length(find(point2cluster == i));
    end
    kdtree_delete(embed_tree); clear embed_tree;
    indA = indCluster';
else
    % load computed kmeans data
    load 'kmeans.mat'
end
%------------- KMEAN --------------------


%------------- Sim NYSTROM --------------
indB = setdiff(1:1000,indA);
sample_K_A = full(sparse_K(indA,indA));
sample_K_B = full(sparse_K(indA,indB));
%---
sample_q = sum(sample_K_A,2);
sample_K_A = sample_K_A./(sample_q*sample_q');
%---
sample_d = sum(sample_K_A,2);
sample_pi = sample_d./sum(sample_d);
sample_dsqrt = sqrt(sample_d);
sample_K_A = sample_K_A./(sample_dsqrt*sample_dsqrt');

[sample_V,sample_E] = eig(sample_K_A);
[sample_V,sample_E] = SortEigen(sample_V,sample_E);

sample_V_left=[];
for i=1:size(sample_V,2)
    sample_V_left(:,i) = sample_V(:,i)./sqrt(sample_pi);
end

sample_Eij = [sample_K_A; sample_K_B'];
s_orig_extended_sample_V_left = sample_Eij * sample_V_left * pinv(sample_E);

s_orig_sample_embedding = s_orig_extended_sample_V_left*(sample_E^(2^t));

ind = [indA indB];
s_orig_sample_embedding(ind,:)=s_orig_sample_embedding(1:length(ind),:);

h=[0,0];
figure;
hold on;
h(1)=plot3(sparse_embedding(:,2),sparse_embedding(:,3),sparse_embedding(:,4), 'b.');
h(2)=plot3(s_orig_sample_embedding(:,2),s_orig_sample_embedding(:,3),s_orig_sample_embedding(:,4), 'ro');
legend(h,{'Eigensolver','Nystrom'});
hold off;
axis([-5,5,-5,5,-5,5]);
fprintf('Simple Nystrom Approx Error:  %.10g\n', ...
    norm(sparse_embedding(:,2:4)-s_orig_sample_embedding(:,2:4)) );
%------------- Sim NYSTROM --------------


%------------- Weighted Nyst ------------
wK_A = full(sparse_K(indA,indA));
wK_B = full(sparse_K(indA,indB));
n = size(wK_A,2);
m = size(wK_B,2);

d1  = sum( [wK_A; wK_B'], 1 );
d2  = sum( wK_B,1 )  + sum( wK_B',1 ) * pinv(wK_A) * wK_B;
d_hat  = 1./[d1 d2]';
wK_A = wK_A.*(d_hat(1:n)*d_hat(1:n)');
wK_B = wK_B.*(d_hat(1:n)*d_hat(n+(1:m))');

Sj = diag(indGroupSize);
Eij = [wK_A; wK_B']*Sj;

dd_wKAz = sum(wK_A*Sj,2);
dd_wKAz_si = diag(1./sqrt(dd_wKAz));

dd_wKAx = sum(Eij,2);
dd_wKAx_si = diag(1./sqrt(dd_wKAx));

wK_A1 = wK_A*Sj;
wK_A2 = dd_wKAz_si*wK_A1*dd_wKAz_si;
[Vd_A Ed_A] = eig(wK_A2);
[Vd_A,Ed_A] = SortEigen(Vd_A,Ed_A);

Vd_A_ex = dd_wKAx_si*Eij*dd_wKAz_si*Vd_A*pinv(Ed_A);

for i=1:size(Vd_A_ex,2)
Vd_A_ex(:,i) = Vd_A_ex(:,i)./norm(Vd_A_ex(:,i));
end

dd_wpi = dd_wKAx ./ sum(dd_wKAx);

Vd_A_left = [];
for i=1:size(Vd_A_ex,2)
Vd_A_left(:,i) = Vd_A_ex(:,i)./sqrt(dd_wpi);
end


Vd_A_embedding = Vd_A_left*(Ed_A^(2^t));

ind = [indA indB];
Vd_A_embedding(ind,:) = Vd_A_embedding(1:length(ind),:);

figure;
h=[0,0];
hold on;
h(1)=plot3(sparse_embedding(:,2),sparse_embedding(:,3),sparse_embedding(:,4), 'b.');
h(2)=plot3(Vd_A_embedding(:,2),Vd_A_embedding(:,3),Vd_A_embedding(:,4),'ro');
hold off;
legend(h,{'Eigensolver','Density-Weighted Nystrom'});
axis([-5,5,-5,5,-5,5]);
fprintf('Density-Weighted Nystrom Approx Error:  %.10g\n', ...
    norm(sparse_embedding(:,2:4)-Vd_A_embedding(:,2:4)) );
%------------- Weighted Nyst ------------


%------------- Alignment ----------------
% ---- eigenmode and signflip ----
tt_embedding=[];
tt_column=[];
tt_sign = [];
UsedIntSet=java.util.HashSet;
UsedIntSet.clear();
CorrespondanceDimension = 4;
EigenSearchSpace = CorrespondanceDimension * 3;
for i=1:CorrespondanceDimension
    largest = 0;
    signchange = 1;
    set = 0;
    for j=1:EigenSearchSpace
        if (~UsedIntSet.contains(j))
            col_vector1 = Vd_A_embedding(:,j);
            col_vector2 = sparse_embedding(:,i);
            dotproduct = (col_vector1/norm(col_vector1))'*(col_vector2/norm(col_vector2));
            if (~isnan(dotproduct) && abs(dotproduct) > largest)
                largest = abs(dotproduct);
                if (dotproduct > 0)
                    signchange = 1;
                else
                    signchange = -1;
                end
                set = j;
            end
        end
    end

    if (set == 0)
        %fprintf('Not enough vectors of extended_V to search\n');
        avg_dotProduct_score = 0;
        score1 = 1e20;
        tt_embedding = [];
        tt_column = [];
        tt_sign = [];
        return;
    end
    tt_embedding(:,i)=signchange*Vd_A_embedding(:,set);		% the whole columns are shuffled
    tt_column(i)=set;
    tt_sign(i) = signchange;
    UsedIntSet.add(set);
end
clear i largest signchange dotproduct set UsedIntSet;
%------------- Alignment ----------------
figure;
h=[0,0];
hold on;
h(1)=plot3(sparse_embedding(:,2),sparse_embedding(:,3),sparse_embedding(:,4), 'b.');
h(2)=plot3(tt_embedding(:,2),tt_embedding(:,3),tt_embedding(:,4), 'ro');
hold off;
legend(h,{'Eigensolver','Sign change and Alignment'});
axis([-5,5,-5,5,-5,5]);
fprintf('Density-Weighted Nystrom + SignChange & Alignment, Approx Error: %.10g\n', ...
    norm(sparse_embedding(:,2:4)-tt_embedding(:,2:4)));



%------------- TPS Nyst ------------
st = tps_func(tt_embedding(indA,2:4), sparse_embedding(indA,2:4));
tps_embedding = zeros(size(tt_embedding(:,2:4)));
for i = 1:length(tt_embedding)
    tps_embedding(i,:) = tps_inplt(st, tt_embedding(i,2:4));
end
figure;
h=[0,0,0,0];
hold on;
h(1)=plot3(sparse_embedding(:,2),sparse_embedding(:,3),sparse_embedding(:,4), 'b.');
h(2)=plot3(tt_embedding(:,2),tt_embedding(:,3),tt_embedding(:,4), 'ro');
h(3)=plot3(tps_embedding(:,1),tps_embedding(:,2),tps_embedding(:,3), 'go');
for i=1:length(indA)
    j = indA(i);
    lines = [tt_embedding(j,2),tt_embedding(j,3),tt_embedding(j,4);
             tps_embedding(j,1)  tps_embedding(j,2) tps_embedding(j,3)];
    h(4)=plot3( lines(:,1), lines(:,2), lines(:,3), ...
        'k-', 'LineWidth', 0.5);
end
hold off;
legend(h,{'Eigensolver','Sign change and Alignment','TPS after alignment','Landmarks Correspondence'});
axis([-5,5,-5,5,-5,5]);
fprintf('TPS Nystrom Approx Error: %.10g\n', ...
    norm(sparse_embedding(:,2:4)-tps_embedding(:,1:3)));
%------------- TPS Nyst ------------

