function [ A, B ] = betaNMFvisualizer( X, fs, shift, beta, nb, it, A, B )
%
% Corded by D. Kitamura (d-kitamura@nii.ac.jp) on 25 Feb. 2016.
%
% This function decomposes input matrix by nonnegative matrix
% factorization (NMF) based on beta-divergence criterion and
% multiplicative update rules, where the estimated bases and 
% activations are visualized.
%
% see also
% http://d-kitamura.sakura.ne.jp/index.html
%
% [syntax]
%   [ A, B ] = betaNMFvisualizer( X, fs, shift )
%   [ A, B ] = betaNMFvisualizer( X, fs, shift, beta )
%   [ A, B ] = betaNMFvisualizer( X, fs, shift, beta, nb )
%   [ A, B ] = betaNMFvisualizer( X, fs, shift, beta, nb, it )
%   [ A, B ] = betaNMFvisualizer( X, fs, shift, beta, nb, it, A )
%   [ A, B ] = betaNMFvisualizer( X, fs, shift, beta, nb, it, A, B )
%
% [inputs]
%        X: observed nonnegative matrix (I x J)
%       fs: sampling frequency of the input signal
%    shift: window shift size of the spectrogram
%     beta: beta value of beta-divergence (default: 2)
%           (0: Itakura-Saito, 1: Generalized KL, 2: Euclidian)
%       nb: number of bases (default: ceil(J/10))
%       it: number of iteration (default: 100)
%        A: Initial basis matrix (I x nb, default: max(rand(I,nb),eps))
%        B: Initial activation matrix (nb x J, default: max(rand(nb,J),eps))
%
% [outputs]
%        A: Decomposed basis matrix (I x nb)
%        B: Decomposed activation matrix (nb x J)
%

% Check errors and set default values -------------------------------------
[I,J] = size(X);
if min(min(X)) < 0 || isreal(X) ~= 1
    error('All the entries of input matrix must be nonnegative.');
end
if (nargin < 2)
    error('Sampling frequency "fs" is required.');
end
if (nargin < 3)
    error('Window shift size "shift" is required.');
end
if (nargin < 4)
    beta = 2;
end
if (nargin < 5)
    nb = ceil(J/10);
end
if (nargin < 6)
    it = 100;
end
if (nargin < 7)
    A = max( rand( I, nb ), eps ); % Initial basis matrix
end
if (nargin < 8)
    B = max( rand( nb, J ), eps ); % Initial activation matrix
end
if size(A,1) ~= I || size(A,2) ~= nb || size(B,1) ~= nb || size(B,2) ~= J
    error('The sizes of input initial matrices are incorrect.\n');
end
% -------------------------------------------------------------------------

% Some initializations ----------------------------------------------------
X = max( X, eps );  % Machine epsilon flooring
E = ones( I, J ); % Ones matrix whose size is the same with observed matrix X
if beta < 1
    phi = 1/( 2 - beta ) ;
elseif beta < 2
    phi = 1 ;
else
    phi = 1/( beta - 1 ) ;
end
% -------------------------------------------------------------------------

% Parameters for figure size ----------------------------------------------
screenSize = get(groot,'ScreenSize'); % Obtaining screen size
Left = floor(screenSize(4)./50); % Distance between left edge of primary display and left edge of figure
Bottom = screenSize(4)./20; % Distance between bottom edge of primary display and bottom edge of figure
Width = floor(screenSize(3)-screenSize(3)./50); % Distance between left and right edges of figure
Height = floor(screenSize(4)-screenSize(4)./20); % Distance between top and bottom edges of figure
handle = figure('OuterPosition',[Left Bottom Width Height]); % Display figure with specified window size
Animation(1:it+1) = struct('cdata',uint8(zeros(ceil(Height),ceil(Width),3)),'colormap',[]); % Preallocation of structure array "Animation"
% -------------------------------------------------------------------------

% Parameters for movie file -----------------------------------------------
movObj = VideoWriter('./output/output.mp4','MPEG-4'); % Declaration of an object for movie file
movObj.FrameRate = 5; % Frame rate of the output movie file
open(movObj); % file opening movie object (must be closed before script stops)
% -------------------------------------------------------------------------

