/*
 * 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.distribution;

import java.util.function.DoubleSupplier;
import org.apache.commons.numbers.gamma.Erf;
import org.apache.commons.numbers.gamma.ErfDifference;
import org.apache.commons.numbers.gamma.Erfcx;
import org.apache.commons.rng.UniformRandomProvider;
import org.apache.commons.rng.sampling.distribution.ZigguratSampler;

/**
 * Implementation of the truncated normal distribution.
 *
 * <p>The probability density function of \( X \) is:
 *
 * <p>\[ f(x;\mu,\sigma,a,b) = \frac{1}{\sigma}\,\frac{\phi(\frac{x - \mu}{\sigma})}{\Phi(\frac{b - \mu}{\sigma}) - \Phi(\frac{a - \mu}{\sigma}) } \]
 *
 * <p>for \( \mu \) mean of the parent normal distribution,
 * \( \sigma \) standard deviation of the parent normal distribution,
 * \( -\infty \le a \lt b \le \infty \) the truncation interval, and
 * \( x \in [a, b] \), where \( \phi \) is the probability
 * density function of the standard normal distribution and \( \Phi \)
 * is its cumulative distribution function.
 *
 * @see <a href="https://en.wikipedia.org/wiki/Truncated_normal_distribution">
 * Truncated normal distribution (Wikipedia)</a>
 */
public final class TruncatedNormalDistribution extends AbstractContinuousDistribution {

    /** The max allowed value for x where (x*x) will not overflow.
     * This is a limit on computation of the moments of the truncated normal
     * as some calculations assume x*x is finite. Value is sqrt(MAX_VALUE). */
    private static final double MAX_X = 0x1.fffffffffffffp511;

    /** The min allowed probability range of the parent normal distribution.
     * Set to 0.0. This may be too low for accurate usage. It is a signal that
     * the truncation is invalid. */
    private static final double MIN_P = 0.0;

    /** sqrt(2). */
    private static final double ROOT2 = Constants.ROOT_TWO;
    /** Normalisation constant 2 / sqrt(2 pi) = sqrt(2 / pi). */
    private static final double ROOT_2_PI = Constants.ROOT_TWO_DIV_PI;
    /** Normalisation constant sqrt(2 pi) / 2 = sqrt(pi / 2). */
    private static final double ROOT_PI_2 = Constants.ROOT_PI_DIV_TWO;

    /**
     * The threshold to switch to a rejection sampler. When the truncated
     * distribution covers more than this fraction of the CDF then rejection
     * sampling will be more efficient than inverse CDF sampling. Performance
     * benchmarks indicate that a normalized Gaussian sampler is up to 10 times
     * faster than inverse transform sampling using a fast random generator. See
     * STATISTICS-55.
     */
    private static final double REJECTION_THRESHOLD = 0.2;

    /** Parent normal distribution. */
    private final NormalDistribution parentNormal;
    /** Lower bound of this distribution. */
    private final double lower;
    /** Upper bound of this distribution. */
    private final double upper;

    /** Stored value of {@code parentNormal.probability(lower, upper)}. This is used to
     * normalise the probability computations. */
    private final double cdfDelta;
    /** log(cdfDelta). */
    private final double logCdfDelta;
    /** Stored value of {@code parentNormal.cumulativeProbability(lower)}. Used to map
     * a probability into the range of the parent normal distribution. */
    private final double cdfAlpha;
    /** Stored value of {@code parentNormal.survivalProbability(upper)}. Used to map
     * a probability into the range of the parent normal distribution. */
    private final double sfBeta;

    /**
     * @param parent Parent distribution.
     * @param z Probability of the parent distribution for {@code [lower, upper]}.
     * @param lower Lower bound (inclusive) of the distribution, can be {@link Double#NEGATIVE_INFINITY}.
     * @param upper Upper bound (inclusive) of the distribution, can be {@link Double#POSITIVE_INFINITY}.
     */
    private TruncatedNormalDistribution(NormalDistribution parent, double z, double lower, double upper) {
        this.parentNormal = parent;
        this.lower = lower;
        this.upper = upper;

        cdfDelta = z;
        logCdfDelta = Math.log(cdfDelta);
        // Used to map the inverse probability.
        cdfAlpha = parentNormal.cumulativeProbability(lower);
        sfBeta = parentNormal.survivalProbability(upper);
    }

