%% Code to perform ICA with signal reconstruction and ordering (ICA+SRO). From 
% the paper: 
% Y. Levin-Schwartz, D. R. Sparta, J. F. Cheer and T. Adali, "Parameter-free
% automated extraction of neuronal signals from calcium imaging data," 
% 2017 IEEE International Conference on Acoustics, Speech and Signal Processing 
%(ICASSP), New Orleans, LA, 2017, pp. 1033-1037.
%
% Code written by Yuri Levin-Schwartz (ylevins1@umbc.edu)

clear
clc
close all

NumRun = 1; % Number of ICA runs (only put in more than one if using ICASSO or other stability analysis methods)
Vid_length = 283; % Number of frames in the video

%% Read in video and preprocess the data
vidObj = VideoReader('BNST3ShockDay2-Shock8-calcium-v1 deinterlaced.mp4');

vidHeight = vidObj.Height;
vidWidth = vidObj.Width;

s = struct('cdata',zeros(vidHeight,vidWidth,3,'uint8'),'colormap',[]);
C=zeros(Vid_length,vidHeight*vidWidth); 
k = 1;
while hasFrame(vidObj)
    s(k).cdata = readFrame(vidObj);
    A = double(rgb2gray(s(k).cdata)); % Convert video to grayscale
    B = medfilt2(A); % Median filter frames
    C(k,:)=zscore(reshape(B,[1,640*640])); % Z-score frames
    k = k+1;
end

%% PCA (and order selection) 
[E1, D1, F1]=svd(C','econ'); % SVD
%[comp_est1,~,~,~] = OrderSelection_E_DS(transpose(C)); % Order selection method 1
[comp_est2, comp_est3, comp_est4] = orderEstimation_HdsvsEig_R(C); % Order selction methods 2,3,4
fprintf('Estimated order : %d\n', comp_est4);
X = (inv(D1(1:comp_est4,1:comp_est4)))*F1(:,1:comp_est4)'*C;

%% ICA-EBM
for u=1:NumRun
W = ICA_EBM(X);

S_hat=W*X;
A=(pinv(F1(:,1:comp_est4)')*(D1(1:comp_est4,1:comp_est4)))*pinv(W); % Back-reconstruct mixing matrix (time courses)

%% Adjust the sign of the components as well as the mixing matrix
for dd=1:comp_est4
    b=sign(skewness(S_hat(dd,:)')); 
    S_hat2(dd,:)=b*S_hat(dd,:); % Adjust spatial maps by sign of the spatial skewness 
    A2(:,dd)=b*A(:,dd); % Adjust mixing matrix by sign of the spatial skewness
%     qwe(dd)=skewness(S_hat2(dd,:));
end
% [~,bb]=sort(qwe,'descend'); 
% sR.signal{1,u}=S_hat2(bb,:); 
% sR.A{1,u}=A2(:,bb); 

end

%% Binarize and keep only the largest object in the binarized image
for abc=1:comp_est4
    S_hat3=im2bw(double(reshape(S_hat2(abc,:),640,640)),1); % Binarize component
    L = bwlabel(S_hat3);
    M=mode(L(find(L))); % Find the largest object in the binarized image
    S_hat4(abc,:)=zeros(1,640*640);
    N=find(L==M); % Select the larget object in the binarized image
    S_hat4(abc,N)=1; 
    for def=1:Vid_length
    A3(def,abc)=mean(uint8(C(def,N))); % Find the average activation in the largest object over time
    end
    sk(abc)=skewness(A3(abc,:));
end

[~,bb]=sort(sk,'descend'); 
S_hat5=S_hat2(bb,:); % Sort by temporal skewness
A4=A2(:,bb); % Sort by temporal skewness

%%%%%%%%%%%%%%Plotting components%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

for jj = 1: floor(comp_est4/5)
figure
subplot(5,2,1), imshow(double(reshape(S_hat5(5*(jj-1)+1,:),640,640)))
subplot(5,2,2), plot(A4(:,1))
subplot(5,2,3), imshow(double(reshape(S_hat5(5*(jj-1)+2,:),640,640)))
subplot(5,2,4), plot(A4(:,2))
subplot(5,2,5), imshow(double(reshape(S_hat5(5*(jj-1)+3,:),640,640)))
subplot(5,2,6), plot(A4(:,3))
subplot(5,2,7), imshow(double(reshape(S_hat5(5*(jj-1)+4,:),640,640)))
subplot(5,2,8), plot(A4(:,4))
subplot(5,2,9), imshow(double(reshape(S_hat5(5*(jj-1)+5,:),640,640)))
subplot(5,2,10), plot(A4(:,5))
end
