//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);
	substack="1-"+sls;
	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("Use Results locations (pt array)","Auto");
	mdefault=mchoices[1]; if(ptarray)mdefault=mchoices[0];

	Dialog.create("Choices");
	Dialog.addChoice("Mode", mchoices, mdefault);
	Dialog.addCheckbox("Z-align?",true);
	Dialog.addCheckbox("XY-align?",true);
	Dialog.addCheckbox("New stack?", false);
	Dialog.addMessage("Just for Auto:");
	Dialog.addString("Substack (for xy z-proj):", substack);
	Dialog.addCheckbox("Set selections?",false);
	Dialog.addCheckbox("Check z-changes before accepting?",false);
 	Dialog.addChoice("Ch:", charray, "2");
	if(!ptarray)Dialog.addMessage("Pt array requires one or a set of points to be measured for each time point");
	Dialog.show();


	mode=Dialog.getChoice();
	doz=Dialog.getCheckbox;
	doxy=Dialog.getCheckbox;
	newstack=Dialog.getCheckbox();
	substack=Dialog.getString();
	selec=Dialog.getCheckbox();
	autoaccept=!Dialog.getCheckbox();
	chchoice=Dialog.getChoice();

	startsub=1; endsub=sls;
	if(substack!="") {
		if(indexOf(substack,"-")==-1){
			startsub=round(substack);
			endsub=round(substack);
		} else {
			substa=split(substack,"-");
			startsub=parseInt(substa[0]);
			endsub=parseInt(substa[1]);
		}
		if(startsub>endsub) {tempsub=startsub; startsub=endsub; endsub=tempsub;}
	}

	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);

	if(mode=="Use Results locations (pt array)" && !ptarray) exit("Pt array requires one or a set of points to be measured for each time point");
	if(mode=="Auto") ptarray=false; else ptarray=true;

	//setBatchMode(true);
	//setBatchMode("exit and display");

	z=newArray(frms);
	x=newArray(frms);
	y=newArray(frms);
	r=newArray(frms);

	accepted=true;
	if(ptarray){
		rsn=nResults/frms;
		xinits=newArray(rsn);
		yinits=newArray(rsn);
		zinits=newArray(rsn);
		isxz=false;
		if(rsn==2){isxz=getBoolean("XZ then XY data?");}
		if(isxz){
			for(i=0;i<frms;i++){
				if(i==0){
					xinits[0]=parseFloat(getResult("X",0))/px;
					xinits[1]=parseFloat(getResult("X",frms))/px;
					yinits[0]=parseFloat(getResult("Y",frms))/py;
					zinits[0]=(parseFloat(getResult("Y",0))+1)/pz;
				}else{
					x[i]=(((parseFloat(getResult("X",0*frms+i))/px)-xinits[0])+((parseFloat(getResult("X",1*frms+i))/px)-xinits[1]))/2;
					y[i]=(parseFloat(getResult("Y",1*frms+i))/py)-yinits[0];
					z[i]=round(((parseFloat(getResult("Y",i))+1)/pz)-zinits[0]);
				}
				r[i]=0; //not coded (yet)
				//print(""+x[i]+"  "+y[i]+"  "+z[i]);
			}
		}else{
			for(i=0;i<frms;i++){
				for(n=0;n<rsn;n++){
					if(i==0){
						xinits[n]=parseFloat(getResult("X",n))/px;
						yinits[n]=parseFloat(getResult("Y",n))/py;
						if(sls>1) zinits[n]=parseFloat(getResult("Slice",n));
					}else{
						x[i]+=(parseFloat(getResult("X",i*rsn+n))/px)-xinits[n];
						y[i]+=(parseFloat(getResult("Y",i*rsn+n))/py)-yinits[n];
						if(sls>1) z[i]+=(parseFloat(getResult("Slice",i*rsn+n)))-zinits[n];
					}
				}
				x[i]/=rsn; y[i]/=rsn; if(sls>1) z[i]/=rsn;
				r[i]=0; //not coded (yet)
			}
		}
	}
	
	zaccepted=false;
	if(doz){
		if(!ptarray){
			z[0]=0;
			ctot=0; dytot=0; dxtot=0;
			if(selec)waitForUser("Place selection for reslice now (currently only xz reslice)");
			for(i=0; i<frms; i++){
				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(bits==24){
					if(chcrgb!="all") {
						for(sl=0;sl<sls;sl++) {Stack.setSlice(sl+1); setMinAndMax(255,255,rgbc);}
					}
					run("8-bit");
				}
				setVoxelSize(1, 1, 1, "pixel");

				selectImage(ResliceID);
				rename("Source");
				SourceID=getImageID();
				if(i==0){
					setLocation(10,10);
					rename("Target");
					getDimensions(Width,Height,temp1,temp2,temp3);
					TargetID=getImageID();
					run("Duplicate...", "title=[ReslicesAdj]");
					setLocation(10,200);
					endim2=getImageID();
					run("Duplicate...", "title=[Reslices]");
					setLocation(10,400);
					endim=getImageID();
				} else {
					c=0;
					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;
					close();
					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+=dy;
					z[i]=round(ctot);
					print("\\Update:Timepoint "+(i+1)+" z"+round(ctot));
		
				}
			}//for frames
			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?");
			}
			selectImage(endim); close; selectImage(endim2); close;
		}//not ptarray
	
		if(accepted){
			zaccepted=true;
			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;
			}
	
			selectImage(IID);
			run("Select None");
			if(newstack){
				showStatus("Duplicating stack...");
				run("Duplicate...","title=["+Zname+"] duplicate");
			} else rename(Zname);
			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);
			 //setBatchMode(false);
		} //z accepted
	}//doz
	stilldo=true;
	if(!accepted && doxy) stilldo=getBoolean("Continue with XY align anyway?");
    if(doxy && stilldo){
		if(doz && newstack) {selectImage(IID); close; selectWindow(Zname);}
		IID=getImageID();
		getDimensions(w,h,chs,sls,frms);
	
		setPasteMode("Copy");
		XYname=getTitle();
		if(endsWith(XYname,"-Zalign")){XYname=substring(XYname,0,indexOf(XYname,"-Zalign"))+"-XYZalign";}
		else XYname=XYname+"-XYalign";
	
		xpos=0;
		ypos=0;
		rotangle=0;

		if(!ptarray){
			x[0]=0; y[0]=0;r[0]=0;
			run("Z Project...", "start="+startsub+" stop="+endsub+" projection=[Max Intensity] all");
			if(selec) {
				waitForUser("You may select a square now to constrain xy-registration");
				if(selectionType()==0)run("Crop");
			}
			getDimensions(wcr,hcr,chscr,slscr,frmscr);
			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");
	
	
			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 "+(wcr-1)+" "+(hcr-1)+" -window targ 0 0 "+(wcr-1)+" "+(hcr-1)
					+" -rigidBody "
					+((wcr-1)/2)+" "+((hcr-1)/2)+" "+((wcr-1)/2)+" "+((hcr-1)/2)+" "
					+((wcr-1)/10)+" "+((hcr-1)/2)+" "+((wcr-1)/10)+" "+((hcr-1)/2)+" "
					+((wcr-1)/(0.9))+" "+((hcr-1)/2)+" "+((wcr-1)/(0.9))+" "+((hcr-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;
				x[i]=xpos;
				y[i]=ypos;
				rotangle+=rotangle;
				r[i]=rotangle;
				if(doz) print("Timepoint "+(i+1)+" x"+xpos+" y"+ypos+" z"+z[i]+" r"+rotangle);
				else print("Timepoint "+(i+1)+" x"+xpos+" y"+ypos+" r"+rotangle);
			}
			selectWindow("targ"); close();
			print("XY Registration concluded");
			if(isOpen("Refined Landmarks")){
				print("[Refined Landmarks]","\\Clear");
				print("[Refined Landmarks]","\\Close");
			}
    	}
    	accepted=true;
		if(!autoaccept && false) { //not implementing right now..
			//code to get zxy image needed above
			selectImage(zxyID);
			zxyID=getImageID();
			setLocation(800,300); selectImage(zID); setLocation(100,300);
			waitForUser("Check xy-projections");
	 		accepted=getBoolean("Accept xy-reg changes?");
	 		selectImage(zxyID); close;
	 		selectImage(zID); close();
		}else {if(!ptarray){selectImage(zID); close();}}
		
		if(accepted){
			concatstr="  title=["+XYname+"]";
			for(i=0; i<frms;i++){
				selectImage(IID); run("Duplicate...", "title=frm"+(i+1)+" duplicate channels=1-"+chs+" slices=1-"+sls+" frames="+(i+1));
				run("Hyperstack to Stack"); 
				if(i>0){
					selectWindow("frm"+(i+1));
					makeRectangle(x[i],y[i],w,h);  
					run("Crop");
					torb="Bottom";
					rorl="Right";
					if(x[i]>0) rorl="Left";
					if(y[i]>0) torb="Top";
					run("Canvas Size...", "width="+w+" height="+h+" position="+torb+"-"+rorl+" zero");
					if(abs(r[i])>0.5)run("Rotate... ", "angle="+r[i]+" grid=0 interpolation=None");
				}
				concatstr=concatstr+" image"+(i+1)+"=frm"+(i+1);
			}
			concatstr=concatstr+" image"+(i+1)+"=[-- None --]";
			run("Concatenate...", concatstr);
			selectWindow(XYname);
			run("Stack to Hyperstack...", "order=xyczt(default) channels="+chs+" slices="+sls+" frames="+frms+" display=Composite");
			if(zaccepted){selectWindow(Zname); close;}
			else if(newstack){selectWindow(Zname); close;}
			if(bits!=24)setTextLuts(textLuts);
			
		}
	}
 //if(doxy && stilldo)
	print("XYZ Registration completed");
	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);
}