//This is useful to correct for z- and xy- drift across timepoints.
//Uses turboreg plugin to align a hyperstack in z and/or in xy.
//The alignment is based on a max projection of the xz reslice for
//the z-alignment, and a regular xy max projection for xy alignment
//Currently only can use one of the channels to align.
//This will not align the slices of a particular individual time point,
//for that use Align Z forEachTimePoint

	title=getTitle();
	IID=getImageID();
	getDimensions(w,h,chs,sls,frms);
	bits=bitDepth();
	if(bits!=24)textLuts=getTextLUTs();

	setPasteMode("Copy");
	Zname=title + "-Zalign";

	if(sls<2 && frms< 2) exit("Requires hyperstack with multi z and frames");
	if(bits==24){
		charray=newArray("all","red","green","blue");
	}else{
		charray=newArray(chs);
		for(i=0;i<chs;i++){charray[i]=toString(i+1);}
	}

	Dialog.create("Choices");
	Dialog.addCheckbox("Z-align?",true);
	Dialog.addCheckbox("XY-align?",true);
 	Dialog.addChoice("Ch:", charray, "2");
	Dialog.addChoice("Mode", newArray("Auto","Manual with suggestion", "Manual"), "Auto");
	Dialog.addCheckbox("Auto bright?", false);
	Dialog.addCheckbox("Auto accept?", true);
	Dialog.addMessage("Manual means line up the two stacks and push Space");
	Dialog.show();

	doz=Dialog.getCheckbox;
	doxy=Dialog.getCheckbox;
	chchoice=Dialog.getChoice();
	mode=Dialog.getChoice();
	autobright=Dialog.getCheckbox();
	autoaccept=Dialog.getCheckbox();
	accepted=true;

	chcrgb="";
	if(chs==1) {chcrgb=chchoice; chchoice="1";}
	
	automode=false;
	manualmode=false;

	if(mode=="Auto") automode=true;
	if(mode=="Manual") manualmode=true;

	if(automode && autoaccept && doz) setBatchMode(true);

     if(doz){
	x=0;
	y=frms;

	zReg=newArray(y);
	zReg[0]=0;

	
	ctot=0;
	dytot=0;
	dxtot=0;
		
	for(i=x; i<y; i++){

		selectImage(IID);
		run("Duplicate...", "title=TimeS"+ (i+1)+" duplicate channels="+chchoice+" slices=1-"+sls+" frames="+(i+1));
		if(bits==24){
			if(chcrgb!="all") {
				if(chcrgb=="red") rgbc=3; else if(chcrgb=="green") rgbc=5; else rgbc=6;
				for(sl=0;sl<sls;sl++) {Stack.setSlice(sl+1); setMinAndMax(255,255,rgbc);}
			}
			run("8-bit");
		}
		run("Properties...", "unit=pixel pixel_width=1 pixel_height=1 voxel_depth=1");

		setLocation(650,300);
		if(i==x){
			OrigID=getImageID();
			if(!automode){
				setSlice(nSlices/2);
				setLocation(50,300);
				if(autobright) setMinAndMax(0,330);
			}
		}
		CompID=getImageID();
		if(!manualmode){
			run("Reslice [/]...", "input=1.000 output=1.000 start=Top avoid");
			ResliceID=getImageID();
			run("Z Project...", "projection=[Max Intensity]");
			rename("Source");
			if(i==x){
				setLocation(10,10);
				rename("Target");
				getDimensions(Width,Height,temp1,temp2,temp3);
				TargetID=getImageID();
			}
			SourceID=getImageID();
			selectImage(ResliceID);
			close();
			selectImage(SourceID);
			if(i==x){
				run("Duplicate...", "title=[ReslicesAdj]");
				setLocation(10,200);
				endim2=getImageID();
				run("Duplicate...", "title=[Reslices]");
				setLocation(10,400);
				endim=getImageID();
			}
		}
		if(automode){
			selectImage(CompID);
			close();
		}

		if(i!=x){
			c=0;
			if(!manualmode){
				run("TurboReg ","-align -window Source 0 0 "+round(Width-1)+" "+round(Height-1)+" -window Target 0 0 "+round(Width-1)+" "+round(Height-1)+" -translation "+round((Width-1)/2)+" "+round((Height-1)/2)+" "+round((Width-1)/2)+" "+round((Height-1)/2)+" -showOutput");
				sy=getResult("sourceY");
				ty=getResult("targetY");
				sx=getResult("sourceX");
				tx=getResult("targetX");
				dy=sy-ty;
				dx=sx-tx;
				dytot+=dy;
				dxtot+=dx;
				c=dy;
				close();
			}
			if(!automode){
				selectImage(OrigID);
				OrigSlice=getSliceNumber();
				selectImage(CompID);
				setSlice(OrigSlice+c);
				if(autobright) setMinAndMax(0,330);
				while(!isKeyDown("space")) wait(50);
				selectImage(OrigID);
				OrigSlice=getSliceNumber();
				close();
				selectImage(CompID);
				CompSlice=getSliceNumber();
				c=CompSlice-OrigSlice;
				OrigID=CompID;
				setLocation(50,300);
			}
			if(!manualmode){
				dytot=dytot-dy+c;
				selectImage(SourceID);
				run("TurboReg ","-transform -window Source "+round(Width)+" "+round(Height)+" -translation "+round(tx+dxtot)+" "+round(ty+dytot)+" "+round(tx)+" "+round(ty)+" -showOutput");				
				copyTo(endim2,true,(i==y));
				selectImage(SourceID);
				run("TurboReg ","-transform -window Source "+round(Width)+" "+round(Height)+" -translation "+round(tx+dxtot)+" "+round(ty)+" "+round(tx)+" "+round(ty)+" -showOutput");
				copyTo(endim,true, (i==y));
				selectImage(TargetID);
				close();
				selectImage(SourceID);
				setLocation(10,10);
				rename("Target");
				TargetID=SourceID;
			}

			ctot+=c;

			zReg[i]=round(ctot);
			print("\\Update:Timepoint "+(i+1)+" z"+round(ctot));

		} 
	}
	

	if(!automode) {selectImage(OrigID); close();}

	if(!manualmode){
		selectImage(TargetID);
		close();
	}


	print("Z Registration concluded");
	if(isOpen("Refined Landmarks")){
		print("[Refined Landmarks]","\\Clear");
		print("[Refined Landmarks]","\\Close");
	}

	accepted=true;
	if(!autoaccept) {
		waitForUser("Check reslices");
 		accepted=getBoolean("Accept z-reg changes?");
	}
	
	if(!manualmode){selectImage(endim); close; selectImage(endim2); close;
}

	if(accepted){
		
		zlength=frms-1;
		noslices=sls;
		slicesmin=1;
		slicesmax=sls;
		Array.getStatistics(zReg, zRegMin, zRegMax, mean, stdDev);
		if(zRegMin<0){
			if(zRegMax<0){
				noslices=sls+zRegMin;
			} else {
				noslices=sls+zRegMin-zRegMax;
				slicesmax=sls-zRegMax;
			}
			slicesmin=1-zRegMin;
		} else {
			noslices=sls-zRegMax;
			slicesmax=noslices;
		}

	     bychs=false;
	     if(bychs){
		run("Hyperstack...", "title=["+Zname+"] type=16-bit display=Composite width="+owidth+" height="+oheight+" channels="+chs+" slices="+noslices+" frames="+frms);
		realendID=getImageID;
		for(k=0;k<frms;k++){
			for(j=0;j<noslices;j++){
				for(i=0;i<chs;i++){
					selectImage(IID); Stack.setPosition(i+1,j+slicesmin+zReg[k],k+1);
					run("Select All"); run("Copy"); mdata=getMetadata("Label");
					selectImage(realendID); Stack.setPosition(i+1,j+1,k+1);
					run("Paste");setMetadata("Label",mdata);
				}
			}
		}
	     } else {   //bychs

		selectImage(IID);
		run("Duplicate...","title=["+Zname+"] duplicate");
		run("Hyperstack to Stack");
 
		selectWindow(Zname);
		if(is("Hyperstack"))run("Stack to Hyperstack...", "order=xyczt(default) channels=1 slices="+(chs*sls*frms)+" frames=1 display=Grayscale");
		selectWindow(Zname);
		n=0;
		for(k=0;k<frms;k++){
			for(j=0;j<sls;j++){
				for(i=0;i<chs;i++){
					if( ((j+1) < slicesmin+zReg[k]) || ((j+1)>slicesmin+zReg[k]+noslices-1)){
						setSlice(n+1+i+(j*chs)+(k*chs*sls)); run("Delete Slice"); n--;
					}
				}
			}
		}
		
		run("Stack to Hyperstack...", "order=xyczt(default) channels="+chs+" slices="+noslices+" frames="+frms+" display=Composite");
		mdata=""+zReg[0]; for(i=1;i<zReg.length;i++) mdata=mdata+" "+zReg[i];
		setMetadata("Info",mdata);
	     }
	     if(bits!=24)setTextLuts(textLuts);
	     if(automode && autoaccept) setBatchMode(false);
	} //accepted
     } //doz
	stilldo=true;
	if(!accepted && doxy) stilldo=getBoolean("Continue with XY align anyway?");
    if(doxy && stilldo){
	
		if(doz) {selectImage(IID); close; selectWindow(Zname);}
		IID=getImageID();
		getDimensions(w,h,chs,sls,frms);
	
		setPasteMode("Copy");
		XYname=getTitle();
		if(!doz) XYname=getTitle()+ "-XYalign";
	
		xpos=0;
		ypos=0;
		rotangle=0;
	
	
	 	wid=w; hei=h; widt=w; heit=h;
		run("Z Project...", "projection=[Max Intensity] all");  zID=getImageID;
		if(bits==24){
			if(chcrgb!="all") {
				if(chcrgb=="red") rgbc=3; else if(chcrgb=="green") rgbc=5; else rgbc=6;
				for(sl=0;sl<frms;sl++) {Stack.setFrame(sl+1); setMinAndMax(255,255,rgbc);}
			}
			run("8-bit");
		}
		run("Duplicate...", "title=targ duplicate channels="+chchoice+" frames=1");
		selectImage(IID); run("Duplicate...", "title=frm1 duplicate channels=1-"+chs+" slices=1-"+sls+" frames=1");
		run("Hyperstack to Stack");
	
		for(i=1; i<frms;i++){
			selectImage(zID);
			run("Duplicate...", "title=sour duplicate channels="+chchoice+" frames="+(i+1));
			run("TurboReg ","-align -window sour 0 0 "+(wid-1)+" "+(hei-1)+" -window targ 0 0 "+(widt-1)+" "+(heit-1)
				+" -rigidBody "
				+((wid-1)/2)+" "+((hei-1)/2)+" "+((widt-1)/2)+" "+((heit-1)/2)+" "
				+((wid-1)/10)+" "+((hei-1)/2)+" "+((widt-1)/10)+" "+((heit-1)/2)+" "
				+((wid-1)/(0.9))+" "+((hei-1)/2)+" "+((widt-1)/(0.9))+" "+((heit-1)/2)
				+" -showOutput");
			close();
			Regdata=getRegData();
			dx1=Regdata[0];
			dy1=Regdata[1];
			angle=Regdata[2];
	
			selectWindow("targ");
			close();
			selectWindow("sour");
			rename("targ");
			xpos+=dx1;
			ypos+=dy1;
			rotangle+=rotangle;
			if(doz) print("\\Update:Timepoint "+(i+1)+" x"+xpos+" y"+ypos+" z"+zReg[i]+" r"+rotangle);
			else print("\\Update:Timepoint "+(i+1)+" x"+xpos+" y"+ypos+" r"+rotangle);
	
			selectImage(IID); run("Duplicate...", "title=frm"+(i+1)+" duplicate channels=1-"+chs+" slices=1-"+sls+" frames="+(i+1));
			run("Hyperstack to Stack"); selectWindow("frm"+(i+1));
			makeRectangle(xpos,ypos,wid, hei);  
			run("Crop");
			torb="Bottom";
			rorl="Right";
			if(xpos>0) rorl="Left";
			if(ypos>0) torb="Top";
			run("Canvas Size...", "width="+wid+" height="+hei+" position="+torb+"-"+rorl+" zero");
			if(abs(rotangle)>0.5)run("Rotate... ", "angle="+rotangle+" grid=0 interpolation=None");
			run("Concatenate...", "stack1=frm"+(i)+" stack2=frm"+(i+1)+" title=frm"+(i+1));
		}
		selectWindow("frm"+i);
		run("Stack to Hyperstack...", "order=xyczt(default) channels="+chs+" slices="+sls+" frames="+frms+" display=Composite");
		rename("XYalign");
		if(bits!=24)setTextLuts(textLuts);
		selectWindow("targ"); close();
		
		accepted=true;
		if(!autoaccept) {
			run("Z Project...", "projection=[Max Intensity] all"); zxyID=getImageID();
			setLocation(800,300); selectImage(zID); setLocation(100,300);
			waitForUser("Check z-projections");
	 		accepted=getBoolean("Accept xy-reg changes?");
		}
		selectImage(zID); close();
		if(!autoaccept){selectImage(zxyID); close;}
		if(accepted) {
			selectWindow("XYalign"); rename(XYname);
			if(doz){ rename(title+"-XYZalign"); selectWindow(Zname); close;}
		} else {selectWindow("XYalign"); close;}
	
	
		print("XY Registration concluded");
		if(isOpen("Refined Landmarks")){
			print("[Refined Landmarks]","\\Clear");
			print("[Refined Landmarks]","\\Close");
		}
	}
 //if(doxy && stilldo)
	wait(100);
	showProgress(0.5);
	showProgress(1.0);
     


