/*
 * 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.rng.examples.sampling;

import java.io.PrintWriter;
import java.util.EnumSet;
import java.util.concurrent.Callable;
import java.io.IOException;
import org.apache.commons.rng.UniformRandomProvider;
import org.apache.commons.rng.simple.RandomSource;

import picocli.CommandLine.Command;
import picocli.CommandLine.Mixin;
import picocli.CommandLine.Option;

import org.apache.commons.rng.sampling.distribution.ZigguratNormalizedGaussianSampler;
import org.apache.commons.rng.sampling.distribution.ZigguratSampler;
import org.apache.commons.rng.sampling.distribution.MarsagliaNormalizedGaussianSampler;
import org.apache.commons.rng.sampling.distribution.StableSampler;
import org.apache.commons.rng.sampling.distribution.TSampler;
import org.apache.commons.rng.sampling.distribution.BoxMullerNormalizedGaussianSampler;
import org.apache.commons.rng.sampling.distribution.ChengBetaSampler;
import org.apache.commons.rng.sampling.distribution.AhrensDieterExponentialSampler;
import org.apache.commons.rng.sampling.distribution.AhrensDieterMarsagliaTsangGammaSampler;
import org.apache.commons.rng.sampling.distribution.InverseTransformParetoSampler;
import org.apache.commons.rng.sampling.distribution.LevySampler;
import org.apache.commons.rng.sampling.distribution.LogNormalSampler;
import org.apache.commons.rng.sampling.distribution.ContinuousUniformSampler;
import org.apache.commons.rng.sampling.distribution.GaussianSampler;
import org.apache.commons.rng.sampling.distribution.ContinuousSampler;

/**
 * Approximation of the probability density by the histogram of the sampler output.
 */
@Command(name = "density",
         description = {"Approximate the probability density of samplers."})
class ProbabilityDensityApproximationCommand implements Callable<Void> {
    /** The standard options. */
    @Mixin
    private StandardOptions reusableOptions;

    /** Number of (equal-width) bins in the histogram. */
    @Option(names = {"-b", "--bins"},
            description = "The number of bins in the histogram (default: ${DEFAULT-VALUE}).")
    private int numBins = 25_000;

    /** Number of samples to be generated. */
    @Option(names = {"-n", "--samples"},
            description = "The number of samples in the histogram (default: ${DEFAULT-VALUE}).")
    private long numSamples = 1_000_000_000;

    /** The samplers. */
    @Option(names = {"-s", "--samplers"},
            split = ",",
            description = {"The samplers (comma-delimited for multiple options).",
                           "Valid values: ${COMPLETION-CANDIDATES}."})
    private EnumSet<Sampler> samplers = EnumSet.noneOf(Sampler.class);

    /** The samplers. */
    @Option(names = {"-r", "--rng"},
            description = {"The source of randomness (default: ${DEFAULT-VALUE})."})
    private RandomSource randomSource = RandomSource.XOR_SHIFT_1024_S_PHI;

    /** Flag to output all samplers. */
    @Option(names = {"-a", "--all"},
            description = "Output all samplers")
    private boolean allSamplers;

    /**
     * The sampler. This enum uses lower case for clarity when matching the distribution name.
     */
    enum Sampler {
        /** The ziggurat gaussian sampler. */
        ZigguratGaussianSampler,
        /** The Marsaglia gaussian sampler. */
        MarsagliaGaussianSampler,
        /** The Box Muller gaussian sampler. */
        BoxMullerGaussianSampler,
        /** The modified ziggurat gaussian sampler. */
        ModifiedZigguratGaussianSampler,
        /** The Cheng beta sampler case 1. */
        ChengBetaSamplerCase1,
        /** The Cheng beta sampler case 2. */
        ChengBetaSamplerCase2,
        /** The Ahrens Dieter exponential sampler. */
        AhrensDieterExponentialSampler,
        /** The modified ziggurat exponential sampler. */
        ModifiedZigguratExponentialSampler,
        /** The Ahrens Dieter Marsaglia Tsang gamma sampler small gamma. */
        AhrensDieterMarsagliaTsangGammaSamplerCase1,
        /** The Ahrens Dieter Marsaglia Tsang gamma sampler large gamma. */
        AhrensDieterMarsagliaTsangGammaSamplerCase2,
        /** The inverse transform pareto sampler. */
        InverseTransformParetoSampler,
        /** The continuous uniform sampler. */
        ContinuousUniformSampler,
        /** The log normal ziggurat gaussian sampler. */
        LogNormalZigguratGaussianSampler,
        /** The log normal Marsaglia gaussian sampler. */
        LogNormalMarsagliaGaussianSampler,
        /** The log normal Box Muller gaussian sampler. */
        LogNormalBoxMullerGaussianSampler,
        /** The log normal modified ziggurat gaussian sampler. */
        LogNormalModifiedZigguratGaussianSampler,
        /** The Levy sampler. */
        LevySampler,
        /** The stable sampler. */
        StableSampler,
        /** The t sampler. */
        TSampler,
    }

