/*
 * 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
 *
 *      https://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.statistics.inference;

import java.util.Collection;
import java.util.Iterator;
import org.apache.commons.numbers.core.Sum;
import org.apache.commons.statistics.distribution.FDistribution;

/**
 * Implements one-way ANOVA (analysis of variance) statistics.
 *
 * <p>Tests for differences between two or more categories of univariate data
 * (for example, the body mass index of accountants, lawyers, doctors and
 * computer programmers). When two categories are given, this is equivalent to
 * the {@link TTest}.
 *
 * <p>This implementation computes the F statistic using the definitional formula:
 *
 * <p>\[ F = \frac{\text{between-group variability}}{\text{within-group variability}} \]
 *
 * @see <a href="https://en.wikipedia.org/wiki/Analysis_of_variance">Analysis of variance (Wikipedia)</a>
 * @see <a href="https://en.wikipedia.org/wiki/F-test#Multiple-comparison_ANOVA_problems">
 * Multiple-comparison ANOVA problems (Wikipedia)</a>
 * @see <a href="https://www.biostathandbook.com/onewayanova.html">
 * McDonald, J.H. 2014. Handbook of Biological Statistics (3rd ed.). Sparky House Publishing, Baltimore, Maryland.
 * One-way anova. pp 145-156.</a>
 * @since 1.1
 */
public final class OneWayAnova {
    /** Default instance. */
    private static final OneWayAnova DEFAULT = new OneWayAnova();

    /**
     * Result for the one-way ANOVA.
     *
     * <p>This class is immutable.
     *
     * @since 1.1
     */
    public static final class Result extends BaseSignificanceResult {
        /** Degrees of freedom in numerator (between groups). */
        private final int dfbg;
        /** Degrees of freedom in denominator (within groups). */
        private final long dfwg;
        /** Mean square between groups. */
        private final double msbg;
        /** Mean square within groups. */
        private final double mswg;
        /** nO value used to partition the variance. */
        private final double nO;

        /**
         * @param dfbg Degrees of freedom in numerator (between groups).
         * @param dfwg Degrees of freedom in denominator (within groups).
         * @param msbg Mean square between groups.
         * @param mswg Mean square within groups.
         * @param nO Factor for partitioning the variance.
         * @param f F statistic
         * @param p P-value.
         */
        Result(int dfbg, long dfwg, double msbg, double mswg, double nO, double f, double p) {
            super(f, p);
            this.dfbg = dfbg;
            this.dfwg = dfwg;
            this.msbg = msbg;
            this.mswg = mswg;
            this.nO = nO;
        }

        /**
         * Gets the degrees of freedom in the numerator (between groups).
         *
         * @return degrees of freedom between groups
         */
        int getDFBG() {
            return dfbg;
        }

        /**
         * Gets the degrees of freedom in the denominator (within groups).
         *
         * @return degrees of freedom within groups
         */
        long getDFWG() {
            return dfwg;
        }

        /**
         * Gets the mean square between groups.
         *
         * @return mean square between groups
         */
        public double getMSBG() {
            return msbg;
        }

        /**
         * Gets the mean square within groups.
         *
         * @return mean square within groups
         */
        public double getMSWG() {
            return mswg;
        }

        /**
         * Gets the variance component between groups.
         *
         * <p>The value is a partitioning of the variance.
         * It is the complement of {@link #getVCWG()}.
         *
         * <p>Partitioning the variance applies only to a model II
         * (random effects) one-way anova. This applies when the
         * groups are random samples from a larger set of groups;
         * partitioning the variance allows comparison of the
         * variation between groups to the variation within groups.
         *
         * <p>If the {@linkplain #getMSBG() MSBG} is less than the
         * {@linkplain #getMSWG() MSWG} this returns 0. Otherwise this
         * creates an estimate of the added variance component
         * between groups as:
         *
         * <p>\[ \text{between-group variance} = A = (\text{MS}_{\text{bg}} - \text{MS}_{\text{wg}}) / n_o \]
         *
         * <p>where \( n_o \) is a number close to, but usually less than,
         * the arithmetic mean of the sample size \(n_i\) of each
         * of the \( a \) groups:
         *
         * <p>\[ n_o = \frac{1}{a-1} \left( \sum_i{n_i} - \frac{\sum_i{n_i^2}}{\sum_i{n_i}} \right) \]
         *
         * <p>The added variance component among groups \( A \) is expressed
         * as a fraction of the total variance components \( A + B \) where
         * \( B \) is the {@linkplain #getMSWG() MSWG}.
         *
         * @return variance component between groups (in [0, 1]).
         */
        public double getVCBG() {
            if (msbg <= mswg) {
                return 0;
            }
            // a is an estimate of the between-group variance
            final double a = (msbg - mswg) / nO;
            final double b = mswg;
            return a / (a + b);
        }

