// (C)2007 Attila Kovacs <attila@submm.caltech.edu>

import java.io.*;
import java.text.*;
import java.util.*;

import nom.tam.fits.*;
import nom.tam.image.*;
import nom.tam.util.*;

public class SimpleCorrelatedRemove {

	static Vector data = new Vector();
	static double[] w;
	static double[] dof;
	static int[] flag;
	static boolean[] isBad;
	static int channels = 0;

	public static void main(String[] args) {
		System.err.println();
		System.err.println("  CorrelatedRemove -- Remove correlations from a Labview time-series.");
		System.err.println("                      2006 Attila Kovacs <attila@submm.caltech.edu> <akovacs@mpifr-bonn.mpg.de>");	
		System.err.println();

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



		int foundChannels = 0;

		try { 
			BufferedReader in = new BufferedReader(new FileReader(args[0]));

			int index = args[0].lastIndexOf('.');
			String fileNameStem = index > 0 ? args[0].substring(0, index) : args[0];
			String fileNameExt = index > 0 ? args[0].substring(index) : "";
			index = args[0].lastIndexOf(System.getProperty("file.separator"));
			String workPath = index > 0 ? args[0].substring(0, index+1) : ""; 
			double samplingFreq = Double.parseDouble(in.readLine());
			String line;

			System.err.println("   Sampled at " + samplingFreq + " Hz.");

			boolean initialized = false;

			System.err.print("   Reading data.");
			while((line = in.readLine()) != null) {
				StringTokenizer tokens = new StringTokenizer(line);
				if(!initialized) {
					System.err.print(" " + tokens.countTokens() + " channels found.");
					foundChannels = tokens.countTokens();
					initialized = true;
				}
				//if(tokens.countTokens() == channels) {
				double[] entry = new double[foundChannels];
				for(int c=0; c<foundChannels && tokens.hasMoreTokens(); c++)
					entry[c] = Double.parseDouble(tokens.nextToken());
				data.add(entry);
			}

			System.err.println(" " + data.size() + " frames read."); 

			channels = foundChannels;
			w = new double[channels];
			dof = new double[channels];
			flag = new int[channels];
			double[] gain = new double[channels];

			Arrays.fill(w, 1.0);
			Arrays.fill(gain, 1.0);

			double[] correlated;
			double[] allGain;
			int[] allChannels;	

			allChannels = new int[channels];
			for(int c=0; c<channels; c++) allChannels[c] = c;

			allGain = new double[allChannels.length];
			isBad = new boolean[allChannels.length];

			Arrays.fill(allGain, 1.0);

			for(int c=0,k=0; c<channels; c++) if(isBad[c]) {
				flag[c] |= 1;
				w[c] = 0.0;
			}



			correlated = new double[data.size()];

			System.err.println("   Analyzing.");

			int gainTurn = 1;
			int weightTurn = 2;

			for(int i=1; i<=7; i++) {
				System.err.print("\tRound " + i + ":");

				Arrays.fill(dof, data.size());

				removeOffsets(allChannels, 0, data.size());

				getCorrelated(allChannels, allGain, correlated);
				if(i >= gainTurn) getGains(allChannels, allGain, correlated, -10.0, 10.0, 1<<2);

				census();

				// Pixel drifts...
				/*
		int deltat = 5000; // 10 s.
		if(i >= 2) for(int t=0; t < data.size(); t+=deltat)
		    removeOffsets(allChannels, t, t+deltat);
				 */

				if(i >= weightTurn) getPixelWeightsAlt();

				census();

				System.err.println();

			}


			System.err.println("   Writing data.");


			try { 
				writeGains(fileNameStem + ".gains.dat", allChannels, allGain); 
				writeCorrelated(fileNameStem + ".correlated.dat", correlated);
			}
			catch(IOException e) {}


			/*
	    System.err.println("   Calculating covariances.");	    
	    double[][] covar = getCovariance();

	    try {
		System.err.println("   Writing covariances.");	    
		writeCovariance(workPath + "covar", covar);
	    }
	    catch(Exception e) {}
			 */

			System.err.println("   Writing timestream.");
			writeData(samplingFreq, fileNameStem + ".clean" + fileNameExt);

			//printData();

			System.err.println();
		}
		catch(IOException e) { System.err.println("  " + e.getMessage()); }
	}

