//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();
	oIID=IID;
	getVoxelSize(px,py,pz,unit);
	getDimensions(w,h,chs,sls,frms);
	substack="1-"+sls;
	bits=bitDepth();
	if(bits!=24)textLuts=getTextLUTs();
	logarray=false;

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

	if(sls<2 && frms< 2) exit("Requires hyperstack with multi 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;}
	else{
		logstr=split(getInfo("log"),"\n");
		ll=logstr.length;
		st=-1;end=-1;
		for(i=0;i<logstr.length;i++){
			if(startsWith(logstr[i],"Timepoint ")){if(st==-1)st=i; else end=i;}
		}
		if(st!=-1){
			logstr=Array.slice(logstr,st,end+1);
			if(logstr.length==(frms-1)) {logarray=true; ptarray=true;}
		}
	}
	mchoices=newArray("Use Results locations (pt array)","Auto");
	mdefault=mchoices[1]; if(ptarray)mdefault=mchoices[0];
	xyzchoice=newArray("Z-align frames","XY-align frames", "XYZ-align frames", "Z-align slices for each frame");
	xyz2=newArray("XY-align frames", "Z-align slices for each frame");
	xyzuse=xyzchoice;
	if(sls==1)xyzuse=xyz2;

	Dialog.create("Choices");
	Dialog.addMessage("XY-adjust stack: "+title);
	Dialog.addChoice("Pt array?", mchoices, mdefault);
	Dialog.addChoice("Mode", xyzuse, xyzuse[0]);
	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();
	xyzmode=Dialog.getChoice();
	doz=(xyzmode==xyzchoice[0] || xyzmode==xyzchoice[2]);
	doxy=(xyzmode==xyzchoice[1] || xyzmode==xyzchoice[2]);
	newstack=Dialog.getCheckbox();
	substack=Dialog.getString();
	selec=Dialog.getCheckbox();
	autoaccept=!Dialog.getCheckbox();
	chchoice=Dialog.getChoice();

	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(xyzmode==xyzchoice[3]){
		//This aligns the slices of each individual time point.
		//This is useful to correct for movement artifact within a timepoint.
		//May be useful to run this first, then Align Hyperstack.

		stackregstr="StackReg ";
		if(indexOf(getInfo("imagej.dir"),"Fiji")>-1) stackregstr="StackReg";

		if(isOpen("AJZALIGNTEMP")) exit("Please close AJZALIGNTEMP");
		//if(chs>1) run("Stack to RGB");
		title=getTitle();

		for(i=1;i<=frms;i++){
			selectWindow(title); i=1;
			run("Duplicate...", "title="+i+" duplicate frames="+i);
			if(chs>1)Stack.setChannel(chchoice);
			run(stackregstr, "transformation=Affine");
			if(i!=1) run("Concatenate...", "stack1=1 stack2="+i+" title=AJZALIGNTEMP");
		}
		rename(title+"-slicealign");
		run("Stack to Hyperstack...", "order=xyczt(default) channels="chs" slices="+sls+" frames="+frms);
		exit();
	}

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

	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(logarray){
		for(i=1;i<frms;i++){
			line=split(logstr[i-1]," ");
			for(j=0;j<line.length;j++){
				if(startsWith(line[j],"x"))x[i]=parseFloat(substring(line[j],1,lengthOf(line[j])));
				if(startsWith(line[j],"y"))y[i]=parseFloat(substring(line[j],1,lengthOf(line[j])));
				if(startsWith(line[j],"z"))z[i]=parseFloat(substring(line[j],1,lengthOf(line[j])));
				if(startsWith(line[j],"r"))r[i]=parseFloat(substring(line[j],1,lengthOf(line[j])));	
			}
		}
	} else 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){
				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;
				r[0]=0;
			for(i=1;i<frms;i++){
				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)");
			tpstr=newArray(1);
			tpstr[0]="TP: 01:x";
			print("\\Clear");
			lstrn=0;
			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=[Source] duplicate channels="+chchoice+" slices=1 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"); }
				setVoxelSize(1, 1, 1, "pixel");
				SourceID=getImageID();
				setLocation(10,200);
				if(i==0){
					setLocation(10,10);
					rename("Target");
					getDimensions(Width,Height,temp1,temp2,temp3);
					TargetID=getImageID();
					run("Duplicate...", "title=[ReslicesAdj]");
					setLocation(10,400);
					endim2=getImageID();
					run("Duplicate...", "title=[Reslices]");
					setLocation(10,600);
					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);
					tpstr[lstrn]=tpstr[lstrn]+" "+IJ.pad(i+1,2)+":"+z[i];
					print("\\Update"+lstrn+":"+tpstr[lstrn]);
					if((i+1)%20==0){lstrn++; tpstr=Array.concat(tpstr,"TP: ");}
				}
			}//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; selectImage(rzid); 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;
			}
			
			if(newstack){
				showStatus("Duplicating stack...");
				selectImage(IID);
				run("Duplicate...","title=["+Zname+"] duplicate");
				IID=getImageID();
			}
			
			if(zRegMin==zRegMax) {showMessage("No Z shift will be done");}
			else{
				selectImage(IID);
				rename(Zname);
				run("Select None");
				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)){
								print("\\Update:ch"+i+"/"+chs+" sl"+j+"/"+sls+" fr"+k+"/"+frms+" n"+n);
								setSlice(n+1+i+(j*chs)+(k*chs*sls)); run("Delete Slice", "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];
				selectWindow(Zname);
				setMetadata("Info",mdata);
				IID=getImageID();
		
			     if(bits!=24)setTextLuts(textLuts);
			     call("java.lang.System.gc");
			}
		} //z accepted
	}//doz
	
	stilldo=true;
	if(!accepted && doxy) stilldo=getBoolean("Continue with XY align anyway?");
    if(doxy && stilldo){
    	selectImage(IID);
		//if(doz && newstack) {
		//	close;
		//	selectImage(oIID);
		//	IID=getImageID();
		//}
		getDimensions(w,h,chs,sls,frms);
    	endsub=minOf(endsub,sls);
	
		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");
			}
			selectImage(zID); close();
    	} //!ptarray

		if(!doz){
				selectImage(IID);
			if(newstack){
				showStatus("Duplicating stack...");
				run("Duplicate...","title=["+XYname+"] duplicate");
				IID=getImageID();
			}else{
				rename(XYname);
			}
		}
			
		concatstr="  title=["+XYname+"]";
		for(i=0; i<frms;i++){
			selectImage(IID);
			if(sls>1){
				run("Duplicate...", "title=frm"+(i+1)+" duplicate channels=1-"+chs+" slices=1-"+sls+" frames=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");
					run("Concatenate...", "  title=["+XYname+"] image1=["+XYname+"] image2=frm"+(i+1));
				} else rename(XYname);
				selectImage(IID); Stack.setFrame(1);
				if(i<(frms-1)){run("Delete Slice", "delete=frame");}else{close();}
			}else{
				if(i>0){
					for(ich=0;ich<chs;ich++){
						Stack.setPosition(ich+1, 1, i+1);
						makeRectangle(0, 0, w, h);
						run("Cut");
						makeRectangle(-x[i], -y[i], w, h);
						run("Paste");
						if(abs(r[i])>0.5)run("Rotate... ", "angle="+r[i]+" grid=0 interpolation=None slice");
					}
				}
			}
		}
		selectWindow(XYname);
		if(sls>1){
			run("Stack to Hyperstack...", "order=xyczt(default) channels="+chs+" slices="+sls+" frames="+frms+" display=Composite");
		}else{
			Stack.setFrame(1);
			run("Select None");
		}
		if(bits!=24)setTextLuts(textLuts);
		call("java.lang.System.gc");
	}
 //if(doxy && stilldo)
	print("XYZ Registration completed");
	wait(100);
	showProgress(0.5);
	wait(100);
	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;
	rotangle=0;
	//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);
	}
	//print("Translation: " + dx1 +" by "+dy1+" rot:"+rotangle);
	return newArray(dx1, dy1, rotangle);
}

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