package jnpf.elnFormula;

import org.junit.Test;

import java.math.BigDecimal;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;

public class FormulaMathFunctionsTest {

    private final FormulaParser parser = new FormulaParser();
    private final FormulaFunctionRegistry functions = new FormulaFunctionRegistry();
    private final FormulaTypeChecker checker = new FormulaTypeChecker(functions);
    private final FormulaEvaluator evaluator = new FormulaEvaluator(functions);

    @Test
    public void recognizesExpandedFunctionSignatures() {
        assertType("LN(x_a1)", FormulaType.DECIMAL);
        assertType("LOG10(x_a1)", FormulaType.DECIMAL);
        assertType("LOG(x_a1, 10)", FormulaType.DECIMAL);
        assertType("EXP(x_a1)", FormulaType.DECIMAL);
        assertType("NTH_ROOT(x_a1, 3)", FormulaType.DECIMAL);
        assertType("COUNT(values_a1)", FormulaType.DECIMAL);
        assertType("MEDIAN(values_a1)", FormulaType.DECIMAL);
        assertType("STDDEV_P(values_a1)", FormulaType.DECIMAL);
        assertType("VAR_S(values_a1)", FormulaType.DECIMAL);
        assertType("VAR_P(values_a1)", FormulaType.DECIMAL);
        assertType("GEOMEAN(values_a1)", FormulaType.DECIMAL);
        assertType("WEIGHTED_MEAN(values_a1, weights_b2)", FormulaType.DECIMAL);
        assertType("SLOPE(values_a1, weights_b2)", FormulaType.DECIMAL);
        assertType("INTERCEPT(values_a1, weights_b2)", FormulaType.DECIMAL);
        assertType("CORREL(values_a1, weights_b2)", FormulaType.DECIMAL);
        assertType("LINEAR_PREDICT(x_a1, values_a1, weights_b2)", FormulaType.DECIMAL);
        assertType("ROUND_HALF_EVEN(x_a1, 2)", FormulaType.DECIMAL);
        assertType("ROUND_SIG(x_a1, 3)", FormulaType.DECIMAL);
        assertType("ROUND_SIG_HALF_EVEN(x_a1, 3)", FormulaType.DECIMAL);
    }

    @Test
    public void rejectsInvalidExpandedFunctionShapes() {
        assertError("LN(values_a1)", "TYPE_MISMATCH");
        assertError("LOG(x_a1)", "FUNCTION_ARGUMENT_COUNT_MISMATCH");
        assertError("COUNT(x_a1)", "TYPE_MISMATCH");
        assertError("WEIGHTED_MEAN(values_a1)", "FUNCTION_ARGUMENT_COUNT_MISMATCH");
        assertError("SLOPE(values_a1, x_a1)", "TYPE_MISMATCH");
        assertError("LINEAR_PREDICT(values_a1, values_a1, weights_b2)", "TYPE_MISMATCH");
        assertError("ROUND_SIG(values_a1, 3)", "TYPE_MISMATCH");
    }

    @Test
    public void evaluatesBigDecimalMathAndIntermediateRounding() {
        assertDecimal("LN(1)", "0");
        assertDecimal("EXP(0)", "1");
        assertClose("LOG10(1000)", "3");
        assertClose("LOG(8, 2)", "3");
        assertDecimal("NTH_ROOT(27, 3)", "3");
        assertDecimal("NTH_ROOT(-27, 3)", "-3");
        assertDecimal("ROUND_HALF_EVEN(2.345, 2)", "2.34");
        assertDecimal("ROUND_SIG(12345, 3)", "12300");
        assertDecimal("ROUND_SIG_HALF_EVEN(1.245, 3)", "1.24");
    }

