package crushtools;
/*******************************************************************************
 * Copyright (c) 2014 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 crushtools.util.AmoebaMinimizer;

import jnum.Util;
import jnum.astro.AstroTime;
import jnum.astro.CelestialCoordinates;
import jnum.astro.EquatorialCoordinates;
import jnum.math.Coordinate2D;
import jnum.math.SphericalCoordinates;
import jnum.math.Vector2D;
import jnum.projection.SphericalProjection;
import nom.tam.fits.*;

public class SkyMapCoadder {
	static String fileName;
	static boolean absoluteFileName = false;
	static Vector<String> options = new Vector<String>();
	
	static SphericalProjection reprojection;
	static boolean weighting = false, scaling = false, offsets = false, dataWeights = false;
	static boolean jackknife = false;
	
	public static void main(String[] args) {		
		Vector<String> history = new Vector<String>();
		
		versionInfo();

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

		Vector<SkyMap> maps = new Vector<SkyMap>();

		AstroTime time = new AstroTime();
		time.now();
		
	
		
		for(int i=0; i<args.length; i++) {
			if(isOption(args[i])) {
				boolean success = option(args[i]);
				if(success) history.add(args[i].substring(1));
			}
			else {
				try {	
					String name = args[i]; 
					SkyMap component = new SkyMap();
					component.verbose = false;
					
					if(weighting | scaling | offsets) {
						System.out.println(" Reading " + name + "\t");	 
						component.read(name);
						if(jackknife) if(Math.random() < 0.5) component.scale(-1.0);
					}
					else component.readHeader(name);
					
					SphericalCoordinates coords = component.projection.getReference();
					
					//component.altaz = false;
					if(coords instanceof CelestialCoordinates) { 
						maps.add(component); 
						history.add("[+] " + name);
					}
					else { System.err.println("WARNING! " + component.fileName + " is not in a celestial frame. Cannot coadd."); }
			

				}
				catch(Exception e) { e.printStackTrace(); }
			} 
		} 
		System.err.println();
			
		if(maps.size() == 0) {
			System.err.println("ERROR! No valid maps to reduce. Exiting...");
			System.exit(1);			
		}
 		
		if(reprojection == null) {
			reprojection = maps.get(0).projection;
			System.err.println("Will use " + reprojection.getFullName() + " (" + reprojection.getFitsID() + ") projection.");
		}
			
		
		SkyMap[] component = maps.toArray(new SkyMap[maps.size()]);
		SkyMap coadded;
			
		if(weighting | scaling | offsets) coadded = fitTogether(component, weighting, scaling, offsets);
		else coadded = addSequential(component);
	
		coadded.history.add("Processed with CRUSH/coadd v." + SkyMap.version + " @ " + time.getFitsTimeStamp()); 
		for(int i=0; i<history.size(); i++) coadded.history.add(" coadd: " + history.get(i));
		
		if(fileName == null) coadded.fileName = SkyMap.workPath + maps.get(0).sourceName + ".coadded.fits";
		else coadded.fileName = absoluteFileName ? fileName : SkyMap.workPath + fileName;
		try { coadded.write(); } catch(Exception e) { e.printStackTrace(); }
	}
	

	
	
	void add(SkyMap[] component, SkyMap coadded) {
		double[] uniform = new double[component.length];
		Arrays.fill(uniform, 1.0);
		add(component, new double[component.length], uniform, coadded);		
	}
	
	// Weights can be defined by ReducedData.weightFactor....
	// Return the dof array...
	static double[] add(SkyMap[] component, double[] scaling, double[] offset, SkyMap coadded) {
		coadded.clear();
		coadded.integrationTime = 0.0;
		
		// add together the maps with the scaling and offsets...
		for(int k=0; k<component.length; k++) {
			SkyMap map = component[k];
			double b = offset[k];
			double a = scaling[k];
			if(a > 0.0) for(int i=0; i<coadded.sizeX; i++) for(int j=0; j<coadded.sizeY; j++) if(map.flag[i][j] == 0) {
				double rms = a * map.rms[i][j];
				double w = map.weightFactor / (rms * rms);
				coadded.signal[i][j] += w * (a * map.signal[i][j] + b);
				coadded.weight[i][j] += w;
				coadded.count[i][j] += map.count[i][j];
			}
			coadded.integrationTime += map.integrationTime;
		}
		
		coadded.normalize();
		
		// Count the dependence
		double[] dof = new double[component.length];
		for(int k=0; k<component.length; k++) {
			SkyMap map = component[k];
			double a = scaling[k];
			if(a > 0.0) for(int i=0; i<coadded.sizeX; i++) for(int j=0; j<coadded.sizeY; j++) if(map.flag[i][j] == 0) {
				double rms = a * map.rms[i][j];
				double coaddedRMS = coadded.rms[i][j];
				dof[k] += 1.0 - coaddedRMS * coaddedRMS * map.weightFactor / (rms * rms);
			}
		}
		
		return dof;
	}
	
	// Weights can be defined by ReducedData.weightFactor....
	static void addWeighting(SkyMap[] component, double[] scaling, double[] offset, SkyMap coadded) {
		double avedw = Double.NaN;
		int i = 1;
		
		while(avedw > 0.01 || i <= 10) {
			int contributors = 0;
			double[] dof = add(component, scaling, offset, coadded);
			double sumw = 0.0;
			
			for(int k=0; k<component.length; k++) {
				if(dof[k] > 0.0) {
					double chi2 = SkyMapCoadder.getChiSquared(coadded, component[k], scaling[k], offset[k]);
					chi2 /= dof[k];
					double weightScale = chi2 > 0.0 ? 1.0 / chi2 : 0.0;
					component[k].weightFactor *= weightScale;
					if(weightScale > 0.0) {
						sumw += component[k].weightFactor;
						contributors++;
					}
				}
				else component[k].weightFactor = 0.0;
			}
			
			// Normalize to average 1 weights...
			if(contributors > 0) {
				double avew = sumw / contributors;
				coadded.weightFactor = avew;
				for(int k=0; k<component.length; k++) component[k].weightFactor /= avew;
			}
			
			i++;
		}
	}
	
	
	static double getChiSquared(SkyMap coadded, SkyMap component, double scaling, double offset) {
		double chi2 = 0.0;
		for(int i=0; i<coadded.sizeX; i++) for(int j=0; j<coadded.sizeY; j++) if(component.flag[i][j] == 0) {
			double dev = (scaling * component.signal[i][j] + offset) - coadded.signal[i][j];
			double rms = scaling * component.rms[i][j];
			double w = component.weightFactor / (rms * rms);
			chi2 += w * dev * dev;
		}
		
		return chi2;
	}
	
	public static SkyMap createCoadded(SkyMap[] component, SphericalProjection projection) {	
		// Create the coadded map
		SkyMap coadded = (SkyMap) component[0].clone();

		// Fill in the coadded header information for the FITS file
		coadded.weightFactor = 1.0;
		coadded.projection = projection;
		coadded.projection.setReference((SphericalCoordinates) component[0].projection.getReference().clone());			
		coadded.creator = "crushtools (coadd)";
		coadded.creatorVersion = SkyMap.version;
		coadded.options = new String[options.size()];
		
		try { coadded.header.addLine(new HeaderCard("JACKKNFE", jackknife, "Is this a randomly jackknifed map?")); }
		catch(HeaderCardException e) {}
		
		for(int i=0; i<options.size(); i++) coadded.options[i] = options.get(i);

		if(!checkCommonGrid(component)) {
			System.err.print("Determining Coadded Map Size... ");
			
			// Make coadded the same resolution as the finest of them all...
			double resolution = component[0].delta;
			for(int i=1; i<component.length; i++) if(component[i].delta < resolution) resolution = component[i].delta;
			
			Vector2D minOffset = new Vector2D(Double.POSITIVE_INFINITY, Double.POSITIVE_INFINITY);
			Vector2D maxOffset = new Vector2D(Double.NEGATIVE_INFINITY, Double.NEGATIVE_INFINITY);

			for(int k=0; k<component.length; k++) {
				SkyMap map = component[k];

				// TODO Should go along edges, not just corners, since projections may 'bulge'
				for(int j=-1; j <= map.sizeY; j += map.sizeY+1) {
					updateReprojectedLimits(coadded.projection, map, 0, j, minOffset, maxOffset);
					updateReprojectedLimits(coadded.projection, map, map.sizeX-1, j, minOffset, maxOffset);
				}

				for(int i=-1; i <= map.sizeX; i+=map.sizeX+1) {
					updateReprojectedLimits(coadded.projection, map, i, 0, minOffset, maxOffset);
					updateReprojectedLimits(coadded.projection, map, i, map.sizeY-1, minOffset, maxOffset);
				}			
			}


			// Make the reference fall on pixel boundaries.
			coadded.i0 = 0.5 - Math.round(minOffset.x()/resolution);
			coadded.j0 = 0.5 - Math.round(minOffset.y()/resolution);	

			double rangeX = maxOffset.x() - minOffset.x();
			double rangeY = maxOffset.y() - minOffset.y();

			int sizeX = (int)Math.ceil(rangeX/resolution) + 2;
			int sizeY = (int)Math.ceil(rangeY/resolution) + 2;

			System.err.println(sizeX + "x" + sizeY);
			
			System.err.println("Creating Storage for Coadded Map...");
			coadded.setSize(sizeX, sizeY);

		}
		else { 
			System.err.println("Creating conformal Coadded Map...");
			coadded.setSize(coadded.sizeX, coadded.sizeY);
		}
		coadded.info();
		
		System.err.println();
		
		return coadded;
	}
		
	private static Coordinate2D reuseOffset = new Vector2D();
	private static SphericalCoordinates reuseCoords = new SphericalCoordinates();
	private static void updateReprojectedLimits(SphericalProjection projection, SkyMap map, int i, int j, Vector2D min, Vector2D max) {
		reuseOffset.set(map.dXofIndex(i), map.dYofIndex(j));		
		map.projection.deproject(reuseOffset, reuseCoords);
		projection.project(reuseCoords, reuseOffset);
		
		min.setX(Math.min(min.x(), reuseOffset.x()));
		min.setY(Math.min(min.y(), reuseOffset.y()));
		max.setX(Math.max(max.x(), reuseOffset.x()));
		max.setY(Math.max(max.y(), reuseOffset.y()));
	}
	
	public static Vector<SkyMap[]> groupConforming(SkyMap[] maps) {
		System.err.print("Grouping together maps with identical grids... ");
		
		Vector<Vector<SkyMap>> grids = new Vector<Vector<SkyMap>>();
		
		for(int i=0; i<maps.length; i++) {
			boolean conforming = false;
			for(int k=0; k<grids.size(); k++) if(maps[i].conforms(grids.get(k).get(0))) {
				grids.get(k).add(maps[i]);
				conforming = true;
				break;
			}
			if(!conforming) {
				Vector<SkyMap> grid = new Vector<SkyMap>();
				grid.add(maps[i]);
				grids.add(grid);
			}
		}
		
		Vector<SkyMap[]> groups = new Vector<SkyMap[]>();
		for(int i=0; i<grids.size(); i++) groups.add(grids.get(i).toArray(new SkyMap[grids.get(i).size()]));
		
		System.err.println(groups.size() + " groups identified.");
		
		return groups;
	}
	
	public static SkyMap addSequential(SkyMap[] maps) {
			
		// Regrid components to the same grid...		
		SkyMap coadded = createCoadded(maps, reprojection);
			
		//Vector<SharcScan> scans = new Vector<SharcScan>();
		Vector<SkyMap[]> groups = groupConforming(maps);
		
		for(int g=0; g<groups.size(); g++) {
			SkyMap[] component = groups.get(g);

			SkyMap groupCoadded = (SkyMap) component[0].clone();
			groupCoadded.setSize(groupCoadded.sizeX, groupCoadded.sizeY);

			for(int k=0; k<component.length; k++) {

				try { 
					System.err.println("Adding " + component[k].fileName + "... ");

					component[k].readData();
					component[k].sanitize();
					
					if(jackknife) if(Math.random() < 0.5) {
						System.err.println("Inverting");
						component[k].scale(-1.0);
					}
					if(dataWeights) component[k].dataWeight();
					else component[k].robustWeight();

					// Fill in the scan number information
					//for(int m=0; m<component[k].scan.length; m++) scans.add(component[k].scan[m]);

					groupCoadded.addDirect(component[k], 1.0);
					
					component[k].empty();
					component[k] = null;
					
					System.runFinalization();
					System.gc();
				}
				catch(Exception e) {
					System.err.println();
					System.err.println("ERROR! Problem encountered while reading image data.");
					System.err.println("       Skipping map...");
				}
			}

			groupCoadded.normalize();
			groupCoadded.sanitize();
			
			System.err.println("Regridding to final grid.");
			SkyMap regrid = (SkyMap) groupCoadded.clone();

			regrid.conformTo(coadded);
			groupCoadded.regridTo(regrid);
			
			System.err.println("Adding to final map.");
			coadded.addDirect(regrid, 1.0);

			regrid.empty();
			groupCoadded.empty();
			
			System.runFinalization();
			System.gc();
			
			System.err.println();
		}
		
		coadded.normalize();
		coadded.sanitize();
		
		//coadded.scan = scans.toArray(new SharcScan[0]);
		
		if(!dataWeights) coadded.robustWeight();
		
		return coadded;		
	}
	
	public static SkyMap fitTogether(SkyMap[] component, boolean weighting, boolean fitScaling, boolean fitOffsets) {
		// Regrid components to the same grid...	
		SkyMap coadded = createCoadded(component, reprojection);
		
		// Fill in the scan number information
		coadded.extraHDUs.clear();
		for(int i=0; i<component.length; i++) coadded.extraHDUs.addAll(component[i].extraHDUs);
		
		System.err.print("Checking if grids are matching... ");
		if(checkCommonGrid(component)) { System.err.println(":-)"); }
		else {
			System.err.println(":-(");
			component = getIdenticalGridMaps(component, coadded);
		}
			
		return fitTogether(component, weighting, fitScaling, fitOffsets, coadded);
	}
	
	public static SkyMap fitTogether(final SkyMap[] component, final boolean weighting, final boolean fitScaling, final boolean fitOffsets, final SkyMap coadded) {
		
		final double[] scaling = new double[component.length];
		final double[] offset = new double[component.length];
	
		for(int i=0; i<component.length; i++) {
			if(dataWeights) component[i].dataWeight();
			else component[i].dataWeight();
		}
		
		class LinearFit extends AmoebaMinimizer {	
			public LinearFit() {

				int nparms = 0;
				if(fitScaling) nparms += component.length;
				if(fitOffsets) nparms += component.length;

				double[] initparm = new double[nparms];
				double[] range = new double[nparms];

				int i=0;
				if(fitScaling) for(int k=0; k<component.length; k++, i++) { initparm[i] = scaling[k]; range[i] = 0.1*scaling[k]; } 
				if(fitOffsets) for(int k=0; k<component.length; k++, i++) { initparm[i] = offset[k]; range[i] = 0.1* (component[k].getMax(component[k].signal) - component[k].getMin(component[k].signal)); }

				verbose = true;

				init(initparm);
				setStartSize(range);
			}

			@Override
			public double evaluate(double[] tryparm) {
				// Do not allow negative scaling... simply use absolute value...

				int n = 0;
				if(fitScaling) for(int k=0; k<component.length; k++) scaling[k] = Math.abs(tryparm[n++]);
				if(fitOffsets) for(int k=0; k<component.length; k++) offset[k] = tryparm[n++]; 

				double sumw = 0.0;
				for(int k=0; k<component.length; k++) sumw += component[k].weightFactor;

				// Renormalize the offsets and scalings to weight mean zero and average one...

				if(fitOffsets) {
					double aveOffset = 0.0;
					for(int k=0; k<component.length; k++) aveOffset += component[k].weightFactor * offset[k];
					if(sumw > 0.0) aveOffset /= sumw;
					for(int k=0; k<component.length; k++) offset[k] -= aveOffset;
				}

				if(fitScaling) {
					double aveScaling = 0.0;
					for(int k=0; k<component.length; k++) aveScaling += component[k].weightFactor * scaling[k];
					if(sumw > 0.0) {
						aveScaling /= sumw;
						for(int k=0; k<component.length; k++) scaling[k] /= aveScaling;
					}
				}

				double[] dof = SkyMapCoadder.add(component, scaling, offset, coadded);

				double sumChi2 = 0.0;
				double sumDOF = 0.0;

				for(int k=0; k<component.length; k++) {
					double chi2 = SkyMapCoadder.getChiSquared(coadded, component[k], scaling[k], offset[k]);

					if(fitOffsets) dof[k]--;
					if(fitScaling) dof[k]--;

					sumChi2 += chi2;
					sumDOF += dof[k];

					chi2 /= dof[k];
				}

				//System.err.println(Util.e3.format(sumChi2/sumDOF));

				return sumChi2/sumDOF;	
			}	
		}
		
			
		Arrays.fill(scaling, 1.0);
		
		if(!(fitScaling || fitOffsets)) {
			if(weighting) {
				System.err.println("Adding maps with derived weighs.");
				addWeighting(component, scaling, offset, coadded);
				System.err.println("Derived Weights: ");
				for(int k=0; k<component.length; k++) System.err.println((k+1) + "\t\t" + Util.f3.format(component[k].weightFactor));
				System.err.println();
				
			}
			else {
				System.err.println("Adding maps directly.");
				add(component, scaling, offset, coadded);
			}
			return coadded;
		}

		System.err.println("Fitting maps together... (this may take a while)");
		System.err.println();
		
		System.err.println("Applying fit to components... ");
		for(int k=0; k<component.length; k++) {
			component[k].scale(scaling[k]);
			component[k].offset(offset[k]);
		}
		
		LinearFit fit = new LinearFit();
		fit.precision = 1e-6;
		fit.maxSteps = 1000;
		fit.minimize();
		
		if(weighting) {
			System.err.println("Deriving component weights.");
			addWeighting(component, scaling, offset, coadded);
		}
		
		System.err.println();
		System.err.println("Summary of fitted parameters:");
		System.err.println();
		System.err.println("Number\t\tScaling\t\tOffset\t\tWeight");
		System.err.println("      \t\t       \t\t" + coadded.nativeUnit.name() + "");
		System.err.println("======================================================");
		
		for(int i=0; i<component.length; i++) 
			System.err.println((i+1) + "\t\t" + Util.f3.format(scaling[i]) 
				+ "\t\t" + Util.e3.format(offset[i]/coadded.nativeUnit.value()) + " \t" + Util.f3.format(component[i].weightFactor));
		
		System.err.println();
		
		if(!dataWeights) coadded.robustWeight();
		
		return coadded;
	}
	
	public static boolean checkCommonGrid(SkyMap[] component) {	
		for(int i=1; i<component.length; i++) if(!component[i].conforms(component[0])) return false;		
		return true;
	}
	
	public static SkyMap regrid(SkyMap map, SkyMap template) {
	
		// Precess coordinates to the same epoch if necessary
		if(template.projection.getReference() instanceof EquatorialCoordinates) {
			EquatorialCoordinates reference = (EquatorialCoordinates) template.projection.getReference();
			EquatorialCoordinates componentRef = (EquatorialCoordinates) map.projection.getReference();
			if(!componentRef.epoch.equals(reference.epoch)) componentRef.precess(reference.epoch);
		}
		
		SkyMap regrid = (SkyMap) map.clone();
		regrid.conformTo(template);
		map.regridTo(regrid);
		
		return regrid;
	}
	 
	public static SkyMap[] getIdenticalGridMaps(SkyMap[] component, SkyMap template) {
		System.err.println("Resampling maps onto common grid.");
		
		SkyMap[] regrid = new SkyMap[component.length];
		
		// Regrid the maps onto a grid equivalent to the final grid
		for(int k=0; k<component.length; k++) regrid[k] = regrid(component[k], template);
		return regrid;
	}



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

		String key = tokens.nextToken();

		options.add(optionString);

		if(key.equalsIgnoreCase("-out")) {
			fileName = Util.getSystemPath(tokens.nextToken());
			absoluteFileName = true;
		}

		else if(key.equalsIgnoreCase("-name")) {
			fileName = Util.getSystemPath(tokens.nextToken());
			absoluteFileName = false;
		}
		
		else if(key.equalsIgnoreCase("-projection")) {
			try { reprojection = SphericalProjection.forName(tokens.nextToken()); }
			catch(Exception e) {
				System.err.println("WARNING! " + e.getMessage());
			}
		}

		else if(key.equalsIgnoreCase("-jackknife")) jackknife = true;
		
		else if(key.equalsIgnoreCase("-offsets")) offsets = true;
		
		else if(key.equalsIgnoreCase("-scaling")) scaling = true;

		else if(key.equalsIgnoreCase("-weighting")) weighting = true;
		
		else if(key.equalsIgnoreCase("-noise")) {
			dataWeights = tokens.nextToken().equalsIgnoreCase("data");
		}

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

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

		return true;
	}

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

	public static void versionInfo() {
		System.out.println("coadd -- Tool for combining 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: coadd [options] <filenames> ...");
		System.out.println();
		System.out.println("Options:");
		System.out.println();
		System.out.println("   -help         Produces this help screen.");
		System.out.println("   -out=         Specify the full 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 the input image path.");
		System.out.println("   -projection=  Specify the spherical projection to use.");
		System.out.println("   -jackknife    Randomly jackknife maps.");
		System.out.println("   -scaling      Fit scaling for each component.");
		System.out.println("   -offsets      Fit offset for each component.");
		System.out.println("   -weighting	 Derive weighting for the components.");
		System.out.println("   -noise=       Estimate noise from 'data' or 'image'. (image)");
		System.out.println();

		System.exit(1);
	}



//LinearFit scaling and offset.


}
