/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.hadoop.hive.ql.udf.generic;

import java.io.DataOutput;
import java.io.IOException;
import java.util.ArrayList;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentTypeException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.parse.SemanticException;
import org.apache.hadoop.hive.serde2.io.DoubleWritable;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StructField;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.DoubleObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.LongObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils;
import org.apache.hadoop.hive.serde2.typeinfo.PrimitiveTypeInfo;
import org.apache.hadoop.hive.serde2.typeinfo.TypeInfo;
import org.apache.hadoop.io.LongWritable;

import edu.uci.ics.hivesterix.runtime.evaluator.BufferSerDeUtil;
import edu.uci.ics.hivesterix.runtime.evaluator.SerializableBuffer;

/**
 * Compute the covariance covar_pop(x, y), using the following one-pass method
 * (ref. "Formulas for Robust, One-Pass Parallel Computation of Covariances and
 * Arbitrary-Order Statistical Moments", Philippe Pebay, Sandia Labs):
 * 
 * Incremental: n : <count> mx_n = mx_(n-1) + [x_n - mx_(n-1)]/n : <xavg> my_n =
 * my_(n-1) + [y_n - my_(n-1)]/n : <yavg> c_n = c_(n-1) + (x_n - mx_(n-1))*(y_n
 * - my_n) : <covariance * n>
 * 
 * Merge: c_X = c_A + c_B + (mx_A - mx_B)*(my_A - my_B)*n_A*n_B/n_X
 * 
 */
@Description(name = "covariance,covar_pop", value = "_FUNC_(x,y) - Returns the population covariance of a set of number pairs", extended = "The function takes as arguments any pair of numeric types and returns a double.\n"
		+ "Any pair with a NULL is ignored. If the function is applied to an empty set, NULL\n"
		+ "will be returned. Otherwise, it computes the following:\n"
		+ "   (SUM(x*y)-SUM(x)*SUM(y)/COUNT(x,y))/COUNT(x,y)\n"
		+ "where neither x nor y is null.")
public class GenericUDAFCovariance extends AbstractGenericUDAFResolver {

	static final Log LOG = LogFactory.getLog(GenericUDAFCovariance.class
			.getName());

	@Override
	public GenericUDAFEvaluator getEvaluator(TypeInfo[] parameters)
			throws SemanticException {
		if (parameters.length != 2) {
			throw new UDFArgumentTypeException(parameters.length - 1,
					"Exactly two arguments are expected.");
		}

		if (parameters[0].getCategory() != ObjectInspector.Category.PRIMITIVE) {
			throw new UDFArgumentTypeException(0,
					"Only primitive type arguments are accepted but "
							+ parameters[0].getTypeName() + " is passed.");
		}

		if (parameters[1].getCategory() != ObjectInspector.Category.PRIMITIVE) {
			throw new UDFArgumentTypeException(1,
					"Only primitive type arguments are accepted but "
							+ parameters[1].getTypeName() + " is passed.");
		}

		switch (((PrimitiveTypeInfo) parameters[0]).getPrimitiveCategory()) {
		case BYTE:
		case SHORT:
		case INT:
		case LONG:
		case FLOAT:
		case DOUBLE:
			switch (((PrimitiveTypeInfo) parameters[1]).getPrimitiveCategory()) {
			case BYTE:
			case SHORT:
			case INT:
			case LONG:
			case FLOAT:
			case DOUBLE:
				return new GenericUDAFCovarianceEvaluator();
			case STRING:
			case BOOLEAN:
			default:
				throw new UDFArgumentTypeException(1,
						"Only numeric or string type arguments are accepted but "
								+ parameters[1].getTypeName() + " is passed.");
			}
		case STRING:
		case BOOLEAN:
		default:
			throw new UDFArgumentTypeException(0,
					"Only numeric or string type arguments are accepted but "
							+ parameters[0].getTypeName() + " is passed.");
		}
	}

