import java.util.*;


public class ReducedCoadder {
    static String version = "13 May 2005";
    
    static String fileName;
    static boolean names = true;
    static boolean celestial = false;
    static boolean pointRMS = false;
    static boolean show = false;
    static boolean offsets = true;
    static boolean weighting = true;

    static int rounds = 5;

    static double scalingFactor = 1.0;
    static double correctingFWHM = Double.NaN;

    static Vector options = new Vector();

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

	if(args.length == 0) usage(); 

	int nmaps=0;
	for(int i=0; i<args.length; i++) if(!isOption(args[i])) nmaps++;
	ReducedData[] maps = new ReducedData[nmaps];
      
	for(int i=0,k=0; i<args.length; i++) {
	    if(isOption(args[i])) option(args[i]);
	    else {
		try {	
		    String name = args[i]; 
		    System.out.print(" Reading " + name + "\t");
		    System.out.flush();		    
		    maps[k] = new ReducedData(name);
		    maps[k].scale(scalingFactor);

		    if(correctingFWHM <= 0.0) maps[k].correctFluxesFor(maps[k].beamFWHM);
		    else maps[k].correctFluxesFor(correctingFWHM);

		    k++;
		}
		catch(Exception e) { e.printStackTrace(); }
	    } 
	}

	if(fileName == null) fileName = Crush.workPath + maps[0].sourceName + ".coadded.fits"; 
	
	ReducedData coadded = add(maps, rounds);