    /**
     * @param sampler Sampler.
     * @param min Right abscissa of the first bin: every sample smaller
     * than that value will increment an additional bin (of infinite width)
     * placed before the first "equal-width" bin.
     * @param max abscissa of the last bin: every sample larger than or
     * equal to that value will increment an additional bin (of infinite
     * width) placed after the last "equal-width" bin.
     * @param outputFile Filename (final name is "pdf.[filename].txt").
     * @throws IOException Signals that an I/O exception has occurred.
     */
    private void createDensity(ContinuousSampler sampler,
                               double min,
                               double max,
                               String outputFile)
        throws IOException {
        final double binSize = (max - min) / numBins;
        final long[] histogram = new long[numBins];

        long belowMin = 0;
        long aboveMax = 0;
        for (long n = 0; n < numSamples; n++) {
            final double r = sampler.sample();

            if (r < min) {
                ++belowMin;
                continue;
            }

            if (r >= max) {
                ++aboveMax;
                continue;
            }

            final int binIndex = (int) ((r - min) / binSize);
            ++histogram[binIndex];
        }

        final double binHalfSize = 0.5 * binSize;
        final double norm = 1 / (binSize * numSamples);

        try (PrintWriter out = new PrintWriter("pdf." + outputFile + ".txt", "UTF-8")) {
            // CHECKSTYLE: stop MultipleStringLiteralsCheck
            out.println("# Sampler: " + sampler);
            out.println("# Number of bins: " + numBins);
            out.println("# Min: " + min + " (fraction of samples below: " + (belowMin / (double) numSamples) + ")");
            out.println("# Max: " + max + " (fraction of samples above: " + (aboveMax / (double) numSamples) + ")");
            out.println("# Bin width: " + binSize);
            out.println("# Histogram normalization factor: " + norm);
            out.println("#");
            out.println("# " + (min - binHalfSize) + " " + (belowMin * norm));
            for (int i = 0; i < numBins; i++) {
                out.println((min + (i + 1) * binSize - binHalfSize) + " " + (histogram[i] * norm));
            }
            out.println("# " + (max + binHalfSize) + " " + (aboveMax * norm));
            // CHECKSTYLE: resume MultipleStringLiteralsCheck
        }
    }

