clear all
close all

nr_sources = 3;
nr_mixtures = 2;
wname = 'sym8';


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%   jednostavno mixanje
%%%%%%%%%%%%%%%%%

SS{1} = double(rgb2gray(imread('4.2.07.tiff')));
SS{2} = double(rgb2gray(imread('4.2.03.tiff')));
SS{3} = double(rgb2gray(imread('4.2.04.tiff')));
% remove mean
SS{1} = SS{1} - mean2(SS{1});
SS{2} = SS{2} - mean2(SS{2});
SS{3} = SS{3} - mean2(SS{3});
% resize image
SS{1} = imresize(SS{1}, 0.3, 'bicubic');
SS{2} = imresize(SS{2}, 0.3, 'bicubic');
SS{3} = imresize(SS{3}, 0.3, 'bicubic');

mixsize = size(SS{1});
wp_level = 3;   % dubina mora biti razumna, 
                %tj. iz par WP koeficijenata ho?emo dovoljan broj uzoraka

AA = abs(randn(nr_mixtures, nr_sources)) % mixing matrix

XX = cell(1,nr_mixtures);
for c1 = 1:nr_mixtures,
    XX{c1} = zeros(mixsize);
    for c2 = 1:nr_sources;
        XX{c1} = XX{c1} + SS{c2}*AA(c1,c2);
    end
end

% mixtures : t -> wp
for c1 = 1 : nr_mixtures,
    X{c1} = wpdec2(XX{c1}, wp_level, wname, 'shannon');
end


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%       Adaptivni izbor rasprsenog podskupa WP 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%za svaki cvor WP-drva pojedinacno napravi se projekcija na sferu pa
%clustering. racuna se distorzija koeficijenata unutar svakog clustera od
%centroida, pa se usrednji. zadrze se koeficijenti onih nekoliko
%cvorova koji su dali najmanju srednju distorziju.
%moguc je slucaj da se zadrze predak i potomci mu, to u principu nije bas ok,
%al koga briga jer radi

N = 3;  %broj cvorova ciji koeficijenati se zadrzavaju
avg_distorsion=[];
for c1 = allnodes(X{1})',
    coeffs = cell(1,nr_mixtures);
    cfs =[];
    for c2 = 1 : nr_mixtures,
        coeffs{c2} = wpcoef(X{c2},c1)';
        coeffs_size = size(coeffs{c2});
        coeffs{c2} = reshape(coeffs{c2}, 1, coeffs_size(1)*coeffs_size(2));
        cfs(:,c2) = coeffs{c2}';
    end
    [centroids, avg_distorsion_tmp] = get_centroids(cfs, nr_sources);
    avg_distorsion = [avg_distorsion, avg_distorsion_tmp];
    
    % malo grafike; prikaz koeficijenata po cvorovima
    %figure, plot(cfs(:,1), cfs(:,2), '.')
    %hold, plot(centroids(:,1), centroids(:,2), 'r*')
    %[depth pos] = ind2depo(4,c1);
    %depth = num2str(depth);
    %pos = num2str(pos);
    %txt = [depth ', ' pos '    avg.dist = ' num2str(avg_distorsion_tmp)];
    %title(txt);
end

