/*******************************************************************************
 * Copyright (c) 2010 Attila Kovacs <attila_kovacs@post.harvard.edu>.
 * All rights reserved. 
 * 
 * This file is part of crush.
 * 
 *     crush 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.
 * 
 *     crush 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 crush.  If not, see <http://www.gnu.org/licenses/>.
 * 
 * Contributors:
 *     Attila Kovacs <attila_kovacs@post.harvard.edu> - initial API and implementation
 ******************************************************************************/
package crush.sourcemodel;

import java.io.*;
import java.util.Collection;

import util.*;
import crush.*;
import crush.array.*;

public class BeamMap<InstrumentType extends Array<?,?>, ScanType extends Scan<? extends InstrumentType, ?>>
		extends SourceMap<InstrumentType, ScanType> {
	
	ScalarMap<InstrumentType, ScanType>[] pixelMap;
	ScalarMap<InstrumentType, ScanType> template;
	
	public BeamMap(InstrumentType instrument) {
		super(instrument);
	}

	@SuppressWarnings("unchecked")
	@Override
	public SourceModel<InstrumentType, ScanType> copy() {
		BeamMap<InstrumentType, ScanType> copy = (BeamMap<InstrumentType, ScanType>) super.copy();
		copy.pixelMap = new ScalarMap[pixelMap.length];
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) copy.pixelMap[p] = (ScalarMap<InstrumentType, ScanType>) pixelMap[p].copy();
		return copy;
	}

	@SuppressWarnings("unchecked")
	@Override
	public void create(Collection<? extends Scan<?,?>> collection) {
		// Set all pixel positions to zero...
		for(Scan<?,?> scan : collection) for(Integration<?,?> integration : scan) 
			for(Pixel pixel : integration.instrument.getMappingPixels()) {
				pixel.getPosition().zero();
				pixel.setIndependent(true);
			}
		
		super.create(collection);
		
		template = new ScalarMap<InstrumentType, ScanType>(instrument);
		template.options = options;
		template.create(collection);
		
		pixelMap = new ScalarMap[instrument.maxPixels() + 1];
	}
	
	@Override
	public void reset() {
		super.reset();
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) pixelMap[p].reset();
	}
	
	@SuppressWarnings("unchecked")
	@Override
	public void add(SourceModel<?, ?> model, double weight) {
		if(!(model instanceof BeamMap<?, ?>)) throw new IllegalArgumentException("ERROR! Cannot add " + model.getClass().getSimpleName() + " to " + getClass().getSimpleName());
		BeamMap<InstrumentType, ScanType> other = (BeamMap<InstrumentType, ScanType>) model;
		
		for(int p=0; p<pixelMap.length; p++) if(other.pixelMap[p] != null) {
			if(pixelMap[p] == null) {
				pixelMap[p] = (ScalarMap<InstrumentType, ScanType>) other.pixelMap[p].copy();
				pixelMap[p].map.weightScale(weight);
			}
			else pixelMap[p].add(other.pixelMap[p], weight);
		}
	}


	@Override
	public long baseFootprint() {
		return pixelMap.length * template.baseFootprint();
	}
	
	
	@Override
	protected void add(Frame exposure, Pixel pixel, MapIndex index, double fGC, double[] sourceGain, double dt, int excludeSamples) {
		int i = pixel.getDataIndex();
		ScalarMap<InstrumentType, ScanType> map = pixelMap[i];
		
		if(map == null) {
			map = (ScalarMap<InstrumentType, ScanType>) template.copy();
			map.id = Integer.toString(i);
			map.standalone();
			pixelMap[i] = map;
		}
			
		pixelMap[i].add(exposure, pixel, index, fGC, sourceGain, dt, excludeSamples);
	}


	@Override
	protected void calcCoupling(Integration<?, ?> integration, Collection<? extends Pixel> pixels, double[] sourceGain) {
		// TODO
	}


	@Override
	public int countPoints() {
		int points = 0;
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) points += pixelMap[p].countPoints();
		return points;
	}

	

	@Override
	public double covariantPoints() {
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) return pixelMap[p].covariantPoints();
		return 1.0;
	}


	@Override
	public void getIndex(Frame exposure, Pixel pixel, CelestialProjector projector, MapIndex index) {
		template.getIndex(exposure, pixel, projector, index);
	}


	@Override
	public double getPixelFootprint() {
		return pixelMap.length * template.getPixelFootprint();
	}


	@Override
	public SphericalProjection getProjection() {
		return template.getProjection();
	}


	@Override
	public boolean isMasked(MapIndex index) {
		return false;
	}


	@Override
	public long pixels() {
		int pixels = 0;
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) pixels += pixelMap[p].pixels();
		return pixels;
	}


	@Override
	public double resolution() {
		return template.resolution();
	}


	@Override
	public void setSize(int sizeX, int sizeY) {
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) pixelMap[p].setSize(sizeX, sizeY);
	}


	@Override
	protected void sync(Frame exposure, Pixel pixel, MapIndex index, double fG,
			double[] sourceGain, double[] usedSourceGain, boolean isMasked) {
		ScalarMap<InstrumentType, ScanType> map = pixelMap[pixel.getDataIndex()];
		if(map != null) map.sync(exposure, pixel, index, fG, sourceGain, usedSourceGain, isMasked);	
	}


	@Override
	public void createBase() {
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) pixelMap[p].createBase();
	}


	@Override
	public void process(Scan<?, ?> scan) {
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) pixelMap[p].process(scan);
	}


	@Override
	public void write(String path) throws Exception {
		if(hasOption("beammap.writemaps")) {
			for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) pixelMap[p].write(path, false);
		}
		calcPixelData(false);
		writePixelData();	
	}

	@Override
	public synchronized void sync() throws InterruptedException {	
		boolean verbose = true;
		for(int p=0; p<pixelMap.length; p++) if(pixelMap[p] != null) {
			if(hasOption("beammap.process")) pixelMap[p].process(verbose);
			verbose = false;
		}
		super.sync();		
	}
	
		
	public void calcPixelData(boolean smooth) {
		double[] peaks = new double[instrument.storeChannels];
		int k = 0;
		
		// Calculate mean rotation angle from sinA and cosA
		// First check if angles are close to +-180 to see if need to branch at zero
		double A0 = scans.get(0).getFirstIntegration().getFirstFrame().getRotation();
		boolean zeroBranch = Math.abs(Math.PI - Math.abs(A0)) < Math.PI/6.0;
		
		double sumA = 0.0;
		int n = 0;
		for(Scan<?,?> scan : scans) for(Integration<?,?> integration : scan) for(Frame exposure : integration) if(exposure != null) {
			double A = exposure.getRotation();
			if(zeroBranch) {
				A += Math.PI;
				A = Math.IEEEremainder(A, 2.0 * Math.PI);
			}
			sumA += A;
			n++;
		}
		double rotation = n > 0 ? sumA /= n : 0.0;
		if(zeroBranch) rotation -= Math.PI;
		
		for(Pixel pixel : scans.get(0).instrument.getMappingPixels()) {
			int i = pixel.getDataIndex();
			ScalarMap<InstrumentType, ScanType> source = pixelMap[i];
			if(source != null) {
				SkyMap map = source.map;
				if(smooth) map.convolveTo(instrument.resolution);
				GaussianSource peak = new GaussianSource(map);
				MapIndex index = map.getS2NImage().indexOfMax();
				peak.center = new Vector2D(index.i, index.j);
				
				peak.finetunePeak();
				// Get the source peak in the pixel.
				peaks[k++] = peak.peak.value;
				
				peak.getLimits(2.0, map.getImageFWHM());
				peak.center = peak.centroid();
				
				// Get the offset position if it makes sense, or set as NaN otherwise...
				pixel.getPosition().set(Double.NaN, Double.NaN);
				//if(instrument instanceof GroundBased) if(!template.isHorizontal()) continue;  
				map.getOffset(peak.center, pixel.getPosition());
				
				// Derotate to array coordinates...
				pixel.getPosition().rotate(-rotation);
				pixel.getPosition().invert();
			}
		}
		if(k == 0) return;
		
		double mean = ArrayUtil.median(peaks, 0, k);
		
		k = 0;
		for(Pixel pixel : scans.get(0).instrument.getMappingPixels()) {
			int i = pixel.getDataIndex();
			ScalarMap<InstrumentType, ScanType> map = pixelMap[i];
			if(map != null) {
				double sourceGain = peaks[k++] / mean;
				for(Channel channel : pixel) channel.coupling *= sourceGain;
			}
		}
	}
	
	
	public void writePixelData() throws IOException {
		String fileName = CRUSH.workPath + File.separator + getDefaultCoreName() + ".rcp";
		PrintStream out = new PrintStream(new FileOutputStream(fileName));
		for(Pixel pixel : scans.get(0).instrument.getMappingPixels()) {
			Vector2D position = pixel.getPosition();
			String positionString = Util.f1.format(position.x / Unit.arcsec) + "\t" + Util.f1.format(position.y / Unit.arcsec);
			
			if(pixelMap[pixel.getDataIndex()] != null) for(Channel channel : pixel) {			
				out.println(pixel.getDataIndex() + "\t" + 
						Util.f3.format(channel.coupling) + "\t" +
						Util.f3.format(channel.gain) + "\t" +
						positionString
				);
			}
		}
		out.flush();
		out.close();
		System.err.println("Written " + fileName);
	}

	@Override
	public String getSourceName() {
		return template.getSourceName();
	}	
}