    /**
     * Program entry point.
     *
     * @throws IOException if failure occurred while writing to files.
     */
    @Override
    public Void call() throws IOException {
        if (allSamplers) {
            samplers = EnumSet.allOf(Sampler.class);
        } else if (samplers.isEmpty()) {
            // CHECKSTYLE: stop regexp
            System.err.println("ERROR: No samplers specified");
            // CHECKSTYLE: resume regexp
            System.exit(1);
        }

        final UniformRandomProvider rng = randomSource.create();

        final double gaussMean = 1;
        final double gaussSigma = 2;
        final double gaussMin = -9;
        final double gaussMax = 11;
        if (samplers.contains(Sampler.ZigguratGaussianSampler)) {
            createDensity(GaussianSampler.of(ZigguratNormalizedGaussianSampler.of(rng),
                                             gaussMean, gaussSigma),
                          gaussMin, gaussMax, "gauss.ziggurat");
        }
        if (samplers.contains(Sampler.MarsagliaGaussianSampler)) {
            createDensity(GaussianSampler.of(MarsagliaNormalizedGaussianSampler.of(rng),
                                             gaussMean, gaussSigma),
                          gaussMin, gaussMax, "gauss.marsaglia");
        }
        if (samplers.contains(Sampler.BoxMullerGaussianSampler)) {
            createDensity(GaussianSampler.of(BoxMullerNormalizedGaussianSampler.of(rng),
                                             gaussMean, gaussSigma),
                          gaussMin, gaussMax, "gauss.boxmuller");
        }
        if (samplers.contains(Sampler.ModifiedZigguratGaussianSampler)) {
            createDensity(GaussianSampler.of(ZigguratSampler.NormalizedGaussian.of(rng),
                                             gaussMean, gaussSigma),
                          gaussMin, gaussMax, "gauss.modified.ziggurat");
        }

        final double betaMin = 0;
        final double betaMax = 1;
        if (samplers.contains(Sampler.ChengBetaSamplerCase1)) {
            final double alphaBeta = 4.3;
            final double betaBeta = 2.1;
            createDensity(ChengBetaSampler.of(rng, alphaBeta, betaBeta),
                          betaMin, betaMax, "beta.case1");
        }
        if (samplers.contains(Sampler.ChengBetaSamplerCase2)) {
            final double alphaBetaAlt = 0.5678;
            final double betaBetaAlt = 0.1234;
            createDensity(ChengBetaSampler.of(rng, alphaBetaAlt, betaBetaAlt),
                          betaMin, betaMax, "beta.case2");
        }

        final double meanExp = 3.45;
        final double expMin = 0;
        final double expMax = 60;
        if (samplers.contains(Sampler.AhrensDieterExponentialSampler)) {
            createDensity(AhrensDieterExponentialSampler.of(rng, meanExp),
                          expMin, expMax, "exp");
        }
        if (samplers.contains(Sampler.ModifiedZigguratExponentialSampler)) {
            createDensity(ZigguratSampler.Exponential.of(rng, meanExp),
                          expMin, expMax, "exp.modified.ziggurat");
        }

        final double gammaMin = 0;
        final double gammaMax1 = 40;
        final double thetaGamma = 3.456;
        if (samplers.contains(Sampler.AhrensDieterMarsagliaTsangGammaSamplerCase1)) {
            final double alphaGammaSmallerThanOne = 0.1234;
            createDensity(AhrensDieterMarsagliaTsangGammaSampler.of(rng, alphaGammaSmallerThanOne, thetaGamma),
                          gammaMin, gammaMax1, "gamma.case1");
        }
        if (samplers.contains(Sampler.AhrensDieterMarsagliaTsangGammaSamplerCase2)) {
            final double alphaGammaLargerThanOne = 2.345;
            final double gammaMax2 = 70;
            createDensity(AhrensDieterMarsagliaTsangGammaSampler.of(rng, alphaGammaLargerThanOne, thetaGamma),
                          gammaMin, gammaMax2, "gamma.case2");
        }

        final double scalePareto = 23.45;
        final double shapePareto = 0.789;
        final double paretoMin = 23;
        final double paretoMax = 400;
        if (samplers.contains(Sampler.InverseTransformParetoSampler)) {
            createDensity(InverseTransformParetoSampler.of(rng, scalePareto, shapePareto),
                          paretoMin, paretoMax, "pareto");
        }

        final double loUniform = -9.876;
        final double hiUniform = 5.432;
        if (samplers.contains(Sampler.ContinuousUniformSampler)) {
            createDensity(ContinuousUniformSampler.of(rng, loUniform, hiUniform),
                          loUniform, hiUniform, "uniform");
        }

        final double scaleLogNormal = 2.345;
        final double shapeLogNormal = 0.1234;
        final double logNormalMin = 5;
        final double logNormalMax = 25;
        if (samplers.contains(Sampler.LogNormalZigguratGaussianSampler)) {
            createDensity(LogNormalSampler.of(ZigguratNormalizedGaussianSampler.of(rng),
                                              scaleLogNormal, shapeLogNormal),
                          logNormalMin, logNormalMax, "lognormal.ziggurat");
        }
        if (samplers.contains(Sampler.LogNormalMarsagliaGaussianSampler)) {
            createDensity(LogNormalSampler.of(MarsagliaNormalizedGaussianSampler.of(rng),
                                              scaleLogNormal, shapeLogNormal),
                          logNormalMin, logNormalMax, "lognormal.marsaglia");
        }
        if (samplers.contains(Sampler.LogNormalBoxMullerGaussianSampler)) {
            createDensity(LogNormalSampler.of(BoxMullerNormalizedGaussianSampler.of(rng),
                                              scaleLogNormal, shapeLogNormal),
                          logNormalMin, logNormalMax, "lognormal.boxmuller");
        }
        if (samplers.contains(Sampler.LogNormalModifiedZigguratGaussianSampler)) {
            createDensity(LogNormalSampler.of(ZigguratSampler.NormalizedGaussian.of(rng),
                                              scaleLogNormal, shapeLogNormal),
                          logNormalMin, logNormalMax, "lognormal.modified.ziggurat");
        }

        if (samplers.contains(Sampler.LevySampler)) {
            final double levyLocation = 1.23;
            final double levyscale = 0.75;
            final double levyMin = levyLocation;
            // Quantile 0 to 0.7 (avoid long tail to infinity)
            final double levyMax = 6.2815;
            createDensity(LevySampler.of(rng, levyLocation, levyscale),
                          levyMin, levyMax, "levy");
        }

        if (samplers.contains(Sampler.StableSampler)) {
            final double stableAlpha = 1.23;
            final double stableBeta = 0.75;
            // Quantiles 0.05 to 0.9 (avoid long tail to infinity)
            final double stableMin = -1.7862;
            final double stableMax = 4.0364;
            createDensity(StableSampler.of(rng, stableAlpha, stableBeta),
                          stableMin, stableMax, "stable");
        }

        if (samplers.contains(Sampler.TSampler)) {
            final double tDegreesOfFreedom = 1.23;
            // Quantiles 0.02 to 0.98 (avoid long tail to infinity)
            final double tMin = -9.9264;
            final double tMax = 9.9264;
            createDensity(TSampler.of(rng, tDegreesOfFreedom),
                          tMin, tMax, "t");
        }

        return null;
    }
}
