% Learning algorithm for EM mixture clustering by sparseFRAME model
% version. No EM alignement.
clear all;
close all;

mex CgetMAX1.c;
mex CgetMAX1_DoG.c;
mex sparseFeatureExtraction.c;
mex MatchingPursuitV2.c;
mex grad_U_sparse.c;
mex filterTiling.c;


dataNames    = {'dogs','cow','butterfly_dragonfly_bug','bike_car_motor_Scooter', 'faces_butterflies_horses','horse_pigeons_eagles_swans_fishes', 'DeerCatWolf'};
gtNumCluster = [2, 2, 3, 4, 3, 5, 3];
sx_list=[100, 100, 100, 60, 100, 60, 100];
sy_list=[110, 110, 110, 80, 110, 70, 100];
numSketchList=[300, 300, 300, 300, 300, 300, 300];
GaborList={[1.4 1 0.7],[1.4 1 0.7],[1.4 1 0.7], [0.7 0.4 0.2],[1.4 1 0.7],[0.7],[1, 0.7]};
DoGList={[],[],[], [], [], [], []};
numEMIterationList=[10, 10, 10, 10, 10, 10, 10];
% set seed for random numbers generation


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
idx = 3; % choose your categoriy
seed = 1;
rng(seed);
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

numEMIteration=numEMIterationList(idx);
numCluster=gtNumCluster(idx);

sx=sx_list(idx); % 70
sy=sy_list(idx); % 90
isSoftClassification=false;
inPath = ['./positiveImages/' dataNames{idx}];
cachePath = './feature';
resultPath = './result';

if ~exist(cachePath,'dir')
    mkdir(cachePath)
else
    rmdir(cachePath,'s')
    mkdir(cachePath)
end
if ~exist(resultPath,'dir')
    mkdir(resultPath)
    mkdir(fullfile(resultPath,'img'));
else
    rmdir(resultPath,'s')
    mkdir(resultPath)
    mkdir(fullfile(resultPath,'img'));
end


nOrient = 16;
nTileRow = 12; %nTileRow \times nTileCol defines the number of paralle chains
nTileCol = 12;
numSample = 3; % how many HMC calls for each learning iteration


sigsq=10;
%nIter = 6; % the number of iterations for learning lambda
nIter = 20; % the number of iterations for learning lambda
lambdaLearningRate = 0.1/sqrt(sigsq);  % 0.1
epsilon = 0.03; 
L= 10;

isLocalNormalize=false; %
isSeparateLocalNormalize=false;
localNormScaleFactor=0.4;

localHalfx=20;
localHalfy=20;
thresholdFactor=0.01;

numSketch = numSketchList(idx); %400 sparsification
locationShiftLimit=5;  % arg-max
orientShiftLimit=1;    % arg-max
numTopFeatureToShow=70; % not used so far, please ignore

IsToRefineAfterConverge=false; 
nIter_AfterConverge=10;

isSaved=1;
isComputelogZ=1;

%% parameters for cpu paralleled computing  
runParallel = 1; % Set variable to know when to run parallel
noWorkers = min(8,numCluster);   %% % Set number of workers, which depends on how many cores in your cpu

%% Step 0: prepare filter, training images,and filter response on images

%% option 1: automatical filters generation
%%parameters for automatically generating scale List
% initialScale = 1.4; %initial scale for Gabor wavelets
% shrinkFactor = sqrt(2.0);
% expandDoG = 20;
% 
% GaborScaleList=[initialScale, initialScale/shrinkFactor, initialScale/(shrinkFactor^2), initialScale/(shrinkFactor^3)];  % scales list for Gabor filters
% DoGScaleList=[initialScale*expandDoG/shrinkFactor];  % scales list for DoG filters

%% option 2: manually filters generation
GaborScaleList=GaborList{idx}; %;  % you can add more scales of Gabor filters here. for this case [0.8, 0.4] also works well
DoGScaleList=DoGList{idx}; 


 
nScaleDoG=length(DoGScaleList);
nScaleGabor=length(GaborScaleList);

filters=[];
for iScale=1:nScaleGabor
    f = MakeFilter(GaborScaleList(iScale),nOrient);
    for i=1:nOrient
     f_r{i} =single(real(f{i}));
     f_i{i} =single(imag(f{i}));
    end
    
    filters = [filters f_r f_i];
end

