/****************************************************
 * Physionet Challenge 2011                         *
 * Entry for event2 and 3                           *
 *                                                  *
 * Author: Lars Johannesen <lj@halloffame.dk>       *
 *                                                  *
 ****************************************************/
package org.physionet.challenge2011;

import java.io.IOException;
import java.io.InputStream;
import java.io.ObjectInputStream;


public class ChallengeEntry {

	public static final String DEBUGTAG = ChallengeEntry.class.toString();
	final static int FS= 500;					//Sampling Frequency
	final static int CH= 12;
	final static int MAX_RT= 220;				//Max expected beats in minutes
	final static int WIN=FS*10;

	// Because FS is fixed:
	// Octave: b = fir1(5, 40/250); 
	final static short LFILT_ORDER = 10;
	double [] LFILT = {
			1.05384029057833e-04,
			1.07692293618465e-02,
			5.43189261717634e-02,
			1.39011974550478e-01,
			2.29255737951395e-01,
			2.68492096154118e-01,
			2.29255737951395e-01,
			1.39011974550478e-01,			
			5.43189261717634e-02,
			1.07692293618465e-02,
			1.05384029057833e-04		
	};
	
	final static short U3_N = 50; // 100ms
	
	final static short MAXQRS=40;
	int[] QRS=new int[MAXQRS*CH]; // Initial candidates
	int[] QRSREAL=new int[MAXQRS]; // True complexes
	int NQRS_REAL=0;
	int[] NQRS=new int[CH];
	final short QRSWIN = 100; // 200 ms
	final short MINLEAD = 6; // QRS complexes must be present in 6 or more leads
	final short QRS_RECONCILE = 80; // 80 ms
	final short LEFT_MOVE=5; // 10 ms	
	final short MAX_LEFT=35; // 70 ms
	
	int[] QRSON=new int[MAXQRS]; // Onsets
	double[] RMSU3 = new double[CH];
	final double RMSU3_THR = 0.7;
	final double RMSU3_THR2 = 0.2;
	final double ONSETTHR=0.15; // When value is below this relative to peak in each lead it is onset
	final double CORTHR=0.70;
	
	//Define Quality values (could also be defined as enum...)
	final static int INIT=0;
	final static int GOOD = 0;
	final static int BAD =  1;
	short[] data=new short[WIN*CH];
	double[] data_filt=new double[WIN*CH];
	double[] data_u3=new double[WIN*CH];
	
	int [] starts = new int[MAXQRS];
	int [] stops = new int[MAXQRS];
	
	final static short PREFIX=150; // 300 ms
	final static short POSTFIX=100; // 200 ms
	final static short MAXLENGTH=2000; // 4000 ms
	double [] medians = new double[MAXLENGTH*CH];
	
	int [] REALQRS_LEAD = new int[CH];
	
	double[] h = new double[MAXQRS];
	double[] A = new double[MAXQRS];
	double[] B = new double[MAXQRS];
	double[] C = new double[MAXQRS];
	double[] D = new double[MAXQRS];
	double[] S = new double[MAXQRS];
	double[] y = new double[MAXQRS];
	double[] x = new double[MAXQRS];
	
	final short PQWIN=3;
	
