using System;
using System.Collections.Generic;

namespace Amphora.RandomComparisons
{
    /// <summary>
    /// The distributions the checks measure their results against.
    /// </summary>
    public static class Stats
    {
        public const double Epsilon = 1e-15;
        public const int MaximumIterations = 200000;

        private static readonly double[] LanczosCoefficients =
        {
            0.99999999999980993, 676.5203681218851, -1259.1392167224028, 771.32342877765313, -176.61502916214059,
            12.507343278686905, -0.13857109526572012, 9.9843695780195716e-6, 1.5056327351493116e-7,
        };

        /// <summary>
        /// Returns the natural log of the gamma function, for arguments of at least one half.
        /// </summary>
        public static double LogGamma(double x)
        {
            if (x < 0.5)
            {
                throw new ArgumentOutOfRangeException(nameof(x), "Only arguments of at least 0.5 are supported.");
            }

            double shifted = x - 1;
            double sum = LanczosCoefficients[0];
            for (int i = 1; i < LanczosCoefficients.Length; i++)
            {
                sum += LanczosCoefficients[i] / (shifted + i);
            }

            double t = shifted + 7.5;
            return 0.5 * Math.Log(2 * Math.PI) + (shifted + 0.5) * Math.Log(t) - t + Math.Log(sum);
        }

        /// <summary>
        /// Returns the regularized lower incomplete gamma function and its complement.  Each is as accurate as it can
        /// be made, since the smaller one is the one that matters in a tail.
        /// </summary>
        public static void IncompleteGamma(double a, double x, out double lower, out double upper)
        {
            if (x <= 0)
            {
                lower = 0;
                upper = 1;
                return;
            }

            double prefix = Math.Exp(-x + a * Math.Log(x) - LogGamma(a));
            if (x < a + 1)
            {
                double ap = a;
                double term = 1 / a;
                double sum = term;
                for (int i = 0; i < MaximumIterations; i++)
                {
                    ap += 1;
                    term *= x / ap;
                    sum += term;
                    if (Math.Abs(term) < Math.Abs(sum) * Epsilon)
                    {
                        break;
                    }
                }

                lower = Math.Min(1, sum * prefix);
                upper = 1 - lower;
                return;
            }

            // Lentz's continued fraction for the upper tail.
            const double tiny = 1e-300;
            double b = x + 1 - a;
            double c = 1 / tiny;
            double d = 1 / b;
            double h = d;
            for (int i = 1; i < MaximumIterations; i++)
            {
                double an = -i * (i - a);
                b += 2;
                d = an * d + b;
                if (Math.Abs(d) < tiny)
                {
                    d = tiny;
                }

                c = b + an / c;
                if (Math.Abs(c) < tiny)
                {
                    c = tiny;
                }

                d = 1 / d;
                double delta = d * c;
                h *= delta;
                if (Math.Abs(delta - 1) < Epsilon)
                {
                    break;
                }
            }

            upper = Math.Min(1, prefix * h);
            lower = 1 - upper;
        }

        /// <summary>
        /// Returns the probability that a fair chi-square with the given degrees of freedom lands at least as far out
        /// as <paramref name="statistic"/>, on whichever side it is.  Small means the statistic is too big or too small
        /// to be chance.
        /// </summary>
        public static double ChiSquareTwoSided(double statistic, double degreesOfFreedom)
        {
            IncompleteGamma(degreesOfFreedom / 2, statistic / 2, out double lower, out double upper);
            return Math.Min(1, 2 * Math.Min(lower, upper));
        }

        /// <summary>
        /// Returns the chance that the smallest of <paramref name="tests"/> independent probabilities is as small as
        /// <paramref name="smallest"/>.  The smallest of several is biased low, so it can't be read as a probability
        /// until it is adjusted.
        /// </summary>
        public static double WorstOf(double smallest, int tests)
        {
            // Below this, 1 - smallest rounds badly, and the first-order term is exact to well past a double's digits.
            if (smallest < 1e-9)
            {
                return tests * smallest;
            }

            return 1 - Math.Pow(1 - smallest, tests);
        }

        /// <summary>
        /// Returns the probability that a standard normal lands at least <paramref name="z"/> from zero.
        /// </summary>
        public static double NormalTwoSided(double z)
        {
            IncompleteGamma(0.5, z * z / 2, out double lower, out double upper);
            return upper;
        }

        /// <summary>
        /// Returns the chance of each count from 0 to <paramref name="counts"/> - 1 under a Poisson distribution.
        /// </summary>
        public static double[] PoissonProbabilities(double mean, int counts)
        {
            var returnMe = new double[counts];
            for (int i = 0; i < counts; i++)
            {
                returnMe[i] = Math.Exp(-mean + i * Math.Log(mean) - LogGamma(i + 1.0));
            }

            return returnMe;
        }

        /// <summary>
        /// Returns the chance that a uniformly random square binary matrix of the given size has exactly this rank.
        /// </summary>
        public static double BinaryRankProbability(int size, int rank)
        {
            double logProbability = -(double)size * size * Math.Log(2);
            for (int i = 0; i < rank; i++)
            {
                double full = Math.Pow(2, size) - Math.Pow(2, i);
                double reduced = Math.Pow(2, rank) - Math.Pow(2, i);
                logProbability += 2 * Math.Log(full) - Math.Log(reduced);
            }

            return Math.Exp(logProbability);
        }

        /// <summary>
        /// Sums observed against expected counts into a chi-square statistic, merging neighboring categories until each
        /// expected count is at least <paramref name="minimumExpected"/>.  Returns the degrees of freedom.
        /// </summary>
        public static double ChiSquareOfCategories(double[] observed, double[] expected, double minimumExpected,
            out double statistic)
        {
            var pooledObserved = new List<double>();
            var pooledExpected = new List<double>();
            double observedPool = 0;
            double expectedPool = 0;
            for (int i = 0; i < observed.Length; i++)
            {
                observedPool += observed[i];
                expectedPool += expected[i];
                if (expectedPool >= minimumExpected)
                {
                    pooledObserved.Add(observedPool);
                    pooledExpected.Add(expectedPool);
                    observedPool = 0;
                    expectedPool = 0;
                }
            }

            // A leftover too small to stand alone joins the last category.
            if (expectedPool > 0 && pooledExpected.Count > 0)
            {
                pooledObserved[pooledObserved.Count - 1] += observedPool;
                pooledExpected[pooledExpected.Count - 1] += expectedPool;
            }

            statistic = 0;
            for (int i = 0; i < pooledExpected.Count; i++)
            {
                double difference = pooledObserved[i] - pooledExpected[i];
                statistic += difference * difference / pooledExpected[i];
            }

            return pooledExpected.Count - 1;
        }
    }
}