	public static void removeOffsets(int[] channel, int fromt, int tot) {
		tot = Math.min(tot, data.size());

		for(int c=0; c<channel.length; c++) {
			double offset = 0.0; 
			for(int t=fromt; t<tot; t++) {
				double[] entry = (double[]) data.get(t);
				offset += entry[channel[c]];
			}
			offset /= tot-fromt;

			if(Math.abs(offset) > 0.98) {
				System.err.print("[S!]");
				flag[channel[c]] |= 1<<16;
			}

			for(int t=fromt; t<tot; t++) {
				double[] entry = (double[]) data.get(t);
				entry[channel[c]] -= offset;
			}
			dof[channel[c]]--;
		}
	}


	public static void getCorrelated(int[] channel, double[] gain, double[] correlated) {
		for(int t=0; t<data.size(); t++) {
			double[] entry = (double[]) data.get(t);
			double sum = 0.0, sumG2 = 0.0;
			for(int c=0; c<channel.length; c++) if(flag[channel[c]] == 0){
				sum += w[channel[c]] * entry[channel[c]] * gain[c];
				sumG2 += w[channel[c]] * gain[c] * gain[c];
			}
			if(sumG2 > 0) {
				double dC = sum / sumG2;
				correlated[t] += dC;
				for(int c=0; c<channel.length; c++) {
					entry[channel[c]] -= gain[c] * dC;
					dof[channel[c]] -= w[channel[c]]*gain[c]*gain[c] / sumG2;
				}
			}
		}
	}

	public static void getGains(int[] channel, double[] gain, double[] correlated, double loLimit, double hiLimit, int flagPattern) {
		double[] dG = new double[gain.length];

		for(int c=0; c<channel.length; c++) {
			double RC = 0.0, C2 = 0.0;

			for(int t=0; t<data.size(); t++) {
				double[] entry = (double[]) data.get(t);
				RC += w[channel[c]] * entry[channel[c]] * correlated[t];
				C2 += w[channel[c]] * correlated[t] * correlated[t];
			}
			if(C2 > 0.0) dG[c] = RC / C2;
		}

		// Normalize gain change...

		/*
	double avedG = 0.0, sumw = 0.0;
	for(int c=0; c<channel.length; c++) if(flag[channel[c]] == 0) { 
	    avedG += w[channel[c]] * gain[c] * dG[c]; 
	    sumw += w[channel[c]] * gain[c] * gain[c];
	}
	if(sumw > 0) avedG /= sumw;
	for(int c=0; c<channel.length; c++) dG[c] -= avedG;
		 */


		double norm = 0.0, sumw = 0.0;
		for(int c=0; c<channel.length; c++) if(flag[channel[c]] == 0) {
			double G1 = gain[c] + dG[c];
			norm += w[channel[c]] * Math.abs(G1); 
			sumw += w[channel[c]];
		}
		if(sumw > 0) norm /= sumw;
		for(int c=0; c<channel.length; c++) dG[c] = (gain[c] + dG[c]) / norm - gain[c];

		// Sync with data.
		for(int c=0; c<channel.length; c++) {
			gain[c] += dG[c];
			for(int t=0; t<data.size(); t++) {
				double[] entry = (double[]) data.get(t);
				entry[channel[c]] -= dG[c] * correlated[t];
			}
			dof[channel[c]]--;
		}

		// Flag/unFlag
		for(int c=0; c<channel.length; c++) {
			if(gain[c] < loLimit || gain[c] > hiLimit) flag[channel[c]] |= flagPattern;
			else flag[channel[c]] &= ~flagPattern;
		}
	}

	public static void census() {
		int good = 0;
		for(int c=0; c<w.length; c++) if(flag[c] == 0 && w[c] > 0.0) good++;

		System.err.print(" " + good);

	}

	public static void getPixelWeights() {
		double sumw = 0.0;
		int n = 0;
		for(int c=0; c<w.length; c++) if(!isBad[c]) {
			if(dof[c] > 10.0) {
				double var=0.0;
				for(int t=0; t<data.size(); t++) {
					double[] entry = (double[]) data.get(t);
					var += entry[c] * entry[c];
				}
				if(var > 0.0) {
					w[c] = 1.0 / (var/dof[c]);
					sumw += w[c];
					n++;
				}
				else w[c] = 0.0;
			}
			else w[c] = 0.0;
		}

		// Flag out channels with unrealistically large weights
		// (rms <~ 1/3 average rms)
		double avew = sumw / n;
		for(int c=0; c<w.length; c++) if(w[c] > 10.0*avew) {
			sumw -= w[c];
			w[c] = 0.0;
		}

		census();

		System.err.print(" " + df.format(Math.sqrt(n/sumw)));
	}