	/**
	 * Evaluate the variance using the algorithm described in
	 * http://en.wikipedia.org/wiki/Algorithms_for_calculating_variance,
	 * presumably by Pébay, Philippe (2008), in "Formulas for Robust, One-Pass
	 * Parallel Computation of Covariances and Arbitrary-Order Statistical
	 * Moments", Technical Report SAND2008-6212, Sandia National Laboratories,
	 * http://infoserve.sandia.gov/sand_doc/2008/086212.pdf
	 * 
	 * Incremental: n : <count> mx_n = mx_(n-1) + [x_n - mx_(n-1)]/n : <xavg>
	 * my_n = my_(n-1) + [y_n - my_(n-1)]/n : <yavg> c_n = c_(n-1) + (x_n -
	 * mx_(n-1))*(y_n - my_n) : <covariance * n>
	 * 
	 * Merge: c_X = c_A + c_B + (mx_A - mx_B)*(my_A - my_B)*n_A*n_B/n_X
	 * 
	 * This one-pass algorithm is stable.
	 * 
	 */
	public static class GenericUDAFCovarianceEvaluator extends
			GenericUDAFEvaluator {

		// For PARTIAL1 and COMPLETE
		private PrimitiveObjectInspector xInputOI;
		private PrimitiveObjectInspector yInputOI;

		// For PARTIAL2 and FINAL
		private StructObjectInspector soi;
		private StructField countField;
		private StructField xavgField;
		private StructField yavgField;
		private StructField covarField;
		private LongObjectInspector countFieldOI;
		private DoubleObjectInspector xavgFieldOI;
		private DoubleObjectInspector yavgFieldOI;
		private DoubleObjectInspector covarFieldOI;

		// For PARTIAL1 and PARTIAL2
		private Object[] partialResult;

		// For FINAL and COMPLETE
		private DoubleWritable result;

		@Override
		public ObjectInspector init(Mode m, ObjectInspector[] parameters)
				throws HiveException {
			super.init(m, parameters);

			// init input
			if (mode == Mode.PARTIAL1 || mode == Mode.COMPLETE) {
				assert (parameters.length == 2);
				xInputOI = (PrimitiveObjectInspector) parameters[0];
				yInputOI = (PrimitiveObjectInspector) parameters[1];
			} else {
				assert (parameters.length == 1);
				soi = (StructObjectInspector) parameters[0];

				countField = soi.getStructFieldRef("count");
				xavgField = soi.getStructFieldRef("xavg");
				yavgField = soi.getStructFieldRef("yavg");
				covarField = soi.getStructFieldRef("covar");

				countFieldOI = (LongObjectInspector) countField
						.getFieldObjectInspector();
				xavgFieldOI = (DoubleObjectInspector) xavgField
						.getFieldObjectInspector();
				yavgFieldOI = (DoubleObjectInspector) yavgField
						.getFieldObjectInspector();
				covarFieldOI = (DoubleObjectInspector) covarField
						.getFieldObjectInspector();
			}

			// init output
			if (mode == Mode.PARTIAL1 || mode == Mode.PARTIAL2) {
				// The output of a partial aggregation is a struct containing
				// a long count, two double averages, and a double covariance.

				ArrayList<ObjectInspector> foi = new ArrayList<ObjectInspector>();

				foi.add(PrimitiveObjectInspectorFactory.writableLongObjectInspector);
				foi.add(PrimitiveObjectInspectorFactory.writableDoubleObjectInspector);
				foi.add(PrimitiveObjectInspectorFactory.writableDoubleObjectInspector);
				foi.add(PrimitiveObjectInspectorFactory.writableDoubleObjectInspector);

				ArrayList<String> fname = new ArrayList<String>();
				fname.add("count");
				fname.add("xavg");
				fname.add("yavg");
				fname.add("covar");

				partialResult = new Object[4];
				partialResult[0] = new LongWritable(0);
				partialResult[1] = new DoubleWritable(0);
				partialResult[2] = new DoubleWritable(0);
				partialResult[3] = new DoubleWritable(0);

				return ObjectInspectorFactory.getStandardStructObjectInspector(
						fname, foi);

			} else {
				setResult(new DoubleWritable(0));
				return PrimitiveObjectInspectorFactory.writableDoubleObjectInspector;
			}
		}

		static class StdAgg implements SerializableBuffer {
			long count; // number n of elements
			double xavg; // average of x elements
			double yavg; // average of y elements
			double covar; // n times the covariance

			@Override
			public void deSerializeAggBuffer(byte[] data, int start, int len) {
				count = BufferSerDeUtil.getLong(data, start);
				start += 8;
				xavg = BufferSerDeUtil.getDouble(data, start);
				start += 8;
				yavg = BufferSerDeUtil.getDouble(data, start);
				start += 8;
				covar = BufferSerDeUtil.getDouble(data, start);
			}

			@Override
			public void serializeAggBuffer(byte[] data, int start, int len) {
				BufferSerDeUtil.writeLong(count, data, start);
				start += 8;
				BufferSerDeUtil.writeDouble(xavg, data, start);
				start += 8;
				BufferSerDeUtil.writeDouble(yavg, data, start);
				start += 8;
				BufferSerDeUtil.writeDouble(covar, data, start);
			}

			@Override
			public void serializeAggBuffer(DataOutput output)
					throws IOException {
				output.writeLong(count);
				output.writeDouble(xavg);
				output.writeDouble(yavg);
				output.writeDouble(covar);
			}
		};

		@Override
		public AggregationBuffer getNewAggregationBuffer() throws HiveException {
			StdAgg result = new StdAgg();
			reset(result);
			return result;
		}

		@Override
		public void reset(AggregationBuffer agg) throws HiveException {
			StdAgg myagg = (StdAgg) agg;
			myagg.count = 0;
			myagg.xavg = 0;
			myagg.yavg = 0;
			myagg.covar = 0;
		}

		@Override
		public void iterate(AggregationBuffer agg, Object[] parameters)
				throws HiveException {
			assert (parameters.length == 2);
			Object px = parameters[0];
			Object py = parameters[1];
			if (px != null && py != null) {
				StdAgg myagg = (StdAgg) agg;
				double vx = PrimitiveObjectInspectorUtils.getDouble(px,
						xInputOI);
				double vy = PrimitiveObjectInspectorUtils.getDouble(py,
						yInputOI);
				myagg.count++;
				myagg.yavg = myagg.yavg + (vy - myagg.yavg) / myagg.count;
				if (myagg.count > 1) {
					myagg.covar += (vx - myagg.xavg) * (vy - myagg.yavg);
				}
				myagg.xavg = myagg.xavg + (vx - myagg.xavg) / myagg.count;
			}
		}

		@Override
		public Object terminatePartial(AggregationBuffer agg)
				throws HiveException {
			StdAgg myagg = (StdAgg) agg;
			((LongWritable) partialResult[0]).set(myagg.count);
			((DoubleWritable) partialResult[1]).set(myagg.xavg);
			((DoubleWritable) partialResult[2]).set(myagg.yavg);
			((DoubleWritable) partialResult[3]).set(myagg.covar);
			return partialResult;
		}

		@Override
		public void merge(AggregationBuffer agg, Object partial)
				throws HiveException {
			if (partial != null) {
				StdAgg myagg = (StdAgg) agg;

				Object partialCount = soi.getStructFieldData(partial,
						countField);
				Object partialXAvg = soi.getStructFieldData(partial, xavgField);
				Object partialYAvg = soi.getStructFieldData(partial, yavgField);
				Object partialCovar = soi.getStructFieldData(partial,
						covarField);

				long nA = myagg.count;
				long nB = countFieldOI.get(partialCount);

				if (nA == 0) {
					// Just copy the information since there is nothing so far
					myagg.count = countFieldOI.get(partialCount);
					myagg.xavg = xavgFieldOI.get(partialXAvg);
					myagg.yavg = yavgFieldOI.get(partialYAvg);
					myagg.covar = covarFieldOI.get(partialCovar);
				}

				if (nA != 0 && nB != 0) {
					// Merge the two partials
					double xavgA = myagg.xavg;
					double yavgA = myagg.yavg;
					double xavgB = xavgFieldOI.get(partialXAvg);
					double yavgB = yavgFieldOI.get(partialYAvg);
					double covarB = covarFieldOI.get(partialCovar);

					myagg.count += nB;
					myagg.xavg = (xavgA * nA + xavgB * nB) / myagg.count;
					myagg.yavg = (yavgA * nA + yavgB * nB) / myagg.count;
					myagg.covar += covarB + (xavgA - xavgB) * (yavgA - yavgB)
							* ((double) (nA * nB) / myagg.count);
				}
			}
		}

		@Override
		public Object terminate(AggregationBuffer agg) throws HiveException {
			StdAgg myagg = (StdAgg) agg;

			if (myagg.count == 0) { // SQL standard - return null for zero
									// elements
				return null;
			} else {
				getResult().set(myagg.covar / (myagg.count));
				return getResult();
			}
		}

		public void setResult(DoubleWritable result) {
			this.result = result;
		}

		public DoubleWritable getResult() {
			return result;
		}
	}

}