	synchronized public int get_result(InputStream iFile, final ECG_MetaData m_MetaData) throws IOException {
		ObjectInputStream in = new ObjectInputStream(iFile);
		try {
			data = (short[])in.readObject();
		} catch (ClassNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		
		
		int [] tmpflag = new int[CH]; 
		for(int i = 0; i < CH; ++i) {
			REALQRS_LEAD[i] = 0;
			tmpflag[i] = 0;
		}
		
		// Step1: Detect LF
		boolean leadfail = true;
		for(int i = 0; i < CH; ++i) {
			for(int j = 1; j < WIN; ++j) {
				if(data[j*CH+i] != data[(j-1)*CH+i]) {
					leadfail = false;
					break;
				}
			}
		
			if(!leadfail) { break; }
		}
		
		if(leadfail) {
			iFile.close();
			return BAD;
		}
		
		// Step2: Calculate HF noise / filter signal for U3
		double allsum=0;
		for(int i = 0; i < CH; ++i) {
			// Transient handling
			data_filt[0*CH+i] = LFILT[0]*data[0*CH+i];
			double diff = (data_filt[0*CH+i] - data_filt[0*CH+i]);
			double cursum = diff*diff;
			
			for(int j=1; j < LFILT_ORDER+1; ++j) {
				data_filt[j*CH+i] = 0;
				for(int k = 0; k < j+1; ++k) {
					data_filt[j*CH+i] = data_filt[j*CH+i] + LFILT[k] * data[(j-k)*CH+i];
				}
				
				diff = data_filt[j*CH+i]-data[j*CH+i];
				cursum = cursum + (diff*diff);
			}
			
			// Filter rest
			for(int j=LFILT_ORDER+1; j < WIN; ++j) {
				data_filt[j*CH+i] = 0;
				for(int k = 0; k < LFILT_ORDER+1; ++k) {
					data_filt[j*CH+i] = data_filt[j*CH+i] + LFILT[k] * data[(j-k)*CH+i];
				}
				
				diff = data_filt[j*CH+i]-data[j*CH+i];
				cursum = cursum + (diff*diff);
			}
			
			allsum = allsum + Math.sqrt(cursum/WIN);
		}
		
		double hfnoise = allsum / CH;
		
		if(hfnoise > 100) {
			iFile.close();
			
			return BAD;
		}
		
		// Step3 : Deploy QRS detector
		for(int i = 0; i < CH; ++i) {
			double sum = 0;
			double cursum = 0;
			
			for(int j = 2; j < U3_N; ++j) {
				sum = sum + Math.pow(data_filt[j*CH+i] - 
							         data_filt[(j-2)*CH+i],2);
			}
			data_u3[0*CH+i] = sum;
			cursum = sum * sum;
			
			for(int j = 1; j < WIN-U3_N/2; ++j) {
				double subterm = 0;
				if(j > U3_N/2) {
					subterm = Math.pow(data_filt[(j-U3_N/2)*CH+i]-
									   data_filt[(j-U3_N/2+2)*CH+i],2);
				}
				
				double addterm = Math.pow(data_filt[(j+U3_N/2)*CH+i]-
						                  data_filt[(j+U3_N/2-2)*CH+i],2);
				
				data_u3[j*CH+i] = data_u3[(j-1)*CH+i] + addterm - subterm;
				
				cursum = cursum + (data_u3[j*CH+i]*data_u3[j*CH+i]);
			}
		
			RMSU3[i] = Math.sqrt(cursum/WIN);	
			NQRS[i] = 0;
			boolean within = false; // Did we recently found a peak ?
			int start = 0; // Start of search for peaks
			int lastloc = 0; // Last location
			int j = 1;
			
			double thr = RMSU3[i] * RMSU3_THR;
			
			while( j < WIN-1) {
				if(data_u3[(j-1)*CH+i] < data_u3[j*CH+i] && data_u3[j*CH+i] > data_u3[(j+1)*CH+i] && data_u3[j*CH+i] > thr) {
					if(!within) {
						within = true;
						start = j;
						lastloc = j;
					} else if(data_u3[i+CH*lastloc] < data_u3[j*CH+i]) {
						lastloc = j;
					}
				}
				
				if(within && (j - start) > U3_N) {
					QRS[CH*NQRS[i]+i] = lastloc;
					NQRS[i] = NQRS[i] + 1;
					
					j = j + QRSWIN;
					within = false;
				} else {
					++j;
				}
			}
		}
		
		// Step4 : Reconcile QRS complexes and determine noise peaks
		int noisepeaks = 0;
		int multinoise = 0;
		NQRS_REAL = 0;
		double tmpsum = 0;
		int tmpn = 0;
		for(int i = 0; i < CH; ++i) {
			for(int iQRS = 0; iQRS < NQRS[i]; ++iQRS) {
				if(QRS[i+CH*iQRS] != -1) {
					tmpsum = QRS[i+CH*iQRS];
					tmpn = 1;
					
					for(int j = 0; j < CH; ++j) {
						if(i != j) {
							for(int jQRS = 0; jQRS < NQRS[j]; ++jQRS) {
								if(QRS[j+CH*jQRS] != -1 && Math.abs(QRS[j+CH*jQRS]-QRS[i+CH*iQRS]) < QRS_RECONCILE ) {
									tmpsum = tmpsum + QRS[j+CH*jQRS];
									tmpn = tmpn + 1;
									QRS[j+CH*jQRS] = -1;
									tmpflag[j] = 1;
									break;
								}
							}
						}
					}

					QRS[i+CH*iQRS] = -1;
					
					if(tmpn == 1) {
						noisepeaks = noisepeaks + tmpn;
					} else if(tmpn > MINLEAD) {
						QRSREAL[NQRS_REAL] = (int) Math.round(tmpsum/tmpn);
						NQRS_REAL = NQRS_REAL+1;
						
						for(int j = 0; j < CH; ++j) {
							if(tmpflag[j] == 1) {
								REALQRS_LEAD[j] = REALQRS_LEAD[j] + 1; 
							}
						}
					} else {
						multinoise = multinoise + tmpn;
					}
					
					for(int j = 0; j < CH; ++j) {
						tmpflag[j] = 0;
					}
				}
			}
		}
		
		// Sort
		int n = NQRS_REAL;
		while (n > 1) {
			int newn = 0;
			for(int i = 0; i < n-1; ++i) {
				if(QRSREAL[i] > QRSREAL[i+1]) {
					int tmp = QRSREAL[i];
					QRSREAL[i] = QRSREAL[i+1];
					QRSREAL[i+1] = tmp;
					newn = i + 1;
				}
			}
			
			n = newn;
		}
		
		// Look for missing beats
		double rr = 0;
		for(int i = 1; i < NQRS_REAL; ++i) {
			rr = rr + QRSREAL[i]-QRSREAL[i-1];
		}
		rr = rr / (NQRS_REAL-1.0);
		
		boolean update = false;
		for(int i = 1; i < NQRS_REAL; ++i) {
			int rra = QRSREAL[i]-QRSREAL[i-1];
			
			// Did we really miss ?
			if(Math.abs(rra - 2.0*rr) < Math.abs(rra - rr)) {
				int start = QRSREAL[i-1] + (int) Math.round(0.39 * Math.sqrt(rr+0.04));
				int stop = QRSREAL[i];
				
				// Find by lead
				for(int j = 0; j < CH; ++j) {
					NQRS[j] = 0;
					
					boolean within = false;
					int startw = 0;
					int lastloc = 0;
					double thr = RMSU3[j]*RMSU3_THR2;
					
					for(int k = start; k < stop; ++k) {
						if(data_u3[(k-1)*CH+j] < data_u3[k*CH+j] && data_u3[k*CH+j] > data_u3[(k+1)*CH+j] && data_u3[k*CH+j] > thr) {
							if(!within) {
								within = true;
								startw = j;
								lastloc = j;
							} else if(data_u3[k+CH*lastloc] < data_u3[k*CH+j]) {
								lastloc = j;
							}
						}
						
						if(within && (j - startw) > U3_N) {
							QRS[CH*NQRS[j]+j] = lastloc;
							NQRS[j] = NQRS[j] + 1;
							break;
						} else {
							++k;
						}
					}
				}
			
				for(int j = 0; j < CH; ++j) {
					tmpflag[j] = 0;
				}
				
				tmpsum = 0;
				tmpn = 0;
				for(int j = 0; j < CH; ++j) {
					for(int jQRS = 0; jQRS < NQRS[j]; ++jQRS) {
						if(QRS[j+CH*jQRS] != -1) {
							tmpsum = QRS[j+CH*jQRS];
							tmpn = 1;
							
							for(int k = 0; k < CH; ++k) {
								if(j != k) {
									for(int kQRS = 0; kQRS < NQRS[k]; ++kQRS) {
										if(QRS[k+CH*kQRS] != -1 && Math.abs(QRS[k+CH*kQRS]-QRS[j+CH*jQRS]) < QRS_RECONCILE) {
											tmpsum = tmpsum + QRS[k+CH*kQRS];
											tmpn = tmpn + 1;
											QRS[k+CH*kQRS] = -1;
											tmpflag[k] = 1;
											break;
										}
									}
								}
							}

							QRS[j+CH*jQRS] = -1;
							
							if(tmpn == 1) {
								noisepeaks = noisepeaks - tmpn;
							} else if(tmpn > MINLEAD) {
								QRSREAL[NQRS_REAL] = (int) Math.round(tmpsum/tmpn);
								NQRS_REAL = NQRS_REAL+1;
								update= true;
								
								for(int l = 0; j < CH; ++j) {
									if(tmpflag[l] == 1) {
										REALQRS_LEAD[l] = REALQRS_LEAD[l] + 1;
									}
								}
							}
						}
					}
				}
			}
		}
		
		// Re-sort list
		if(update) {
			n = NQRS_REAL;
			while (n > 1) {
				int newn = 0;
				for(int j = 0; j < n-1; ++j) {
					if(QRSREAL[j] > QRSREAL[j+1]) {
						int tmp = QRSREAL[j];
						QRSREAL[j] = QRSREAL[j+1];
						QRSREAL[j+1] = tmp;
						newn = j + 1;
					}
				}
		
				n = newn;
			}
		}
		
		// Step5: Valid NQRS ? 
		if(NQRS_REAL < 5 || NQRS_REAL > 37) {
			iFile.close();
			return BAD;
		}
		
		int nonleads = 0;
		for(int i = 0; i < CH; ++i) {
			if(REALQRS_LEAD[i] < 0.5*NQRS_REAL) {
				nonleads = nonleads + 1;
			}
		}
			
		// Step6: Not too leads that are not contributing ?
		if(nonleads > 2) {
			iFile.close();
			return BAD;
		}
		
		// Step7: Detect PQ knots
		for(int i = 0; i < NQRS_REAL; ++i) {
			allsum = 0;
			
			for(int j = 0; j < CH; ++j) {
				int onset = QRSREAL[i];
				double thrval = data_u3[j+CH*onset]*ONSETTHR;
				int monset = onset;
				
				while(onset > 0 && data_u3[j+CH*onset] > thrval) {
					if(data_u3[j+CH*onset] < data_u3[j+CH*onset]) {
						monset = onset;
					}
					
					if(QRSREAL[i] - onset > MAX_LEFT) {
						onset = monset;
						break;
					}
					
					--onset;
				}
				
				allsum = allsum + onset;
			}
			
			QRSON[i] = (int) (Math.round(allsum/CH)-LEFT_MOVE);
		}
		
		// Step8: Estimate BW noise and filt data
		allsum = 0;
		for(int i = 0; i < CH; ++i) {
			double cursum = 0;
			int NP=0;
			
			// Calculate x,y
			int ji = 0;
			while(ji < NQRS_REAL) {
				if(QRSON[ji]-PQWIN > 0 && QRSON[ji] + PQWIN < WIN) {
					x[NP] = QRSON[ji];
				
					double tmpsum2 = 0;
					for(int k = QRSON[ji]-PQWIN; k <= QRSON[ji]+PQWIN; ++k) {
						tmpsum2 = tmpsum2 + data_filt[i+CH*k];
					}
				
					y[NP] = tmpsum2 / (PQWIN*2+1);
					NP = NP+1;
				}
				++ji;
			}
			
			// Calculate  h
			for(int j = 0; j < NP-1; ++j) {
				h[j] = x[j+1]-x[j];
			}
			
			// Calculate A,B,C,D
			for(int j = 1; j < NP-1; ++j) {
				int k = j - 1;
				D[k] = 2 * (h[j-1]+h[j]);
				A[k] = h[j+1];
				B[k] = h[j];
				C[k] = 6 * ((y[j+1]-y[j])/h[j] - (y[j]-y[j-1])/h[j-1]);
			}
			
			// Solve matrix
			for(int j = 1; j < NP-1; ++j) {
				double m = 0;
				if(D[j-1] != 0) {
					m = B[j] / D[j-1];
				}
				
				D[j] = D[j] - m * A[j-1];
				C[j] = C[j] - m * C[j-1];
			}
			
			if(D[NP-1] == 0) {
				C[NP-1] = 0;
			} else {
				C[NP-1] = C[NP-1] / D[NP-1];
			}
			
			for(int j = NP-2; j >= 0; --j) {
				if(D[j] == 0) {
					C[j] = 0;
				} else {
					C[j] = (C[j]-A[j]*C[j+1]) / D[j];
				}
			}
			
			// Calculate S
			for(int j = 1; j < NP-2; ++j) {
				int k = j - 1;
				S[j] = C[k]; 
			}
			S[0] = 0;
			S[NP-1] = 0;
			
			// Calculate A,B,C,D
			for(int j = 0; j < NP-1; ++j) {
				A[j] = (S[j+1]-S[j]) / (6*h[j]);
				B[j] = S[j]/2.0;
				C[j] = (y[j+1]-y[j]) / h[j] - (2.0 * h[j] * S[j] + h[j] * S[j+1])/6.0;
				D[j] = y[j];
			}
			
			int START = (int) x[0];	
			int STOP = (int) x[NP-1];
			
			for(int num = 0; num < NP-1; ++num) {
				for(int j = (int) x[num]; j < x[num+1]; ++j) {
					double bw = A[num] * Math.pow(j-x[num],3) + B[num] * Math.pow(j-x[num],2) + C[num] * (j-x[num]) + D[num];
					data_filt[i+CH*j] = data_filt[i+CH*j] - bw;
					
					cursum = cursum + Math.pow(bw,2);
 				}
			}
			
			// Correct BW
			
			allsum = allsum + Math.sqrt(cursum/(STOP-START));
		}// End bw correct
	
		double lf_noise = Math.sqrt(allsum / CH);
		
		if(lf_noise > 35) {
			iFile.close();
			return BAD;
		}
		
		// Calculate median beat length
		double lengths = 0;
		int NBEATS = 0;
		for(int i = 1; i < NQRS_REAL-1; ++i) {
			if(QRSREAL[i]-PREFIX > 0) {
				starts[NBEATS] = QRSREAL[i] - PREFIX;
				stops[NBEATS] = QRSREAL[i+1] - POSTFIX;
				lengths = lengths + (stops[NBEATS] - starts[NBEATS]);
				NBEATS = NBEATS + 1;
			}
		}
		lengths = lengths / (NBEATS);
		
		if(lengths > MAXLENGTH) {
			lengths = MAXLENGTH;
		}
		
		// Calculate median beat by lead
		for(int i = 0; i < CH; ++i) {			
			// For each sample in median
			for(int j = 0; j < lengths; ++j) {
				double val = 0;
				int n3 = 0;
				
				// For each beat belonging to it
				for(int k = 0; k < NBEATS; ++k) {
					// Is the sample within ?
					if(j < stops[k]-starts[k]) {
						int curpos = starts[k]+j;
						val = val + data_filt[curpos*CH+i];
						n3 = n3 + 1;
					}
				}
				
				medians[j*CH+i] = val/n3;
			}
		}
		
		// Calculate RMS / SD in beats with good correlation
		double allRMSsum = 0;
		double allSDsum = 0;
		int allN = 0;
		for(int i = 0; i < CH; ++i) {
			double ystd = 0;
			double ymean = 0;
			for(int j = 0; j < lengths; ++j) {
				ymean = ymean + medians[j*CH+i];
				ystd = ystd + (medians[j*CH+i]*medians[j*CH+i]);
			}
			
			for(int j = 0; j < NBEATS; ++j) {
				int N = (int) lengths;
				if(N > stops[j] - starts[j]) {
					N = stops[j] - starts[j];
				}
				N = N - 1;
				
				double cur_ystd = Math.sqrt((ystd - (ymean*ymean)/N) / (N-1));
				double cur_ymean = ymean / N;
				
				double xmean = 0;
				double xstd = 0;
				double xyprod = 0;
				
				for(int k=0; k < N; ++k) {
					xmean = xmean + data_filt[(starts[j]+k)*CH+i];
					xstd = xstd + (data_filt[(starts[j]+k)*CH+i]*data_filt[(starts[j]+k)*CH+i]);
					xyprod = xyprod + data_filt[(starts[j]+k)*CH+i] * medians[k*CH+i];
				}
				
				xstd = Math.sqrt((xstd - (xmean*xmean)/N)/(N-1));
				xmean = xmean / N;  
				
				double corr = (xyprod - N*xmean*cur_ymean)/((N-1)*xstd*cur_ystd);
				
				if(corr > CORTHR) {
					double residuum = data_filt[(starts[j]+0)*CH+i] - medians[0*CH+i];
					double resmean = residuum;
					double resstd = residuum * residuum;
					double diffsum = 0;
					
					for(int k = 1; k < N; ++k) {
						double nresiduum = data_filt[(starts[j]+k)*CH+i] - medians[k*CH+i];
						
						resmean = resmean + nresiduum;
						resstd = resstd + (nresiduum*nresiduum);
						diffsum = diffsum + (nresiduum-residuum)*(nresiduum-residuum);
						
						residuum = nresiduum;
					}
					
					resstd = Math.sqrt((resstd-(resmean*resmean)/N)/(N-1));
					diffsum = Math.sqrt(diffsum / (N-1));
					
					allRMSsum = allRMSsum + diffsum;
					allSDsum = allSDsum + resstd;
					
					allN = allN + 1;
				} // End cor
			} // End foreach beat
		}// End calc RMS/SD
		
		allRMSsum = allRMSsum / allN;
		allSDsum = allSDsum / allN;
		
		if(allSDsum > 50 || allRMSsum > 50) {		
			iFile.close();
			return BAD;
		}
		
		iFile.close();
		return GOOD;
	}
}