    /**
     * Creates a truncated normal distribution.
     *
     * <p>Note that the {@code mean} and {@code sd} is of the parent normal distribution,
     * and not the true mean and standard deviation of the truncated normal distribution.
     * The {@code lower} and {@code upper} bounds define the truncation of the parent
     * normal distribution.
     *
     * @param mean Mean for the parent distribution.
     * @param sd Standard deviation for the parent distribution.
     * @param lower Lower bound (inclusive) of the distribution, can be {@link Double#NEGATIVE_INFINITY}.
     * @param upper Upper bound (inclusive) of the distribution, can be {@link Double#POSITIVE_INFINITY}.
     * @return the distribution
     * @throws IllegalArgumentException if {@code sd <= 0}; if {@code lower >= upper}; or if
     * the truncation covers no probability range in the parent distribution.
     */
    public static TruncatedNormalDistribution of(double mean, double sd, double lower, double upper) {
        if (sd <= 0) {
            throw new DistributionException(DistributionException.NOT_STRICTLY_POSITIVE, sd);
        }
        if (lower >= upper) {
            throw new DistributionException(DistributionException.INVALID_RANGE_LOW_GTE_HIGH, lower, upper);
        }

        // Use an instance for the parent normal distribution to maximise accuracy
        // in range computations using the error function
        final NormalDistribution parent = NormalDistribution.of(mean, sd);

        // If there is no computable range then raise an exception.
        final double z = parent.probability(lower, upper);
        if (z <= MIN_P) {
            // Map the bounds to a standard normal distribution for the message
            final double a = (lower - mean) / sd;
            final double b = (upper - mean) / sd;
            throw new DistributionException(
                "Excess truncation of standard normal : CDF(%s, %s) = %s", a, b, z);
        }

        // Here we have a meaningful truncation. Note that excess truncation may not be optimal.
        // For example truncation close to zero where the PDF is constant can be approximated
        // using a uniform distribution.

        return new TruncatedNormalDistribution(parent, z, lower, upper);
    }

    /**
     * Gets the mean for the parent distribution.
     *
     * <p>Note that the mean is of the parent normal distribution,
     * and not the true mean of the truncated normal distribution.
     * This is the {@code mean} parameter used to construct the truncated distribution.
     *
     * @return the parent mean.
     * @see #getMean
     * @since 1.3
     */
    public double getParentMean() {
        return parentNormal.getMean();
    }

    /**
     * Gets the standard deviation for the parent distribution.
     *
     * <p>Note that the standard deviation (SD) is of the parent normal distribution,
     * and not the true standard deviation of the truncated normal distribution.
     * This is the {@code sd} parameter used to construct the truncated distribution.
     *
     * @return the parent standard deviation.
     * @since 1.3
     */
    public double getParentStandardDeviation() {
        return parentNormal.getStandardDeviation();
    }

    /** {@inheritDoc} */
    @Override
    public double density(double x) {
        if (x < lower || x > upper) {
            return 0;
        }
        return parentNormal.density(x) / cdfDelta;
    }

    /** {@inheritDoc} */
    @Override
    public double probability(double x0, double x1) {
        if (x0 > x1) {
            throw new DistributionException(DistributionException.INVALID_RANGE_LOW_GT_HIGH,
                                            x0, x1);
        }
        return parentNormal.probability(clipToRange(x0), clipToRange(x1)) / cdfDelta;
    }

    /** {@inheritDoc} */
    @Override
    public double logDensity(double x) {
        if (x < lower || x > upper) {
            return Double.NEGATIVE_INFINITY;
        }
        return parentNormal.logDensity(x) - logCdfDelta;
    }

    /** {@inheritDoc} */
    @Override
    public double cumulativeProbability(double x) {
        if (x <= lower) {
            return 0;
        } else if (x >= upper) {
            return 1;
        }
        return parentNormal.probability(lower, x) / cdfDelta;
    }

    /** {@inheritDoc} */
    @Override
    public double survivalProbability(double x) {
        if (x <= lower) {
            return 1;
        } else if (x >= upper) {
            return 0;
        }
        return parentNormal.probability(x, upper) / cdfDelta;
    }

