%% Covariance domain modelling
%==============================================================================
%
%	Written by Krishnaprasad N R
%	E-mail: k.namburramamohan@tudelft.nl
%	
%   References
%   [1] Krishnaprasad Nambur Ramamohan, Sundeep Prabhakar Chepuri, Daniel
%   Fernandez Comesana and Geert Leus,
%   "Blind Calibration Sparse Arrays for DOA estimation with Analog and One-Bit measurements", ICASSP 2019
%
%==============================================================================
clc;
clear all; close all;

%% Script Settings
sensType   = 'APS';         % APS or AVS array
covarType  = logical(1);    % 1 - Ideal  or 0 - Sample covariance matirix
noiseType  = logical(1);    % 1 - With   or 0 - without Noise
noiseVar   = logical(1);    % 1 - Known  or 0 - Unknown noise variance
arrayType  = logical(1);    % 1 - Sparse or 0 - Uniform linear array

%% Settings for generation of Measurement data of Linear array

SNR         = 0;              % [dB]
N           = 500;            % # of time snapshots
D           = 4;              % # of sources    
R           = 2;              % # of references
d           = 0.5;            % # Spacing between Sensors in terms of wavelength
theta       = 1:0.1:180;      % Predefined grid
grid_size   = length(theta);  % Length of the grid

if arrayType     
    %pos = [ 0, 1, 2, 3, 6, 10];
    pos = [ 0, 1, 2, 3, 6, 9];
    
    M = pos(end)+1;             % # of Possible Sensors
    selMat = eye(M);            % Selection matrix for creating sparse arrays

    K  = length( pos );
    Ga = selMat( pos + 1, :);    
else % ULA
   
   M = 10;
   K = M;
   Ga = selMat;
   
end

% Define target DOAs
doaRange  =  1/sqrt(2);   % in cosine space
theta_tar = acosd( linspace( doaRange, -doaRange, D) );
theta_tar = round(theta_tar(1:D));

%% Generation of measurement data without gain and phase uncertainitites

% Dictionary steering matrix
Agrid = gen_a( M, d, theta, sensType);

% True steering matrix
Amat = gen_a( M, d, theta_tar, sensType );

% Generation of the source signals
sig2_s = 1; % Source Power
Source = sqrt(sig2_s / 2).*( randn( D, N ) + 1j* randn( D, N ) );

% Based on type of Sensors select number of channels
if sensType == 'AVS'

    M = 3 * M;
    K = 3 * K;
    Ga = kron(eye(3), Ga);
    
end
Agrid = Ga * Agrid;
Amat  = Ga * Amat;

% Generation of the noise signals
sig2_n = sig2_s / (10^(SNR / 10)); % Noise Power
Ns = sqrt(sig2_n / 2) .* (randn(K,N) + 1j* randn(K,N));

% Measurement data
if noiseType
    
    Xmat = Amat * Source + Ns;
 
else
    
    Xmat = Amat * Source ;
    
end

%% Generation of measurement data with gain and phase uncertainities

% Amplitude errors

vec_amp = [ ones( R, 1);  ( ones( K - R, 1 ) -  0.6 * randn( K - R, 1 ) )];
vec_amp = abs(vec_amp);
Gerr = diag( vec_amp );

% Phase errors

vec_pha = [ zeros( R, 1); 0.8*randn( K - R, 1)];
Perr = diag ( exp( 1j * vec_pha ) );

% Calibration error matrix
Dcal = Gerr * Perr;

% Generation of measurement data with calibration errors
Ymat =  Dcal * Xmat;


