path = getDir("select the experiment folder");
sdqdir = path + "/stardist quantif/";
File.makeDirectory(sdqdir);


workflow();


function workflow() {
	do {
		Dialog.createNonBlocking("Particles segmentation and quantification workflow");
	    Dialog.addMessage("Open a single color stack and a RoiSet delineating neurites\nthen tick the validate checkbox to continue");
	    Dialog.addCheckbox("threshold before segmentation", true);
	    Dialog.addCheckbox("Validate", false);
		Dialog.show();
		threshold_before = Dialog.getCheckbox();
	    validate = Dialog.getCheckbox();
	} while (validate == false || nImages == 0 || roiManager("count") == 0);

	print("All results saved in " + sdqdir);

	neuritesCount = roiManager("count");
	neuritesList = newArray(neuritesCount); //keep ROI names to reopen it afterwards
	savedParticleRois = newArray(neuritesCount); //keep RoiSets names to reassign them to the right channel afterwards	
	thresholds = newArray(neuritesCount);
	
	exportNeuritesRois(neuritesList);

	for (i = 0; i < neuritesCount; i++) {
		neuriteRoi = neuritesList[i] + ".roi";
		imgName_threshold = processOne(neuriteRoi);
		imgName = imgName_threshold[0];
		threshold = toString(imgName_threshold[1]);
		thresholds[i] = imgName + " " + threshold;
		print(imgName + " threshold = " + threshold);
		savedParticleRois[i] = saveParticleRoiSet(imgName);
		measureRois(imgName);
	}
	saveThresholds();
	updateAllRoisToHS();
	assembleParticleRois();
}


function exportNeuritesRois(neuritesList){
	for (i = 0; i < lengthOf(neuritesList); i++) {
		roiManager('select', i);
		neuritesList[i] = Roi.getName();
	    roiManager("save", sdqdir + neuritesList[i] + ".roi");
	}
	roiManager("deselect");
	roiManager("delete");
}

function processOne(neuriteRoiName){
	neuriteRoiPath = sdqdir + neuriteRoiName;
	roiManager("Open", neuriteRoiPath);
	roiManager("select", 0);
	stackname = getTitle();
    run("Duplicate...", "ignore use");
	roiManager("deselect");
    roiManager("select", 0);
    dupname = getTitle();
    //run("Duplicate...", "ignore use");
    resetMinAndMax();
    
    // uncomment next section to add first thresholding
    ///*
    if (threshold_before){
	    setAutoThreshold("Otsu dark");    
	    waitForUser("set threshold then click ok");
	    getThreshold(lower, upper);
	    print(dupname + " thresholds " + lower + " " +upper);
	    selectImage(dupname);
		setOption("BlackBackground", true);
		run("Convert to Mask");
		run("16-bit");
		run("Enhance Contrast...", "saturated=0 normalize");
		imageCalculator("AND create", stackname, dupname);
		close(dupname);
		rename(dupname);
		run("Enhance Contrast", "saturated=0.35");
    }
		//*/
	imgName = dupname;
	
	selectImage(imgName);
    roiManager("deselect");
    roiManager("select", 0);
    
    validate = false;
//    threshold = 0.5;
    threshold = 0.4;
    
    do {
   		Dialog.createNonBlocking("Adjust stardist threshold");
	    Dialog.addNumber("Threshold", threshold);
	    Dialog.addCheckbox("Validate", false);
	    Dialog.addMessage("Once satisfied with the segmentation\ntick the validate checkbox to continue");
	    if(roiManager("count") >= 1) {
		    Dialog.show();
	    }
	    validate = Dialog.getCheckbox();
	    threshold = Dialog.getNumber();
	    
    //stardist thresholds : 0.04 = all particles, 0.67 = strong particles; norm bottom=1 percentileTop=99.9 nmsThreshold=0.33
    //if nonorm, threshold = 0.002, percBott=1 percTop=100, nmsThres=0.25
		if (validate == false ) {
		run("Command From Macro","command=[de.csbdresden.stardist.StarDist2D], args=['input':'" + imgName +"', 'modelChoice':'Versatile (fluorescent nuclei)', 'normalizeInput':'true', 'percentileBottom':'0.1', 'percentileTop':'99.8', 'probThresh':'"+toString(threshold)+ "', 'nmsThresh':'0.33', 'outputType':'ROI Manager', 'nTiles':'1', 'excludeBoundary':'2', 'roiPosition':'Automatic', 'verbose':'true', 'showCsbdeepProgress':'true', 'showProbAndDist':'false'], process=[false]");
		//note that after Stardist execution, all ROIs in manager before its execution will be erased
		
		roiManager("Open", neuriteRoiPath); //add neurite to particle list
		roiManager("show all");
		
		}
		else {
			roiManager("show none");
			//removes AIS contour
			count=roiManager("count");
			if (count > 1) {
				roiManager("select", count-1);
				roiManager("delete");
			}
		}
    } while(validate == false);
	if (count >1) {
		filterParticlesInNeurite(imgName, neuriteRoiPath);
	}
	close(imgName);//don't remove !
	return newArray(imgName, threshold);
}

