% alignment  by sparse FRAME (generative bossting) 
 close all;
 clear;


mex Ccopy.c; % copy around detect location
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;
mex sparseFRAME_SUM2_logZ_partial.c;
mex mexc_TemplateAffineTransform.cpp;


% Full List of folders of alignment tasks
%
% (1) maple
% (2) lion
% (3) flower
% (4) wolf
% (5) cougar
% (6) hummingbird
% (7) goose
% 

for task_id = 2 % choose your categoriy ID shown in above, for all of them, please use 1:7 

    for seed = 1
        
        %%%%%%%%%%%%%%%%%%%%%%   
        %%% The algorithmic tuning parameters are set by function "getParametersForExp" based on different tasks, while others shown in this file are fixed.   
        %%%%%%%%%%%%%%%%%%%%%% 
        para = getParametersForExp(task_id);
        disp(['we will start the clustering task ' num2str(task_id) ' : ' para.foldName]);
        
        %%%% CPU paralleled for computing SUM2, MAX2, SUM3, and MAX3
        noWorkers=4;    % Set number of workers, which depends on how many cores in your cpu
        runParallel=1;  % Set variable to know when to run parallel
        
        % set seed for random numbers generation
        rng(seed);        
              
        resizeByScalingFactor = para.resizeByScalingFactor; % if true, resize all training images without changing image aspect ratio, otherwise, resize them into the same size as the template size (sx, sy)
        isToUseNewFilterBankAtLastIteration = para.isToUseNewFilterBankAtLastIteration;
       % resizeFactor = para.resizeFactor; % parameter about scaling of the training images
        
        %%%% EM clustering parameters
        numEMIteration = para.numEMIteration;        
        %%%%%
        
        categoryName=para.foldName;
        %%%%% parameters for multiple selection (cell size) or deformable parts %%%%%%%%%%%%%%%%%%%
        nPartCol=para.nPartCol;
        nPartRow=para.nPartRow;
        part_sx=para.part_sx;
        part_sy=para.part_sy;
        
        gradient_threshold_scale=0.8;     %  we use addaptive threshold in the multiple selection.  The threshold is equal to (gradient_threshold_scale * maximum gradient).
         
        %%%%% parameters for large template
        sx=nPartRow*part_sx;   % template size x
        sy=nPartCol*part_sy;   % template size y
        rotateShiftLimit = para.rotateShiftLimit;   % template rotation  from -rotateShiftLimit to rotateShiftLimit, eg. (1)-2:2 if rotateShiftLimit=2 (2)0 is without rotation
        rotationRange = -rotateShiftLimit:rotateShiftLimit;
        numRotate = length(rotationRange);
        RatioDisplacementSUM2 = para.ratioDisplacementSUM2;   % default=0. Compute all values in SUM3 map
        
        nOrient = 16;
        locationShiftLimit=para.locationShiftLimit;  % arg-max
        orientShiftLimit=para.orientShiftLimit;    % arg-max
        %%%%%%
        
        
        %%%%%%% parameters about resolution
        numResolution = para.numResolutionEnlarge + para.numResolutionShrink +1;
        originalResolution = para.numResolutionShrink +1 ;    % original resolution is the one at which the imresize factor = 1.  the middle one is by default
        scaleStepSize = para.scaleStepSize;  
        %%%%%%%
        
               
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        interval = 1; % 5 control how many iterations we need to wait to select next wavlet
        numWavelet = para.numWavelet;   % default 300, 370
        %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
        
        
        %%%%% parameters for Gibbs sampler
        sigsq=10;
        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=-25:3:25;  % range of c
        lambdaLearningRate = 0.1/sqrt(sigsq);  % 0.1
        %%%%
        
        %%%%%%%%%%%%%% multiple chains
        useMultiChain=true;
        nTileRow = 8; %nTileRow \times nTileCol defines the number of paralle chains
        nTileCol = 8;
        if useMultiChain==false
            nTileRow = 1;
            nTileCol = 1;
        end
        %%%%%%%%%%%%%
        
        %%% local normalization parameters
        isLocalNormalize= para.isLocalNormalize; %false; %
        isSeparate=false;
        localNormScaleFactor=2; %0.5, 1, 2, 3, or 4, we suggest 2;
        thresholdFactor=0.01;
        %%%%%%%%%%%%%%%%%%%%%%%%%%
        
        inPath = ['./positiveImages/' para.foldName];
        cachePath = './feature';
        resultPath = ['./result_task_' num2str(task_id) '_seed_' num2str(seed) '_' categoryName];
        
        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
        
        
        %% Step 0: prepare filter, training images,and filter response on images
        GaborScaleList = para.GaborScaleList;
        DoGScaleList = para.DoGScaleList;
        
        if isLocalNormalize
            DoGScaleList=[];
        end
        
        nScaleDoG=length(DoGScaleList);
        nScaleGabor=length(GaborScaleList);
        
        
        %%% create filter bank
        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));
        Corr = CorrFilterFrame(filters); %calculate correlation among filters
        
        %%%%%%%% optional: for drawing sketch in the learning process
        [allFilterR, ~, filterSymbol] = MakeFilterBank(GaborScaleList, 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
        %%%%%%%%%%
        
        files = dir(fullfile(inPath,'*.jpg'));
        numImage=length(files);
        disp(['start filtering']);
        
        
        for iImg = 1:numImage
            copyfile(fullfile(inPath,files(iImg).name),fullfile(resultPath,'img',files(iImg).name));
            imageOriginal = imread(fullfile(inPath,files(iImg).name));
                        
            if resizeByScalingFactor % resize the training images by a scaling factor
                  [x,y,~]=size(imageOriginal); 
                  if x<y
                     resizeFactor = sx/x;                     
                  else
                     resizeFactor = sy/y;   
                  end
                  img = imresize(imageOriginal, resizeFactor);
            else   % resize the training images into bounding box or template (sx, sy)
                  img = imresize(imageOriginal, [sx, sy]);
            end
           
            if ndims(img)==3
                img=rgb2gray(img);
            end
            img = im2single(img);
                        
            allSizex = zeros(1, numResolution);
            allSizey = zeros(1, numResolution);
            ImageMultiResolution = cell(1, numResolution);
            
            for resolution=1:numResolution
                resizeFactor = 1.0 + (resolution - originalResolution)*scaleStepSize;
                img2 = imresize(img, resizeFactor, 'nearest');  % images at multiple resolutions
                
                img2 = img2-mean(img2(:));
                img2 = img2/std(img2(:))*sqrt(sigsq);
                
                ImageMultiResolution{resolution} = img2;
                
                [sizex, sizey] = size(ImageMultiResolution{resolution});
                allSizex(resolution) = sizex;
                allSizey(resolution) = sizey;
            end
            
            % compute MAX1 map for images in different resolutions
            disp(['======> start filtering and maxing image ' num2str(iImg)]);
            tic
            [SUM1mapFind, MAX1mapFind] = applyfilterBank_MultiResolution_sparseV4(ImageMultiResolution, filters, halfFilterSizes, nOrient,...
                locationShiftLimit,orientShiftLimit,isLocalNormalize,isSeparate,localNormScaleFactor,thresholdFactor,nScaleGabor,nScaleDoG, sqrt(sigsq));  % if using local normalization, the last parameter is important.
            
            mapName = fullfile(cachePath,['SUMMAXmap-image' num2str(iImg) '.mat']);
            save(mapName, 'imageOriginal', 'ImageMultiResolution','SUM1mapFind', 'MAX1mapFind','allSizex', 'allSizey');
            
            disp(['filtering time: ' num2str(toc) ' seconds']);
            
        end
        
        
        %% Prepare variables for EM
        
        %%%%%%% initialization of alignment
        model=struct('rHat',[],'template',[],'logZ',[], 'S2T', []);  % structure to store information of cluster
         
        %%% initialize the observed statistics by setting zeros
        for iFilter = 1:numFilter
            model.rHat{iFilter}=zeros(sx, sy,'single');
        end
        
        t=0; % index of image in the cluster, as well as the number of images in cluster
        for iImg = 1:numImage
            tic
            
            imageLoaded = load(fullfile(cachePath,['SUMMAXmap-image' num2str(iImg)]));
            
            % we initialize the alignment by cropping patch form original resolution image, with 0 level rotation and center (img_x/2. img_y/2)
            rot_init = 0;    % 0 level
            ind_init = originalResolution;  % original resolution
            Fx_init = floor(imageLoaded.allSizex(originalResolution)/2);  % center x
            Fy_init = floor(imageLoaded.allSizey(originalResolution)/2);   % center y
            
            cropedImage=single(zeros(sx, sy));
            Ccopy(cropedImage, single(imageLoaded.ImageMultiResolution{ind_init}), Fx_init, Fy_init, floor(sx/2), floor(sy/2), sx, sy, imageLoaded.allSizex(ind_init), imageLoaded.allSizey(ind_init), -rot_init*pi/nOrient);
         %   Ccopy(cropedImage, single(imageLoaded.ImageMultiResolution{Mind}), MFx, MFy, floor(sx/2), floor(sy/2), sx, sy, imageLoaded.allSizex(Mind), imageLoaded.allSizey(Mind), -Mrot*pi/nOrient);

            % optinal: output cropped images for iteration 0
            savingFolder0=fullfile(resultPath,['iteration0/detectedCropped/']);
            if ~exist(savingFolder0)
                mkdir(savingFolder0);
            end
            
            gLow = min(cropedImage(:));
            gHigh = max(cropedImage(:));
            img_tem = (cropedImage-gLow)/(gHigh-gLow);
            imwrite(img_tem,fullfile(savingFolder0,['detected-img-' num2str(iImg,'%04d') '.png']));
            
            cropedImage = cropedImage - mean(cropedImage(:));
            cropedImage = cropedImage/std(cropedImage(:))*sqrt(sigsq);
            
            % compute feature map to learn
            [~, MAX1] = applyfilterBank_MultiResolution_sparseV4({cropedImage}, filters, halfFilterSizes, nOrient, locationShiftLimit,orientShiftLimit,...
                isLocalNormalize,isSeparate,localNormScaleFactor,thresholdFactor,nScaleGabor,nScaleDoG, sqrt(sigsq));  % if using local normalization, the last parameter is important.
            
            % sum over the observed statistics (within cluster)
            for iFilter = 1:numFilter
                model.rHat{iFilter}=model.rHat{iFilter}+MAX1{iFilter}/numImage;
            end
            
            disp(['cropping time for image ' num2str(t) ': ' num2str(toc) ' seconds']);
        end% iImage
        
        it = 0; % used in fuction ShowResults.m
        ShowResults;
        
        %% EM iteration
        for it = 1 : numEMIteration
            
            %%%% M-step
            disp(['M-step of iteration ' num2str(it)]);
            
            %% learn model for each cluster
            disp('Learning: learning sparse FRAME model:');
            
            savingFolder=fullfile(resultPath,['iteration' num2str(it) '/']);
            if ~exist(savingFolder)
                mkdir(savingFolder);
            end
            
            clusterSavingFolder=fullfile(savingFolder,'model/');
            if ~exist(clusterSavingFolder)
                mkdir(clusterSavingFolder);
            end
            
            [template, currSample, logZ]=sparseFRAMElearnGibbs_multipleSelection(filters, nScaleGabor, nScaleDoG, nOrient, filterSymbol, half, model.rHat, ...
                sx, sy, halfFilterSizes, locationShiftLimit, nTileRow, nTileCol, lambdaLearningRate, numFilter, numWavelet, interval, Corr, threshold_corrBB, c_val_list, ...
                lower_bound_rand, upper_bound_rand, nPartCol, nPartRow, part_sx, part_sy, gradient_threshold_scale, clusterSavingFolder);
            
            model.template=template;
            model.currSample=currSample;
            model.logZ=single(logZ);
            
            tic            
            [model.S2T] = transformTemplate(sx, sy, rotateShiftLimit, nOrient, model.template, nScaleGabor);
            disp(['transforming the learned template : ' num2str(toc) ' seconds']);
            
            %%%% E-step
            disp(['E-step of iteration ' num2str(it)]);
            
            detectedCroppedSavingFolder=fullfile(savingFolder, 'detectedCropped/');
            if ~exist(detectedCroppedSavingFolder)
                mkdir(detectedCroppedSavingFolder);
            end
            
            boundingBoxSavingFolder=fullfile(savingFolder, 'boundingBox/');
            if ~exist(boundingBoxSavingFolder)
                mkdir(boundingBoxSavingFolder);
            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(numImage/noWorkers);
%             Assignments_left=mod(numImage,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];
%             
          
 %           spmd
                
 %              for iImg= startIndex(labindex):endIndex(labindex)
                for iImg = 1 : numImage
                    mapName = fullfile(cachePath,['SUMMAXmap-image' num2str(iImg)]);
                    imageLoaded = load(mapName);
                    [~] = detectObject(imageLoaded, model, iImg, numFilter, sx, sy, nOrient, rotateShiftLimit, detectedCroppedSavingFolder, boundingBoxSavingFolder, RatioDisplacementSUM2);      
                end
                
%            end
            
            %%%%%%%%%%%%%%%%%
            
            ShowResults;  % show the member before learning   
            
            %%
            disp('Cropping images and preparing feature maps for next iteration of learning:');               
           
          
            if isToUseNewFilterBankAtLastIteration  &&  it == numEMIteration            
                % re-learn the aligned images without using local normalization or with a new fiter bank in order to get better visualizations
                learnAtLastIteration;                
            else    
                
                %%% initialize the observed statistics by setting zeros
                for iFilter = 1:numFilter
                    model.rHat{iFilter}=zeros(sx, sy,'single');
                end
            
                for iImg = 1:numImage
                    tic      
                    % load cropped images for training
                    imageLoaded = load(fullfile(detectedCroppedSavingFolder,['detected-img-' num2str(iImg,'%04d')]), 'cropedImage');

                    cropedImage = imageLoaded.cropedImage;

                    cropedImage = cropedImage - mean(cropedImage(:));
                    cropedImage = cropedImage/std(cropedImage(:))*sqrt(sigsq);


                    % compute feature map to learn
                    [~,MAX1] = applyfilterBank_MultiResolution_sparseV4({cropedImage}, filters, halfFilterSizes, nOrient, locationShiftLimit,orientShiftLimit,...
                        isLocalNormalize,isSeparate,localNormScaleFactor,thresholdFactor,nScaleGabor,nScaleDoG, sqrt(sigsq));  % if using local normalization, the last parameter is important.


                    % sum over the observed statistics (within cluster)
                    for iFilter = 1:numFilter
                        model.rHat{iFilter}=model.rHat{iFilter}+MAX1{iFilter}/numImage;
                    end

                    disp(['cropping time for image ' num2str(iImg) ': ' num2str(toc) ' seconds']);
                end% iImage

            end
            
            
        end %it        
            
    end
end