function W = jbss_sos(x, tau, cost_selection)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% JBSS using second order statistics 
% For a general description of the algorithm and its relationship with others,
% see http://mlsp.umbc.edu/jointBSS_introduction.html
% Input:    x: the mixture of dimensionality N*T*K
%
%           tau: the time-delays, e.g. [0] or [0:5] or [0,3,7,12]
%           By default, set tau to [0]
%           Always assume tau(1)=0
%
%           cost_selection: 
%           In case of 'cross', only the cross cost is minimized by using
%           sequential orthogonal procrustes
%           In case of 'whole', a gradient search algorithm is used
%           starting from the initial guess provided by sequential
%           orthogonal procrustes
%
% Output:   W: demixing matrix
%
% References:
% [1] X.-L. Li, T. Adali, & M. Anderson, "Joint Blind Source Separation by Generalized Joint Diagonalization of Cumulant Matrices," Signal Process., 2011, 91, 2314-2322
% [2] X.-L. Li, M. Anderson, & T. Adali, "Second and Higher-Order Correlation Analysis of Multiset Multidimensional Variables by Joint Diagonalization," Lecture Notes in Computer Science: Independent Component Analysis and Blind Signal Separation, Latent Variable Analysis and Signal Separation, Springer Berlin / Heidelberg, 2010, 6365, 197-204
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if nargin==0
   help jbss_sos
   disp('Executing test program.')
   demo_sos
   return
end

[N,T,K] = size(x);
if nargin < 2
    tau = [0];
    cost_selection = 'cross';
end
if nargin < 3
    cost_selection = 'cross';
end
tau(1) = 0;
tau = abs(tau);
L = length(tau);

% first remove mean
for k = 1 : K
    for n = 1 : N
        x(n,:,k) = x(n,:,k) - mean(x(n,:,k));
    end
end

% second whitening the mixtures
P = zeros(N,N,K);
for k = 1 : K
    R = x(:,:,k)*x(:,:,k)'/T;
    P(:,:,k) = inv(sqrtm(R));
    x(:,:,k) = P(:,:,k)*x(:,:,k);
end

% third construct the correlation matrix at L time-delays 
C = zeros(N,N,K,K,L);
for k1 = 1 : K
    for k2 = 1 : K
        for el = 1 : L
            C(:,:,k1,k2,el) = x(:,1:T-tau(el),k1)*x(:,tau(el)+1:T,k2)'/T;
        end
    end
end

% joint diagonalization
W0 = zeros(N,N,K);
for k = 1 : K
    W0(:,:,k) = eye(N);
end

switch cost_selection
    
    case 'whole'
        W0 = diag_cross( C, W0, 1e-8 );
        W = diag_whole( C, W0);
    
    otherwise
        if L == 1
            W = sym_diag_cross( C, W0, 1e-8 );
        else
            W = diag_cross( C, W0, 1e-8 );
        end
        
end

% recovery the whitening 
for k = 1 : K
    W(:,:,k) = W(:,:,k) * P(:,:,k);
end

return

function demo_sos

N = 10;  % number of sources
K = 10;  % number of groups
T = 1000;   % sample size

% generate the mixtures
J = sqrt(-1);
S = zeros(N,T,K);
for n = 1 : N
    
    temp1 = randn(K,T)+J*randn(K,T);  
    temp  = zeros(K,T);
    B = randn(K,K,3)+J*randn(K,K,3);
    for p = 0 : 2
        for t = 3 : T
            temp(:,t) = temp(:,t) + B(:,:,p+1)*temp1(:,t-p); % introduce nonwhiteness and spatial correlation
        end
    end

    for k = 1 : K
        S(n,:,k) = temp(k,:);
        S(n,:,k) = S(n,:,k) - mean(S(n,:,k));
        S(n,:,k) = S(n,:,k)/sqrt(var(S(n,:,k)));
    end
end
A = randn(N,N,K)+J*randn(N,N,K);
x = zeros(N,T,K);
for k = 1 : K
    x(:,:,k) = A(:,:,k)*S(:,:,k);
end


% separation
W = jbss_sos(x);
%W = jbss_sos(x,[0:2]); 
%W = jbss_sos(x,[0:2],'whole');
% show results
figure;
T1 = zeros(N,N);
for k = 1 : K
    Tk = W(:,:,k)*A(:,:,k);
    Tk = abs(Tk);
    for n = 1 : N
        Tk(n,:) = Tk(n,:)/max(abs(Tk(n,:)));
    end
    T1 = T1 + Tk/K;
end
P=zeros(N); 
[~,imax]=max(T1);
ind=sub2ind([N N],1:N,imax);
P(ind)=1;
T1=P*T1;