        /**
         * Gets the variance component within groups.
         *
         * <p>The value is a partitioning of the variance.
         * It is the complement of {@link #getVCBG()}. See
         * that method for details.
         *
         * @return variance component within groups (in [0, 1]).
         */
        public double getVCWG() {
            if (msbg <= mswg) {
                return 1;
            }
            final double a = (msbg - mswg) / nO;
            final double b = mswg;
            return b / (a + b);
        }
    }

    /** Private constructor. */
    private OneWayAnova() {
        // Do nothing
    }

    /**
     * Return an instance using the default options.
     *
     * @return default instance
     */
    public static OneWayAnova withDefaults() {
        return DEFAULT;
    }

    /**
     * Computes the F statistic for an ANOVA test for a collection of category data,
     * evaluating the null hypothesis that there is no difference among the means of
     * the data categories.
     *
     * <p>Special cases:
     * <ul>
     * <li>If the value in each category is the same (no variance within groups) but different
     * between groups, the f-value is {@linkplain Double#POSITIVE_INFINITY infinity}.</li>
     * <li>If the value in every group is the same (no variance within or between groups),
     * the f-value is {@link Double#NaN NaN}.</li>
     * </ul>
     *
     * @param data Category summary data.
     * @return F statistic
     * @throws IllegalArgumentException if the number of categories is less than
     * two; a contained category does not have at least one value; or all
     * categories have only one value (zero degrees of freedom within groups)
     */
    public double statistic(Collection<double[]> data) {
        final double[] f = new double[1];
        aov(data, f);
        return f[0];
    }

    /**
     * Performs an ANOVA test for a collection of category data,
     * evaluating the null hypothesis that there is no difference among the means of
     * the data categories.
     *
     * <p>Special cases:
     * <ul>
     * <li>If the value in each category is the same (no variance within groups) but different
     * between groups, the f-value is {@linkplain Double#POSITIVE_INFINITY infinity} and the p-value is zero.</li>
     * <li>If the value in every group is the same (no variance within or between groups),
     * the f-value and p-value are {@link Double#NaN NaN}.</li>
     * </ul>
     *
     * @param data Category summary data.
     * @return test result
     * @throws IllegalArgumentException if the number of categories is less than
     * two; a contained category does not have at least one value; or all
     * categories have only one value (zero degrees of freedom within groups)
     */
    public Result test(Collection<double[]> data) {
        return aov(data, null);
    }

