package ij.plugin;
import ij.*;
import ij.gui.*;
import ij.process.*;
import ij.measure.Calibration;
import java.awt.*;

/** This plugin implements the Edit/Selection/Straighten command. */
public class Straightener implements PlugIn {
	static boolean processStack;

 	public void run(String arg) {
		ImagePlus imp = IJ.getImage();
		Roi roi = imp.getRoi();
		boolean rotatedRectangle = roi!=null && (roi instanceof RotatedRectRoi);
		if (roi==null || !(roi.isLine()||rotatedRectangle)) {
			IJ.error("Straightener", "Line, or rotated rectangle, selection required");
			return;
		}
		if (rotatedRectangle) {
			IJ.run(imp, "Duplicate...", " ");
			return;
		}
		if (!imp.lock()) return;
		int width = (int)Math.round(roi.getStrokeWidth());
		int originalWidth = width;
		boolean isMacro = IJ.macroRunning() && Macro.getOptions()!=null;
		int stackSize = imp.getStackSize();
		if (stackSize==1) processStack = false;
		String newTitle = WindowManager.getUniqueName(imp.getTitle());
		if (width<=1 || isMacro || stackSize>1) {
			if (width<=1) width = 20;
			GenericDialog gd = new GenericDialog("Straightener");
			gd.addStringField("Title:", newTitle, 15);
			gd.addNumericField("Line Width:", width, 0, 3, "pixels");
			if (stackSize>1)
				gd.addCheckbox("Process Entire Stack", processStack);
			gd.showDialog();
			if (gd.wasCanceled()) {imp.unlock(); return;}
			newTitle = gd.getNextString();
			width = (int)gd.getNextNumber();
			Line.setWidth(width);
			if (stackSize>1)
				processStack = gd.getNextBoolean();
		}
		roi = (Roi)imp.getRoi().clone();
		int type = roi.getType();
		if (type==Roi.FREELINE)
			IJ.run(imp, "Fit Spline", "");
		ImageProcessor ip2 = null;
		ImagePlus imp2 = null;
		if (processStack) {
			ImageStack stack2 = straightenStack(imp, roi, width);
			imp2 = new ImagePlus(newTitle, stack2);
		} else {
			ip2 = straighten(imp, roi, width);
			imp2 = new ImagePlus(newTitle, ip2);
		}
		imp.unlock();
		if (imp2==null)
			return;
		Calibration cal = imp.getCalibration();
		if (cal.pixelWidth==cal.pixelHeight)
			imp2.setCalibration(cal);
		imp2.show();
		if (isMacro) Line.setWidth(originalWidth);
	}

	public ImageProcessor straighten(ImagePlus imp, Roi roi, int width) {
		ImageProcessor ip2;
		if (imp.getBitDepth()==24 && roi.getType()!=Roi.LINE)
			ip2 = straightenRGB(imp, width);
		else if (imp.isComposite() && ((CompositeImage)imp).getMode()==IJ.COMPOSITE) {
			if (roi.getType()==Roi.LINE)
				ip2 = rotateCompositeLine(imp, width);
			else
				ip2 = straightenComposite(imp, width);
		} else if (roi.getType()==Roi.LINE)
			ip2 = rotateLine(imp, width);
		else
			ip2 = straightenLine(imp, width);
		return ip2;
	}

	public ImageStack straightenStack(ImagePlus imp, Roi roi, int width) {
		int current = imp.getCurrentSlice();
		int n = imp.getStackSize();
		ImageStack stack2 = null;
		for (int i=1; i<=n; i++) {
			IJ.showProgress(i, n);
			imp.setSlice(i);
			ImageProcessor ip2 = straighten(imp, roi, width);
			if (stack2==null)
				stack2 = new ImageStack(ip2.getWidth(), ip2.getHeight());
			stack2.addSlice(null, ip2);
		}
		imp.setSlice(current);
		return stack2;
	}

