//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();
	getVoxelSize(px,py,pz,unit);
	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);}
	}
	ptarray=false;
	if(nResults>0 && nResults%frms==0) ptarray=true;
	mchoices=newArray("Auto","Manual with suggestion", "Manual");
	if(ptarray) mchoices=newArray("Use Results locations (pt array)","Auto","Manual with suggestion", "Manual");

	Dialog.create("Choices");
	Dialog.addCheckbox("Z-align?",true);
	Dialog.addCheckbox("XY-align?",true);
 	Dialog.addChoice("Ch:", charray, "2");
	Dialog.addChoice("Mode", mchoices, mchoices[0]);
	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"; if(chcrgb=="red") rgbc=3; else if(chcrgb=="green") rgbc=5; else if(chcrgb=="blue") rgbc=6;}
	
chchoice=parseInt(chchoice);
	automode=false;
	manualmode=false;

	if(mode=="Auto"||mode=="Manual with suggestion") automode=true;
	if(mode=="Manual"||mode=="Manual with suggestion") manualmode=true;
	if(mode!="Use Results locations (pt array)") ptarray=false;

	//if(!manualmode && autoaccept && doz) setBatchMode(true);
	//setBatchMode("exit and display");


	x=newArray(frms);
	y=newArray(frms);
	z=newArray(frms);
	if(ptarray){
		rsn=nResults/frms;
		for(i=0;i<frms;i++){
			for(n=0;n<rsn;n++){
				x[i]+=parseFloat(getResult("X",i*rsn+n))/px;
				y[i]+=parseFloat(getResult("Y",i*rsn+n))/py;
				if(sls>1) z[i]+=parseFloat(getResult("Slice",i*rsn+n));
			}
			x[i]/=rsn; y[i]/=rsn; if(sls>1) z[i]/=rsn;
			if(i>0){x[i]-=x[0]; y[i]-=y[0]; z[i]-=z[0];}
		}
		x[0]=0; y[0]=0; z[0]=0;
	}

	if(doz){
		if(!ptarray){
			z[0]=0;
			ctot=0; dytot=0;
 dxtot=0;
		
			for(i=0; i<frms; i++){
				if(automode){
					if(i==0){
						run("Reslice [/]...", "input=1.000 output=1.000 start=Top avoid");
						rid=getImageID();
						run("Z Project...", "projection=[Max Intensity] all");
						rzid=getImageID();
						selectImage(rid); close();
						selectImage(rzid); Stack.setChannel(chchoice);
					}
					selectImage(rzid);
					run("Duplicate...", "title=TimeS"+ (i+1)+" duplicate channels="+chchoice+" slices=1 frames="+(i+1));
					ResliceID=getImageID();
				}
				if(manualmode){
					selectImage(IID);
					run("Duplicate...", "title=TimeS"+ (i+1)+" duplicate channels="+chchoice+" slices=1-"+sls+" frames="+(i+1));
				}
				if(bits==24){
					if(chcrgb!="all") {
						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==0){
					if(manualmode){
						OrigID=getImageID();
						setSlice(nSlices/2);
						setLocation(50,300);
						if(autobright) setMinAndMax(0,330);
						CompID=getImageID();
					}
				}
				if(automode){
					selectImage(ResliceID);
					rename("Source");
					if(i==0){
						setLocation(10,10);
						rename("Target");
						getDimensions(Width,Height,temp1,temp2,temp3);
						TargetID=getImageID();
					}
					SourceID=getImageID();
					if(i==0){
						run("Duplicate...", "title=[ReslicesAdj]");
						setLocation(10,200);
						endim2=getImageID();
						run("Duplicate...", "title=[Reslices]");
						setLocation(10,400);
						endim=getImageID();
					}
				}
		
				if(i!=0){
					c=0;
					if(automode){
						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(manualmode){
						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(automode){
						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==frms));
						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==frms));
						selectImage(TargetID);
						close();
						selectImage(SourceID);
						setLocation(10,10);
						rename("Target");
						TargetID=SourceID;
					}
		
					ctot+=c;
		
					z[i]=round(ctot);
					print("\\Update:Timepoint "+(i+1)+" z"+round(ctot));
		
		} 
			}
			if(manualmode) {selectImage(OrigID); close();}
			if(automode){
				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(automode){selectImage(endim); close; selectImage(endim2); close;}
		} else {accepted=true; print("Using pts for Z Registration...");}

	
		if(accepted){
			zlength=frms-1;
			noslices=sls;
			slicesmin=1;
			slicesmax=sls;
			Array.getStatistics(z, 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+z[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("Select None");
				showStatus("Duplicating stack...");
				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;
				showStatus("Deleting slices...");
				for(k=0;k<frms;k++){
					for(j=0;j<sls;j++){
						for(i=0;i<chs;i++){
							if( ((j+1) < slicesmin+z[k]) || ((j+1)>slicesmin+z[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=""+z[0]; for(i=1;i<z.length;i++) mdata=mdata+" "+z[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;

		if(!ptarray){
			run("Z Project...", "projection=[Max Intensity] all");  zID=getImageID;
			if(bits==24){
				if(chcrgb!="all") {
					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++){
			if(!ptarray){
				selectImage(zID);
				run("Duplicate...", "title=sour duplicate channels="+chchoice+" frames="+(i+1));
				run("TurboReg ","-align -window sour 0 0 "+(w-1)+" "+(h-1)+" -window targ 0 0 "+(w-1)+" "+(h-1)
					+" -rigidBody "
					+((w-1)/2)+" "+((h-1)/2)+" "+((w-1)/2)+" "+((h-1)/2)+" "
					+((w-1)/10)+" "+((h-1)/2)+" "+((w-1)/10)+" "+((h-1)/2)+" "
					+((w-1)/(0.9))+" "+((h-1)/2)+" "+((w-1)/(0.9))+" "+((h-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"+z[i]+" r"+rotangle);
				else print("\\Update:Timepoint "+(i+1)+" x"+xpos+" y"+ypos+" r"+rotangle);
			}else{xpos=x[i]; ypos=y[i];}
	
			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,w,h);  
			run("Crop");
			torb="Bottom";
			rorl="Right";
			if(xpos>0) rorl="Left";
			if(ypos>0) torb="Top";
			run("Canvas Size...", "width="+w+" height="+h+" 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);
		if(!ptarray){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?");
		}
		if(!ptarray){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);
}