    /** {@inheritDoc} */
    @Override
    public double inverseCumulativeProbability(double p) {
        ArgumentUtils.checkProbability(p);
        // Exact bound
        if (p == 0) {
            return lower;
        } else if (p == 1) {
            return upper;
        }
        // Linearly map p to the range [lower, upper]
        final double x = parentNormal.inverseCumulativeProbability(cdfAlpha + p * cdfDelta);
        return clipToRange(x);
    }

    /** {@inheritDoc} */
    @Override
    public double inverseSurvivalProbability(double p) {
        ArgumentUtils.checkProbability(p);
        // Exact bound
        if (p == 1) {
            return lower;
        } else if (p == 0) {
            return upper;
        }
        // Linearly map p to the range [lower, upper]
        final double x = parentNormal.inverseSurvivalProbability(sfBeta + p * cdfDelta);
        return clipToRange(x);
    }

    /** {@inheritDoc} */
    @Override
    public Sampler createSampler(UniformRandomProvider rng) {
        // Map the bounds to a standard normal distribution
        final double u = parentNormal.getMean();
        final double s = parentNormal.getStandardDeviation();
        final double a = (lower - u) / s;
        final double b = (upper - u) / s;
        // If the truncation covers a reasonable amount of the normal distribution
        // then a rejection sampler can be used.
        double threshold = REJECTION_THRESHOLD;
        // If the truncation is entirely in the upper or lower half then adjust the
        // threshold as twice the samples can be used
        if (a >= 0 || b <= 0) {
            threshold *= 0.5;
        }

        if (cdfDelta > threshold) {
            // Create the rejection sampler
            final ZigguratSampler.NormalizedGaussian sampler = ZigguratSampler.NormalizedGaussian.of(rng);
            final DoubleSupplier gen;
            // Use mirroring if possible
            if (a >= 0) {
                // Return the upper-half of the Gaussian
                gen = () -> Math.abs(sampler.sample());
            } else if (b <= 0) {
                // Return the lower-half of the Gaussian
                gen = () -> -Math.abs(sampler.sample());
            } else {
                // Return the full range of the Gaussian
                gen = sampler::sample;
            }
            // Sample in [a, b] using rejection
            return () -> {
                double x = gen.getAsDouble();
                while (x < a || x > b) {
                    x = gen.getAsDouble();
                }
                // Avoid floating-point error when mapping back
                return clipToRange(u + x * s);
            };
        }

        // Default to an inverse CDF sampler
        return super.createSampler(rng);
    }

    /**
     * {@inheritDoc}
     *
     * <p>Represents the true mean of the truncated normal distribution rather
     * than the parent normal distribution mean.
     *
     * <p>For \( \mu \) mean of the parent normal distribution,
     * \( \sigma \) standard deviation of the parent normal distribution, and
     * \( a \lt b \) the truncation interval of the parent normal distribution, the mean is:
     *
     * <p>\[ \mu + \frac{\phi(a)-\phi(b)}{\Phi(b) - \Phi(a)}\sigma \]
     *
     * <p>where \( \phi \) is the probability density function of the standard normal distribution
     * and \( \Phi \) is its cumulative distribution function.
     */
    @Override
    public double getMean() {
        final double u = parentNormal.getMean();
        final double s = parentNormal.getStandardDeviation();
        final double a = (lower - u) / s;
        final double b = (upper - u) / s;
        return u + moment1(a, b) * s;
    }

    /**
     * {@inheritDoc}
     *
     * <p>Represents the true variance of the truncated normal distribution rather
     * than the parent normal distribution variance.
     *
     * <p>For \( \mu \) mean of the parent normal distribution,
     * \( \sigma \) standard deviation of the parent normal distribution, and
     * \( a \lt b \) the truncation interval of the parent normal distribution, the variance is:
     *
     * <p>\[ \sigma^2 \left[1 + \frac{a\phi(a)-b\phi(b)}{\Phi(b) - \Phi(a)} -
     *       \left( \frac{\phi(a)-\phi(b)}{\Phi(b) - \Phi(a)} \right)^2 \right] \]
     *
     * <p>where \( \phi \) is the probability density function of the standard normal distribution
     * and \( \Phi \) is its cumulative distribution function.
     */
    @Override
    public double getVariance() {
        final double u = parentNormal.getMean();
        final double s = parentNormal.getStandardDeviation();
        final double a = (lower - u) / s;
        final double b = (upper - u) / s;
        return variance(a, b) * s * s;
    }