function getRegData(){

	sourceX0 = getResult("sourceX", 0); // First line of the table.
	sourceY0 = getResult("sourceY", 0);
	targetX0 = getResult("targetX", 0);
	targetY0 = getResult("targetY", 0);
	dx1 = sourceX0 - targetX0;
	dy1 = sourceY0 - targetY0;
	//print("\\Update:Translation [pixel]: " + dx1 +" in X and "+dy1+" in Y");
	//translation = sqrt(dx1 * dx1 + dy1 * dy1); // Amount of translation, in pixel units.
	if(nResults>1){
		sourceX1 = getResult("sourceX", 1); // Second line of the table.
		sourceY1 = getResult("sourceY", 1);
		targetX1 = getResult("targetX", 1);
		targetY1 = getResult("targetY", 1);
		sourceX2 = getResult("sourceX", 2); // Third line of the table.
		sourceY2 = getResult("sourceY", 2);
		targetX2 = getResult("targetX", 2);
		targetY2 = getResult("targetY", 2);
		dx = sourceX2 - sourceX1;
		dy = sourceY2 - sourceY1;
		sourceAngle = atan2(dy, dx);
		dx = targetX2 - targetX1;
		dy = targetY2 - targetY1;
		targetAngle = atan2(dy, dx);
		rotation = targetAngle - sourceAngle; // Amount of rotation, in radian units.
		//print("Amount of rotation [degree]: " + (rotation * 180.0 / PI));
		rotangle=(rotation * 180.0 / PI);
		return newArray(dx1, dy1, rotangle);
	} else return newArray(dx1,dy1);
}