for iScale=1:nScaleDoG

    f0 = single(dog(DoGScaleList(iScale),0));
    filters = [filters f0];

end

numFilter = length(filters);
halfFilterSizes = zeros(size(filters));
for iF = 1:numFilter
   halfFilterSizes(iF)=(size(filters{iF},1)-1)/2;
end
overAllArea = sum((sx-2*halfFilterSizes).*(sy-2*halfFilterSizes));


files = dir(fullfile(inPath,'*.jpg'));
numImage=length(files);
disp(['start filtering']); tic

imgCell = cell(0);
for iImg = 1:numImage
    copyfile(fullfile(inPath,files(iImg).name),fullfile(resultPath,'img',files(iImg).name));
    img = imread(fullfile(inPath,files(iImg).name));
    img = imresize(img,[sx, sy]);
    if ndims(img)==3
       img=rgb2gray(img);
    end
    img = im2single(img);
    img = img-mean(img(:));
    img = img/std(img(:))*sqrt(sigsq);
    
    imgCell{iImg}=img;  % saving all the raw images
    
    
    S1 = cell(numFilter,1);
    M1 = cell(numFilter,1);
    
    % SUM1
    for iFilter = 1:numFilter
        Y = filter2(filters{iFilter},img);
        S1{iFilter} = abs(single(Y));
        M1{iFilter} = zeros(sx,sy,'single');
    end
    
    
    if isLocalNormalize
       
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% Local Normalization
    if isSeparate        
      for iScale=1:nScaleGabor  
        h = halfFilterSizes(1+(iScale-1)*(2*nOrient)); 
        [S1(1+(iScale-1)*(2*nOrient) : nOrient+(iScale-1)*(2*nOrient))]= ...
            LocalNormalize(S1(1+(iScale-1)*(2*nOrient):nOrient+(iScale-1)*(2*nOrient)),[],h,round(localNormScaleFactor*h),round(localNormScaleFactor*h),thresholdFactor);
        [S1(nOrient+1+(iScale-1)*(2*nOrient): 2*nOrient+(iScale-1)*(2*nOrient))] =...
            LocalNormalize(S1(nOrient+1+(iScale-1)*(2*nOrient): 2*nOrient+(iScale-1)*(2*nOrient)),[],h,round(localNormScaleFactor*h),round(localNormScaleFactor*h),thresholdFactor);
      end
      % for DoG
      for iScale=1:nScaleDoG  
          ind  = 2*nOrient*nScaleGabor+iScale;
          h = halfFilterSizes(ind);       
          [S1(ind)]=LocalNormalize(S1(ind),[],h,round(localNormScaleFactor*h),round(localNormScaleFactor*h),thresholdFactor);                
      end    
    else        
        for iScale=1:nScaleGabor
           h = halfFilterSizes(1+(iScale-1)*(2*nOrient)); 
           [S1(1+(iScale-1)*(2*nOrient):nOrient+(iScale-1)*(2*nOrient)),S1(nOrient+1+(iScale-1)*(2*nOrient): 2*nOrient+(iScale-1)*(2*nOrient))]=...
               LocalNormalize(S1(1+(iScale-1)*(2*nOrient):nOrient+(iScale-1)*(2*nOrient)),S1(nOrient+1+(iScale-1)*(2*nOrient): 2*nOrient+(iScale-1)*(2*nOrient)),h,round(localNormScaleFactor*h),round(localNormScaleFactor*h),thresholdFactor);
        end
        % for DoG
        for iScale=1:nScaleDoG 
           ind  = 2*nOrient*nScaleGabor+iScale;
           h = halfFilterSizes(ind);       
           [S1(ind)]=LocalNormalize(S1(ind),[],h,round(localNormScaleFactor*h),round(localNormScaleFactor*h),thresholdFactor);                
        end        
    end
 
    end
    
    %arg-Max for MAX1map
    % Gabor
    for iScale=1:nScaleGabor
      CgetMAX1(1,sx,sy,nOrient,locationShiftLimit,orientShiftLimit,1,S1(1+(iScale-1)*(2*nOrient):nOrient+(iScale-1)*(2*nOrient)),M1(1+(iScale-1)*(2*nOrient):nOrient+(iScale-1)*(2*nOrient)));
      CgetMAX1(1,sx,sy,nOrient,locationShiftLimit,orientShiftLimit,1,S1(nOrient+1+(iScale-1)*(2*nOrient): 2*nOrient+(iScale-1)*(2*nOrient)),M1(nOrient+1+(iScale-1)*(2*nOrient): 2*nOrient+(iScale-1)*(2*nOrient))); 
 
    end
    % DoG
    for iScale=1:nScaleDoG
       ind  = 2*nOrient*nScaleGabor+iScale;       
       CgetMAX1_DoG(1,sx,sy,locationShiftLimit,1,S1(ind),M1(ind)); 
  
    end
    
    mapName = fullfile(cachePath,['SUMMAXmap-image' num2str(iImg) '.mat']);
    current_file_name=files(iImg).name;
    save(mapName, 'M1','current_file_name');  % only save the MAX1
