% Code for basis segmentation, described in: 
%
% Ella Bingham, Aristides Gionis, Niina Haiminen, Heli Hiisilä, Heikki
% Mannila, Evimaria Terzi, "Segmentation and dimensionality reduction". 2006
% SIAM Conference on Data Mining, April 20-22, 2006, Bethesda, Maryland, USA.
%
% Code written by Heli Hiisilä and Ella Bingham, 2005, with contributions
% from the co-authors of the paper.

function [output_struct] = BasisSegmentation(X,k,m,p);

%% X - input data, has rows as dimensions and columns as observations
%% k - amount of segments 
%% m - amount of basis vectors
%% p = 1 (median) p = 2 (mean)

%% output_struct - contains the output variables:
%% error_kk, error1, error2, error3, boundaries_kk, boundaries1,
%% boundaries2, boundaries3, levels_kk, levels1, levels2, levels3.
%% Here "kk" refers to plain (k,k)-segmentation, "1" to SEG-PCA, "2" to
%% SEG-PCA-DP, and "3" to PCA-SEG. "error" gives the error of
%% reconstructing the data, "boundaries" gives the segment boundaries, and
%% "levels" gives the segment means.
   
% Make the data zero mean and unit variance. The rows of X
% correspond to dimensions, and the columns to observations.
[d,n] = size(X);
for i=1:d
  X(i,:) = X(i,:)-mean(X(i,:));
  X(i,:) = X(i,:)/std(X(i,:));
end;

%%% Plain (k,k)-segmentation
[boundaries_kk,levels_kk] = dynprog(X,k,p);
%%% Error of (k,k)-segmentation
p=2;
error_kk = 0;
for j = 1 : k
 X_ = X(:,[boundaries_kk(j):boundaries_kk(j+1)-1]);
 error_kk = error_kk + Lperror(X_,representative(X_,p),p);
end;


%%% Algorithm SEG-PCA

disp('Computing SEG-PCA...')
%%% Let's make the U matrix with samples in the k segments
%%% replaced by the mean of the segment
[foo,N] = size(boundaries_kk);
[d,n] = size(X);
for i = 1 : N-1
 if(i == 1)
   U = (mean(X(:,[boundaries_kk(i):boundaries_kk(i+1)-1])')'*ones(1,(boundaries_kk(i+1)-boundaries_kk(i))));
 else 
   U = [U,(mean(X(:,[boundaries_kk(i):boundaries_kk(i+1)-1])')'*ones(1,(boundaries_kk(i+1)-boundaries_kk(i))))];
 end;
end;
% And let's then calculate the PCA basis matrix V
[V_U_pca,D] = eig(cov(U')); 
[d,d] = size(V_U_pca);
V_U_pca = V_U_pca(:,[d-m+1:d]); 
% Reconstruction error
error1 =  Error(X,V_U_pca,pinv(V_U_pca),p,boundaries_kk);
% Segment boundaries and mean levels are the same as in the original
% (k,k) segmentation
boundaries1 = boundaries_kk;
levels1 = levels_kk;


%%% Algorithm SEG-PCA-DP

disp('Computing SEG-PCA-DP...')
% Continuing from SEG-PCA by dynamic programming with the basis for the
% original data X
[boundaries2,levels2] = dynprog_(X,V_U_pca,k,p);
% Reconstruction error
error2 =  Error(X,V_U_pca,pinv(V_U_pca),p,boundaries2);


%%% Algorithm PCA-SEG

disp('Computing PCA-SEG...')
% Compute PCA on the original data, and project the data onto the PCA basis
[Vpca,D] = eig(cov(X')); 
[d,d] = size(Vpca);
Vpca = Vpca(:,[d-m+1:d]);
Apca = Vpca'*X;
% Segment the projected data
[boundaries3,levels3] = dynprog(Apca,k,p);
% Reconstruction error
error3 =  Error(X,Vpca,pinv(Vpca),p,boundaries3); 


%%% Visualize the segmentations of the algorithms
ShowBasis(X,boundaries_kk,boundaries2,boundaries3,Apca)

%%% Visualize the segment-wise coefficients of the algorithms
ShowCoefficients(X,boundaries_kk,boundaries2,boundaries3,V_U_pca,Vpca);

%%% Visualize the basis vectors of the algorithms
figure, imagesc(-V_U_pca(:,end:-1:end-m+1)); colormap(gray); colorbar;
title('Basis vectors of algorithms SEG-PCA and SEG-PCA-DP');
figure, imagesc(-Vpca(:,end:-1:end-m+1)); colormap(gray); colorbar;
title('Basis vectors of algorithm PCA-SEG');

%%% Output variables
output_struct = struct('error_kk',error_kk,'error1',error1,'error2',error2,'error3',error3,'boundaries_kk',boundaries_kk, 'boundaries1',boundaries1,'boundaries2',boundaries2,'boundaries3',boundaries3,'levels_kk',levels_kk,'levels1',levels1,'levels2',levels2,'levels3',levels3);