	public static void getPixelWeightsAlt() {
		double sumw = 0.0;
		int n = 0;
		for(int c=0; c<w.length; c++) if(!isBad[c]) {
			if(dof[c] > 10.0) {
				double var=0.0;
				double[] last = (double[]) data.get(0);
				for(int t=1; t<data.size(); t++) {
					double[] entry = (double[]) data.get(t);

					double value = entry[c] - last[c];
					var += value * value / 2.0;
					last = entry;
				}
				if(var > 0.0) {
					w[c] = 1.0 / (var/(dof[c]-1));
					sumw += w[c];
					n++;
				}
				else w[c] = 0.0;
			}
			else w[c] = 0.0;
		}

		// Flag out channels with unrealistically large weights
		// (rms <~ 1/3 average rms)
		double avew = sumw / n;
		for(int c=0; c<w.length; c++) if(w[c] > 10.0*avew) {
			sumw -= w[c];
			w[c] = 0.0;
		}

		census();

		System.err.print(" " + df.format(Math.sqrt(n/sumw)));
	}

	static DecimalFormat df = new DecimalFormat("0.000E0");    

	public static void printData() {

		for(int t=0; t<data.size(); t++) {
			double[] entry = (double[]) data.get(t);
			for(int c=0; c<320; c++) System.out.print((flag[c] != 0 ? 0.0 : df.format(entry[c])) + "\t");
			System.out.println();
		}
	}


	public static void writeData(double freq, String fileName) throws IOException {
		DecimalFormat e5 = new DecimalFormat("0.00000E0");
		PrintStream out = new PrintStream(new BufferedOutputStream(new FileOutputStream(fileName), 1000000));
		out.println(df.format(freq));
		for(int t=0; t<data.size(); t++) {
			double[] entry = (double[]) data.get(t);
			for(int c=0; c<channels; c++)
				out.print((isBad[c] || flag[c] != 0 ? 0.0 : e5.format(entry[c])) + "\t");
			out.println();
		}
		out.flush();
		out.close();
	}

	public static void writeGains(String name, int[] channel, double[] gain) throws IOException {
		PrintWriter out = new PrintWriter(new FileOutputStream(name));
		for(int c=0; c<gain.length; c++) out.println((c+1) + "\t" + (channel[c]+1) + "\t" 
				+ df.format(gain[c]) + "\t"
				+ df.format(w[channel[c]]) +"\t" 
				+ flag[channel[c]]);
		out.close();
	}

	public static void writeCorrelated(String name, double[] correlated) throws IOException {
		PrintWriter out = new PrintWriter(new FileOutputStream(name));
		out.println("33.333");
		for(int t=0; t<correlated.length; t++) out.println(df.format(correlated[t]));
		out.close();
	}


	public static double[][] getCovariance() {
		double[][] covar = new double[channels][channels];

		for(int t=0; t<data.size(); t++) {
			double[] entry = (double[]) data.get(t);

			for(int c1=0; c1<entry.length; c1++) for(int c2=c1+1; c2<entry.length; c2++)
				covar[c1][c2] += entry[c1] * entry[c2];
		}

		for(int c1=0; c1<covar.length; c1++) for(int c2=c1+1; c2<covar.length; c2++) {
			covar[c1][c2] *= Math.sqrt(w[c1] * w[c2])/data.size();
			if(covar[c1][c2] == 0.0 || flag[c1] != 0 || flag[c2] != 0) 
				covar[c1][c2] = Double.NaN;
			//covar[c1][c2] = 0.0;

			covar[c2][c1] = covar[c1][c2];
		}
		return covar;
	}


	public static void writeCovariance(String name, double[][] covar) throws IOException, FitsException {
		Fits fits = new Fits();
		BasicHDU hdu = Fits.makeHDU(covar);
		fits.addHDU(hdu);
		fits.write(new BufferedDataOutputStream(new FileOutputStream(name + ".fits")));	
	}

	public static void usage() {

		System.out.println("  Usage: correlated <laboca-data-file>");
		System.out.println();
		System.out.println("       <data-file>    The Labview-style ASCII time series file.");

	}
}
