%% Accumulate the SUM2 maps and compute MAX2 scores

%% generate rotated templates (currently only deal with one scale)
count = 0;
rotationRange = -3:3; % only local rotation
for templateScaleInd = 0:0 % large scale change
    for rotation = rotationRange 
        for rowScale = 2.^[0] % small scale change and aspect ratio change
            for colScale = 2.^[0]
                count = count + 1;
                templateTransform{count} = [templateScaleInd rowScale colScale rotation];
            end
        end
    end
end
nTransform = count;

% the actual transformation computation
TransformedTemplate = cell(nTransform,1);

maxSquareSize = 0;

learnedTemplateName = ['working/learnedTemplate' '_' num2str(startx) '_' num2str(starty)];
load(learnedTemplateName); 
selectedScale = zeros(1,length(selectedx));
numElement = length(selectedx);
% center the template
selectedx = selectedx - partSize/2;
selectedy = selectedy - partSize/2;
for iT = 1:nTransform
    templateScaleInd = templateTransform{iT}(1);
    rowScale = templateTransform{iT}(2);
    colScale = templateTransform{iT}(3);
    rotation = templateTransform{iT}(4);
    [tmpSelectedRow tmpSelectedCol tmpSelectedOri tmpSelectedScale] = ...
        mexc_TemplateAffineTransform( templateScaleInd, rowScale,...
        colScale, rotation, single(selectedx), single(selectedy),...
        single(selectedOrient), single(selectedScale), numOrient );
    TransformedTemplate{iT}.selectedRow = tmpSelectedRow;
    TransformedTemplate{iT}.selectedCol = tmpSelectedCol;
    TransformedTemplate{iT}.selectedOri = tmpSelectedOri;
    TransformedTemplate{iT}.selectedScale = tmpSelectedScale;
    TransformedTemplate{iT}.selectedLambda = single(selectedlambda);
    newSelectedLogZ = zeros(1,numElement);
    TransformedTemplate{iT}.selectedLogZ = single(selectedLogZ);
    TransformedTemplate{iT}.subsampleM1 = subsample;
    TransformedTemplate{iT}.sym = displayTemplate( ...
        tmpSelectedRow,tmpSelectedCol,tmpSelectedOri, tmpSelectedScale, ones(numElement,1),...
        allSymbol,numOrient,subsample);
    if size(TransformedTemplate{iT}.sym,1) > maxSquareSize
        maxSquareSize = size(TransformedTemplate{iT}.sym,1);
    end
end

% display the transformed templates

tmpTT = cell(1,nTransform);
for iT = 1:nTransform
    % inscribe into a fixed-size square
    tmpTT{iT} = zeros(maxSquareSize,maxSquareSize);
    squareSize = size(TransformedTemplate{iT}.sym,1);
    margin = floor((maxSquareSize - squareSize)/2);
    tmpTT{iT}(margin+(1:squareSize),margin+(1:squareSize)) = -single(TransformedTemplate{iT}.sym);
end
im = displayImages(tmpTT,4*2,50,50);
imwrite(im,sprintf('output/transformedTemplates%d_%d.png',startx,starty));
im_TransformedTemplates = im; % for illustration


% -------------------------------------------------
% template rotation ends
% -------------------------------------------------