% Prameters of axis ranges for plotting -----------------------------------
FreqMin = 1; % Minimum frequency (Hz) for basis plotting
FreqMax = 8000; % Maximum frequency (Hz) for basis plotting
PowMin = -80; % Minimum power (dB) for basis plotting
PowMax = 50; % Maximum power (dB) for basis plotting
TimeMin = shift/fs; % Minimum time-frame (s) for activation plotting
TimeMax = (shift/fs)*J; % Maximum time frame (s) for activation plotting
AmpMin = 0; % Minimum amplitude for activation plotting
AmpMax = 50; % Maximum amplitude for activation plotting
% -------------------------------------------------------------------------

% Plotting initial basis and activation matrices --------------------------
for b = 1 : nb
    %%% Basis plotting %%%
    subaxis(nb,2,2*b-1,'Margin',0.05); plot(linspace(fs/(2*I),fs/2,I),10*log10(A(:,b)));
    axis([FreqMin FreqMax PowMin PowMax]);
    title(sprintf('Basis %d',b)); xlabel('Frequency [Hz]'); ylabel('Power');
    %%% Activation plotting %%%
    subaxis(nb,2,2*b,'Margin',0.05); plot(linspace(shift/fs,(shift/fs)*J,J),B(b,:));
    axis([TimeMin TimeMax AmpMin max(B(b,:))]);
    title(sprintf('Activation %d',b)); xlabel('Time [s]'); ylabel('Amplitude');
end
Animation(1) = getframe(handle); % Obtaining (capturing) current figure
writeVideo(movObj,Animation(1)); % Writing current video frame
% -------------------------------------------------------------------------

% NMF update and plotting initial basis and activation matrices -----------
fprintf('NMF iteration: ');
for i = 1 : it
    [ A, B ] = multiplicativeUpdate_local( X, A, B, E, beta, phi ); % Update rules
    for b = 1 : nb
        %%% Basis plotting %%%
        subaxis(nb,2,2*b-1,'Margin',0.05); plot(linspace(fs/(2*I),fs/2,I),10*log10(A(:,b)));
        axis([FreqMin FreqMax PowMin PowMax]);
        title(sprintf('Basis %d',b)); xlabel('Frequency [Hz]'); ylabel('Power');
        %%% Activation plotting %%%
        subaxis(nb,2,2*b,'Margin',0.05); plot(linspace(shift/fs,(shift/fs)*J,J),B(b,:));
        axis([TimeMin TimeMax AmpMin max(B(b,:))]);
        title(sprintf('Activation %d',b)); xlabel('Time [s]'); ylabel('Amplitude');
    end
    Animation(i+1) = getframe(handle); % Obtaining (capturing) current figure
    writeVideo(movObj,Animation(i+1)); % Writing current video frame
    fprintf('%d ',i);
end
close(movObj);
% -------------------------------------------------------------------------

scale = 1./sum(A); % Normalization scale
A = A .* repmat( scale, size(A,1), 1 ); % Normalizing each basis
B = B ./ repmat( scale', 1, size(B,2) ); % Normalizing each activation
fprintf('\n');
end

%%% subfunction for update multiplivate update %%%%
function [ A, B ] = multiplicativeUpdate_local( X, A, B, E, beta, phi )
if beta == 2 % Update rules based on EUC distance
    A = A .* (( X*B.' )./( A*(B*B.') ));
    A = max(A,eps);
    B = B .* (( A.'*X )./( A.'*A*B ));
    B = max(B,eps);
elseif beta == 1 % Update rules based on KL divergence
    A = A .* (( (X./(A*B))*B' )./( E*B' ));
    A = max(A,eps);
    B = B .* (( A'*(X./(A*B)) )./( A'*E ));
    B = max(B,eps);
elseif beta == 0 % Update rules based on IS divergence
    AB = A*B;
    A = A .* sqrt(( (X.*((AB).^(-2)))*B' )./( ((AB).^(-1))*B' ));
    A = max(A,eps);
    AB = A*B;
    B = B .* sqrt(( A'*(X.*((AB).^(-2))) )./( A'*((AB).^(-1)) ));
    B = max(B,eps);
else % Update rules based on beta-divergence
    AB = A*B;
    A = A .* ((( (X.*((AB).^(beta-2)))*B' )./( ((AB).^(beta-1))*B' )).^(phi));
    A = max(A,eps);
    AB = A*B;
    B = B .* ((( A'*(X.*((AB).^(beta-2))) )./( A'*((AB).^(beta-1)) )).^(phi));
    B = max(B,eps);
end
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% EOF %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%