package crushtools;
/*******************************************************************************
 * Copyright (c) 2015 Attila Kovacs <attila[AT]sigmyne.com>.
 * All rights reserved. 
 * 
 * This file is part of crushtools.
 * 
 *     crushtools is free software: you can redistribute it and/or modify
 *     it under the terms of the GNU General Public License as published by
 *     the Free Software Foundation, either version 3 of the License, or
 *     (at your option) any later version.
 * 
 *     crushtools is distributed in the hope that it will be useful,
 *     but WITHOUT ANY WARRANTY; without even the implied warranty of
 *     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *     GNU General Public License for more details.
 * 
 *     You should have received a copy of the GNU General Public License
 *     along with crushtools.  If not, see <http://www.gnu.org/licenses/>.
 * 
 * Contributors:
 *     Attila Kovacs <attila[AT]sigmyne.com> - initial API and implementation
 ******************************************************************************/

import java.util.*;

import jnum.Unit;
import jnum.Util;
import jnum.astro.AstroTime;
import jnum.math.Coordinate2D;
import jnum.math.SphericalCoordinates;
import jnum.math.Vector2D;
import jnum.util.BufferedRandom;
import nom.tam.fits.*;

public class ImageTool {
	static SkyMap image;

	public static void main(String args[]) {
		versionInfo();

		if(args.length == 0) usage();
		else if(args[0].equalsIgnoreCase("-help")) usage();

		AstroTime time = new AstroTime();
		time.now();
		
		try {	    
			image = new SkyMap(args[args.length-1]);

			for(int i=0; i<args.length-1; i++) {
				option(args[i]);
				//boolean success = option(args[i]);
				//if(success) image.history.add(" imagetool: " + args[i].substring(1));
			}
			
			System.out.println(" Updated image information: \n");
			
			image.info();
			image.history.add("Processed with CRUSH/imagetool v." + SkyMap.version + " @ " + time.getFitsTimeStamp()); 
			   
			image.write();
		}
		catch(Exception e) { e.printStackTrace(); }
	}