imagesc(1:N,1:N,T1)
colormap('bone')
title('the joint global matrix')
colorbar
disp('Ideally image is identity matrix.')
return

function W = diag_cross( T, W0, Tol )
% joint diagonalization algorithm minimizing the cross cost by sequential
% orthogonal procrustes
% input: T: target matrices of dimension N*N*K*K*L 
%        W0: initial guess
%        Tol: stopping condition
% output: W: demixing matrix of demension N*N*K  
N = size(T,1);
K = size(T,3);
L = size(T,5);
D = zeros(N,K,K,L);

if nargin < 3
    Tol = 1e-10;
else if nargin < 2
        W0 = [];
        Tol = 1e-10;
    end
end
if isempty(W0)
    W0 = zeros(N,N,K);
    for k = 1 : K
        W0(:,:,k) = randn(N);
        W0(:,:,k) = sqrtm(W0(:,:,k)*W0(:,:,k)')\W0(:,:,k);
    end
end
W = W0;

A = zeros(2*L*N*(K-1), N, K);
for p = 1 : K
    cnt = 0;
    for q = 1 : K
        if q == p
            continue;
        end
        for el = 1 : L
            A(cnt*N+1:cnt*N+N, :, p) = T(:,:,p,q,el)';
            cnt = cnt + 1;
            A(cnt*N+1:cnt*N+N, :, p) = T(:,:,q,p,el);
            cnt = cnt + 1;
        end
    end
end
            

max_iter = 1000;
cost = zeros(max_iter,1);
last_W = W;
%tol = 1;
for iter = 1 : max_iter
  
    cost(iter) = 0;
    for p = 1 : K
        Wp = W(:,:,p);
        for q = 1 : K
            if q == p
                continue;
            end
            Wqh = W(:,:,q)';
            for el = 1 : L
                D(:,p,q,el) = diag( Wp*T(:,:,p,q,el)*Wqh );
                %B = Wp*T(:,:,p,q,el)*Wqh;
                %B = B - diag(diag(B));
                %cost(iter) = cost(iter) + norm( B ,'fro')^2;
            end
        end
    end

    %A = zeros(2*L*N*(K-1), N);
    B = zeros(2*L*N*(K-1), N);
    randp = randperm(K);
    for p0 = 1 : K
        p = randp(p0);
        
        cnt = 0;
        for q = 1 : K
            if q == p
                continue;
            end
            for el = 1 : L
                %A(cnt*N+1:cnt*N+N, :) = T(:,:,p,q,el)';
                B(cnt*N+1:cnt*N+N, :) = W(:,:,q)'*diag(D(:,p,q,el)');
                cnt = cnt + 1;
                %A(cnt*N+1:cnt*N+N, :) = T(:,:,q,p,el);
                B(cnt*N+1:cnt*N+N, :) = W(:,:,q)'*diag(D(:,q,p,el));
                cnt = cnt + 1;
            end
        end
        [U,~,V] = svd(B'*A(:,:,p));
        W(:,:,p) = U*V';
        
    end
             
    tol = 0;
    for k = 1 : K
        tol = max(tol, max(1-diag(W(:,:,k)*last_W(:,:,k)')));
    end
    if tol < Tol
        break;
    else
        last_W = W;
    end
    
end
%figure;
%plot(log(cost(1:iter)))
return

function W = diag_whole( T, W0 )
% joint diagonalization algorithm minimizing the whole cost
% input: T: target matrices of dimension N*N*K*K*L 
%        W0: initial guess for W
% output: W: demixing matrix of demension N*N*K  
% N = size(T,1);
% K = size(T,3);
% L = size(T,5);

W = W0;
mu = 1/20;
W = iva_diag_kernel( T, W, mu );
return

function W = iva_diag_kernel( T, W, mu )
% the actual gradient search algorithm 
% input: T: target matrices of dimension N*N*K*K*L 
%        W: initial guess
%        mu: step size
% output: W: demixing matrix of demension N*N*K  
N = size(T,1);
K = size(T,3);
L = size(T,5);

max_iter = 1000;
cost = zeros(max_iter,1);
% grad = zeros(N,N,K);
min_cost = inf;
best_W = W;
last_W = W;
cost_increase_counter = 0;
% tol = 1;
for iter = 1 : max_iter
  
    cost(iter) = 0;
    for p = 1 : K
        for q = 1 : K
            for el = 1 : L
                B = W(:,:,p)*T(:,:,p,q,el)*W(:,:,q)';
                B = B - diag(diag(B));
                cost(iter) = cost(iter) + norm( B ,'fro')^2;
            end
        end
    end
    
    grad = zeros(N,N,K);  
    for p = 1 : K
        
            for el = 1 : L
                
                for q = 1 : K
                    B = W(:,:,p)*T(:,:,p,q,el)*W(:,:,q)';
                    B = B - diag(diag(B));
                    grad(:,:,p) = grad(:,:,p) + B*W(:,:,q)*T(:,:,p,q,el)';
                    
                    B = W(:,:,p)*T(:,:,q,p,el)'*W(:,:,q)';
                    B = B - diag(diag(B));
                    grad(:,:,p) = grad(:,:,p) + B*W(:,:,q)*T(:,:,q,p,el);
                end
            end
            grad(:,:,p) = grad(:,:,p) - W(:,:,p)*grad(:,:,p)'*W(:,:,p);
            grad(:,:,p) = grad(:,:,p)/norm(grad(:,:,p));    % normalization 

            W(:,:,p) = W(:,:,p) - mu*grad(:,:,p);
            W(:,:,p) = sqrtm(W(:,:,p)*W(:,:,p)')\W(:,:,p); %inv(sqrtm(W(:,:,p)*W(:,:,p)'))*W(:,:,p);

    end
    
    if cost(iter) < min_cost
        min_cost = cost(iter);
        best_W = last_W;
    end
        
        if cost(iter) > min_cost
            cost_increase_counter = cost_increase_counter + 1;
            if cost_increase_counter > 1
                cost_increase_counter = 0;
                mu = 0.5*mu;
                W = best_W;
                last_W = W;
                continue;
            end
        end
    
    tol = 0;
    for k = 1 : K
        tol = max(tol, max(1-diag(W(:,:,k)*last_W(:,:,k)')));
    end
    if tol < 1e-5
        break;
    else
        last_W = W;
    end
    
end
W = best_W;
% figure;
% plot(log(cost(1:iter)))
return

function W = sym_diag_cross( T, W0, Tol )
% joint diagonalization algorithm minimizing the cross cost by sequential
% orthogonal procrustes
% input: T: target matrices of dimension N*N*K*K*L 
%        W0: initial guess
%        Tol: stopping condition
% output: W: demixing matrix of demension N*N*K
%
% assume T is symmetric, i.e., T(:,:,k1,k2,:) = T(:,:,k2,k1,:)' and only
% operate on T(:,:,k1,k2,:), k2 > k1
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
N = size(T,1);
K = size(T,3);
L = size(T,5);
D = zeros(N,K,K,L);

if nargin < 3
    Tol = 1e-10;
else if nargin < 2
        W0 = [];
        Tol = 1e-10;
    end
end
if isempty(W0)
    W0 = zeros(N,N,K);
    for k = 1 : K
        W0(:,:,k) = randn(N);
        W0(:,:,k) = sqrtm(W0(:,:,k)*W0(:,:,k)')\W0(:,:,k); %inv(sqrtm(W0(:,:,k)*W0(:,:,k)'))*W0(:,:,k);
    end
end
W = W0;           

max_iter = 1000;
cost = zeros(max_iter,1);
last_W = W;
% tol = 1;
for iter = 1 : max_iter
  
    cost(iter) = 0;
    for p = 1 : K
        Wp = W(:,:,p);
        for q = p+1 : K
            Wqh = W(:,:,q)';
            for el = 1 : L
                D(:,p,q,el) = diag( Wp*T(:,:,p,q,el)*Wqh );
                %B = Wp*T(:,:,p,q,el)*Wqh;
                %B = B - diag(diag(B));
                %cost(iter) = cost(iter) + norm( B ,'fro')^2;
            end
        end
    end

    A = zeros(L*N*(K-1), N);
    B = zeros(L*N*(K-1), N);
    randp = randperm(K);
    for p0 = 1 : K
        p = randp(p0);
        
        cnt = 0;
        for q = p+1 : K
            for el = 1 : L
                A(cnt*N+1:cnt*N+N, :) = T(:,:,p,q,el)';
                B(cnt*N+1:cnt*N+N, :) = W(:,:,q)'*diag(D(:,p,q,el)');
                cnt = cnt + 1;
            end
        end
        for q = 1 : p-1
            for el = 1 : L
                A(cnt*N+1:cnt*N+N, :) = T(:,:,q,p,el);
                B(cnt*N+1:cnt*N+N, :) = W(:,:,q)'*diag(D(:,q,p,el));
                cnt = cnt + 1;
            end
        end
        [U,~,V] = svd(B'*A);
        W(:,:,p) = U*V';
        
    end
             
    tol = 0;
    for k = 1 : K
        tol = max(tol, max(1-diag(W(:,:,k)*last_W(:,:,k)')));
    end
    if tol < Tol
        break;
    else
        last_W = W;
    end
    
end
%figure;
%plot(log(cost(1:iter)))
return