% main function for the sparse-FRAME model for general images 
% (1) sparsification by Forward selection 
% (2) fitting FRAME model by mean arg-MAX1map 


clear;
close all;

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


categoryName = 'hummingbird';  % tiger

inPath = ['Image/' categoryName '/'];
Synfolder='./synthesiedImage/';
if ~exist(Synfolder)
   mkdir(Synfolder);
end

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 = 14;
 
 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

else

  %% option 2: manually filters generation
  GarborScaleList=[1.2 0.7];
  DoGScaleList=[10];

end

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

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

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

localHalfx=20;
localHalfy=20;
thresholdFactor=0.01;
sigsq=10;

isLocalNormalize=false;
isSeparate=false;  % it matters, if isV2=false;
isToGenerateMovie=false;

% lambdaLearningRate = 0.01/sqrt(sigsq);
% epsilon = 0.009; 
% L= 10;
% % 
lambdaLearningRate = 0.1/sqrt(sigsq);  % 0.1
%epsilon = 0.03; %0.03; 
L= 10;
stepsize =1;
% multiple chains
useMultiChain=true;
nTileRow = 6; %nTileRow \times nTileCol defines the number of paralle chains
nTileCol = 6;
numSample = 3; % how many HMC calls for each learning iteration


%% Step 0: prepare filter and training images

nScaleGabor=length(GarborScaleList);

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

%nGaborFilter=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


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

if useMultiChain 
  prevSample = single(randn(sx*nTileRow,sy*nTileCol));  
else
  prevSample = single(randn(size(img)));
end

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));

% % % % [template, numSelected]= filtersSelectionV2(imgCell, GarborScaleList, DoGScaleList, nOrient, numSketch, locationShiftLimit, orientShiftLimit, sx, sy);
% % % % %numSketch=30;
% % % % template.selectedLambdas=single(zeros(1,numSelected));
% % % % template.numSelected = single(numSelected);



countMap = cell(numFilter,1);
for iFilter = 1:numFilter
    countMap{iFilter}=zeros(sx,sy);
end
countTotal=0;


for iter = 1:nIteration  
    
     disp( [ 'iteration: ' num2str(iter)]);  
     
     if mod(iter-1,interval)==0
        %%% forward selection
        [x_added, y_added, filter_added]=forwardSelection(filters, halfFilterSize, locationShiftLimit, rHat, prevSample, 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;
            
            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
            
            disp(['filter: ' num2str(filter_added), ' x: ' num2str(x_added), ' y: ', num2str(y_added)]);            
            
        end        
             
     end  
    
    
     tic 
     if useMultiChain         
      % the field of enlargeTemplate stores a big tiled template for multiple chains HMC sampling      
      %[rModel_selected, currSample]=multiChainHMC_sparse(template, filters, prevSample, epsilon, L, numSample, nTileRow, nTileCol);  
      [rModel_selected, currSample]=multiChainLangevin_sparse(template, filters, prevSample, stepsize, L, numSample, nTileRow, nTileCol);  
     else      
      
      %currSample=HMC_sparse(template, filters, prevSample, epsilon, L);
      currSample=Langevin_sparse(template, filters, prevSample, stepsize, L);
      [rModel_selected]= sparseFeatureExtraction(single(currSample),filters, template, true);
     end     
     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;
     
     prevSample=currSample;
           
     if mod(iter-1,interval)==0
       % save synthesied image
       gLow = min(currSample(:));
       gHigh = max(currSample(:));
       disp([ 'min: ' num2str(gLow) ' max: ' num2str(gHigh)]);
       disp([ 'SSD: ' num2str(SSD(iter))]);
       img = (currSample-gLow)/(gHigh-gLow);
       imwrite(img,[Synfolder num2str(iter,'%04d') '.png']);
     end
end


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


curr_projected = zeros(nTileRow*sx, nTileCol*sy);
currSample2 = zeros(nTileRow*sx, nTileCol*sy);

index=0;
for j=1:nTileCol
   for i=1:nTileRow
     index=index+1;
     crop=currSample(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']);
     currSample2(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']);
     curr_projected(1+(i-1)*sx: sx+(i-1)*sx, 1+(j-1)*sy: sy+(j-1)*sy) = img_projected; 
   end
end
imwrite(currSample2,[EachChainSynFolder 'raw.png']);
imwrite(curr_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');