	public static boolean option(String optionString) {
		StringTokenizer tokens = new StringTokenizer(optionString, "=,");

		String key = tokens.nextToken();

		if(key.equalsIgnoreCase("-crop") || key.equalsIgnoreCase("-clip")) {
			if(!tokens.hasMoreTokens()) image.autoCrop();
			else {
				double dXmin = Double.parseDouble(tokens.nextToken()) * Unit.arcsec;
				double dYmin = Double.parseDouble(tokens.nextToken()) * Unit.arcsec;
				double dXmax = Double.parseDouble(tokens.nextToken()) * Unit.arcsec;
				double dYmax = Double.parseDouble(tokens.nextToken()) * Unit.arcsec;

				SphericalCoordinates coords = image.projection.getReference();
				
				// For equatorial images X offsets are RA offsets...
				if(coords.isReverseLongitude()) {
					double temp = dXmin;
					dXmin = -dXmax;
					dXmax = -temp;				
				}
				
				image.crop(dXmin, dYmin, dXmax, dYmax);
			}
		}
		else if(key.equalsIgnoreCase("-pxlcrop")) {       
		    image.crop(
		            Integer.parseInt(tokens.nextToken()), Integer.parseInt(tokens.nextToken()), 
		            Integer.parseInt(tokens.nextToken()), Integer.parseInt(tokens.nextToken())
		    );       
		}
		else if(key.equalsIgnoreCase("-out")) {
			image.fileName = Util.getSystemPath(tokens.nextToken());
		}
		else if(key.equalsIgnoreCase("-pick")) {
			try { image.writeImage(tokens.nextToken()); }
			catch(Exception e) { e.printStackTrace(); }
			System.exit(0);
		}
		else if(key.equalsIgnoreCase("-replace")) {
			StringTokenizer values = new StringTokenizer(tokens.nextToken(), ":= \t");
			String plane = values.nextToken().toLowerCase();
			String fileName = values.nextToken();
			
		
			try {
			    Fits fits = new Fits(fileName);
				BasicHDU<?> hdu = fits.getHDU(0);
				if(plane.equals("flux")) {
					System.err.println("Replacing flux plane.");
					image.signal = image.readDoubleImage(hdu, image.nativeUnit.value());
					image.makeFlagFrom(image.signal);
				}
				else if(plane.equals("weight")) {
					System.err.println("Replacing weight plane.");
					image.weight = image.readDoubleImage(hdu, 1.0 / (image.nativeUnit.value() * image.nativeUnit.value()));
					image.calc();
				}
				else if(plane.equals("time")) {
					System.err.println("Replacing integration-time plane.");
					image.count = image.readDoubleImage(hdu, Unit.s);
					
				}
				else if(plane.equals("flag")) {
					System.err.println("Replacing flag plane.");
					try { image.flag = image.readIntImage(hdu); }
					catch(IllegalArgumentException e) {
						image.makeFlagFrom(image.readDoubleImage(hdu, 1.0));
					}
				}
				else {
					System.err.println("Unknown image plane: " + plane);
					fits.close();
					return false;
				}
				image.addProcessingHistory("repaced " + plane + " from " + fileName) ;
				fits.close();
			}
			catch(Exception e) { e.printStackTrace(); }
		}
		else if(key.equalsIgnoreCase("-flag")) {
			try {
				String fileName = tokens.nextToken();
				Fits fits = new Fits(fileName);
				BasicHDU<?> hdu = fits.getHDU(0);
				try { image.flag = image.readIntImage(hdu); }
				catch(IllegalArgumentException e) {
					image.makeFlagFrom(image.readDoubleImage(hdu, 1.0));
				}	
				image.addProcessingHistory("new flags from " + fileName) ;
				System.err.println("Applying new flag image.");
				fits.close();
			}
			catch(Exception e) { e.printStackTrace(); }
		}
		else if(key.equalsIgnoreCase("-name")) {
			image.fileName = SkyMap.workPath + Util.getSystemPath(tokens.nextToken());
		}
		else if(key.equalsIgnoreCase("-correct")) {
			String value = tokens.nextToken();
			if(value.equalsIgnoreCase("point")) image.correct(image.instrumentFWHM);
			else if(value.equalsIgnoreCase("none")) image.uncorrect();
			else image.correct(Double.parseDouble(value) * Unit.arcsec);
		}	
		else if(key.equalsIgnoreCase("-regrid")) {
			image.regrid(Double.parseDouble(tokens.nextToken()) * Unit.arcsec);
		}
		else if(key.equalsIgnoreCase("-convolve")) {
			image.convolve(Double.parseDouble(tokens.nextToken()) * Unit.arcsec);
		}
		else if(key.equalsIgnoreCase("-smooth")) {
			image.convolveTo(Double.parseDouble(tokens.nextToken()) * Unit.arcsec);
		}
		else if(key.equalsIgnoreCase("-extFilter")) {
			double fwhm = Double.parseDouble(tokens.nextToken()) * Unit.arcsec;
			double blanking = tokens.hasMoreTokens() ? Double.parseDouble(tokens.nextToken()) : Double.NaN;
			image.filterAbove(fwhm, blanking);
		}

		else if(key.equalsIgnoreCase("-deconvolve")) {
			double[][] beam = null;
			double replacementFWHM = image.instrumentFWHM / 2.0;
			
			if(tokens.hasMoreTokens()) {
				
				StringTokenizer specs = new StringTokenizer(tokens.nextToken(), ":");
				String beamSpec = specs.nextToken(); 

				try { beam = image.getBeam(Double.parseDouble(beamSpec) * Unit.arcsec); }

				catch(NumberFormatException e)  {
					try {
						SkyMap beamImage = new SkyMap(beamSpec);
						beamImage.scale(1.0 / beamImage.getMax(beamImage.signal));
						beam = beamImage.signal;
					}
					catch(Exception e2) { System.err.println("ERROR! Cannot open beam FITS file " + beamSpec); }		
				}		

				if(specs.hasMoreTokens()) replacementFWHM = Double.parseDouble(specs.nextToken()) * Unit.arcsec;

				if(beam != null) image.clean(beam, 0.1, replacementFWHM);
			}
			else image.clean();
		}
		else if(key.equalsIgnoreCase("-minexp")) {
			image.exposureClip(Double.parseDouble(tokens.nextToken()));
		}
		else if(key.equalsIgnoreCase("-maxnoise")) {
			image.rmsClip(Double.parseDouble(tokens.nextToken()));
		}
		else if(key.equalsIgnoreCase("-s2nclip")) {
			image.s2nClip(Double.parseDouble(tokens.nextToken()));
		}
		else if(key.equalsIgnoreCase("-growflags")) {
			double radius = image.getImageFWHM();
			int pattern = ~0;
			if(tokens.hasMoreTokens()) radius = Double.parseDouble(tokens.nextToken()) * Unit.arcsec;
			if(tokens.hasMoreTokens()) pattern = Integer.decode(tokens.nextToken());
			image.growFlags(radius, pattern);			
		}
		else if(key.equalsIgnoreCase("-rmsscale")) {
			image.rmsScale(Double.parseDouble(tokens.nextToken()));
		}
		else if(key.equalsIgnoreCase("-flatnoise")) {
			image.setNoise(Double.parseDouble(tokens.nextToken()) * image.nativeUnit.value());
		}
		else if(key.equalsIgnoreCase("-scale")) {
			image.scale(Double.parseDouble(tokens.nextToken()));
		}
		else if(key.equalsIgnoreCase("-noise")) {
			String value = tokens.nextToken();
			if(value.equalsIgnoreCase("data")) image.dataWeight();
			else if(value.equalsIgnoreCase("image")) image.robustWeight();
			else return false;
		}
		else if(key.equalsIgnoreCase("-offset")) {
			image.offset(Double.parseDouble(tokens.nextToken()) * image.nativeUnit.value());
		}
		else if(key.equalsIgnoreCase("-origin")) {
			String x = tokens.nextToken();
			String y = tokens.nextToken();
			try {
				Vector2D offset = new Vector2D(Double.parseDouble(x), Double.parseDouble(y));
				offset.scale(Unit.arcsec);
				SphericalCoordinates coords = image.projection.getReference();
				
				if(coords.isReverseLongitude()) offset.invertX();
				if(coords.isReverseLatitude()) offset.invertY();
				
				image.projection.getReference().addOffset(offset);
				image.addProcessingHistory("reference point moved by "
						+ Util.F3.format(x) + "," + Util.f3.format(y) + " arcsec."
				);
				
			}
			catch(NumberFormatException e) {
				SphericalCoordinates origin = image.projection.getReference().copy();
				origin.parse(x + " " + y);
				Coordinate2D iRef = image.projection.getProjected(origin);
				image.i0 = iRef.x() / image.delta;
				image.j0 = iRef.y() / image.delta;
				image.projection.getReference().copy(origin);
				image.addProcessingHistory("reference set to " + origin.toString(2));
			}
		}
		else if(key.equalsIgnoreCase("-random")) {
			Random random = new BufferedRandom();
			image.smoothFWHM = image.delta / Math.sqrt(1.14);
			image.extFilterFWHM = Double.NaN;
			for(int i=0; i<image.sizeX; i++) for(int j=0; j<image.sizeY; j++) image.signal[i][j] = random.nextGaussian() / Math.sqrt(image.weight[i][j]);
			image.calc();
			image.addProcessingHistory("new random image.");
		}
		else if(key.equalsIgnoreCase("-masks")) {
			StringTokenizer actions = new StringTokenizer(tokens.nextToken(), ", \t");
			
			while(actions.hasMoreTokens()) {
				String action = actions.nextToken();
				if(action.equalsIgnoreCase("add")) image.addSources();
				else if(action.equalsIgnoreCase("remove")) image.removeSources();
				else if(action.equalsIgnoreCase("pointadd")) image.addPointSources();
				else if(action.equalsIgnoreCase("pointremove")) image.removePointSources();
				else if(action.equalsIgnoreCase("flag")) image.flagRegions();
				else if(action.equalsIgnoreCase("xflag")) image.flagOutsideRegions();
				else if(action.equalsIgnoreCase("forget")) image.clearRegions();
				else if(action.equalsIgnoreCase("image")) {
					image.extFilterFWHM = Double.NaN;
					for(int i=0; i<image.sizeX; i++) for(int j=0; j<image.sizeY; j++) image.signal[i][j] = 0.0;
					image.addSources();
					image.calcS2N();
					image.addProcessingHistory("new image from source mask.");
				}
				else if(action.equalsIgnoreCase("stack")) {
					image = image.stack();
					int iExt = image.fileName.lastIndexOf(".");
					if(iExt < 0) iExt = image.fileName.length();
					image.fileName = image.fileName.substring(0, iExt) + ".stacked" + image.fileName.substring(iExt);
					image.regions.clear();
					image.polygons.clear();
				}
				else if(action.equalsIgnoreCase("match")) {
					double pointingRMS = image.delta;
					if(action.length() > "match:".length()) {
						pointingRMS = Double.parseDouble(action.substring("match:".length())) * Unit.arcsec;
					}
					image.matchPeaks(pointingRMS);
					System.exit(0);
				}
				else if(action.toLowerCase().startsWith("extract")) {
					double pointingRMS = image.delta;
					if(action.length() > "extract:".length()) {
						pointingRMS = Double.parseDouble(action.substring("extract:".length())) * Unit.arcsec;
					}
					image.extractRegions(pointingRMS);
					System.exit(0);
				}
			}
		}
		else if(key.equalsIgnoreCase("-shift")) {
			SphericalCoordinates coords = image.projection.getReference();
			double dx = Double.parseDouble(tokens.nextToken());
			double dy = Double.parseDouble(tokens.nextToken());
			image.i0 += (coords.isReverseLongitude() ? 1 : -1) * dx * Unit.arcsec / image.delta;
			image.j0 += (coords.isReverseLatitude() ? 1 : -1) * dy * Unit.arcsec / image.delta;
			image.addProcessingHistory("pointing adjusted by " + dx + "," + dy + " arcsec.");
		}
		else if(key.equalsIgnoreCase("-unit")) {
			String unitName = tokens.nextToken();
			image.setUnit(unitName);
			image.addProcessingHistory("signal unit changed to " + unitName);
		}
		else if(key.equalsIgnoreCase("-source")) {
			image.sourceName = tokens.nextToken();
			image.addProcessingHistory("source name changed to " + image.sourceName);
		}
		else if(key.equalsIgnoreCase("-mask")) {
			String maskName = Util.getSystemPath(tokens.nextToken());
		    try { image.readMask(maskName, false); }
		    catch(Exception e) {
		    	try { image.readMask(SkyMap.workPath + maskName, false); }
		    	catch(Exception e2) { System.err.println("ERROR!: " + e2.getMessage()); }
		    }
		}
		else if(key.equalsIgnoreCase("-polygon")) {
			String fileName = Util.getSystemPath(tokens.nextToken());
		    try { image.addPolygon(fileName); }
		    catch(Exception e) {
		    	try { image.addPolygon(SkyMap.workPath + fileName); }
		    	catch(Exception e2) { System.err.println("ERROR!: " + e2.getMessage()); }
		    }
		}
		else if(key.equalsIgnoreCase("-changeKey")) {
			StringTokenizer args = new StringTokenizer(tokens.nextToken(), ":");
			if(args.countTokens() < 2) return false;

			String fitsKey = args.nextToken();
			String fitsValue = args.nextToken();

			System.out.println("Changing value of FITS key " + fitsKey + " to " + fitsValue);

			try { 
				image.changeKey(fitsKey, fitsValue); 
			}
			catch(Exception e) { System.err.println("ERROR!: " + e.getMessage()); }
		}
		else if(key.equalsIgnoreCase("-printHeader")) {
			System.out.println("\nFITS Header:\n"); 
			image.printHeader(); 
			System.exit(0);
		}
		else if(key.equalsIgnoreCase("-depth")) {
			System.err.println("Average depth " + Util.e2.format(image.getTypicalRMS()/image.nativeUnit.value()) + " " + image.nativeUnit.name());
		}
		else if(key.equalsIgnoreCase("-beam")) {
			image.instrumentFWHM = Double.parseDouble(tokens.nextToken()) * Unit.arcsec;
			image.addProcessingHistory("instrument beam set to " + (image.instrumentFWHM / Unit.arcsec) + " arcsec FWHM.");
		}
		else if(key.equalsIgnoreCase("-level")) {
			image.robustLevel();
		}
		else if(key.equalsIgnoreCase("-profile")) {
			image.profile(true);
		}
		else if(key.equalsIgnoreCase("-azimuthal")) {
		    double resolution = tokens.hasMoreTokens() ? Double.parseDouble(tokens.nextToken()) * Unit.deg : 15.0 * Unit.deg;
            image.angularProfile(0.0, resolution, true);
        }
		
			
		else if(key.equalsIgnoreCase("-help")) usage();
		else return false;

		return true;

	}
	