function filterParticlesInNeurite(imgName, neuriteRoiPath) {
	particlesCount = roiManager("count");
	particlesIn = newArray(particlesCount -1);
	roiManager("Open", neuriteRoiPath); //add neurite to particle list
	neuriteRoiIndex = roiManager("count")-1;
	roiManager("select", neuriteRoiIndex);

	for (j = 0; j < particlesCount; j++) {
		particlesIn[j] = false; //default initialization
		roiManager("select", newArray(neuriteRoiIndex, j));
		roiManager("and");
		overlap = getValue("selection.size");
		//print(j + " " + overlap);
		
		//temporary create a roi to check if the particle is IN the neurite or on its boundary, removes the roi afterwards		
		if (overlap > 0) { 
			roiManager("add");
			overlapRoiIndex = neuriteRoiIndex+1;
			roiManager("select", overlapRoiIndex);
			roiManager("measure");
			overlapArea = getResult("Area");
			roiManager("deselect");
			roiManager("select", j);
			roiManager("measure");
			particleArea = getResult("Area");
			roiManager("deselect");
			roiManager("select", neuriteRoiIndex+1);
			roiManager("delete");

			if (overlapArea > particleArea/2){
				particlesIn[j] = true;
			}
		}
	}
		
	roiManager("deselect");
	roiManager("Select", neuriteRoiIndex);
	
	neuriteRoiPrefix = substring(Roi.getName(), 0, 5);
	roiManager("deselect");

	for (j = particlesCount-1 ; j >= 0 ; j--) {
		roiManager("select", j);
		
		if (particlesIn[j] == false) {
			roiManager("delete");
		}
		else {
			name = RoiManager.getName(j);
			roiManager("rename", neuriteRoiPrefix + name);
		}
	}
	pcount = roiManager("count") -1;
	roiManager("select", pcount);
	getStatistics(neuriteArea);
	print("pcount : " + pcount + ", ROI area : " + neuriteArea + ", density : " + (pcount/neuriteArea));
}

function saveParticleRoiSet(imgName) {
	roiManager("deselect");
	
	roiSetName = "" + getShortImgName(imgName) + " RoiSet.zip";
	roiManager("save", sdqdir + roiSetName);
	return roiSetName;
}

function getShortImgName(imgName) {
	tokens = split(imgName, "_");
	shortImgName = tokens[0] + "_" + tokens[1] + "_" + tokens[2];
	return shortImgName;
}

function measureRois(imgName){
	
	roiManager("deselect");
	roiManager("multi-measure");
	saveAs("Results", sdqdir + getShortImgName(imgName) + " Particles.csv");

	close("Results");
	roiManager("delete");
}

function updateAllRoisToHS () {
	close("*");
	close("Roi Manager");
	path3C = File.openDialog("open the 3C stack");
	open(path3C);
	
	for (i = 0; i < lengthOf(savedParticleRois); i++) {
		roiManager("Open", sdqdir + savedParticleRois[i]);
		
		tokens = split(savedParticleRois[i], "_ ");
		imgId = String.join(newArray(tokens[0], tokens[1], tokens[2]), "_"); //Cxx_Nxx_cond
		print(imgId);
		
		for (j = 1; j <= nSlices; j++) {
    		Stack.setSlice(j);
			curImg = getInfo("slice.label");
			//print(curImg);
			tokens = split(curImg, "_");
			curImgId = String.join(newArray(tokens[0], tokens[1], tokens[2]), "_");
			//print(curImgId);
			if (curImgId == imgId) {
				Stack.getPosition(c, s, f);
				roiCount = roiManager("count");
				for (k = 0 ; k < roiCount ; k++) {
					roiManager("select", k);
					Stack.setChannel(c);
					Stack.setSlice(s);
					Stack.setFrame(f);
					roiManager("update");
				}

				if (roiManager("count") > 0) {
					roiManager("save", sdqdir + imgId + " RoiSet 3C.zip");
					roiManager("deselect");
					roiManager("delete");					
				}
			}
			//break;
		}
	}	
}

function assembleParticleRois(){
	for (i = 0; i < lengthOf(savedParticleRois); i++) {
		filename3C = savedParticleRois[i].replace("\\.zip", " 3C.zip");
		roiManager("Open", sdqdir + filename3C);
	}
	roiManager("deselect");
	roiManager("sort");
	roiManager("save", sdqdir + " Stack RoiSet 3C.zip");
}

function saveThresholds() {
	selectWindow("Log");
	saveAs("text", sdqdir + "thresholds.txt");
	//close("Log");
}


function afterWardsMeasureRois(){
	path = getDir("select the experiment folder");
	sdqdir = path + "/stardist quantif/";
	
	curImg = getInfo("slice.label");
	roiManager("Open", sdqdir + savedParticleRois[i]);
		
	roiManager("deselect");
	roiManager("multi-measure");
	saveAs("Results", sdqdir + getShortImgName(imgName) + " Particles.csv");

	close("Results");
	roiManager("delete");
}