function addArray(item, array){
	temparray=newArray(array.length+1);
	for(i=0; i<array.length; i++) temparray[i]=array[i];
	temparray[i]=item;
	return temparray;
}

function copyTo(imageid,isTr,isLast){
	if(isTr){
		if(bits==16){run("16-bit"); run("Divide...", "stack value=16");}
		else if(bits==8 || bits==24){run("8-bit");}
		else{exit("Don't know what to do with "+bits+" bits");}
	}
	run("Select All");
	run("Copy");
	if(isTr)close;
	selectImage(imageid);
	setSlice(nSlices); run("Add Slice"); run("Select None");
	run("Paste");
	resetMinAndMax;
	if(isLast){ setSlice(1); resetMinAndMax;}

}

function getTextLUTs(){
	if(bitDepth ==24) exit("Does not work on RGB image");
	Stack.getPosition(ch,sl,frm);
	getDimensions(w,h,chs,sls,frms);
	textLuts=newArray(chs);
	for(i=0;i<chs;i++){
		if(Stack.isHyperstack)Stack.setChannel(i+1);
		getLut(reds,greens,blues);
		max=maxOf(reds[128],maxOf(greens[128],blues[128]));
		if(max==reds[128]) {
			result="Red";
			if(max==greens[128]) result="Yellow";
			if(max==blues[128]) result="Magenta";
			if(max==greens[128] && max==blues[128]) result="Grays";
		}
		else if(max==greens[128]) {
			result="Green";
			if(max==blues[128]) result="Cyan";
		}
		else if(max==blues[128]) result="Blue";
		textLuts[i]=""+result;
	}
	Stack.setPosition(ch,sl,frm);
	return textLuts;
}

function setTextLuts(textLuts){
	if(bitDepth ==24) exit("Does not work on RGB image");
	Stack.getPosition(ch,sl,frm);
	getDimensions(w,h,chs,sls,frms);
	for(i=0;i<chs;i++){
		if(Stack.isHyperstack)Stack.setChannel(i+1);
		if(i<textLuts.length) run(textLuts[i]);
	}
	Stack.setPosition(ch,sl,frm);
}