% Learning sparse FRAME models by generalized Gibbs sampler and generative epsilon-boosting

clear;
close all;

mex CgetMAX1.c;
mex CgetMAX1_DoG.c;
mex filterTiling.c;
mex sparseFeatureExtraction.c;
mex ComputeCorrBB.c;
mex I_plus_cB.c;
mex gibbs_v3.c;
mex sparseFiltering.c;
mex mexc_updateCorrelationTable.c;
mex drawSketch.c;

categoryName = 'cat';  

inPath = ['positiveImage/' categoryName '/'];
Synfolder='./synthesiedImage/';
if ~exist(Synfolder)
   mkdir(Synfolder);
else
   rmdir(Synfolder,'s');
   mkdir(Synfolder);
end

templateFolder='./sketchTemplate/';
if ~exist(templateFolder)
   mkdir(templateFolder);
else
   rmdir(templateFolder,'s');
   mkdir(templateFolder);
end

nOrient = 16;
locationShiftLimit=2;  % local max
orientShiftLimit=1;    % local max

%%%%%%%%%%%%%%%%%%%
sx=100;
sy=100;
%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
interval = 5; % 5 control how many iterations we need to wait to select next wavlet 
numWavelet = 300;   % default 300
nIteration = numWavelet * interval * 2;  % number of wavelet we will selected is equal to nIteration/interval
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


isLocalNormalize=true;   % No DoG will be used, if isLocalNormalize=true
isSeparate=false;  % it matters, if isV2=false;
localNormScaleFactor=2;
thresholdFactor=0.01;
sigsq=10;
isToGenerateMovie=false;


isOption1=true;

if isOption1
 %% option 1: automatical filters generation
 % parameters for automatically generating scale List
 initialScale = 1.4; %initial scale for Gabor wavelets
 shrinkFactor = sqrt(2.0);
 expandDoG = 13.5;
 
 GarborScaleList=[initialScale, initialScale/shrinkFactor, initialScale/(shrinkFactor^2), initialScale/(shrinkFactor^3)];  % scales list for Gabor filters
 DoGScaleList=[initialScale*expandDoG,  initialScale*expandDoG/shrinkFactor];  % scales list for DoG filters;
  
 if isLocalNormalize    % No DoG will be used, if isLocalNormalize=true
    DoGScaleList=[];
 end    
  
else

  %% option 2: manually filters generation
  GarborScaleList=[1.2 0.7];
  DoGScaleList=[10];
  
  if isLocalNormalize    % No DoG will be used, if isLocalNormalize=true
    DoGScaleList=[];
  end   

end

%%%%% parameters for Gibbs sampler
threshold_corrBB = 0;    % threshold for correlation matrix  0.01
lower_bound_rand = 0.001;   % lower bound of random number
upper_bound_rand = 0.999;   % upper bound of random number
c_val_list=-30:2:30;  % range of c
lambdaLearningRate = 0.1/sqrt(sigsq);  % 0.1

% multiple chains
useMultiChain=true;
nTileRow = 6; %nTileRow \times nTileCol defines the number of paralle chains
nTileCol = 6;
if useMultiChain==false
   nTileRow = 1;
   nTileCol = 1;
end

%% Step 0: prepare filter and training images
tic;
nScaleGabor=length(GarborScaleList);

outPath = ['./results/'];
if ~exist(outPath)
   mkdir(outPath)
end