%% Evaluation of Covariance Matrix
if covarType 
    
    Sp = diag(ones(D,1)); % Sourece Powers
    
    if noiseType 
    
        Rx = Amat * Sp * Amat' + sig2_n*eye(K);
        Ry = Dcal * (Amat * Sp * Amat' + sig2_n * eye(K)) *  Dcal';
        
    else
        
        Rx = Amat * Sp * Amat';
        Ry = Dcal * (Amat * Sp * Amat') * Dcal';
        
    end
else % Sample Covariance
    
    Rx = (1 / N) * (Xmat * Xmat'); 
    Ry = (1 / N) * (Ymat * Ymat'); 
    
end

%% Joint Estimation of the Calibration parameters and DOAs

% Difference co-array manifold matrix 
ACo = kr( conj( Agrid ), Agrid );

% Vectroize the measurment covariance matrix with Calibration terms
ry = Ry(:);

% Generation of mapping matrix as described in Eqn(21)
if noiseVar
  
    GCo = [ diag( ry ) -ACo ];
    rn = sig2_n * vec( eye( K ) ); % Vec Noise covariance matrix

else
    
    GCo = [ diag( ry ) -ACo -vec( eye( K ) ) ];
    
end

% Defining the optimization problem in CVX

Mco = K^2;
cvx_solver mosek
cvx_begin sdp
  if noiseVar
    variable wCo( Mco + grid_size ) complex
  else
    variable wCo( Mco + grid_size + 1 ) complex
  end
  B = reshape( wCo( 1:Mco ), [ K, K ] );
  b = wCo( 1:K );
  B = (1/2) * ( B + B' );
  if noiseVar
      if noiseType 
        minimize ( norm ( rn - GCo * wCo )) 
      else 
        minimize ( norm ( GCo * wCo ) )
      end
  else 
    minimize ( norm ( GCo * wCo ) )
  end
  subject to
    norm( wCo( (Mco+1):(Mco+grid_size) ), 1 ) <= D;
    [B  b; b' 1] >= 0;
    B( 1:R, 1:R ) == ones( R, R );       
cvx_end

% Check if problem was successfully solved
disp(['Problem is ' cvx_status])
if ~strfind(cvx_status,'Solved')
  return
end

[U,S,~] = svd(B);
bcor = U(:,1);
bcor = bcor./bcor(1);

%% Plotting results
fSize = 24;
figure('position',[-200, -200, 1000, 1000]),
stem( abs( wCo ) )
xlim([1 length(wCo)])
ylabel('$|\hat\gamma|$','FontSize',fSize,'Interpreter','Latex');
xlabel('Index','FontSize',fSize,'Interpreter','Latex');
set(gca,'FontSize',fSize,'FontName','Times New Roman');

fSize = 24;
figure('position',[-200, -200, 1000, 1000]),
stem(1./(abs(wCo(1:K))))
hold on
stem(1./(abs(bcor)))
plot(vec_amp,'*')
legend('JE-Covar Modeling -Opt', 'JE-Covar Modeling -SVD','Expected','FontSize',16,'Interpreter','Latex')
xlabel('Channel index','FontSize',fSize,'Interpreter','Latex')
ylabel('$\hat\psi$','FontSize',fSize,'Interpreter','Latex');
set(gca,'FontSize',fSize,'FontName','Times New Roman');

fSize = 24;
figure('position',[-200, -200, 1000, 1000]),
stem(-angle(wCo(1:K)))
hold on
stem(-angle(bcor))
plot(vec_pha,'*')
legend('JE-Covar Modeling -Opt', 'JE-Covar Modeling -SVD','Expected','FontSize',16,'Interpreter','Latex')
xlabel('Channel index','FontSize',fSize,'Interpreter','Latex')
ylabel('$\hat\phi (^{\circ})$','FontSize',fSize,'Interpreter','Latex');
set(gca,'FontSize',fSize,'FontName','Times New Roman');

aux = logical(sum(logical(theta' == theta_tar),2));
fSize = 24;
figure('position',[-200, -200, 1000, 1000]),
stem(theta, abs( wCo( ( K^2+1 ) : ( K^2+grid_size ) ) ) )
hold on 
plot(theta_tar,0.5*max(abs(wCo(K^2+1:(K^2+grid_size)))),'*')
xlim([1 180]);
% legend('JE-Covar Modeling -Opt', 'JE-Covar Modeling -SVD','Expected','FontSize','16','Interpreter','Latex')
xlabel('$\hat\theta$','FontSize',fSize,'Interpreter','Latex')
ylabel('DOA Spectra','FontSize',fSize,'Interpreter','Latex');
set(gca,'FontSize',fSize,'FontName','Times New Roman');

%% Compensating Gain and Phase errors

gain_est = diag( ( abs( bcor ) ) );
phase_est = diag( exp( 1j*angle( bcor ) ) );

DcalEst = gain_est * phase_est;

RyM = DcalEst * Ry * DcalEst';

%% Spatial Smoothing Music 
thetaFine = 1:0.05:180;
if arrayType  
    P_U  = ssmusic(RyM, M, d, D, thetaFine, sensType, pos);
    P_U1 = ssmusic(Ry,  M, d, D, thetaFine, sensType, pos);
    P_U2 = ssmusic(Rx,  M, d, D, thetaFine, sensType, pos);

    
    % Figure for plotting
    fSize = 16;
    figure('position',[-200, -200, 1000, 400]),
    semilogy(thetaFine, abs(P_U./max(P_U)),'b','LineWidth',2)
    hold on
    semilogy(thetaFine, (abs(P_U1./max(P_U1))),'k','LineWidth',1,'LineStyle', '-.')
    semilogy(thetaFine, (abs(P_U2./max(P_U2))),'m','LineWidth',1,'LineStyle', '--')
    xlim([15 165])
    xlabel('$\theta$', 'Interpreter', 'latex')
    ylabel('Normalized MUSIC Spectrum', 'Interpreter', 'latex')
    legend({'Post Sens Err Est', 'With Sens Err','Without Sens Err'},'FontSize',16,'Interpreter', 'latex','Location', 'southeast')
    set(gca,'XGrid','on','GridAlpha',1,'FontSize',fSize,'FontName','Times New Roman');

end

%% Coventional Music and MVDR on the Uniform Linear array
if ~arrayType  
    % MVDR spectrum evaluation
%     mvdrWSensErrs = mvdrSpectrum(Ry,  M, d, thetaFine, sensType);
%     mvdrWoSensErrs = mvdrSpectrum(Rx,  M, d, thetaFine, sensType);
%     mvdrAfCorrSensErrs = mvdrSpectrum(RyM, M, d, thetaFine, sensType);

    % MUSIC spectrum evaluation
    musicWSensErrs      = musicSpectrum(Ry,  M, d, D, thetaFine, sensType);
    musicWoSensErrs     = musicSpectrum(Rx,  M, d, D, thetaFine, sensType);
    musicAfCorrSensErrs = musicSpectrum(RyM, M, d, D, thetaFine, sensType);

    % Plot Music and MVDR spectrum
%     fSize = 24;
%     figure('position',[-200, -200, 1000, 1000]),
%     plot(thetaFine,abs(musicAfCorrSensErrs)./max(abs(musicAfCorrSensErrs)))    
%     hold on 
%     plot(thetaFine,abs(musicWSensErrs)./max(abs(musicWSensErrs)))
%     plot(thetaFine,abs(musicWoSensErrs)./max(abs(musicWoSensErrs)))    
%     stem(theta_tar,ones(1,D),'k')
%     xlim([1 180])
%     legend({'Post Est of Sensor Errors','With Sensor Errors','With No Sensor Errors','Location of DOAs'},'Interpreter','Latex','FontSize',16);
%     xlabel('$\hat\theta$','FontSize',fSize,'Interpreter','Latex');
%     ylabel('Music Spectra','FontSize',fSize,'Interpreter','Latex');
%     set(gca,'FontSize',fSize,'FontName','Times New Roman');
    
    
    fSize = 16;
    figure('position',[-200, -200, 1000, 400]),
    semilogy(thetaFine, abs(musicAfCorrSensErrs)./max(abs(musicAfCorrSensErrs)),'b','LineWidth',2)
    hold on
    semilogy(thetaFine, abs(musicWSensErrs)./max(abs(musicWSensErrs)),'k','LineWidth',1,'LineStyle', '-.')
    semilogy(thetaFine, abs(musicWoSensErrs)./max(abs(musicWoSensErrs)),'m','LineWidth',1,'LineStyle', '--')
    xlim([30 150])
    xticks(theta_tar)
    xlabel('$\theta$', 'Interpreter', 'latex')
    ylabel('Normalized MUSIC Spectrum', 'Interpreter', 'latex')
    legend({'Post Sens Err Est', 'With Sens Err','Without Sens Err','Location of DOAs'},'FontSize',16,'Interpreter', 'latex','Location', 'southeast')
    set(gca,'XGrid','on','GridAlpha',1,'FontSize',fSize,'FontName','Times New Roman');
end

%% Function Defn

function ssm = ssmusic(Ry, M, d, D, theta, type, pos)
    [~, x_U2] = sample_covariance_to_difference_coarray( pos, Ry, 'x_U' );
    R_tilde2 = toeplitz(x_U2((length(x_U2)+1)/2:end), x_U2((length(x_U2)+1)/2:-1:1));
    R_tilde2 = 1/2*(R_tilde2 + R_tilde2');
    ssm = musicSpectrum(R_tilde2, M, d, D, theta, type);
end

function mvdr = mvdrSpectrum(Ry, M, d, theta, type)
              
    grid_size = length(theta);    
    if type == 'AVS'
        M = M/3;
    end
    
    % Dictionary steering matrix
    Agrid = gen_a(M,d,theta,type);
    invRy = pinv(Ry);
    for i = 1:grid_size
        mvdr(i) = 1./(Agrid(:,i)' * invRy * Agrid(:,i));
    end
end


function music = musicSpectrum(Ry, M, d, D, theta, type)
    
    grid_size = length(theta);    
    if type == 'AVS'
        M = M/3;
    end
    
    % Dictionary steering matrix
    Agrid = gen_a(M,d,theta,type);
    [uy,~,~] = svd(Ry);
    Uy = uy(:,D+1:end);
    for i = 1:grid_size
        music(i) = 1./(Agrid(:,i)' * (Uy*Uy') * Agrid(:,i));
    end

end


function varargout = sample_covariance_to_difference_coarray( S, R_S_tilde, output_type )
%SAMPLE_COVARIANCE_TO_DIFFERENCE_COARRAY:
%   Convert the sample covariance matrix R_S_tilde to quantities on the
%   difference coarray.
%==============================================================================
%
%	Written by Chun-Lin Liu
%	E-mail: cl.liu@caltech.edu
%	Project website:
%	http://systems.caltech.edu/dsp/students/clliu/coarray.html
%	Last revised on July 13, 2017
%
%   If you use this function, please cite reference [1] or [2].
%
%   Input:
%
%       1) S: The discrete integer set for sensor locations
%       2) R_S_tilde: The sample covariance matrix defined on S
%       3) output_type: The quantities on the difference coarray [2].
%           This parameter is case-insensitive.
%           Please see [2] for the definitions of D, U, and V.
%
%           a) 'x_D': The sample autocorrelation vector on the difference
%           coarray D.
%           b) 'x_U': The sample autocorrelation vector on the ULA segment
%           of the difference coarray (set U).
%           c) 'x_V': The sample autocorrelation vector on the set V. Here
%           the coarray interpolation method is nuclear norm minimization.
%               Note: cvx package (http://cvxr.com/cvx/) is required here.
%           d) 'R_D': Augmented covariance matrix based on x_D.
%           e) 'R_U': Augmented covariance matrix based on x_U.
%           f) 'R_V': Augmented covariance matrix based on x_V.
%               Note: cvx package (http://cvxr.com/cvx/) is required here.
%
%   Output:
%
%       1) If there is only one output parameter, then it means either
%           x_D, x_U, x_V, R_D, R_U, or R_V.
%       2) If there are two output variables,
%           the first argument is either D, U, or V;
%           the second argument is either x_D, x_U, x_V, R_D, R_U, or R_V.
%
%   References
%   [1] C.-L. Liu and P. P. Vaidyanathan, “Remarks on the Spatial Smoothing Step in Coarray MUSIC,�? IEEE Signal Processing Letters, vol. 22, no. 9, pp. 1438-1442, Sep. 2015. 
%   [2] C.-L. Liu, P. P. Vaidyanathan and P. Pal, “Coprime Coarray Interpolation for DOA Estimation via Nuclear Norm Minimization,�? in Proc. of 2016 IEEE International Symposium on Circuits and Systems (ISCAS 2016), pp. 2639-2642, Montreal, Canada, May 2016. 
%
%==============================================================================

switch (upper(output_type))
    case 'X_D'
        
        [n1, n2] = ndgrid(S);
        [D, ~] = weight_function( S, 'D' );
        LEN_D = length(D);
        
        x_D = zeros( LEN_D, 1 );
        for mm = 1 : LEN_D
            % Select samples with the same lag
            data = R_S_tilde( n1 - n2 == D(mm) );
            % ====================
            % Average over all the samples with the same lag.
            % Definition 3 in [1]
            x_D(mm) = mean( data );
            % ====================
        end
        output_1 = D;
        output_2 = x_D;
        
    case 'X_U'
        
        [n1, n2] = ndgrid(S);
        [U, ~] = weight_function( S, 'U' );
        LEN_U = length(U);
        
        x_U = zeros( LEN_U, 1 );
        for mm = 1 : LEN_U
            % Select samples with the same lag
            data = R_S_tilde( n1 - n2 == U(mm) );
            % ====================
            % Average over all the samples with the same lag.
            % Definition 3 in [1]
            x_U(mm) = mean( data );
            % ====================
        end
        output_1 = U;
        output_2 = x_U;
        
    case 'X_V'
        
        % Nuclear norm minimization
        [V, R_V] = sample_covariance_to_difference_coarray( S, R_S_tilde, 'R_V' );
        
        output_1 = V;
        output_2 = [R_V(1, end:-1:2).'; R_V(:, 1)];
        
    case 'R_D'
        
        [D, x_D] = sample_covariance_to_difference_coarray( S, R_S_tilde, 'x_D' );
        LEN_D = length(D);
        
        output_1 = D;
        output_2 = toeplitz(x_D((LEN_D+1)/2:end), x_D((LEN_D+1)/2:-1:1));
        
    case 'R_U'
        
        [U, x_U] = sample_covariance_to_difference_coarray( S, R_S_tilde, 'x_U' );
        LEN_U = length(U);
        
        output_1 = U;
        output_2 = toeplitz(x_U((LEN_U+1)/2:end), x_U((LEN_U+1)/2:-1:1));
        
    case 'R_V'
        
        % Sample covariance matrix on the difference coarray
        [D, x_D] = sample_covariance_to_difference_coarray( S, R_S_tilde, 'x_D' );
        
        % Nuclear norm minimization
        [V, ~] = weight_function( S, 'V' );
        LEN_V = length(V);
        Vp = V((1+LEN_V)/2:end);
        LEN_Vp = length(Vp);
        [~, index_S4p_pos, ~] = intersect( Vp, D(D >= 0) );
        cvx_begin % quiet
            variable R_V(LEN_Vp, LEN_Vp) hermitian toeplitz
            minimize( norm_nuc(R_V) )
            subject to
                R_V(index_S4p_pos, 1) == x_D(D >= 0);
        cvx_end
        
        output_1 = V;
        if ( strcmp(cvx_status, 'Infeasible') )
            output_2 = zeros(size(R_V));
        else
            output_2 = R_V;
        end
        
    otherwise
        error('output_type is invalid!')
end


if (nargout == 1)
    varargout{1} = output_2;
elseif (nargout == 2)
    varargout{1} = output_1;
    varargout{2} = output_2;
else
    error('Output arguments are not supported!')
end

end


function a = gen_a(M,Delta,theta, type)

Delta = (0:M-1)'*Delta;
theta(1,:) = theta;
a = exp(1i*2*pi*Delta*cosd(theta)) ;


if type =='AVS'   
% AVS 
h = [ones(1,length(theta)) ; cosd(theta) ;sind(theta)];
a = kr(h,a);
end

end


function [D, w] = weight_function( S, coarray_set )
%WEIGHT_FUNCTION: Returns the weight function of sensor array geometry S
%==============================================================================
%
%	Written by Chun-Lin Liu
%	E-mail: cl.liu@caltech.edu
%	Project website: 
%	http://systems.caltech.edu/dsp/students/clliu/coarray.html
%	Last revised on July 13, 2017
%
%   If you use this function, please cite related references [1-6].
%
%   The definition of weight function w(m) can be found in [1-6].
%   Input:
%       1) S: The discrete integer set
%       2) coarray_set: specify the output grid D. It has three options
%           'D': w is defined over D
%           'U': defined over U, the central maximum contiguous ULA portion
%           'V': defined over V, including the holes
%           The definitions of D, U, and V can be found in [6].
%
%   Output:
%       1) D: Output grid
%       2) w: weight function on D
%
%   Examples:
%       [D, w] = weight_function([0; 1; 4], 'D');
%       We obtain
%           D = [-4; -3; -1; 0; 1; 3; 4];
%           w = [ 1;  1;  1; 3; 1; 1; 1];
%
%       [D, w] = weight_function([0; 1; 4], 'U');
%       We obtain
%           D = [-1; 0; 1];
%           w = [ 1; 3; 1];
%
%       [D, w] = weight_function([0; 1; 4], 'V');
%       We obtain
%           D = [-4; -3; -2; -1; 0; 1; 2; 3; 4];
%           w = [ 1;  1;  0;  1; 3; 1; 0; 1; 1];
%
%	References
%	Journal
%   [1] C.-L. Liu and P. P. Vaidyanathan, “Remarks on the Spatial Smoothing Step in Coarray MUSIC,” IEEE Signal Processing Letters, vol. 22, no. 9, pp. 1438-1442, Sep. 2015.
%	[2] C.-L. Liu and P. P. Vaidyanathan, “Super Nested Arrays: Linear Sparse Arrays with Reduced Mutual Coupling - Part I: Fundamentals,” IEEE Transactions on Signal Processing, vol. 64, no. 15, pp. 3997-4012, Aug. 2016. 
%	[3] C.-L. Liu and P. P. Vaidyanathan, “Super Nested Arrays: Linear Sparse Arrays with Reduced Mutual Coupling - Part II: High-Order Extensions,” IEEE Transactions on Signal Processing, vol. 64, no. 16, pp. 4203-4217, Aug. 2016. 
%	Conference
%	[4] C.-L. Liu and P. P. Vaidyanathan, “Super Nested Arrays: Sparse Arrays with Less Mutual Coupling than Nested Arrays,” in Proc. of 2016 IEEE International Conference on Acoustics Speech and Signal Processing (ICASSP 2016), pp. 2976-2980, Shanghai, China, Mar. 2016.
%	[5] C.-L. Liu and P. P. Vaidyanathan, “High Order Super Nested Arrays,” in Proc. of the Ninth IEEE Sensor Array and Multichannel Signal Processing Workshop (SAM 2016), Rio de Janeiro, Brazil, Jul. 2016.
%	[6] C.-L. Liu, P. P. Vaidyanathan and P. Pal, “Coprime Coarray Interpolation for DOA Estimation via Nuclear Norm Minimization,” in Proc. of 2016 IEEE International Symposium on Circuits and Systems (ISCAS 2016), pp. 2639-2642, Montreal, Canada, May 2016. 
% 
%==============================================================================

    % S is a vector
    if ( ~isvector(S) ||  ~isreal(S) )
        error('The parameter S has to be a a real vector')
    end
    
    S = unique(S);
    [n1, n2] = ndgrid(S);
    n1_n2_mat = n1 - n2;
    n1_n2_vec = n1_n2_mat(:);
    
    D_set = unique(n1_n2_vec, 'sorted');
    w_set = histc(n1_n2_vec, D_set);
    
    D_MAX = max(D_set); % The maximum element in the difference coarray
    
    
    switch(coarray_set)
        case 'D'
            D = D_set;
            w = w_set;
        case 'U'
            % Locate the maximum contiguous set
            D_full = (-D_MAX : D_MAX).';
            H = setdiff(D_full, D_set); % Holes
            if (isempty(H))
                % D is hole-free
                D = D_set; w = w_set;
            else
                N_max = min(abs(H)) - 1;
                D = (-N_max : N_max).';
                w = w_set(D_set >= -N_max & D_set <= N_max);
            end
        case 'V'
            % Locate the central maximum contiguous set
            D_full = (-D_MAX : D_MAX).';
            [~, ia, ~] = intersect(D_full, D_set);
            if (length(ia) == length(D_full))
                % D is hole-free
                D = D_set; w = w_set;
            else
                D = D_full;
                w = zeros(size(D));
                w(ia) = w_set;
            end
        otherwise
            error('The parameter coarray_set is either ''D'', ''U'', or ''V''')
    end


end