    @Test
    public void rejectsMathDomainAndRoundingBoundaryViolations() {
        assertFormulaError("LN(0)", "LOG_DOMAIN_ERROR");
        assertFormulaError("LOG10(-1)", "LOG_DOMAIN_ERROR");
        assertFormulaError("LOG(10, 1)", "INVALID_LOG_BASE");
        assertFormulaError("LOG(10, -2)", "INVALID_LOG_BASE");
        assertFormulaError("EXP(1001)", "EXPONENT_OUT_OF_RANGE");
        assertFormulaError("EXP(-1001)", "EXPONENT_OUT_OF_RANGE");
        assertFormulaError("NTH_ROOT(16, 1)", "INVALID_ROOT_DEGREE");
        assertFormulaError("NTH_ROOT(16, 2.5)", "INVALID_ROOT_DEGREE");
        assertFormulaError("NTH_ROOT(-16, 2)", "NEGATIVE_EVEN_ROOT");
        assertFormulaError("ROUND_SIG(1.23, 0)", "INVALID_SIGNIFICANT_DIGITS");
        assertFormulaError("ROUND_SIG(1.23, 35)", "INVALID_SIGNIFICANT_DIGITS");
        assertFormulaError("ROUND_SIG(1.23, 1.5)", "INVALID_SIGNIFICANT_DIGITS");
    }

    @Test
    public void evaluatesStatisticsAndLinearCalibration() {
        Map<String, FormulaValue> values = statisticValues();
        assertDecimal("COUNT(empty_d4)", values, "0");
        assertDecimal("MEDIAN(values_a1)", values, "2");
        assertDecimal("MEDIAN(even_values_f6)", values, "2.5");
        assertDecimal("STDDEV_P(two_values_g7)", values, "1");
        assertDecimal("VAR_S(values_a1)", values, "1");
        assertClose("VAR_P(values_a1)", values, "0.6666666666666666666666666666666667");
        assertDecimal("GEOMEAN(geomean_h8)", values, "2");
        assertDecimal("WEIGHTED_MEAN(weight_values_i9, weights_c3)", values, "17.5");
        assertDecimal("SLOPE(x_values_a1, y_values_b2)", values, "2");
        assertDecimal("INTERCEPT(x_values_a1, y_values_b2)", values, "1");
        assertDecimal("CORREL(x_values_a1, y_values_b2)", values, "1");
        assertDecimal("LINEAR_PREDICT(4, x_values_a1, y_values_b2)", values, "9");
    }

    @Test
    public void rejectsInvalidStatisticAndRegressionInputs() {
        Map<String, FormulaValue> values = statisticValues();
        assertFormulaError("MEDIAN(empty_d4)", values, "INSUFFICIENT_SAMPLE");
        assertFormulaError("STDDEV_P(empty_d4)", values, "INSUFFICIENT_SAMPLE");
        assertFormulaError("VAR_S(one_value_e5)", values, "INSUFFICIENT_SAMPLE");
        assertFormulaError("VAR_P(empty_d4)", values, "INSUFFICIENT_SAMPLE");
        assertFormulaError("GEOMEAN(negative_values_j0)", values, "GEOMEAN_DOMAIN_ERROR");
        assertFormulaError("WEIGHTED_MEAN(values_a1, one_value_e5)", values,
                "LIST_LENGTH_MISMATCH");
        assertFormulaError("WEIGHTED_MEAN(two_values_g7, negative_weights_k1)", values,
                "INVALID_WEIGHT");
        assertFormulaError("WEIGHTED_MEAN(two_values_g7, zero_weights_l2)", values,
                "INVALID_WEIGHT");
        assertFormulaError("SLOPE(x_values_a1, one_value_e5)", values,
                "LIST_LENGTH_MISMATCH");
        assertFormulaError("SLOPE(constant_values_m3, y_values_b2)", values, "ZERO_VARIANCE");
        assertFormulaError("CORREL(x_values_a1, constant_values_m3)", values, "ZERO_VARIANCE");
    }