filters=[];
for iScale=1:nScaleGabor

    f = MakeFilter(GarborScaleList(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
numGabor = nScaleGabor*nOrient*2;

nScaleDoG=length(DoGScaleList);
for iScale=1:nScaleDoG

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

end

numFilter = length(filters);

halfFilterSize = zeros(size(filters));
for iF = 1:numFilter
   halfFilterSize(iF)=(size(filters{iF},1)-1)/2;
end
Corr = CorrFilterFrame(filters); %calculate correlation among filters
disp(['finished creating filters bank: ' num2str(toc) ' seconds']);


imgCell = cell(0);
files = dir([inPath '/*.jpg']);
for iImg = 1:length(files)
    img = imread(fullfile(inPath,files(iImg).name));
    img = imresize(img,[sx, sy],'nearest');
    im=im2double(rgb2gray(img));
    im=im-mean(im(:));
    im=im/std(im(:))*sqrt(sigsq);
    imgCell{iImg}=im;
end
%% Step 1: compute training sample averages

img = rgb2gray(img);
rHat = cell(numFilter,1);
for iFilter = 1:numFilter
    rHat{iFilter}=zeros(sx,sy);
end
tic;
for iImg = 1:length(files)
    S1 = cell(numFilter,1);
    M1 = cell(numFilter,1);
    
    % SUM1
    for iFilter = 1:numFilter
                
        Y = filter2(filters{iFilter},imgCell{iImg});
        S1{iFilter} = abs(single(Y));
        M1{iFilter} = zeros(sx,sy,'single');
    end
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% Local Normalization   
    if isLocalNormalize        
    
      if isSeparate        
        for iScale=1:nScaleGabor  
           h = halfFilterSize(1+(iScale-1)*(2*nOrient)); 
           [S1(1+(iScale-1)*(2*nOrient) : nOrient+(iScale-1)*(2*nOrient))]= LocalNormalizeV3(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))] = LocalNormalizeV3(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 = halfFilterSize(ind);       
           [S1(ind)]=LocalNormalizeV3(S1(ind),[],h,round(localNormScaleFactor*h),round(localNormScaleFactor*h),thresholdFactor);
        end        
      else
          
        for iScale=1:nScaleGabor
           h = halfFilterSize(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))]= LocalNormalizeV3(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 = halfFilterSize(ind);       
           [S1(ind)]=LocalNormalizeV3(S1(ind),[],h,round(localNormScaleFactor*h),round(localNormScaleFactor*h),thresholdFactor);                
        end       
        
      end
      
      %%%%%%%
      for iFilter = 1:numFilter 
        S1{iFilter }=S1{iFilter}*sqrt(sigsq);    
      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
    
    for iFilter = 1:numFilter
        rHat{iFilter}= rHat{iFilter}+double(M1{iFilter});
    end
end
for iFilter = 1:numFilter
    rHat{iFilter}=rHat{iFilter}/length(files);
end
disp(['finished filtering: '  num2str(toc) ' seconds']);


%% Step 2: optimize the exponential model
% our model parameter

I_syn = single(randn(sx*nTileRow,sy*nTileCol));  


SSD=zeros(nIteration,1);
nValidFeature= sum((sx-2*halfFilterSize-2*locationShiftLimit).*(sy-2*halfFilterSize-2*locationShiftLimit));

% initialize the template
filterSelected=[];
xSelected=[];
ySelected=[];
selectedLambdas=[];

template=struct('selectedFilter',single(filterSelected),'selectedRow',single(xSelected),'selectedCol',single(ySelected),...
    'selectedLambdas',single(selectedLambdas),'numSelected',single(0));

countMap = cell(numFilter,1); % this map separate sin part and cos part, and count them separately.
for iFilter = 1:numFilter
    countMap{iFilter}=zeros(sx,sy);
end
countTotal=0;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
countMap_sketch = cell(nScaleGabor*nOrient+nScaleDoG,1);  % for drawing sketches, this map combine sin part and cos part, and only count them once.
for iFilter = 1:nScaleGabor*nOrient+nScaleDoG
    countMap_sketch{iFilter}=zeros(sx,sy);
end
filterList_sketch = [];
rowList_sketch=[];
colList_sketch=[];

[allFilterR, ~, allSymbol] = MakeFilterBank(GarborScaleList, DoGScaleList, nOrient); % filter bank for function "drawSketch"
half = zeros(1, nOrient*nScaleGabor+nScaleDoG); 
for (i=1:nOrient*nScaleGabor+nScaleDoG)
    half(i) = (size(allFilterR{1, i}, 1)-1)/2;
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

I_squared_sum=zeros(1, nTileCol*nTileRow,'single');
index=0;
for j=1:nTileCol
    for i=1:nTileRow
       index=index+1;
       I_crop=I_syn(1+(i-1)*sx:sx+(i-1)*sx, 1+(j-1)*sy:sy+(j-1)*sy);  
       I_squared_sum(index) = single(sum(sum(I_crop.^2)));  % I.^2
    end
end
rHat_selected=[];
response_IB=[];
corrBB=zeros(numWavelet*2, numWavelet*2, 'single');

nStroke=0;
for iter = 1:nIteration  
    
     disp( [ 'iteration: ' num2str(iter)]);  
     
     if mod(iter-1,interval)==0     % we add one basis every few steps. The interval is set by users.  
        %%% forward selection
        [x_added, y_added, filter_added]=forwardSelection(filters, halfFilterSize, locationShiftLimit, rHat, I_syn, nTileRow, nTileCol, sx, sy);    
       
        
        if countMap{filter_added}(x_added, y_added)==0
                       
            % record the selected wavelet
            countMap{filter_added}(x_added, y_added)=1;
            countTotal = countTotal+1;
            
            % add the information to template
            template.selectedFilter = [template.selectedFilter filter_added];
            template.selectedRow = [template.selectedRow x_added];
            template.selectedCol = [template.selectedCol y_added];
            template.selectedLambdas = [template.selectedLambdas 0];
            template.numSelected = template.numSelected +1;
            
            
            % for drawing sketch template, we create a template combining sin part and cos part together
            if filter_added <= numGabor; %Gabor
                 Scale = floor( ( filter_added-1)/ (2*nOrient)) + 1;
                 Orient = mod(filter_added - 1, nOrient) + 1;
                 filter_added_ForSketch = Orient + (Scale-1) * nOrient;
            else  % DoG
                 filter_added_ForSketch = filter_added - nScaleGabor*nOrient;                
            end
            if countMap_sketch{filter_added_ForSketch}(x_added, y_added)==0
                    countMap_sketch{filter_added_ForSketch}(x_added, y_added)=1;
                    filterList_sketch = [filterList_sketch filter_added_ForSketch];
                    rowList_sketch=[rowList_sketch x_added];
                    colList_sketch=[colList_sketch y_added];                
            
                    sketchTemplate = zeros(sx, sy); 
                    drawSketch( nOrient * nScaleGabor+nScaleDoG, half, sx, sy, allSymbol(1, :),  sketchTemplate , ...
                       double(filterList_sketch), double(rowList_sketch), double(colList_sketch), length(filterList_sketch));
                    sketchTemplate = -sketchTemplate;
                    gLow = min(sketchTemplate(:));
                    gHigh = max(sketchTemplate(:));
                    sketchTemplate = (sketchTemplate-gLow)/(gHigh-gLow); 
                    nStroke= nStroke+1;
                    imwrite(sketchTemplate ,[templateFolder 'sketch' num2str(nStroke,'%04d') '.png']); 
                              
            end
%             if useMultiChain  
%                template.enlargedTemplate = enlargingTemplate(nTileRow, nTileCol, sx, sy, template); 
%             end
            
            % update rHat_selected for computation
%             rHat_selected = zeros(1,template.numSelected);
%             for i=1:template.numSelected
%                  rHat_selected(i)=rHat{template.selectedFilter(i)}(template.selectedRow(i),template.selectedCol(i));
%             end
            rHat_selected=[rHat_selected rHat{filter_added}(x_added,y_added)];            
            
            mexc_updateCorrelationTable(corrBB, Corr, single(halfFilterSize), template);                        
           
          %  corrBB=ComputeCorrBB(Corr,template,single(halfFilterSize));   % update corrBB           
            
            nonzero_corrBB_List=cell(1,template.numSelected);  % neighborhood system (i-th cell contains the indices of all bases whose correlations with B_i are not equal to 0.)
            corrBB2=corrBB(1:template.numSelected, 1:template.numSelected); % crop the valid part from corrBB

            for i=1:template.numSelected
                  nonzero_corrBB_List{i}=single(find(corrBB2(i,:)>threshold_corrBB | corrBB2(i,:)<-threshold_corrBB));
            end
            %%%%
            %response_IB=zeros(nTileCol*nTileRow,template.numSelected,'single');
            response_added=zeros(nTileCol*nTileRow, 1,'single');
            index=0;
            for j=1:nTileCol
                for i=1:nTileRow
                    index=index+1;
                    I_crop=I_syn(1+(i-1)*sx:sx+(i-1)*sx, 1+(j-1)*sy:sy+(j-1)*sy);                     
                    response_added(index) = sparseFiltering(I_crop, filters, filter_added, x_added, y_added);                          
                end
            end
            response_IB=[response_IB, response_added];
            
            disp(['filter: ' num2str(filter_added), ' x: ' num2str(x_added), ' y: ', num2str(y_added)]);            
            
        end        
             
     end  
    
    
     tic      
     % multi-chain sampling
     index=0;     
     for j=1:nTileCol
        for i=1:nTileRow
            index=index+1;
            I_syn_single_chain=I_syn(1+(i-1)*sx:sx+(i-1)*sx, 1+(j-1)*sy:sy+(j-1)*sy);             
            response_IB_single_chain = response_IB(index,:);
            I_squared_sum_single_chain = I_squared_sum(index);
            
            % gibbs sampling
            [I_syn(1+(i-1)*sx:sx+(i-1)*sx, 1+(j-1)*sy:sy+(j-1)*sy), I_squared_sum(index), response_IB(index,:)]= ...
                Gibbs_sparseV5(template, filters, I_syn_single_chain, corrBB(1:template.numSelected, 1:template.numSelected), single(c_val_list), ...
                I_squared_sum_single_chain, response_IB_single_chain, nonzero_corrBB_List, lower_bound_rand, upper_bound_rand);        
        end
     end   
       
     rModel_selected=mean(abs(response_IB),1);   % |<I, B>|        
     
     disp(['one iteration learning time: ' num2str(toc) ' seconds']);

     %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%5   
     % compute gradient and do graidnet ascent
       
     gradientF_selected = rHat_selected-rModel_selected;
    
     SSD(iter)=mean(abs(gradientF_selected));
            
     template.selectedLambdas = template.selectedLambdas + lambdaLearningRate * gradientF_selected;
     
                
     if mod(iter-1,interval)==0
       % save synthesied image
       gLow = min(I_syn(:));
       gHigh = max(I_syn(:));
       disp([ 'min: ' num2str(gLow) ' max: ' num2str(gHigh)]);
       disp([ 'SSD: ' num2str(SSD(iter))]);
       img = (I_syn-gLow)/(gHigh-gLow);
       imwrite(img,[Synfolder num2str(iter,'%04d') '.png']);
     end
     
     
end

save([outPath,'/model.mat'],'template','filters','GarborScaleList','DoGScaleList','sx','sy','locationShiftLimit','orientShiftLimit',...
    'nOrient','isLocalNormalize','isSeparate','sigsq','categoryName','localNormScaleFactor','thresholdFactor');


if isToGenerateMovie
filename='movie.gif';
ii=0;
for i=1:nIteration
    if mod(i-1,interval)==0
      str=[Synfolder num2str(i,'%04d') '.png'];
      Img=imread(str);
      ii=ii+1;
      figure(ii)
      imshow(Img,'border','tight','initialmagnification','fit');
      set(gcf,'Position',[50,50,sy*nTileCol,sx*nTileRow]);     % the dimension of the gif should be lareger than this default setting.
      axis normal;
      frame=getframe(ii);
      im=frame2im(frame);
      [I,map]=rgb2ind(im,256);
      k=ii-0;
      close;
      if k==1;
        imwrite(I,map,filename,'gif','Loopcount',inf,'DelayTime',0.15);
      else
        imwrite(I,map,filename,'gif','WriteMode','append','DelayTime',0.15);
      end
    end
end
end

%% solution of least squares min |I_syn - B_selected * C|^2
%% C = inv(B'B)B'I;  I_reconstruct = B*C = B* inv(B'B)B' I_syn, we compute and store R=B inv(B'B)B' as reconstruction matrix
tic;
B=[];
for i=1:template.numSelected    
   Bi = filterTiling(filters{template.selectedFilter(i)}, sx, sy, template.selectedRow(i), template.selectedCol(i));
   B = [B,Bi(:)];
end
R=B*(pinv(B'*B)*(B'));

disp(['preparing the reconstruction matrix for least square: ' num2str(toc) ' seconds']);
%%%% save each chain of synthesized images respectively in the last iteration
EachChainSynFolder=[Synfolder 'EachChainSyn/'];
if exist(EachChainSynFolder)
    rmdir(EachChainSynFolder,'s'); 
    mkdir(EachChainSynFolder);
else
    mkdir(EachChainSynFolder);
end
%%%% reconstruct each chain of synthesized images respectively in the last iteration
EachChainReconstructionFolder=[Synfolder 'EachChainReconSyn/'];
if exist(EachChainReconstructionFolder)
    rmdir(EachChainReconstructionFolder,'s'); 
    mkdir(EachChainReconstructionFolder);
else
    mkdir(EachChainReconstructionFolder);
end


I_multi_chain_projected = zeros(nTileRow*sx, nTileCol*sy);
I_multi_chain = zeros(nTileRow*sx, nTileCol*sy);

index=0;
for j=1:nTileCol
   for i=1:nTileRow
     index=index+1;
     crop=I_syn(1+(i-1)*sx:sx+(i-1)*sx, 1+(j-1)*sy:sy+(j-1)*sy);
     gLow = min(crop(:));
     gHigh = max(crop(:));
     img = (crop-gLow)/(gHigh-gLow);       
     imwrite(img,[EachChainSynFolder 'original_chain-' num2str(index,'%04d') '.png']);
     I_multi_chain(1+(i-1)*sx: sx+(i-1)*sx, 1+(j-1)*sy: sy+(j-1)*sy) = img; 
     %% project the synthesized images onto the domain specified by the selected basis
     img_projected = reshape(R*crop(:),sx,sy);  % R = B*(inv(B'*B)*(B'));
     gLow = min(img_projected(:));
     gHigh = max(img_projected(:));
     img_projected = (img_projected-gLow)/(gHigh-gLow);       
     imwrite(img_projected,[EachChainReconstructionFolder 'reconstruct_chain-' num2str(index,'%04d') '.png']);
     I_multi_chain_projected(1+(i-1)*sx: sx+(i-1)*sx, 1+(j-1)*sy: sy+(j-1)*sy) = img_projected; 
   end
end
imwrite(I_multi_chain,[EachChainSynFolder 'raw.png']);
imwrite(I_multi_chain_projected,[EachChainReconstructionFolder 'reconstruct.png']);


% showing how many lambdas selected.
A=template.selectedLambdas;
numPosLambdas=length(A(A>0));
numNegLambdas=length(A(A<0));
numZeroLambdas=length(A(A==0));
disp([ num2str(numPosLambdas) ' positive lambdas left,' num2str(numNegLambdas) ' negative lambdas left and ' ...
    num2str(numZeroLambdas) ' zero lambdas left.']);

 

h = figure; 
plot(1:nIteration,SSD);
title('Convergence');
xlabel('Iterations');
ylabel('SSD');
saveas(h,[Synfolder 'plot.png'],'png');