	public ImageProcessor straightenLine(ImagePlus imp, int width) {
		Roi tempRoi = imp.getRoi();
		if (tempRoi == null) return null;	//roi may have changed asynchronously
		if (tempRoi instanceof Line) {
			FloatPolygon fp = ((Line)tempRoi).getFloatPoints();
			tempRoi = new PolygonRoi(fp.xpoints, fp.ypoints, 2, Roi.POLYLINE);
		} else if (!(tempRoi instanceof PolygonRoi))
			return null;
		PolygonRoi roi = (PolygonRoi)tempRoi;
		if (roi==null)
			return null;
		if (roi.getState()==Roi.CONSTRUCTING)
			roi.exitConstructingMode();
		if (roi.isSplineFit())
			roi.removeSplineFit();
		int type = roi.getType();
		int n = roi.getNCoordinates();
		double len = roi.getLength();
		roi.fitSplineForStraightening();
		if (roi.getNCoordinates()<2)
			return null;
		FloatPolygon p = roi.getFloatPolygon();
		n = p.npoints;
		ImageProcessor ip = imp.getProcessor();
		ImageProcessor ip2 = new FloatProcessor(n, width);
		//ImageProcessor distances = null;
		//if (IJ.debugMode)  distances = new FloatProcessor(n, 1);
		float[] pixels = (float[])ip2.getPixels();
		double x1, y1;
		// the following will be taken as the previous point; extrapolate back one pixel
		double x2 = p.xpoints[0]-(p.xpoints[1]-p.xpoints[0]);
		double y2 = p.ypoints[0]-(p.ypoints[1]-p.ypoints[0]);
		if (width<=1)
			ip2.putPixelValue(0, 0, ip.getInterpolatedValue(x2, y2));
		for (int i=0; i<n; i++) {
			if (!processStack&&(i%10)==0) IJ.showProgress(i, n);
			x1=x2; y1=y2;
			x2 = p.xpoints[i];
			y2 = p.ypoints[i];
			//if (distances!=null) distances.putPixelValue(i, 0, (float)Math.sqrt((x2-x1)*(x2-x1)+(y2-y1)*(y2-y1)));
			if (width==1) {
				ip2.putPixelValue(i, 0, ip.getInterpolatedValue(x2, y2));
				continue;
			}
			double dx = x2-x1;
			double dy = y1-y2;
            double length = (float)Math.sqrt(dx*dx+dy*dy);
            dx /= length;
            dy /= length;
			//IJ.log(i+"  "+x2+"  "+dy+"  "+(dy*width/2f)+"   "+y2+"  "+dx+"   "+(dx*width/2f));
			double x = x2-dy*(width-1)/2.0;
			double y = y2-dx*(width-1)/2.0;
			int j = 0;
			int n2 = width;
			do {
				ip2.putPixelValue(i, j++, ip.getInterpolatedValue(x, y));;
				//ip.drawDot((int)x, (int)y);
				x += dy;
				y += dx;
			} while (--n2>0);
		}
		if (!processStack) IJ.showProgress(n, n);
		if (type==Roi.FREELINE)
			roi.removeSplineFit();
		else
			imp.draw();
		if (imp.getBitDepth()!=24) {
			ip2.setColorModel(ip.getColorModel());
			ip2.resetMinAndMax();
		}
		return ip2;
	}

	public ImageProcessor rotateLine(ImagePlus imp, int width) {
		Roi roi = imp.getRoi();
		if (roi==null || roi.getType()!=Roi.LINE)
			throw new IllegalArgumentException("Straight line selection expected");
		ImageProcessor ip2 = imp.getBitDepth()==24?straightenRGB(imp, width):straightenLine(imp, width);
		return ip2;
	}

	ImageProcessor straightenRGB(ImagePlus imp, int width) {
		int w=imp.getWidth(), h=imp.getHeight();
		int size = w*h;
		byte[] r = new byte[size];
        byte[] g = new byte[size];
        byte[] b = new byte[size];
		ColorProcessor cp = (ColorProcessor)imp.getProcessor();
		cp.getRGB(r, g, b);
		Roi roi = imp.getRoi();
		if (roi == null) return null;
        ImagePlus imp2 = new ImagePlus("red", new ByteProcessor(w, h, r, null));
        imp2.setRoi((Roi)roi.clone());
        ImageProcessor red = straightenLine(imp2, width);
        if (red==null) return null;
        imp2 = new ImagePlus("green", new ByteProcessor(w, h, g, null));
        imp2.setRoi((Roi)roi.clone());
        ImageProcessor green = straightenLine(imp2, width);
        if (green==null) return null;
        imp2 = new ImagePlus("blue", new ByteProcessor(w, h, b, null));
        imp2.setRoi((Roi)roi.clone());
        ImageProcessor blue = straightenLine(imp2, width);
        if (blue==null) return null;
        ColorProcessor cp2 = new ColorProcessor(red.getWidth(), red.getHeight());
        red = red.convertToByte(false);
        green = green.convertToByte(false);
        blue = blue.convertToByte(false);
        cp2.setRGB((byte[])red.getPixels(), (byte[])green.getPixels(), (byte[])blue.getPixels());
        imp.setRoi(imp2.getRoi());
        return cp2;
 	}

	ImageProcessor straightenComposite(ImagePlus imp, int width) {
		Image img = imp.getImage();
		ImagePlus imp2 = new ImagePlus("temp", new ColorProcessor(img));
		imp2.setRoi(imp.getRoi());
		ImageProcessor ip2 = straightenRGB(imp2, width);
        imp.setRoi(imp2.getRoi());
        return ip2;
	}

	ImageProcessor rotateCompositeLine(ImagePlus imp, int width) {
		Image img = imp.getImage();
		ImagePlus imp2 = new ImagePlus("temp", new ColorProcessor(img));
		imp2.setRoi(imp.getRoi());
		ImageProcessor ip2 = rotateLine(imp2, width);
		return ip2;
	}

}