    /**
     * {@inheritDoc}
     *
     * <p>The lower bound of the support is equal to the lower bound parameter
     * of the distribution.
     */
    @Override
    public double getSupportLowerBound() {
        return lower;
    }

    /**
     * {@inheritDoc}
     *
     * <p>The upper bound of the support is equal to the upper bound parameter
     * of the distribution.
     */
    @Override
    public double getSupportUpperBound() {
        return upper;
    }

    /**
     * Clip the value to the range [lower, upper].
     * This is used to handle floating-point error at the support bound.
     *
     * @param x Value x
     * @return x clipped to the range
     */
    private double clipToRange(double x) {
        return clip(x, lower, upper);
    }

    /**
     * Clip the value to the range [lower, upper].
     *
     * @param x Value x
     * @param lower Lower bound (inclusive)
     * @param upper Upper bound (inclusive)
     * @return x clipped to the range
     */
    private static double clip(double x, double lower, double upper) {
        if (x <= lower) {
            return lower;
        }
        return x < upper ? x : upper;
    }

    // Calculation of variance and mean can suffer from cancellation.
    //
    // Use formulas from Jorge Fernandez-de-Cossio-Diaz adapted under the
    // terms of the MIT "Expat" License (see NOTICE and LICENSE).
    //
    // These formulas use the complementary error function
    //   erfcx(z) = erfc(z) * exp(z^2)
    // This avoids computation of exp terms for the Gaussian PDF and then
    // dividing by the error functions erf or erfc:
    //   exp(-0.5*x*x) / erfc(x / sqrt(2)) == 1 / erfcx(x / sqrt(2))
    // At large z the erfcx function is computable but exp(-0.5*z*z) and
    // erfc(z) are zero. Use of these formulas allows computation of the
    // mean and variance for the usable range of the truncated distribution
    // (cdf(a, b) != 0). The variance is not accurate when it approaches
    // machine epsilon (2^-52) at extremely narrow truncations and the
    // computation -> 0.
    //
    // See: https://github.com/cossio/TruncatedNormal.jl

    /**
     * Compute the first moment (mean) of the truncated standard normal distribution.
     *
     * <p>Assumes {@code a <= b}.
     *
     * @param a Lower bound
     * @param b Upper bound
     * @return the first moment
     */
    static double moment1(double a, double b) {
        // Assume a <= b
        if (a == b) {
            return a;
        }
        if (Math.abs(a) > Math.abs(b)) {
            // Subtract from zero to avoid generating -0.0
            return 0 - moment1(-b, -a);
        }

        // Here:
        // |a| <= |b|
        // a < b
        // 0 < b

        if (a <= -MAX_X) {
            // No truncation
            return 0;
        }
        if (b >= MAX_X) {
            // One-sided truncation
            return ROOT_2_PI / Erfcx.value(a / ROOT2);
        }

        // pdf = exp(-0.5*x*x) / sqrt(2*pi)
        // cdf = erfc(-x/sqrt(2)) / 2
        // Compute:
        // -(pdf(b) - pdf(a)) / cdf(b, a)
        // Note:
        // exp(-0.5*b*b) - exp(-0.5*a*a)
        // Use cancellation of powers:
        // exp(-0.5*(b*b-a*a)) * exp(-0.5*a*a) - exp(-0.5*a*a)
        // expm1(-0.5*(b*b-a*a)) * exp(-0.5*a*a)

        // dx = -0.5*(b*b-a*a)
        final double dx = 0.5 * (b + a) * (b - a);
        final double m;
        if (a <= 0) {
            // Opposite signs
            m = ROOT_2_PI * -Math.expm1(-dx) * Math.exp(-0.5 * a * a) / ErfDifference.value(a / ROOT2, b / ROOT2);
        } else {
            final double z = Math.exp(-dx) * Erfcx.value(b / ROOT2) - Erfcx.value(a / ROOT2);
            if (z == 0) {
                // Occurs when a and b have large magnitudes and are very close
                return (a + b) * 0.5;
            }
            m = ROOT_2_PI * Math.expm1(-dx) / z;
        }

        // Clip to the range
        return clip(m, a, b);
    }