end
disp(['filtering time: ' num2str(toc) ' seconds']);

%% Prepare variables for EM
prob = zeros(1, numCluster);
SUM2scoreAll = zeros(numImage, numCluster);
dataWeightAll = rand(numImage,numCluster);
dataWeightAll = dataWeightAll./repmat(sum(dataWeightAll,2),1,numCluster);

if isSoftClassification==false
    [maxs,ind]=max(dataWeightAll,[],2);  % find max in each row
    dataWeightAll=zeros(size(dataWeightAll));  % set zeros
    dataWeightAll(sub2ind(size(dataWeightAll),1:length(ind),ind'))=1;  % hard classification
end

clusters=struct('imageIndex',cell(numCluster,1),'rHat',[],'template',[],'logZ',[],'sampleImages',[],'mixtureWeight',[]);
    for iImg=1:size(dataWeightAll,1)
      [~, ind]=max(dataWeightAll(iImg,:));
      clusters(ind).imageIndex=[clusters(ind).imageIndex,iImg];
    end
it = 0;
ShowAssignment;

%save dataWeightAll dataWeightAll;
%load dataWeightAll; % remove this line if you do not want strict reproducibility

IsConverged=false;

%% EM iteration
for (it = 1 : numEMIteration)
    
    old_dataWeightAll=dataWeightAll;
    disp(['M-step of iteration ' num2str(it)]);
    % initialize the folders and aggregate the rHat map;
    for c = 1:numCluster
        savingFolder=fullfile(resultPath,['iteration' num2str(it) '/']);
        if ~exist(savingFolder)
            mkdir(savingFolder);
        end
        sumWeight = sum(dataWeightAll(:,c));
        clusters(c).mixtureWeight=sumWeight;
        clusters(c).rHat = cell(numFilter,1);
        for iFilter = 1:numFilter
            clusters(c).rHat{iFilter}=zeros(sx,sy,'single');
        end
    end
    for iImg = 1:numImage
        mapName = fullfile(cachePath,['SUMMAXmap-image' num2str(iImg)]);
        featureMap=load(mapName);
        for iCluster = 1:numCluster
            if dataWeightAll(iImg,iCluster)~=0
                for iFilter = 1:numFilter
                    clusters(iCluster).rHat{iFilter}= clusters(iCluster).rHat{iFilter}+ featureMap.M1{iFilter}.*(dataWeightAll(iImg,iCluster)/clusters(iCluster).mixtureWeight);
                end
            end
        end
    end

    
%%% cpu paralleled learning     
     
      if ~runParallel
        noWorkers = 1;   
      end      
      %% Set matlabpool as needed
      isOpen = matlabpool('size') > 0;
      isOpenCorr = matlabpool('size') == noWorkers;
      if runParallel && ~isOpenCorr,
          matlabpool close force
          matlabpool(noWorkers)
      end
      if ~runParallel && isOpen,
          matlabpool close force
      end
 
      % assign tasks to the paralleled workers
      numAssignmentList1=ones(1,noWorkers)*floor(numCluster/noWorkers);
      Assignments_left=mod(numCluster,noWorkers);
      numAssignmentList2=[ones(1,Assignments_left), zeros(1,noWorkers-Assignments_left)];
      numAssignmentList=numAssignmentList1 + numAssignmentList2;
      endIndex=cumsum(numAssignmentList);
      startIndex=[1,endIndex(1,1:end-1)+1];

      ticID=tic;  
      temp_result=cell(numCluster,3);  % colum 1 is template, colum 2 is currSample, colum 3 is logZ, 
      spmd
          
       temp_result=codistributed(temp_result, codistributor1d(1,numAssignmentList));
      
       for iCluster= startIndex(labindex):endIndex(labindex)
          
             
          savingFolder=fullfile(resultPath,['iteration' num2str(it) ],['cluster' num2str(iCluster) '/']);
          if ~exist(savingFolder)
              mkdir(savingFolder);
          end
          
          training=cell(0);
          imageIndexSet=clusters(iCluster).imageIndex;
          for iI=1:length(imageIndexSet)
              training{iI}=double(imgCell{imageIndexSet(iI)});
          end
          
          disp(['start filter selection']);
          [template, deformedTemplate]= filtersSelection_em(training, GaborScaleList, DoGScaleList, nOrient, numSketch, numTopFeatureToShow,...
                locationShiftLimit, orientShiftLimit, sx, sy, savingFolder);
          template.selectedLambdas=single(zeros(1,template.numSelected));
          
          disp(['start learning Frame model']);
          [template,currSample,logZ]=sparseFRAMElearn(template, nIter, filters, clusters(iCluster).rHat, sx, sy, halfFilterSizes, ...
               locationShiftLimit, nTileRow,nTileCol,epsilon,L,lambdaLearningRate,numSample, isSaved,savingFolder);
     
          temp_result(iCluster,1)={template};
          temp_result(iCluster,2)={currSample};
          temp_result(iCluster,3)={logZ};      
        
       end
       
      end
      temp_result=gather(temp_result);
      
      for ic=1:numCluster
         clusters(ic).template=temp_result{ic,1};
         clusters(ic).currSample=temp_result{ic,2};
         clusters(ic).logZ=temp_result{ic,3};
      end
      
      disp(['learning time of FRAME by GPU for the ' num2str(numCluster) ' clusters takes ' num2str(toc(ticID)) ' seconds']);
      
      if matlabpool('size') > 0;
            matlabpool close
      end
%%%%%

    %% E-step
    disp(['E-step of iteration ' num2str(it)]);
    for iImg = 1 : numImage
        mapName = fullfile(cachePath,['SUMMAXmap-image' num2str(iImg)]);
        featureMap=load(mapName);
        for iCluster= 1:numCluster
           
            template=clusters(iCluster).template;
            M1_selected = zeros(1,template.numSelected);
            for i=1:template.numSelected
                  M1_selected(i)=featureMap.M1{template.selectedFilter(i)}(template.selectedRow(i),template.selectedCol(i));
            end
            SUM2score=template.selectedLambdas*M1_selected';    

            SUM2scoreAll(iImg, iCluster)= SUM2score - clusters(iCluster).logZ;
        end % iCluster
    end
    if isSoftClassification
        for (c = 1 : numCluster)
            prob(c) = sum(dataWeightAll(:, c))/numImage;
        end
        for (c = 1 : numCluster)
            for (img = 1 : numImage)
                dataWeightAll(img, c) = prob(c)/sum(prob(1, :).*exp(SUM2scoreAll(img, :)-SUM2scoreAll(img, c)));
            end
        end
    else
        dataWeightAll=zeros(size(dataWeightAll));  % set zeros
        [maxs,ind]=max(SUM2scoreAll,[],2);  % find max in each row
        dataWeightAll(sub2ind(size(SUM2scoreAll),1:length(ind),ind'))=1;  % hard classification
    end
    
    save(fullfile(resultPath,['iteration' num2str(it)], 'dataWeightAll.mat'),'dataWeightAll');
    save(fullfile(resultPath,['iteration' num2str(it)], 'mixtureModel.mat'),'clusters','numImage','filters');
    
    %%%%% generating html
    for iCluster = 1:numCluster
      clusters(iCluster).imageIndex=[];
    end
    for iImg=1:size(dataWeightAll,1)
      [~, ind]=max(dataWeightAll(iImg,:));
      clusters(ind).imageIndex=[clusters(ind).imageIndex,iImg];
    end
    ShowAssignment;
    
    if (norm(dataWeightAll(:)-old_dataWeightAll(:),1)==0 && ~isSoftClassification)
       IsConverged = true; 
    end
    
    if (IsToRefineAfterConverge && IsConverged)  % after the the assignment does not change, we refine the model by using more iterations in updating the Lambdas
        nIter = nIter_AfterConverge;
    end
    
    
end %it

[purity, entropy] = evaluateClustering(inPath, 'gt.txt', dataWeightAll);