aveMAX2score = 0;
rowTranslationPerImage = zeros(numImage,1);
colTranslationPerImage = zeros(numImage,1);
rotationPerImage = zeros(numImage,1);
for img = 1:numImage % for each image

    multipleResolutionImageName = ['working/multipleResolutionImage' num2str(img)];
    load(multipleResolutionImageName); 
    
      storeExponentialModelName = ['storedExponentialModel'];   
      load(storeExponentialModelName);

      SUM1MAX1mapName = ['working/SUM1MAX1map' 'image' num2str(img)]; 
      
      load(SUM1MAX1mapName, 'SUM1map', 'MAX1map', 'ARGMAX1map', 'M1RowShift', 'M1ColShift', 'M1OriShifted');
      
		MAX2score = -1000;
		for iRotation = 1:length(TransformedTemplate) % for each rotation
			subsampleS2 = 1;
			SUM2map = mexc_ComputeSUM2( numOrient, MAX1map, TransformedTemplate(iRotation), subsampleS2 );
			if rotationRange(iRotation)==0
				im_SUM2map = zeros(size(SUM2map{1}));
				top = startx + partSize/2 - partLocRange;
				bottom = startx + partSize/2 + partLocRange;
				left = starty + partSize/2 - partLocRange;
				right = starty + partSize/2 + partLocRange;
				im_SUM2map( (top:bottom),(left:right) ) = SUM2map{1}((top:bottom),(left:right));
			end

			%% compute MAX2 scores for all resolutions in a local area
			% modified: search in a local area only, and in a single scale
			top = startx + partSize/2 - partLocRange;
			bottom = startx + partSize/2 + partLocRange;
			left = starty + partSize/2 - partLocRange;
			right = starty + partSize/2 + partLocRange;
			top = max(1,top); left = max(1,left);
			bottom = min(bottom,templateSize(1));
			right = min(right,templateSize(2));
			tmp = SUM2map{1}((top:bottom),(left:right));
			[m idx] = max(tmp(:));
			if m > MAX2score
				MAX2score = m;
				Fy = ceil(idx/(bottom-top+1));
				Fx = idx - (Fy-1) * (bottom-top+1);
				Fx = Fx + top - 1;
				Fy = Fy + left - 1;
				Fx = Fx - floor(partSize/2);
				Fy = Fy - floor(partSize/2);
				Fo = iRotation;
			end
        end
        
        rowTranslationPerImage(img) = Fx + floor(partSize/2);
        colTranslationPerImage(img) = Fy + floor(partSize/2);
        rotationPerImage(img) = rotationRange(Fo);

		% only deal with single resolution
		translatedTemplate = MGetBack( Fx+floor(partSize/2),Fy+floor(partSize/2), MAX1map, ARGMAX1map,...
		  M1RowShift, M1ColShift, M1OriShifted, TransformedTemplate{Fo},...
		  allSymbol,numOrient );
		
        srcIm = inscribe(templateSize,J0,255);
        im_Overlay = srcIm + translatedTemplate*100;
		if img == starting || img == 1 || img == 2
			towrite = im_Overlay;
			towrite = uint8(255 * (towrite-min(towrite(:)))/(max(towrite(:))-min(towrite(:))));
			imwrite(towrite,sprintf('output/%s_part_detection_image%d_part%d_%d.png',category,img,startx,starty)); % overlaid sketch at optimal resolution
		end

        %% Crop detected window for learning
        denseX = -floor(partSize/2) + (1:partSize);
        denseY = denseX;
        count = 0;
        inRow = zeros(length(denseX)*length(denseY),1,'single');
        inCol = zeros(length(denseX)*length(denseY),1,'single');
        for y = denseY
            for x = denseX
                count = count+1;
                inRow(count) = x;
                inCol(count) = y;
            end
        end
		
		tScale = 0; rScale = 1; cScale = 1; inO = zeros(numel(inRow),1,'single'); inS = zeros(numel(inRow),1,'single');
		rotation = rotationRange(Fo);
		[outRow, outCol] = ...
            mexc_TemplateAffineTransform(tScale,rScale,cScale,rotation,inRow,inCol,inO,inS,numOrient);
            
        outRow = outRow * reflection; % deal with reflection

        % crop S1 map
        tScale = 0; destHeight = partSize; destWidth = partSize; nScale = 1; reflection = 1;
        SUM1mapLearn(img,:) = mexc_CropInstance( SUM1map,...
            Fx+floor(partSize/2),...
            Fy+floor(partSize/2),...
            rotation,tScale,reflection,outRow,outCol,numOrient,numScale, partSize, partSize );

        %% Crop detected image patch for visualization
        cropped(img) = mexc_CropInstance( {single(srcIm)},...
            Fx+floor(partSize/2),...
            Fy+floor(partSize/2),...
            0,0,reflection,outRow,outCol,1,1, partSize, partSize );

        aveMAX2score = MAX2score + aveMAX2score;
   
    %% illustration
    if GraphicsIsOn
        clf('reset')
        subplot(4,3,1:3);
        text0 = ['Iter ' num2str(it) '. #E step. Using the current template#to detect instances. '...;
            'Detection is done#by a local search for the best match.'];
        imshow(text2im(text0,700,80,[1 1 1]));
        subplot(4,3,4); % original image
        imshow(inscribe(templateSize,J0,0),[]);
        title(sprintf('image %d',img));
        pause(0.5);
        subplot(4,3,5); % SUM1 map
        imshow(min(6,SUM1map{1}),[]);
        title('SUM1 map 1st orientation');
        pause(0.5);
        subplot(4,3,7:9); % rotated templates
        imshow(im_TransformedTemplates,[]);
        title('rotated current templates');
        pause(0.5);
        subplot(4,3,10); % S2 map
        imshow(im_SUM2map,[]);
        title('SUM2 map for non-rotated template');
        pause(0.5);
        subplot(4,3,11); % template matching overlay
        imshow(im_Overlay,[]);
        title(['matched ' num2str(rotationRange(Fo))]);
        pause(0.5);
        subplot(4,3,12); % cropped image patch
        imshow(cropped{img},[]);
        title('cropped/rotated patch');
        pause(0.5);
    end
    
    % do not save for now
    % save( sprintf('working/SUM2MAX2map_part%d_image%d.mat',count_PartLoc,img), 'SUM2map' ); 
end

aveMAX2score = aveMAX2score / numImage;
save( sprintf('working/LearnedPart_%d_%d.mat',startx,starty), 'aveMAX2score', 'TransformedTemplate',...
    'rowTranslationPerImage','colTranslationPerImage','rotationPerImage');