    @Test
    public void preservesStatisticsForLargeNearbyDecimals() {
        Map<String, FormulaValue> values = statisticValues();

        assertDecimal("VAR_P(large_nearby_n4)", values, "0.25");
        assertDecimal("SLOPE(large_nearby_n4, large_y_o5)", values, "2");
        assertDecimal("CORREL(large_nearby_n4, large_y_o5)", values, "1");
        assertDecimal("LINEAR_PREDICT(10000000000000000000000000000000000, "
                + "large_nearby_n4, large_y_o5)", values, "5");
        assertDecimal("CORREL(large_nearby_n4, large_negative_y_t0)", values, "-1");
        assertDecimal("CORREL(values_a1, partial_correlation_y_u1)", values, "0.5");
    }

    @Test
    public void roundsLargeAveragesOnlyAfterAggregation() {
        Map<String, FormulaValue> values = statisticValues();
        String expected = "9999999999999999999999999999999998";

        assertDecimal("MEAN(large_nearby_n4)", values, expected);
        assertDecimal("MEDIAN(large_nearby_n4)", values, expected);
        assertDecimal("WEIGHTED_MEAN(large_nearby_n4, equal_weights_s9)", values, expected);
    }

    @Test
    public void validatesEveryGeometricMeanItemBeforeReturningZero() {
        assertFormulaError("GEOMEAN(zero_then_negative_p6)", statisticValues(),
                "GEOMEAN_DOMAIN_ERROR");
    }

    @Test
    public void keepsCorrelationInsidePearsonRangeAfterRounding() {
        assertDecimal("CORREL(correlation_x_q7, correlation_y_r8)", statisticValues(), "1");
    }

    private void assertType(String expression, FormulaType expected) {
        assertEquals(expected, checker.check(parser.parse(expression), slots()));
    }

    private void assertError(String expression, String expectedCode) {
        try {
            checker.check(parser.parse(expression), slots());
            fail("expected " + expectedCode);
        } catch (FormulaException ex) {
            assertEquals(expectedCode, ex.getCode());
        }
    }

    private void assertDecimal(String expression, String expected) {
        assertDecimal(expression, Collections.emptyMap(), expected);
    }

    private void assertDecimal(String expression, Map<String, FormulaValue> values,
                               String expected) {
        FormulaAst.Node ast = parser.parse(expression);
        assertEquals(FormulaType.DECIMAL, checker.check(ast, slots()));
        assertEquals(expected, evaluator.evaluate(ast, values).toPlainString());
    }

    private void assertClose(String expression, String expected) {
        assertClose(expression, Collections.emptyMap(), expected);
    }

    private void assertClose(String expression, Map<String, FormulaValue> values,
                             String expected) {
        FormulaAst.Node ast = parser.parse(expression);
        checker.check(ast, slots());
        BigDecimal actual = evaluator.evaluate(ast, values).getDecimal();
        BigDecimal difference = actual.subtract(new BigDecimal(expected)).abs();
        if (difference.compareTo(new BigDecimal("1E-30")) > 0) {
            fail(expression + " expected approximately " + expected + " but was " + actual);
        }
    }

    private void assertFormulaError(String expression, String expectedCode) {
        assertFormulaError(expression, Collections.emptyMap(), expectedCode);
    }

    private void assertFormulaError(String expression, Map<String, FormulaValue> values,
                                    String expectedCode) {
        try {
            FormulaAst.Node ast = parser.parse(expression);
            checker.check(ast, slots());
            evaluator.evaluate(ast, values);
            fail("expected " + expectedCode);
        } catch (FormulaException ex) {
            assertEquals(expectedCode, ex.getCode());
        }
    }