    /**
     * Compute the second moment of the truncated standard normal distribution.
     *
     * <p>Assumes {@code a <= b}.
     *
     * @param a Lower bound
     * @param b Upper bound
     * @return the first moment
     */
    private static double moment2(double a, double b) {
        // Assume a < b.
        // a == b is handled in the variance method
        if (Math.abs(a) > Math.abs(b)) {
            return moment2(-b, -a);
        }

        // Here:
        // |a| <= |b|
        // a < b
        // 0 < b

        if (a <= -MAX_X) {
            // No truncation
            return 1;
        }
        if (b >= MAX_X) {
            // One-sided truncation.
            // For a -> inf : moment2 -> a*a
            // This occurs when erfcx(z) is approximated by (1/sqrt(pi)) / z and terms
            // cancel. z > 6.71e7, a > 9.49e7
            return 1 + ROOT_2_PI * a / Erfcx.value(a / ROOT2);
        }

        // pdf = exp(-0.5*x*x) / sqrt(2*pi)
        // cdf = erfc(-x/sqrt(2)) / 2
        // Compute:
        // 1 - (b*pdf(b) - a*pdf(a)) / cdf(b, a)
        // = (cdf(b, a) - b*pdf(b) -a*pdf(a)) / cdf(b, a)

        // Note:
        // For z -> 0:
        //   sqrt(pi / 2) * erf(z / sqrt(2)) -> z
        //   z * Math.exp(-0.5 * z * z) -> z
        // Both computations below have cancellation as b -> 0 and the
        // second moment is not computable as the fraction P/Q
        // since P < ulp(Q). This always occurs when b < MIN_X
        // if MIN_X is set at the point where
        //   exp(-0.5 * z * z) / sqrt(2 pi) == 1 / sqrt(2 pi).
        // This is JDK dependent due to variations in Math.exp.
        // For b < MIN_X the second moment can be approximated using
        // a uniform distribution: (b^3 - a^3) / (3b - 3a).
        // In practice it also occurs when b > MIN_X since any a < MIN_X
        // is effectively zero for part of the computation. A
        // threshold to transition to a uniform distribution
        // approximation is a compromise. Also note it will not
        // correct computation when (b-a) is small and is far from 0.
        // Thus the second moment is left to be inaccurate for
        // small ranges (b-a) and the variance -> 0 when the true
        // variance is close to or below machine epsilon.

        double m;

        if (a <= 0) {
            // Opposite signs
            final double ea = ROOT_PI_2 * Erf.value(a / ROOT2);
            final double eb = ROOT_PI_2 * Erf.value(b / ROOT2);
            final double fa = ea - a * Math.exp(-0.5 * a * a);
            final double fb = eb - b * Math.exp(-0.5 * b * b);
            // Assume fb >= fa && eb >= ea
            // If fb <= fa this is a tiny range around 0
            m = (fb - fa) / (eb - ea);
            // Clip to the range
            m = clip(m, 0, 1);
        } else {
            final double dx = 0.5 * (b + a) * (b - a);
            final double ex = Math.exp(-dx);
            final double ea = ROOT_PI_2 * Erfcx.value(a / ROOT2);
            final double eb = ROOT_PI_2 * Erfcx.value(b / ROOT2);
            final double fa = ea + a;
            final double fb = eb + b;
            m = (fa - fb * ex) / (ea - eb * ex);
            // Clip to the range
            m = clip(m, a * a, b * b);
        }
        return m;
    }

    /**
     * Compute the variance of the truncated standard normal distribution.
     *
     * <p>Assumes {@code a <= b}.
     *
     * @param a Lower bound
     * @param b Upper bound
     * @return the first moment
     */
    static double variance(double a, double b) {
        if (a == b) {
            return 0;
        }

        final double m1 = moment1(a, b);
        double m2 = moment2(a, b);
        // variance = m2 - m1*m1
        // rearrange x^2 - y^2 as (x-y)(x+y)
        m2 = Math.sqrt(m2);
        final double variance = (m2 - m1) * (m2 + m1);

        // Detect floating-point error.
        if (variance >= 1) {
            // Note:
            // Extreme truncations in the tails can compute a variance above 1,
            // for example if m2 is infinite: m2 - m1*m1 > 1
            // Detect no truncation as the terms a and b lie far either side of zero;
            // otherwise return 0 to indicate very small unknown variance.
            return a < -1 && b > 1 ? 1 : 0;
        } else if (variance <= 0) {
            // Floating-point error can create negative variance so return 0.
            return 0;
        }

        return variance;
    }
}