# Humboldt-Universität zu Berlin - Faculty of Mathematics and Natural Sciences - Strukturforschung / Elektronenmikroskopie

///////////////////////////////////////////////////////////
// This function finds the position of a small image within
// a bigger one:
///////////////////////////////////////////////////////////
// The images are passed by reference (images always are):
image findSmallInBig(image big,image small) {

number Nx,Ny,Nx2,Ny2,Nx3,Ny3;
Image realChi2Map

GetSize(big,Nx,Ny);
GetSize(small,Nx2,Ny2);

// allocate memory for the temporary image:
timg1 := ComplexImage("reciprocal space Image", 8, Nx, Ny)

// create a mask that has the same size as the small image:

// compute the FFT of mask:
// compute the FFT of big^2 and save it in timg1:
timg1 = big*big;
T_fft_C2C(timg1,timg1)

// assign a variable name to the sub-region of the mask that we want to copy
// the small image to:

timg1 = big
T_fft_C2C(timg1,timg1);
T_ifft_C2C(chi2Map,chi2Map);

// shift the center of the chi2-map to the center of the image:
chi2Map := T_shiftImageCenterComplex(chi2Map)
// scale and convert the chi2Map to a real image:
// realChi2Map = real(chi2Map)/(Nx*Ny) + sum(small*small)
realChi2Map = real(chi2Map)
// showimage(realChi2Map)

// free the allocated space again:
deleteImage(chi2Map)
deleteImage(timg1)

return realChi2Map;
}
///////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////
// This function simulates an image stack for testing this script
image simulateStack() {
number sx = 512,sy = 512;
number noise = 10;

Image stack = IntegerImage("Stack",2,1,sx,sy,3);
RealImage img1 = RealImage("Image 1",4,sx,sy);

img1 = (cos(2*Pi()*25*icol/sx)*sin(2*Pi()*15*irow/sy));
img1 = img1*sin(2*Pi()*2.73*sqrt((icol/sx-0.5)*(icol/sx-0.5)+(irow/sy-0.5)*(irow/sy-0.5)));
img1 = img1*img1;

// RealImage img2 = T_shiftImage(img1,shiftX,shiftY); // shifts with wrapping edges
RealImage img2 = T_shiftImageNoWrap(img1,20,50); // shifts without wrapping edges

// add some noise to images:
stack[0,0,0,sx,sy,1] = 1000*(img1+PoissonRandom(noise*img1));
stack[0,0,1,sx,sy,2] = 1000*(img2+PoissonRandom(noise*img2));
img2 = T_shiftImageNoWrap(img1,10,-30); // shifts without wrapping edges
stack[0,0,2,sx,sy,3] = 1000*(img2+PoissonRandom(noise*img2));

setname(stack,"Simulated Stack");
showimage(stack);
deleteimage(img1);
deleteimage(img2);
return stack
}
//////////////////////////////////////////////////////////

number width,height,sx,sy,sz,z,w1,w2,h1,h2;
number top,left,bottom,right
number top2,left2,bottom2,right2
number mval,mx,my;
RealImage stack,imgBig,imgSmall,imgSum,imgShifted,chi2Map,imgTotal,imgScale,imgScaleOld
ROI subRegion

stack = getFrontImage();
// stack = simulateStack();

stack.Get3DSize(sx,sy,sz);

// exit this script, if the stack is just a simple image:
if (sz<2) throw("This stack only has one image - no alignment necessary!\n")

ImageDisplay imgDisp = getFrontImage().ImageGetImageDisplay(0)
number ROIcount = ImageDisplayCountROIs(imgDisp);

if (ROIcount == 0) {
subRegion = CreateROI();
Result("No region of interest has been selected - will use central region of image 2\n")
ROISetRectangle(subRegion,floor(0.25*sy),floor(0.25*sx),ceil(0.75*sy),ceil(0.75*sx));
}
subRegion = imgDisp.ImageDisplayGetROI(0)
subRegion.ROIGetRectangle(top, left, bottom, right)
// result("region: "+top+" , "+left+" , "+bottom+" , "+right+"\n")

imgTotal = stack[0,0,0,sx,sy,1];
imgScaleOld = exprsize(sx,sy,1)

// initialize the big image with the top image of the stack:
// loop through the remaining images of the stack:
for (z=1;z<sz;z++) {

imgSmall := stack[left,top,z,right,bottom,z+1]
imgShifted := stack[0,0,z,sx,sy,z+1]
imgBig := imgTotal

chi2Map := findSmallInBig(imgBig,imgSmall)

// find the relative image shift:
mval = min(chi2Map,mx,my)
chi2Map.getSize(width,height)

mx = mx-floor(width/2);
my = my-floor(height/2);

result("Shift of image "+z+": "+(mx+left)+", "+(my+top)+" (val="+mval+")\n");

///////////////////////////////////////////////////////////////////////
// We must now add the 2 images according to the relative image shift:

imgShifted.getSize(w1,h1);
imgBig.getSize(w2,h2);

// In order to add the 2 images, we must make sure that we keep track of where
// the edges are
top2 = min(-my-top,0)
left2 = min(-mx-left,0)
bottom2 = max(h1+abs(my+top),h2-top2)
right2 = max(w1+abs(mx+left),w2-left2)
// The sum image must have the maximum size:
imgSum := RealImage("Sum image",8,right2,bottom2)
imgScale := RealImage("Sum image",8,right2,bottom2)

// result("top: "+(-top2)+" , "+(-left2)+" , "+(h2-top2)+" , "+(w2-left2)+"("+right2+", "+bottom2+")\n")
// add both image to the sum image
imgSum[-top2,-left2,h2-top2,w2-left2] += imgBig*imgScaleOld
imgScale[-top2,-left2,h2-top2,w2-left2] += imgScaleOld

imgSum[-top2-my-top,-left2-mx-left,-top2-my+h1-top,-left2-mx+w1-left] += imgShifted
imgScale[-top2-my-top,-left2-mx-left,-top2-my+h1-top,-left2-mx+w1-left] += 1

deleteImage(imgTotal)
deleteImage(imgScaleOld)
imgTotal = tert(imgScale>0,imgSum/imgScale,0)
imgScaleOld = imgScale;

}
//////////////////////////////////////////////////////
setname(imgTotal,"Sum of "+getname(stack))
showimage(imgTotal)
setname(imgScale,"Overlap image")
showimage(imgScale)