% zadrzavanje koeficijenata sa najmanjom srednjom distorzijom
cfs = [];
coeffs = cell(1, nr_mixtures);
[avg_distorsion, node_ind] = sort(avg_distorsion);
for c1 = 1 : nr_mixtures,
    cfs_tmp = [];
    for c2 =  1 : N, 
        coeffs{c1} = wpcoef(X{c1},node_ind(c2)-1)';
        coeffs_size = size(coeffs{c1});
        coeffs{c1} = reshape(coeffs{c1}, 1, coeffs_size(1)*coeffs_size(2));
        cfs_tmp = [cfs_tmp; coeffs{c1}'];
    end
    cfs = [cfs cfs_tmp];
end

%za sparse dio mix signala potraziti centroide;
[centroids, avg_distorsion_tmp] = get_centroids(cfs, nr_sources);

figure, plot(cfs(:,1), cfs(:,2), '.')
hold, plot(centroids(:,1), centroids(:,2), 'r*')


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%      Separacija
%%%%%%%%%%%%%%%%%%

AA
AArec = centroids'

if nr_sources < nr_mixtures,    
    invAA_JADE = jadeR(cfs', nr_sources);
    invAA_rec = pinv(AArec);
end

if nr_sources == nr_mixtures,    
    invAA_JADE = jadeR(cfs', nr_sources);
    invAA_rec = inv(AArec);
end

if nr_sources <= nr_mixtures,
    for c1 = 1:nr_sources,
        SSrec{c1} = zeros(mixsize);
        for c2 = 1:nr_mixtures;
            SSrec{c1} = SSrec{c1} + XX{c2}*invAA_JADE(c1,c2);
            %invAA_rec za geom. umjesto invAA_JADE za JADE
        end
    end
end

if nr_sources >nr_mixtures,     % JADE plache da ne mozhe
    termcoeffs = cell(1, nr_sources);
    termnodes = leaves(X{1}); %indeksi terminalnih cvorova
    quasi_clusts = cell(1, nr_sources);

    for c1 = 1:nr_sources,
        Xrec{c1} = cfs2wpt('sym8', mixsize, leaves(X{1}), 4);   %kreiraj prazno WP drvo
    end

    for c1 = termnodes',
        c1
        termcfs = [];
        for c2 = 1 : nr_mixtures,
            termcoeffs{c2} = wpcoef(X{c2}, c1);
            termcoeffs_size = size(termcoeffs{c2});  % suvisno racunati svaki put...mozda bude optimizirano...jednom...
            termcfs = [termcfs reshape(termcoeffs{c2}, 1, termcoeffs_size(1)*termcoeffs_size(2))'];
        end
       %figure, plot(termcfs(:,1), termcfs(:,2), '.');
   
       %trebam projicirane koeficijente. projiciram ih po drugi put -> %optimizacija
        space_size = size(termcfs);
        for c2 = 1 : space_size(1),
            termcfs_proj(c2,:) = double(abs(termcfs(c2,:)./norm(termcfs(c2,:), 'fro')));
        end
        [val, quasi_clust_ind] = min(distfcm(centroids, termcfs_proj));   %nadji pripadnost kvaziklasterima
        for c2 = 1:nr_sources,
            %zadrzavamo koeficijente koje izbacujemo, pripadaju izvoru c1
            %(kvaziklaster). ostalo ostaje
            row_ind = find(quasi_clust_ind==c2);
            termcoeff_for_Xrec = zeros(1, termcoeffs_size(1)*termcoeffs_size(1));
            termcoeff_for_Xrec(row_ind) = termcfs(row_ind,2);  % npr. iz mjesavine 2
            termcoeff_for_Xrec = reshape(termcoeff_for_Xrec, termcoeffs_size(2), termcoeffs_size(2));   % natrag u formu WP koeficijenta
            Xrec{c2} = write(Xrec{c2}, 'cfs', c1, termcoeff_for_Xrec);    %spremamo ga u odgv. WP drvo
            SSrec{c2} = wprec2(Xrec{c2});
        end  
    end  
end   


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                   Rezultati
%%%%%%%%%%%%%%%%%%%%%%%%

%sources
figure, imagesc(SS{1}), colormap(gray);
figure, imagesc(SS{2}), colormap(gray);
figure, imagesc(SS{3}), colormap(gray);
% mixes
figure, imagesc(XX{1}), colormap(gray);
figure, imagesc(XX{2}), colormap(gray);
%figure, imagesc(XX{3}), colormap(gray);
% recovered
figure, imagesc(SSrec{1}), colormap(gray);
figure, imagesc(SSrec{2}), colormap(gray);
figure, imagesc(SSrec{3}), colormap(gray);

