/*
 * 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.commons.math.stat.correlation;

import org.apache.commons.math.TestUtils;
import org.apache.commons.math.distribution.TDistribution;
import org.apache.commons.math.distribution.TDistributionImpl;
import org.apache.commons.math.linear.RealMatrix;
import org.apache.commons.math.linear.BlockRealMatrix;

import junit.framework.TestCase;

public class PearsonsCorrelationTest extends TestCase {

	protected final double[] longleyData = new double[] { 60323, 83.0, 234289, 2356, 1590, 107608, 1947, 61122, 88.5,
			259426, 2325, 1456, 108632, 1948, 60171, 88.2, 258054, 3682, 1616, 109773, 1949, 61187, 89.5, 284599, 3351,
			1650, 110929, 1950, 63221, 96.2, 328975, 2099, 3099, 112075, 1951, 63639, 98.1, 346999, 1932, 3594, 113270,
			1952, 64989, 99.0, 365385, 1870, 3547, 115094, 1953, 63761, 100.0, 363112, 3578, 3350, 116219, 1954, 66019,
			101.2, 397469, 2904, 3048, 117388, 1955, 67857, 104.6, 419180, 2822, 2857, 118734, 1956, 68169, 108.4,
			442769, 2936, 2798, 120445, 1957, 66513, 110.8, 444546, 4681, 2637, 121950, 1958, 68655, 112.6, 482704,
			3813, 2552, 123366, 1959, 69564, 114.2, 502601, 3931, 2514, 125368, 1960, 69331, 115.7, 518173, 4806, 2572,
			127852, 1961, 70551, 116.9, 554894, 4007, 2827, 130081, 1962 };

	protected final double[] swissData = new double[] { 80.2, 17.0, 15, 12, 9.96, 83.1, 45.1, 6, 9, 84.84, 92.5, 39.7,
			5, 5, 93.40, 85.8, 36.5, 12, 7, 33.77, 76.9, 43.5, 17, 15, 5.16, 76.1, 35.3, 9, 7, 90.57, 83.8, 70.2, 16, 7,
			92.85, 92.4, 67.8, 14, 8, 97.16, 82.4, 53.3, 12, 7, 97.67, 82.9, 45.2, 16, 13, 91.38, 87.1, 64.5, 14, 6,
			98.61, 64.1, 62.0, 21, 12, 8.52, 66.9, 67.5, 14, 7, 2.27, 68.9, 60.7, 19, 12, 4.43, 61.7, 69.3, 22, 5, 2.82,
			68.3, 72.6, 18, 2, 24.20, 71.7, 34.0, 17, 8, 3.30, 55.7, 19.4, 26, 28, 12.11, 54.3, 15.2, 31, 20, 2.15,
			65.1, 73.0, 19, 9, 2.84, 65.5, 59.8, 22, 10, 5.23, 65.0, 55.1, 14, 3, 4.52, 56.6, 50.9, 22, 12, 15.14, 57.4,
			54.1, 20, 6, 4.20, 72.5, 71.2, 12, 1, 2.40, 74.2, 58.1, 14, 8, 5.23, 72.0, 63.5, 6, 3, 2.56, 60.5, 60.8, 16,
			10, 7.72, 58.3, 26.8, 25, 19, 18.46, 65.4, 49.5, 15, 8, 6.10, 75.5, 85.9, 3, 2, 99.71, 69.3, 84.9, 7, 6,
			99.68, 77.3, 89.7, 5, 2, 100.00, 70.5, 78.2, 12, 6, 98.96, 79.4, 64.9, 7, 3, 98.22, 65.0, 75.9, 9, 9, 99.06,
			92.2, 84.6, 3, 3, 99.46, 79.3, 63.1, 13, 13, 96.83, 70.4, 38.4, 26, 12, 5.62, 65.7, 7.7, 29, 11, 13.79,
			72.7, 16.7, 22, 13, 11.22, 64.4, 17.6, 35, 32, 16.92, 77.6, 37.6, 15, 7, 4.97, 67.6, 18.7, 25, 7, 8.65,
			35.0, 1.2, 37, 53, 42.34, 44.7, 46.6, 16, 29, 50.43, 42.8, 27.7, 22, 29, 58.33 };

	/**
	 * Test Longley dataset against R.
	 */
	public void testLongly() throws Exception {
		RealMatrix matrix = createRealMatrix(longleyData, 16, 7);
		PearsonsCorrelation corrInstance = new PearsonsCorrelation(matrix);
		RealMatrix correlationMatrix = corrInstance.getCorrelationMatrix();
		double[] rData = new double[] { 1.000000000000000, 0.9708985250610560, 0.9835516111796693, 0.5024980838759942,
				0.4573073999764817, 0.960390571594376, 0.9713294591921188, 0.970898525061056, 1.0000000000000000,
				0.9915891780247822, 0.6206333925590966, 0.4647441876006747, 0.979163432977498, 0.9911491900672053,
				0.983551611179669, 0.9915891780247822, 1.0000000000000000, 0.6042609398895580, 0.4464367918926265,
				0.991090069458478, 0.9952734837647849, 0.502498083875994, 0.6206333925590966, 0.6042609398895580,
				1.0000000000000000, -0.1774206295018783, 0.686551516365312, 0.6682566045621746, 0.457307399976482,
				0.4647441876006747, 0.4464367918926265, -0.1774206295018783, 1.0000000000000000, 0.364416267189032,
				0.4172451498349454, 0.960390571594376, 0.9791634329774981, 0.9910900694584777, 0.6865515163653120,
				0.3644162671890320, 1.000000000000000, 0.9939528462329257, 0.971329459192119, 0.9911491900672053,
				0.9952734837647849, 0.6682566045621746, 0.4172451498349454, 0.993952846232926, 1.0000000000000000 };
		TestUtils.assertEquals("correlation matrix", createRealMatrix(rData, 7, 7), correlationMatrix, 10E-15);

		double[] rPvalues = new double[] { 4.38904690369668e-10, 8.36353208910623e-12, 7.8159700933611e-14,
				0.0472894097790304, 0.01030636128354301, 0.01316878049026582, 0.0749178049642416, 0.06971758330341182,
				0.0830166169296545, 0.510948586323452, 3.693245043123738e-09, 4.327782576751815e-11,
				1.167954621905665e-13, 0.00331028281967516, 0.1652293725106684, 3.95834476307755e-10,
				1.114663916723657e-13, 1.332267629550188e-15, 0.00466039138541463, 0.1078477071581498,
				7.771561172376096e-15 };
		RealMatrix rPMatrix = createLowerTriangularRealMatrix(rPvalues, 7);
		fillUpper(rPMatrix, 0d);
		TestUtils.assertEquals("correlation p values", rPMatrix, corrInstance.getCorrelationPValues(), 10E-15);
	}

	/**
	 * Test R Swiss fertility dataset against R.
	 */
	public void testSwissFertility() throws Exception {
		RealMatrix matrix = createRealMatrix(swissData, 47, 5);
		PearsonsCorrelation corrInstance = new PearsonsCorrelation(matrix);
		RealMatrix correlationMatrix = corrInstance.getCorrelationMatrix();
		double[] rData = new double[] { 1.0000000000000000, 0.3530791836199747, -0.6458827064572875,
				-0.6637888570350691, 0.4636847006517939, 0.3530791836199747, 1.0000000000000000, -0.6865422086171366,
				-0.6395225189483201, 0.4010950530487398, -0.6458827064572875, -0.6865422086171366, 1.0000000000000000,
				0.6984152962884830, -0.5727418060641666, -0.6637888570350691, -0.6395225189483201, 0.6984152962884830,
				1.0000000000000000, -0.1538589170909148, 0.4636847006517939, 0.4010950530487398, -0.5727418060641666,
				-0.1538589170909148, 1.0000000000000000 };
		TestUtils.assertEquals("correlation matrix", createRealMatrix(rData, 5, 5), correlationMatrix, 10E-15);

		double[] rPvalues = new double[] { 0.01491720061472623, 9.45043734069043e-07, 9.95151527133974e-08,
				3.658616965962355e-07, 1.304590105694471e-06, 4.811397236181847e-08, 0.001028523190118147,
				0.005204433539191644, 2.588307925380906e-05, 0.301807756132683 };
		RealMatrix rPMatrix = createLowerTriangularRealMatrix(rPvalues, 5);
		fillUpper(rPMatrix, 0d);
		TestUtils.assertEquals("correlation p values", rPMatrix, corrInstance.getCorrelationPValues(), 10E-15);
	}

	/**
	 * Constant column
	 */
	public void testConstant() {
		double[] noVariance = new double[] { 1, 1, 1, 1 };
		double[] values = new double[] { 1, 2, 3, 4 };
		assertTrue(Double.isNaN(new PearsonsCorrelation().correlation(noVariance, values)));
	}

	/**
	 * Insufficient data
	 */

	public void testInsufficientData() {
		double[] one = new double[] { 1 };
		double[] two = new double[] { 2 };
		try {
			new PearsonsCorrelation().correlation(one, two);
			fail("Expecting IllegalArgumentException");
		} catch (IllegalArgumentException ex) {
			// Expected
		}
		RealMatrix matrix = new BlockRealMatrix(new double[][] { { 0 }, { 1 } });
		try {
			new PearsonsCorrelation(matrix);
			fail("Expecting IllegalArgumentException");
		} catch (IllegalArgumentException ex) {
			// Expected
		}
	}

	/**
	 * Verify that direct t-tests using standard error estimates are consistent
	 * with reported p-values
	 */
	public void testStdErrorConsistency() throws Exception {
		TDistribution tDistribution = new TDistributionImpl(45);
		RealMatrix matrix = createRealMatrix(swissData, 47, 5);
		PearsonsCorrelation corrInstance = new PearsonsCorrelation(matrix);
		RealMatrix rValues = corrInstance.getCorrelationMatrix();
		RealMatrix pValues = corrInstance.getCorrelationPValues();
		RealMatrix stdErrors = corrInstance.getCorrelationStandardErrors();
		for (int i = 0; i < 5; i++) {
			for (int j = 0; j < i; j++) {
				double t = Math.abs(rValues.getEntry(i, j)) / stdErrors.getEntry(i, j);
				double p = 2 * (1 - tDistribution.cumulativeProbability(t));
				assertEquals(p, pValues.getEntry(i, j), 10E-15);
			}
		}
	}

	/**
	 * Verify that creating correlation from covariance gives same results as
	 * direct computation from the original matrix
	 */
	public void testCovarianceConsistency() throws Exception {
		RealMatrix matrix = createRealMatrix(longleyData, 16, 7);
		PearsonsCorrelation corrInstance = new PearsonsCorrelation(matrix);
		Covariance covInstance = new Covariance(matrix);
		PearsonsCorrelation corrFromCovInstance = new PearsonsCorrelation(covInstance);
		TestUtils.assertEquals("correlation values", corrInstance.getCorrelationMatrix(),
				corrFromCovInstance.getCorrelationMatrix(), 10E-15);
		TestUtils.assertEquals("p values", corrInstance.getCorrelationPValues(),
				corrFromCovInstance.getCorrelationPValues(), 10E-15);
		TestUtils.assertEquals("standard errors", corrInstance.getCorrelationStandardErrors(),
				corrFromCovInstance.getCorrelationStandardErrors(), 10E-15);

		PearsonsCorrelation corrFromCovInstance2 = new PearsonsCorrelation(covInstance.getCovarianceMatrix(), 16);
		TestUtils.assertEquals("correlation values", corrInstance.getCorrelationMatrix(),
				corrFromCovInstance2.getCorrelationMatrix(), 10E-15);
		TestUtils.assertEquals("p values", corrInstance.getCorrelationPValues(),
				corrFromCovInstance2.getCorrelationPValues(), 10E-15);
		TestUtils.assertEquals("standard errors", corrInstance.getCorrelationStandardErrors(),
				corrFromCovInstance2.getCorrelationStandardErrors(), 10E-15);
	}

	public void testConsistency() {
		RealMatrix matrix = createRealMatrix(longleyData, 16, 7);
		PearsonsCorrelation corrInstance = new PearsonsCorrelation(matrix);
		double[][] data = matrix.getData();
		double[] x = matrix.getColumn(0);
		double[] y = matrix.getColumn(1);
		assertEquals(new PearsonsCorrelation().correlation(x, y), corrInstance.getCorrelationMatrix().getEntry(0, 1),
				Double.MIN_VALUE);
		TestUtils.assertEquals("Correlation matrix", corrInstance.getCorrelationMatrix(),
				new PearsonsCorrelation().computeCorrelationMatrix(data), Double.MIN_VALUE);
	}

	protected RealMatrix createRealMatrix(double[] data, int nRows, int nCols) {
		double[][] matrixData = new double[nRows][nCols];
		int ptr = 0;
		for (int i = 0; i < nRows; i++) {
			System.arraycopy(data, ptr, matrixData[i], 0, nCols);
			ptr += nCols;
		}
		return new BlockRealMatrix(matrixData);
	}

	protected RealMatrix createLowerTriangularRealMatrix(double[] data, int dimension) {
		int ptr = 0;
		RealMatrix result = new BlockRealMatrix(dimension, dimension);
		for (int i = 1; i < dimension; i++) {
			for (int j = 0; j < i; j++) {
				result.setEntry(i, j, data[ptr]);
				ptr++;
			}
		}
		return result;
	}

	protected void fillUpper(RealMatrix matrix, double diagonalValue) {
		int dimension = matrix.getColumnDimension();
		for (int i = 0; i < dimension; i++) {
			matrix.setEntry(i, i, diagonalValue);
			for (int j = i + 1; j < dimension; j++) {
				matrix.setEntry(i, j, matrix.getEntry(j, i));
			}
		}
	}
}