	coadded.fileName = fileName;
	if(show) coadded.show();
	try { coadded.write(); } catch(Exception e) { e.printStackTrace(); }
    }

    public static ReducedData[] getIdenticalGridMaps(ReducedData[] map) {
	double minRA, maxRA, minDEC, maxDEC;

	int N = map.length;
	
	double RA0 = map[0].RA;
	double DEC0 = map[0].DEC;
	double minDelta = map[0].delta;
	
	// Regrid all maps to the same delta
	for(int i=1; i<N; i++) if(map[i].delta < minDelta) minDelta = map[i].delta;
	for(int i=0; i<N; i++) if(map[i].delta > minDelta) map[i].regrid(minDelta);
	
	minRA  = map[0].getRAofIndex(-1); maxRA  = map[0].getRAofIndex(map[0].sizeX);
	minDEC = map[0].getDECofIndex(-1); maxDEC = map[0].getDECofIndex(map[0].sizeY);

	for(int i=1; i<map.length; i++) {
	    double mapMinRA = map[i].getRAofIndex(-1);
	    double mapMinDEC = map[i].getDECofIndex(-1);
	    double mapMaxRA = map[i].getRAofIndex(map[i].sizeX);
	    double mapMaxDEC = map[i].getDECofIndex(map[i].sizeY);

	    if(mapMinRA  < minRA ) minRA  = mapMinRA;
	    if(mapMinDEC < minDEC) minDEC = mapMinDEC;
	    if(mapMaxRA  > maxRA) maxRA  = mapMaxRA;
	    if(mapMaxDEC > maxDEC) maxDEC = mapMaxDEC;
	}

	// Make the coadded a pixel wider on all edges
	// and readjust the reference to the new edge;
	int sizeRA =  (int) Math.ceil((maxRA  - minRA ) * Math.cos(DEC0) / minDelta);
	int sizeDEC = (int) Math.ceil((maxDEC - minDEC) / minDelta);	

	System.out.println("Size: " + sizeRA + " x " + sizeDEC);

	ReducedData[] regrid = new ReducedData[N];
	// Regrid the maps onto a grid equivalent to the final grid
	for(int k=0; k<N; k++) {	
	    regrid[k] = (ReducedData) map[k].clone();
	    regrid[k].setSize(sizeRA, sizeDEC);

	    regrid[k].delta = minDelta;
	    regrid[k].dRA0 = (minRA - regrid[k].RA) * Math.cos(regrid[k].DEC);
	    regrid[k].dDEC0 = (minDEC - regrid[k].DEC);
	    regrid[k].i0 = 0;
	    regrid[k].j0 = 0;

	    for(int i=0; i<map[k].sizeX; i++) for(int j=0; j<map[k].sizeY; j++) if(map[k].flag[i][j] == 0) {
		    double aRA = map[k].getRAofIndex(i);
		    double aDEC = map[k].getDECofIndex(j);

		    regrid[k].addPointAt(aRA, aDEC, map[k].clean[i][j], map[k].weight[i][j], map[k].count[i][j]);
		}
		
	    regrid[k].normalize();
	}

	System.out.println("\n\nCoadded Map is:");
	regrid[0].info();

	return regrid;
    }

    public static ReducedData add(ReducedData[] map, int iter) {
	ReducedData[] regrid = getIdenticalGridMaps(map);

	int N = map.length;

	// create the map weights array
	double[] mapWeight = new double[map.length];
	for(int i=0; i<N; i++) mapWeight[i] = 1.0;

	// Create the coadded map
	ReducedData coadded = (ReducedData) regrid[0].clone();
	int sizeRA = coadded.sizeX;
	int sizeDEC = coadded.sizeY;
	coadded.setSize(sizeRA, sizeDEC);
	coadded.header = null;

	// Fill in the coadded header information for the FITS file
	coadded.creator = "coadd";
	coadded.creatorVersion = Crush.version;
	coadded.options = new String[options.size()];
	for(int i=0; i<options.size(); i++) coadded.options[i] = (String) options.get(i);
	
	// Fill in the scan number information
	int nScans = 0;
	for(int i=0; i<N; i++) nScans += map[i].scan.length;
	coadded.scan = new SharcScan[nScans];
	for(int i=0,j=0; i<N; i++) for(int k=0; k<map[i].scan.length; k++,j++)
	    coadded.scan[j] = map[i].scan[k];
	

	double[] offset = new double[N];
	

	if(!weighting) iter=1;

	// Here comes the iterative coadding with smart weights adjusting
	for(int ii=0; ii<iter; ii++) {
	    // System.out.println("\tIteration " + (ii+1));

	    if(offsets) {
		OffsetFit fit = new OffsetFit(regrid, mapWeight);
		fit.minimize();
		fit.apply();
	    }
	    	    
	    // Now just add them all together
	    coadded.clear();

	    for(int i=0; i<sizeRA; i++) for(int j=0; j<sizeDEC; j++) {	
		for(int k=0; k<N; k++) if(regrid[k].flag[i][j] == 0) {
		    double w = mapWeight[k] * regrid[k].weight[i][j];
		    coadded.clean[i][j] += w * regrid[k].clean[i][j];
		    coadded.count[i][j] += regrid[k].count[i][j];
		    coadded.weight[i][j] += w;
		    coadded.flag[i][j] = 0;
		}
	    }

	    coadded.normalize();
	    
	    // Re-adjust the map weights -> 1/chi^2
	    for(int k=0; k<N; k++) {
		double chisq = 0.0;
		int n=0;
		for(int i=0; i<sizeRA; i++) for(int j=0; j<sizeDEC; j++) 
		    if(regrid[k].flag[i][j] == 0 && coadded.flag[i][j] == 0) {
			double dev = regrid[k].clean[i][j] - coadded.clean[i][j];
			chisq += dev * dev * regrid[k].weight[i][j];
			n++;
		    }
		
		if(chisq > 0.0) mapWeight[k] = n / chisq;
		
		// Do not allow weights above 1. They would lead to RMS estimate
		// smaller than justified by the original maps
		if(mapWeight[k] > 1.0) mapWeight[k] = 1.0;
	    }			    
	}

	if(weighting) for(int i=0; i<N; i++) System.out.println( " weight " + i + " = " + mapWeight[i]); 
	coadded.map = coadded.clean;
	 
	
	if(pointRMS) for(int i=0; i<coadded.sizeX; i++) for(int j=0; j<coadded.sizeY; j++) {
	    double var = 0.0, sumw = 0.0;
	    int n=0;
	    for(int k=0; k<N; k++) {
		double w = mapWeight[k] * regrid[k].rms[i][j];
		double value = regrid[k].clean[i][j] - coadded.clean[i][j];
		var += w * value*value;
		sumw += w;
		n++;
	    }
	    coadded.rms[i][j] = Math.sqrt(var/sumw/(n-1));
	}

	return coadded;		
    }

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

	String key = tokens.nextToken();

	options.add(optionString);
	
	if(key.equalsIgnoreCase("-out")) fileName = tokens.nextToken();	    

	else if(key.equalsIgnoreCase("-name")) fileName = Crush.workPath + tokens.nextToken();

	else if(key.equalsIgnoreCase("-rounds")) rounds = Integer.parseInt(tokens.nextToken());	    
	
	else if(key.equalsIgnoreCase("-celestial")) celestial = true;

	else if(key.equalsIgnoreCase("-pointrms")) pointRMS = true; 
	
	else if(key.equalsIgnoreCase("-show")) show = true; 

	else if(key.equalsIgnoreCase("-nooffset")) offsets = false;

	else if(key.equalsIgnoreCase("-noweight")) weighting = false;

	else if(key.equalsIgnoreCase("-direct")) { offsets = false; weighting = false; }

	else if(key.equalsIgnoreCase("-scale")) scalingFactor = Double.parseDouble(tokens.nextToken());

	else if(key.equalsIgnoreCase("-correct")) {
	    String value = tokens.nextToken();
	    if(value.equalsIgnoreCase("none")) correctingFWHM = Double.NaN;
	    else if(value.equalsIgnoreCase("point")) correctingFWHM = 0.0;
	    else correctingFWHM = Double.parseDouble(value) * Unit.ArcSECOND;
	}

	else if(key.equalsIgnoreCase("-help")) usage();

	else System.out.println("WARNING! Unrecognized option: " + key);
    
    }

    public static boolean isOption(String argument) {
	return argument.charAt(0) == '-';
    }

    public static void versionInfo() {
	System.out.println("coadd -- Coaddition and Manipulation Utility for SHARC2 maps");
	System.out.println("         Author: Attila Kovacs <attila@submm.caltech.edu");
	System.out.println("         Version: " + version);
	System.out.println();
    }

    public static void usage() {	
	System.out.println("usage: coadd [options] scanNo...");
	System.out.println();
	System.out.println("Options:");
	System.out.println();
	System.out.println("   -help      Produces this help screen.");
	System.out.println("   -show      Display coadded data when done.");
	System.out.println("   -out=      Specify the output file name. The default value is");
	System.out.println("              [source name].coadded.fits.");
	System.out.println("   -name=     Specify the output file name relative to output path.");
	System.out.println("   -rounds=   The number of iterations in weighted coadding (5)");
	System.out.println("   -celestial Coadd assuming all maps are centered on the same source.");
	System.out.println("   -pointrms  Calculate the rms as the scatter per point.");
	System.out.println("   -noweight  Do not do iterative weighting. Use map variance directly.");
	System.out.println("   -nooffset  Do not solve for map offsets.");
	System.out.println("   -direct    Add maps directly. Same as -nooffset and -noweight together.");
	System.out.println("   -scale=    Scale the following maps by the specified value (1.0).");
	System.out.println("   -correct=  Correct fluxes for FWHM source size in arcsec.");

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

}