	public static void clipRMS(double level) {
		for(int i=0; i<image.sizeX; i++) for(int j=0; j<image.sizeY; j++)
			if(Math.abs(image.rms[i][j]) > level) image.flag[i][j] = 1;
	}

	public static void versionInfo() {
		System.out.println("imagetool -- Image Manipulation Utility for CRUSH maps and more.");
		System.out.println("             Author: Attila Kovacs <attila[AT]sigmyne.com>");
		System.out.println("             Version: " + SkyMap.version);
		System.out.println();
	}

	public static void usage() {
		System.out.println("Usage: imagetool [options] <filename>");
		System.out.println();
		System.out.println("Options:");
		System.out.println();

		optionsSummary();

		System.out.println();	
		System.exit(0);
	}

	public static void optionsSummary() {
		System.out.println("  [Output Options]");
		System.out.println("    -out=          Full output file name with path. Default overwrites input.");
		System.out.println("    -name=         Output file name relative to input path.");
		System.out.println("    -pick=         Writes a single image (flux|s2n|rms|time) then exits.");
		System.out.println();
		System.out.println("  [Clipping/Cropping]");
		System.out.println("    -crop          Automatically crop the image at its edges.");
		System.out.println("    -crop=         Crop at offsets dX1,dY1,dX2,dY2 (in arcsec).");
		System.out.println("    -pxlCrop=      Crop by pixel indices imin,jmin,imax,jmax (0-based).");
		System.out.println("    -s2nClip=      Flag pixels below the specified significance level.");
		System.out.println("    -minExp=       Flag pixels with smaller realtive exposures.");
		System.out.println("    -maxNoise=     Flag pixels with higher relative noise levels.");
		System.out.println("    -growFlags=    Specify a radius (in arcsec) by which flagged areas are");
		System.out.println("                   increased.");
		System.out.println();
		System.out.println("  [Smoothing/Filtering]");
		System.out.println("    -convolve=     Convolve with the specified FWHM beam size (arcsec). ");
		System.out.println("    -smooth=       Smooth map to the specified FWHM beam size (arcsec). ");
		System.out.println("    -extFilter=    Filter extended structures above this FWHM (arcsec). ");
		System.out.println("    -deconvolve[=] Deconvolve image. Optional argument is either:");
		System.out.println("                     BeamFWHM[:replacementFWHM]");
		System.out.println("                     <BeamFITSFile>[:replacementFWHM]");
		System.out.println();
		System.out.println("  [Masking Options]");
		System.out.println("    -mask=         Specify a region mask file.");
		System.out.println("    -masks=        Perform tasks with mask files. Defined tasks are:");
		System.out.println("                        forget                Delete existing masks.");
		System.out.println("                        add/remove            Add/Remove as Gaussian sources.");
		System.out.println("                        pointadd/pointremove  Add/Remove as point sources.");
		System.out.println("                        extract               Extract counterparts to catalog.");
		System.out.println("                        flag                  Flag circular regions.");
		System.out.println("                        xflag                 Flag outside regions.");
		System.out.println("                        stack                 Stack sources from mask.");
		System.out.println("                        match                 Match sources to mask.");
		System.out.println("    -flag=         Specify a FITS image for image flagging.");
		System.out.println("    -profile       Obtain radial profiles around circular regions.");
		System.out.println("    -azimuthal     Obtain azimuthal profiles around circular regions.");
		System.out.println();
		System.out.println("  [FITS Header]");
		System.out.println("    -source=       Change the source name.");
		System.out.println("    -Jy=           Specify the voltage response for a 1 Jy point source.");	
		System.out.println("    -shift=        Shift map alignment by the comma separated offsets (arcsec).");
		System.out.println("    -changeKey=    Change an entry in the FITS header (key:value)");
		System.out.println("    -printHeader   Print the FITS image header.");
		System.out.println();
		System.out.println("  [Miscellaneous]");
		System.out.println("    -noise=        Select the source of the noise estimate. ('data' or 'image')");
		System.out.println("    -regrid=       Regrid image to the specified grid size (arcsec).");
		System.out.println("    -correct=      Apply flux corrections to the specified FWHM source(s).");
		System.out.println("    -scale=        Scale the map and noise with the specified factor.");
		System.out.println("    -rmsScale=     Scale the map noise with the specified factor.");
		System.out.println("    -offset=       Add a specified offset (in map units).");
		System.out.println("    -origin=       Comma separated offsets or coordinates of the new origin.");
		System.out.println("    -unit=         Specify the output map unit.");		
		System.out.println("    -random        Creates an unsmoothed random map based on the supplied map.");
		System.out.println("    -flatnoise=    Use the specified flat noise rms for the entire image.");
		System.out.println("    -replace=      <plane>:<file> to replace an image <plane> ('flux','weight',");
		System.out.println("                   'time' or 'rms') with the FITS image from <file>.");
		System.out.println("    -help          Provides this help screen.");
		System.out.println();
	}
}


