/*
    * 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.statistics.inference;
  
  import java.util.EnumSet;
  import java.util.Objects;
  import org.apache.commons.statistics.descriptive.DoubleStatistics;
  import org.apache.commons.statistics.descriptive.Statistic;
  import org.apache.commons.statistics.distribution.TDistribution;
  
  /**
   * Implements Student's t-test statistics.
   *
   * <p>Tests can be:
   * <ul>
   * <li>One-sample or two-sample
   * <li>One-sided or two-sided
   * <li>Paired or unpaired (for two-sample tests)
   * <li>Homoscedastic (equal variance assumption) or heteroscedastic (for two sample tests)
   * </ul>
   *
   * <p>Input to tests can be either {@code double[]} arrays or the mean, variance, and size
   * of the sample.
   *
   * @see <a href="https://en.wikipedia.org/wiki/Student%27s_t-test">Student&#39;s t-test (Wikipedia)</a>
   * @since 1.1
   */
  public final class TTest {
      /** Default instance. */
      private static final TTest DEFAULT = new TTest(AlternativeHypothesis.TWO_SIDED, false, 0);
  
      /** Alternative hypothesis. */
      private final AlternativeHypothesis alternative;
      /** Assume the two samples have the same population variance. */
      private final boolean equalVariances;
      /** The true value of the mean (or difference in means for a two sample test). */
      private final double mu;
  
      /**
       * Result for the t-test.
       *
       * <p>This class is immutable.
       */
      public static final class Result extends BaseSignificanceResult {
          /** Degrees of freedom. */
          private final double degreesOfFreedom;
  
          /**
           * Create an instance.
           *
           * @param statistic Test statistic.
           * @param degreesOfFreedom Degrees of freedom.
           * @param p Result p-value.
           */
          Result(double statistic, double degreesOfFreedom, double p) {
              super(statistic, p);
              this.degreesOfFreedom = degreesOfFreedom;
          }
  
          /**
           * Gets the degrees of freedom.
           *
           * @return the degrees of freedom
           */
          public double getDegreesOfFreedom() {
              return degreesOfFreedom;
          }
      }
  
      /**
       * @param alternative Alternative hypothesis.
       * @param equalVariances Assume the two samples have the same population variance.
       * @param mu true value of the mean (or difference in means for a two sample test).
       */
      private TTest(AlternativeHypothesis alternative, boolean equalVariances, double mu) {
          this.alternative = alternative;
          this.equalVariances = equalVariances;
          this.mu = mu;
      }
  
      /**
       * Return an instance using the default options.
       *
       * <ul>
       * <li>{@link AlternativeHypothesis#TWO_SIDED}
      * <li>{@link DataDispersion#HETEROSCEDASTIC}
      * <li>{@linkplain #withMu(double) mu = 0}
      * </ul>
      *
      * @return default instance
      */
     public static TTest withDefaults() {
         return DEFAULT;
     }
 
     /**
      * Return an instance with the configured alternative hypothesis.
      *
      * @param v Value.
      * @return an instance
      */
     public TTest with(AlternativeHypothesis v) {
         return new TTest(Objects.requireNonNull(v), equalVariances, mu);
     }
 
     /**
      * Return an instance with the configured assumption on the data dispersion.
      *
      * <p>Applies to the two-sample independent t-test.
      * The statistic can compare the means without the assumption of equal
      * sub-population variances (heteroscedastic); otherwise the means are compared
      * under the assumption of equal sub-population variances (homoscedastic).
      *
      * @param v Value.
      * @return an instance
      * @see #test(double[], double[])
      * @see #test(double, double, long, double, double, long)
      */
     public TTest with(DataDispersion v) {
         return new TTest(alternative, Objects.requireNonNull(v) == DataDispersion.HOMOSCEDASTIC, mu);
     }
 
     /**
      * Return an instance with the configured {@code mu}.
      *
      * <p>For the one-sample test this is the expected mean.
      *
      * <p>For the two-sample test this is the expected difference between the means.
      *
      * @param v Value.
      * @return an instance
      * @throws IllegalArgumentException if the value is not finite
      */
     public TTest withMu(double v) {
         return new TTest(alternative, equalVariances, Arguments.checkFinite(v));
     }
 
     /**
      * Computes a one-sample t statistic comparing the mean of the dataset to {@code mu}.
      *
      * <p>The returned t-statistic is:
      *
      * <p>\[ t = \frac{m - \mu}{ \sqrt{ \frac{v}{n} } } \]
      *
      * @param m Sample mean.
      * @param v Sample variance.
      * @param n Sample size.
      * @return t statistic
      * @throws IllegalArgumentException if the number of samples is {@code < 2}; or the
      * variance is negative
      * @see #withMu(double)
      */
     public double statistic(double m, double v, long n) {
         Arguments.checkNonNegative(v);
         checkSampleSize(n);
         return computeT(m - mu, v, n);
     }
 
     /**
      * Computes a one-sample t statistic comparing the mean of the sample to {@code mu}.
      *
      * @param x Sample values.
      * @return t statistic
      * @throws IllegalArgumentException if the number of samples is {@code < 2}
      * @see #statistic(double, double, long)
      * @see #withMu(double)
      */
     public double statistic(double[] x) {
         final long n = checkSampleSize(x.length);
         final DoubleStatistics s = DoubleStatistics.of(
             EnumSet.of(Statistic.MEAN, Statistic.VARIANCE), x);
         final double m = s.getAsDouble(Statistic.MEAN);
         final double v = s.getAsDouble(Statistic.VARIANCE);
         return computeT(m - mu, v, n);
     }
 
     /**
      * Computes a paired two-sample t-statistic on related samples comparing the mean difference
      * between the samples to {@code mu}.
      *
      * <p>The t-statistic returned is functionally equivalent to what would be returned by computing
      * the one-sample t-statistic {@link #statistic(double[])}, with
      * the sample array consisting of the (signed) differences between corresponding
      * entries in {@code x} and {@code y}.
      *
      * @param x First sample values.
      * @param y Second sample values.
      * @return t statistic
      * @throws IllegalArgumentException if the number of samples is {@code < 2}; or the
      * the size of the samples is not equal
      * @see #withMu(double)
      */
     public double pairedStatistic(double[] x, double[] y) {
         final long n = checkSampleSize(x.length);
         final double m = StatisticUtils.meanDifference(x, y);
         final double v = StatisticUtils.varianceDifference(x, y, m);
         return computeT(m - mu, v, n);
     }
 
     /**
      * Computes a two-sample t statistic on independent samples comparing the difference in means
      * of the datasets to {@code mu}.
      *
      * <p>Use the {@link DataDispersion} to control the computation of the variance.
      *
      * <p>The heteroscedastic t-statistic is:
      *
      * <p>\[ t = \frac{m1 - m2 - \mu}{ \sqrt{ \frac{v_1}{n_1} + \frac{v_2}{n_2} } } \]
      *
      * <p>The homoscedastic t-statistic is:
      *
      * <p>\[ t = \frac{m1 - m2 - \mu}{ \sqrt{ v (\frac{1}{n_1} + \frac{1}{n_2}) } } \]
      *
      * <p>where \( v \) is the pooled variance estimate:
      *
      * <p>\[ v = \frac{(n_1-1)v_1 + (n_2-1)v_2}{n_1 + n_2 - 2} \]
      *
      * @param m1 First sample mean.
      * @param v1 First sample variance.
      * @param n1 First sample size.
      * @param m2 Second sample mean.
      * @param v2 Second sample variance.
      * @param n2 Second sample size.
      * @return t statistic
      * @throws IllegalArgumentException if the number of samples in either dataset is
      * {@code < 2}; or the variances are negative.
      * @see #withMu(double)
      * @see #with(DataDispersion)
      */
     public double statistic(double m1, double v1, long n1,
                             double m2, double v2, long n2) {
         Arguments.checkNonNegative(v1);
         Arguments.checkNonNegative(v2);
         checkSampleSize(n1);
         checkSampleSize(n2);
         return equalVariances ?
             computeHomoscedasticT(mu, m1, v1, n1, m2, v2, n2) :
             computeT(mu, m1, v1, n1, m2, v2, n2);
     }
 
     /**
      * Computes a two-sample t statistic on independent samples comparing the difference
      * in means of the samples to {@code mu}.
      *
      * <p>Use the {@link DataDispersion} to control the computation of the variance.
      *
      * @param x First sample values.
      * @param y Second sample values.
      * @return t statistic
      * @throws IllegalArgumentException if the number of samples in either dataset is {@code < 2}
      * @see #withMu(double)
      * @see #with(DataDispersion)
      */
     public double statistic(double[] x, double[] y) {
         final long n1 = checkSampleSize(x.length);
         final long n2 = checkSampleSize(y.length);
         final DoubleStatistics.Builder b = DoubleStatistics.builder(Statistic.MEAN, Statistic.VARIANCE);
         final DoubleStatistics s1 = b.build(x);
         final double m1 = s1.getAsDouble(Statistic.MEAN);
         final double v1 = s1.getAsDouble(Statistic.VARIANCE);
         final DoubleStatistics s2 = b.build(y);
         final double m2 = s2.getAsDouble(Statistic.MEAN);
         final double v2 = s2.getAsDouble(Statistic.VARIANCE);
         return equalVariances ?
             computeHomoscedasticT(mu, m1, v1, n1, m2, v2, n2) :
             computeT(mu, m1, v1, n1, m2, v2, n2);
     }
 
     /**
      * Perform a one-sample t-test comparing the mean of the dataset to {@code mu}.
      *
      * <p>Degrees of freedom are \( v = n - 1 \).
      *
      * @param m Sample mean.
      * @param v Sample variance.
      * @param n Sample size.
      * @return test result
      * @throws IllegalArgumentException if the number of samples is {@code < 2}; or the
      * variance is negative
      * @see #statistic(double, double, long)
      */
     public Result test(double m, double v, long n) {
         final double t = statistic(m, v, n);
         final double df = n - 1.0;
         final double p = computeP(t, df);
         return new Result(t, df, p);
     }
 
     /**
      * Performs a one-sample t-test comparing the mean of the sample to {@code mu}.
      *
      * <p>Degrees of freedom are \( v = n - 1 \).
      *
      * @param sample Sample values.
      * @return the test result
      * @throws IllegalArgumentException if the number of samples is {@code < 2}; or the
      * the size of the samples is not equal
      * @see #statistic(double[])
      */
     public Result test(double[] sample) {
         final double t = statistic(sample);
         final double df = sample.length - 1.0;
         final double p = computeP(t, df);
         return new Result(t, df, p);
     }
 
     /**
      * Performs a paired two-sample t-test on related samples comparing the mean difference between
      * the samples to {@code mu}.
      *
      * <p>The test is functionally equivalent to what would be returned by computing
      * the one-sample t-test {@link #test(double[])}, with
      * the sample array consisting of the (signed) differences between corresponding
      * entries in {@code x} and {@code y}.
      *
      * @param x First sample values.
      * @param y Second sample values.
      * @return the test result
      * @throws IllegalArgumentException if the number of samples is {@code < 2}; or the
      * the size of the samples is not equal
      * @see #pairedStatistic(double[], double[])
      */
     public Result pairedTest(double[] x, double[] y) {
         final double t = pairedStatistic(x, y);
         final double df = x.length - 1.0;
         final double p = computeP(t, df);
         return new Result(t, df, p);
     }
 
     /**
      * Performs a two-sample t-test on independent samples comparing the difference in means of the
      * datasets to {@code mu}.
      *
      * <p>Use the {@link DataDispersion} to control the computation of the variance.
      *
      * <p>The heteroscedastic degrees of freedom are estimated using the
      * Welch-Satterthwaite approximation:
      *
      * <p>\[ v = \frac{ (\frac{v_1}{n_1} + \frac{v_2}{n_2})^2 }
      *                { \frac{(v_1/n_1)^2}{n_1-1} + \frac{(v_2/n_2)^2}{n_2-1} } \]
      *
      * <p>The homoscedastic degrees of freedom are \( v = n_1 + n_2 - 2 \).
      *
      * @param m1 First sample mean.
      * @param v1 First sample variance.
      * @param n1 First sample size.
      * @param m2 Second sample mean.
      * @param v2 Second sample variance.
      * @param n2 Second sample size.
      * @return test result
      * @throws IllegalArgumentException if the number of samples in either dataset is
      * {@code < 2}; or the variances are negative.
      * @see #statistic(double, double, long, double, double, long)
      */
     public Result test(double m1, double v1, long n1,
                        double m2, double v2, long n2) {
         final double t = statistic(m1, v1, n1, m2, v2, n2);
         final double df = equalVariances ?
                 -2.0 + n1 + n2 :
                 computeDf(v1, n1, v2, n2);
         final double p = computeP(t, df);
         return new Result(t, df, p);
     }
 
     /**
      * Performs a two-sample t-test on independent samples comparing the difference in means of
      * the samples to {@code mu}.
      *
      * <p>Use the {@link DataDispersion} to control the computation of the variance.
      *
      * @param x First sample values.
      * @param y Second sample values.
      * @return the test result
      * @throws IllegalArgumentException if the number of samples in either dataset
      * is {@code < 2}
      * @see #statistic(double[], double[])
      * @see #test(double, double, long, double, double, long)
      */
     public Result test(double[] x, double[] y) {
         // Here we do not call statistic(double[], double[]) because the degreesOfFreedom
         // requires the variance. So repeat the computation and compute p.
         final long n1 = checkSampleSize(x.length);
         final long n2 = checkSampleSize(y.length);
         final DoubleStatistics.Builder b = DoubleStatistics.builder(Statistic.MEAN, Statistic.VARIANCE);
         final DoubleStatistics s1 = b.build(x);
         final double m1 = s1.getAsDouble(Statistic.MEAN);
         final double v1 = s1.getAsDouble(Statistic.VARIANCE);
         final DoubleStatistics s2 = b.build(y);
         final double m2 = s2.getAsDouble(Statistic.MEAN);
         final double v2 = s2.getAsDouble(Statistic.VARIANCE);
         final double t;
         final double df;
         if (equalVariances) {
             t = computeHomoscedasticT(mu, m1, v1, n1, m2, v2, n2);
             df = -2.0 + n1 + n2;
         } else {
             t = computeT(mu, m1, v1, n1, m2, v2, n2);
             df = computeDf(v1, n1, v2, n2);
         }
         final double p = computeP(t, df);
         return new Result(t, df, p);
     }
 
     /**
      * Computes t statistic for one-sample t-test.
      *
      * @param m Sample mean.
      * @param v Sample variance.
      * @param n Sample size.
      * @return t test statistic
      */
     private static double computeT(double m, double v, long n) {
         return m / Math.sqrt(v / n);
     }
 
     /**
      * Computes t statistic for two-sample t-test without the assumption of equal
      * samples sizes or sub-population variances.
      *
      * @param mu Expected difference between means.
      * @param m1 First sample mean.
      * @param v1 First sample variance.
      * @param n1 First sample size.
      * @param m2 Second sample mean.
      * @param v2 Second sample variance.
      * @param n2 Second sample size.
      * @return t test statistic
      */
     private static double computeT(double mu,
                                    double m1, double v1, long n1,
                                    double m2, double v2, long n2)  {
         return (m1 - m2 - mu) / Math.sqrt((v1 / n1) + (v2 / n2));
     }
 
     /**
      * Computes approximate degrees of freedom for two-sample t-test without the
      * assumption of equal samples sizes or sub-population variances.
      *
      * @param v1 First sample variance.
      * @param n1 First sample size.
      * @param v2 Second sample variance.
      * @param n2 Second sample size.
      * @return approximate degrees of freedom
      */
     private static double computeDf(double v1, long n1,
                                     double v2, long n2) {
         // Sample sizes are specified as a double to avoid integer overflow
         final double d1 = n1;
         final double d2 = n2;
         return (((v1 / d1) + (v2 / d2)) * ((v1 / d1) + (v2 / d2))) /
             ((v1 * v1) / (d1 * d1 * (n1 - 1)) + (v2 * v2) / (d2 * d2 * (n2 - 1)));
     }
 
     /**
      * Computes t statistic for two-sample t-test under the hypothesis of equal
      * sub-population variances.
      *
      * @param mu Expected difference between means.
      * @param m1 First sample mean.
      * @param v1 First sample variance.
      * @param n1 First sample size.
      * @param m2 Second sample mean.
      * @param v2 Second sample variance.
      * @param n2 Second sample size.
      * @return t test statistic
      */
     private static double computeHomoscedasticT(double mu,
                                                 double m1, double v1, long n1,
                                                 double m2, double v2, long n2)  {
         final double pooledVariance = ((n1 - 1) * v1 + (n2 - 1) * v2) / (-2.0 + n1 + n2);
         return (m1 - m2 - mu) / Math.sqrt(pooledVariance * (1.0 / n1 + 1.0 / n2));
     }
 
     /**
      * Computes p-value for the specified t statistic.
      *
      * @param t T statistic.
      * @param degreesOfFreedom Degrees of freedom.
      * @return p-value for t-test
      */
     private double computeP(double t, double degreesOfFreedom) {
         if (alternative == AlternativeHypothesis.LESS_THAN) {
             return TDistribution.of(degreesOfFreedom).cumulativeProbability(t);
         }
         if (alternative == AlternativeHypothesis.GREATER_THAN) {
             return TDistribution.of(degreesOfFreedom).survivalProbability(t);
         }
         // two-sided
         return 2.0 * TDistribution.of(degreesOfFreedom).survivalProbability(Math.abs(t));
     }
 
     /**
      * Check sample data size.
      *
      * @param n Data size.
      * @return the sample size
      * @throws IllegalArgumentException if the number of samples {@code < 2}
      */
     private static long checkSampleSize(long n) {
         if (n <= 1) {
             throw new InferenceException(InferenceException.TWO_VALUES_REQUIRED, n);
         }
         return n;
     }
 }