class OffsetFit extends AmoebaMinimizer {
    int N;
    ReducedData[] map;
    double[] mapWeight;
    
    public OffsetFit(ReducedData[] data, double[] weight) {
	map = data;
	mapWeight = weight;
	N = map.length;
	
	double[] initparm = new double[N];		
	init(initparm);
	
	for(int i=0; i<N; i++) p[i+1][i] += 0.1 * (map[i].getMax() - map[i].getMin());		
    }
    
    public double minimize() {
	double min = super.minimize();
	
	double ave = 0.0;
	for(int k=0; k<N; k++) ave += parm[k];
	ave /= N;

	// Remove the offsets from the maps
	for(int k=0; k<N; k++) 
	    for(int i=0; i<map[k].sizeX; i++) for(int j=0; j<map[k].sizeY; j++) 
		map[k].clean[i][j] -= parm[k] - ave;
	
	return min;
    }

    public double evaluate(double[] tryparm) {
	double chiSquared = 0.0;
	double sumw = 0.0;
	int n = 0;
	
	for(int k=0; k<N; k++) for(int m=k+1; m<N; m++) {
	    for(int i=0; i<map[k].sizeX; i++) for(int j=0; j<map[k].sizeY; j++) 
		if(map[k].flag[i][j] == 0) if(map[m].flag[i][j] == 0) {
		    double wk = mapWeight[k] / (map[k].rms[i][j] * map[k].rms[i][j]);
		    double wm = mapWeight[m] / (map[m].rms[i][j] * map[m].rms[i][j]);
		    
		    if(wk > 0.0 && wm > 0.0) {    
			double var = (1.0/wk + 1.0/wm);
		       
			double value = map[k].clean[i][j] - parm[k];
			value -= map[m].clean[i][j] - parm[m];
			
			chiSquared += value*value / var;
			n++;
		    }
		}
	}
	
	return chiSquared/(n-N+1);
    }	
    
    public void apply() {
	for(int k=0; k<N; k++) for(int m=k+1; m<N; m++) {
	    for(int i=0; i<map[k].sizeX; i++) for(int j=0; j<map[k].sizeY; j++) 
		map[k].clean[i][j] -= parm[k];
	    map[k].calcS2N();
	}
    }
}