    private static Map<String, FormulaValue> statisticValues() {
        Map<String, FormulaValue> result = new LinkedHashMap<>();
        result.put("values_a1", list("1", "2", "3"));
        result.put("x_values_a1", list("1", "2", "3"));
        result.put("y_values_b2", list("3", "5", "7"));
        result.put("weights_c3", list("1", "3"));
        result.put("empty_d4", list());
        result.put("one_value_e5", list("1"));
        result.put("even_values_f6", list("1", "2", "3", "4"));
        result.put("two_values_g7", list("1", "3"));
        result.put("geomean_h8", list("1", "4"));
        result.put("weight_values_i9", list("10", "20"));
        result.put("negative_values_j0", list("-1", "4"));
        result.put("negative_weights_k1", list("1", "-1"));
        result.put("zero_weights_l2", list("0", "0"));
        result.put("constant_values_m3", list("2", "2", "2"));
        result.put("large_nearby_n4", list(
                "9999999999999999999999999999999998",
                "9999999999999999999999999999999999"));
        result.put("large_y_o5", list("1", "3"));
        result.put("zero_then_negative_p6", list("0", "-1"));
        String[] correlationX = {"4056.4", "1871.7", "-0.0163297", "-0.0085862", "7.1395"};
        result.put("correlation_x_q7", list(correlationX));
        result.put("correlation_y_r8", scaledList("1.234567890123456789", correlationX));
        result.put("equal_weights_s9", list("1", "1"));
        result.put("large_negative_y_t0", list("3", "1"));
        result.put("partial_correlation_y_u1", list("0", "-1", "1"));
        return result;
    }

    private static FormulaValue list(String... values) {
        FormulaValue[] items = new FormulaValue[values.length];
        for (int i = 0; i < values.length; i++) {
            items[i] = FormulaValue.decimal(values[i]);
        }
        return FormulaValue.decimalList(items);
    }

    private static FormulaValue scaledList(String factor, String... values) {
        BigDecimal multiplier = new BigDecimal(factor);
        FormulaValue[] items = new FormulaValue[values.length];
        for (int i = 0; i < values.length; i++) {
            items[i] = FormulaValue.decimal(new BigDecimal(values[i]).multiply(multiplier));
        }
        return FormulaValue.decimalList(items);
    }

    private static Map<String, FormulaType> slots() {
        Map<String, FormulaType> result = new LinkedHashMap<>();
        result.put("x_a1", FormulaType.DECIMAL);
        result.put("values_a1", FormulaType.DECIMAL_LIST);
        result.put("weights_b2", FormulaType.DECIMAL_LIST);
        result.put("x_values_a1", FormulaType.DECIMAL_LIST);
        result.put("y_values_b2", FormulaType.DECIMAL_LIST);
        result.put("weights_c3", FormulaType.DECIMAL_LIST);
        result.put("empty_d4", FormulaType.DECIMAL_LIST);
        result.put("one_value_e5", FormulaType.DECIMAL_LIST);
        result.put("even_values_f6", FormulaType.DECIMAL_LIST);
        result.put("two_values_g7", FormulaType.DECIMAL_LIST);
        result.put("geomean_h8", FormulaType.DECIMAL_LIST);
        result.put("weight_values_i9", FormulaType.DECIMAL_LIST);
        result.put("negative_values_j0", FormulaType.DECIMAL_LIST);
        result.put("negative_weights_k1", FormulaType.DECIMAL_LIST);
        result.put("zero_weights_l2", FormulaType.DECIMAL_LIST);
        result.put("constant_values_m3", FormulaType.DECIMAL_LIST);
        result.put("large_nearby_n4", FormulaType.DECIMAL_LIST);
        result.put("large_y_o5", FormulaType.DECIMAL_LIST);
        result.put("zero_then_negative_p6", FormulaType.DECIMAL_LIST);
        result.put("correlation_x_q7", FormulaType.DECIMAL_LIST);
        result.put("correlation_y_r8", FormulaType.DECIMAL_LIST);
        result.put("equal_weights_s9", FormulaType.DECIMAL_LIST);
        result.put("large_negative_y_t0", FormulaType.DECIMAL_LIST);
        result.put("partial_correlation_y_u1", FormulaType.DECIMAL_LIST);
        return result;
    }
}