% main function for the fast sparse-FRAME model for general images 
% (1) sparsification by Matching Pursuit and arg-max handling deformation 
% (2) fitting FRAME model by mean arg-MAX1map 
% (3) sparse structure avoids the convolution, which makes computation fast
% (4) least squares to project the synthesized images into sparse domain (C implementation)
% (5) orthogonal matching pursuit is used in sparsification
% fast sparse-FRAME_MPA (MPA: matching pursuit algorithm)

clear;
close all;

mex CgetMAX1.c;
mex CgetMAX1_DoG.c;
mex OrthogonalMatchingPursuit.c;
mex sparseFeatureExtraction.c;
mex grad_U_sparse.c
mex ComputeProjection.c

category='flamingo';
inPath = ['Image/' category '/'];
Synfolder='./synthesizedImage/';
if ~exist(Synfolder)
   mkdir(Synfolder);
end

%% 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;  % 10 if local normalization

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
%DoGScaleList=[];
%% option 2: manually filters generation
% GarborScaleList=[0.4, 0.8, 0.2];
% DoGScaleList=[10];

nOrient = 15;
locationShiftLimit=3;  % arg-max
orientShiftLimit=1;    % arg-max
sx=127;
sy=100;
numSketch=300;  % 300


localHalfx=20;
localHalfy=20;
thresholdFactor=0.01;
scaleFactor=0.6;
nIteration=100;
sigsq=10;

isLocalNormalizeFrame=false;
isV2=false;  % it matters, if isLocalNormalize=true;  %% V1 is Wenze's version, V2 is Zhangzhang's version
isSeparate=false;  % it matters, if isV2=false;
isToGenerateMovie=false;

isLocalNormalizeFilterSelection=false;

% lambdaLearningRate = 0.01/sqrt(sigsq);
% epsilon = 0.009; 
% L= 10;
% % 
lambdaLearningRate = 0.1/sqrt(sigsq);  % 0.1
epsilon = 0.03; 
L= 10;

% 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

% OMP or MP
isOrthogonal = true;
subsamplingOMP = 1;
%% 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

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
FRAMEfilterBankInfo.Corr = Corr;
FRAMEfilterBankInfo.halfFilterSize = halfFilterSize;
FRAMEfilterBankInfo.filters = filters;

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 isLocalNormalizeFrame
    
    if isV2  % V2 use Separate local normalization
      for iScale=1:nScaleGabor 
        h = halfFilterSize(1+(iScale-1)*(2*nOrient));
        S1(1+(iScale-1)*(2*nOrient) : nOrient+(iScale-1)*(2*nOrient))= LocalNormalizeV2(S1(1+(iScale-1)*(2*nOrient):nOrient+(iScale-1)*(2*nOrient)),h,localHalfx,localHalfy,thresholdFactor);
        S1(nOrient+1+(iScale-1)*(2*nOrient): 2*nOrient+(iScale-1)*(2*nOrient)) = LocalNormalizeV2(S1(nOrient+1+(iScale-1)*(2*nOrient): 2*nOrient+(iScale-1)*(2*nOrient)),h,localHalfx,localHalfy,thresholdFactor);
      end
      % DoG
      for iScale=1:nScaleDoG
         ind  = 2*nOrient*nScaleGabor+iScale;
         h = halfFilterSize(ind);       
         S1(ind)=LocalNormalizeV2(S1(ind),h,localHalfx,localHalfy,thresholdFactor);     
      end
      
    else
    
    %%%%% V1 
    if isSeparate        
      for iScale=1:nScaleGabor  
        h = halfFilterSize(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,localHalfx,localHalfy,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,localHalfx,localHalfy,thresholdFactor);
      end
      % for DoG
      for iScale=1:nScaleDoG  
          ind  = 2*nOrient*nScaleGabor+iScale;
          h = halfFilterSize(ind);       
          [S1(ind)]=LocalNormalize(S1(ind),[],h,localHalfx,localHalfy,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))]= 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,localHalfx,localHalfy,thresholdFactor);
        end
        % for DoG
        for iScale=1:nScaleDoG 
           ind  = 2*nOrient*nScaleGabor+iScale;
           h = halfFilterSize(ind);       
           [S1(ind)]=LocalNormalize(S1(ind),[],h,localHalfx,localHalfy,thresholdFactor);                
        end        
    end
    %%%%%% end of V1
    
    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));

%% filter selection
ticID=tic;
%[indEliminated, numSelected]= filtersSelection(imgCell, GarborScaleList, DoGScaleList, nOrient, numSketch, locationShiftLimit, orientShiftLimit, sx, sy);

[template, numSelected]= OrthogonalFiltersSelection(imgCell, GarborScaleList, DoGScaleList, nOrient, numSketch, locationShiftLimit, ...
    orientShiftLimit, sx, sy, isOrthogonal,subsamplingOMP,isLocalNormalizeFilterSelection,scaleFactor,thresholdFactor,FRAMEfilterBankInfo);
template.selectedLambdas=single(zeros(1,numSelected));
template.numSelected = single(numSelected);
disp(['finished filters selection: '  num2str(toc(ticID)) ' seconds']);


if useMultiChain
    % the field of enlargeTemplate stores a big tiled template for multiple chains HMC sampling 
    template.enlargedTemplate = enlargingTemplate(nTileRow, nTileCol, sx, sy, template);
 
end

rHat_selected = zeros(1,numSelected);
for i=1:numSelected
   rHat_selected(i)=rHat{template.selectedFilter(i)}(template.selectedRow(i),template.selectedCol(i));
end


for iter = 1:nIteration
    disp( [ 'iteration: ' num2str(iter)]);
    tic  
  
     if useMultiChain    
      [rModel_selected, currSample]=multiChainHMC_sparse(template, filters, prevSample, epsilon, L, numSample, nTileRow, nTileCol);  
     
     else        
      
      currSample=HMC_sparse(template, filters, prevSample, epsilon, 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;
           
    % 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


%% solution of least squares min |I_syn - B_selected * C|^2


%%%% 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 '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'));
     [invBB, BI, coefficient, img_projected]=ComputeProjection(Corr,template,single(halfFilterSize),filters,crop);
     curr_projected(1+(i-1)*sx:sx+(i-1)*sx, 1+(j-1)*sy:sy+(j-1)*sy)=img_projected;
     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']);
%%%%%%

if isToGenerateMovie
filename='movie.gif';
for i=1:nIteration
    str=[ num2str(i,'%04d') '.png'];
    Img=imread(str);
    figure(i)
    imshow(Img,'border','tight','initialmagnification','fit');
    set(gcf,'Position',[50,50,130,130]);     % the dimension of the gif should be lareger than this default setting.
    axis normal;
    frame=getframe(i);
    im=frame2im(frame);
    [I,map]=rgb2ind(im,256);
    k=i-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

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