    /**
     * Performs an ANOVA test for a collection of category data, evaluating the null
     * hypothesis that there is no difference among the means of the data categories.
     *
     * <p>This is a utility method to allow computation of the F statistic without
     * the p-value or partitioning of the variance. If the {@code statistic} is not null
     * the method will record the F statistic in the array and return null.
     *
     * @param data Category summary data.
     * @param statistic Result for the F statistic (or null).
     * @return test result (or null)
     * @throws IllegalArgumentException if the number of categories is less than two; a
     * contained category does not have at least one value; or all categories have only
     * one value (zero degrees of freedom within groups)
     */
    private static Result aov(Collection<double[]> data, double[] statistic) {
        Arguments.checkCategoriesRequiredSize(data.size(), 2);
        long n = 0;
        for (final double[] array : data) {
            n += array.length;
            Arguments.checkValuesRequiredSize(array.length, 1);
        }
        final long dfwg = n - data.size();
        if (dfwg == 0) {
            throw new InferenceException(InferenceException.ZERO, "Degrees of freedom within groups");
        }

        // wg = within group
        // bg = between group

        // F = Var(bg) / Var(wg)
        // Var = SS / df
        // SStotal = sum((x - u)^2) = sum(x^2) - sum(x)^2/n
        //         = SSwg + SSbg
        // Some cancellation of terms reduces the computation to 3 sums:
        // SSwg = [ sum(x^2) - sum(x)^2/n ] - [ sum_g { sum(sum(x^2) - sum(x)^2/n) } ]
        // SSbg = SStotal - SSwg
        //      = sum_g { sum(x)^2/n) } - sum(x)^2/n
        // SSwg = SStotal - SSbg
        //      = sum(x^2) - sum_g { sum(x)^2/n) }

        // Stabilize the computation by shifting all to a common mean of zero.
        // This minimise the magnitude of x^2 terms.
        // The terms sum(x)^2/n -> 0. Included them to capture the round-off.
        final double m = StatisticUtils.mean(data);
        final Sum sxx = Sum.create();
        final Sum sx = Sum.create();
        final Sum sg = Sum.create();
        // Track if each group had the same value
        boolean eachSame = true;
        for (final double[] array : data) {
            eachSame = eachSame && allMatch(array[0], array);
            final Sum s = Sum.create();
            for (final double v : array) {
                final double x = v - m;
                s.add(x);
                // sum(x)
                sx.add(x);
                // sum(x^2)
                sxx.add(x * x);
            }
            // Create the negative sum so we can subtract it via 'add'
            // -sum_g { sum(x)^2/n) }
            sg.add(-pow2(s.getAsDouble()) / array.length);
        }

        // Note: SS terms should not be negative given:
        // SS = sum((x - u)^2)
        // This can happen due to floating-point error in sum(x^2) - sum(x)^2/n
        final double sswg = Math.max(0, sxx.add(sg).getAsDouble());
        // Flip the sign back
        final double ssbg = Math.max(0, -sg.add(pow2(sx.getAsDouble()) / n).getAsDouble());
        final int dfbg = data.size() - 1;
        // Handle edge-cases:
        // Note: 0 / 0 -> NaN : x / 0 -> inf
        // These are documented results and should output p=NaN or 0.
        // This result will occur naturally.
        // However the SS totals may not be 0.0 so we correct these cases.
        final boolean allSame = eachSame && allMatch(data);
        final double msbg = allSame ? 0 : ssbg / dfbg;
        final double mswg = eachSame ? 0 : sswg / dfwg;
        final double f = msbg / mswg;
        if (statistic != null) {
            statistic[0] = f;
            return null;
        }
        final double p = FDistribution.of(dfbg, dfwg).survivalProbability(f);

        // Support partitioning the variance
        // ni = size of each of the groups
        // nO=(1/(a−1))*(sum(ni)−(sum(ni^2)/sum(ni))
        final double nO = (n - data.stream()
                .mapToDouble(x -> pow2(x.length)).sum() / n) / dfbg;

        return new Result(dfbg, dfwg, msbg, mswg, nO, f, p);
    }

    /**
     * Return true if all values in the array match the specified value.
     *
     * @param v Value.
     * @param a Array.
     * @return true if all match
     */
    private static boolean allMatch(double v, double[] a) {
        for (final double w : a) {
            if (v != w) {
                return false;
            }
        }
        return true;
    }

    /**
     * Return true if all values in the arrays match.
     *
     * <p>Assumes that there are at least two arrays and that each array has the same
     * value throughout. Thus only the first element in each array is checked.
     *
     * @param data Arrays.
     * @return true if all match
     */
    private static boolean allMatch(Collection<double[]> data) {
        final Iterator<double[]> iter = data.iterator();
        final double v = iter.next()[0];
        while (iter.hasNext()) {
            if (iter.next()[0] != v) {
                return false;
            }
        }
        return true;
    }

    /**
     * Compute {@code x^2}.
     *
     * @param x Value.
     * @return {@code x^2}
     */
    private static double pow2(double x) {
        return x * x;
    }
}
