using System;
using System.Numerics;
using System.Threading;
using FastRng.Distributions;

namespace FastRng;

/// <summary>
/// ShapeFitter is a rejection sampler, cf. https://en.wikipedia.org/wiki/Rejection_sampling
/// </summary>
public sealed class ShapeFitter<TNum> where TNum : IFloatingPointIeee754<TNum>, IDivisionOperators<TNum, TNum, TNum>
{
    private readonly TNum[] probabilities;
    private readonly IRandom<TNum> rng;
    private readonly TNum max;
    private readonly TNum sampleSize;
    private readonly IDistribution<TNum> uniform;

    /// <summary>
    /// Creates a shape fitter instance.
    /// </summary>
    /// <param name="shapeFunction">The function which describes the desired shape.</param>
    /// <param name="rng">The random number generator instance to use.</param>
    /// <param name="sampleSize">The number of sampling steps to sample the given function.</param>
    public ShapeFitter(Func<TNum, TNum> shapeFunction, IRandom<TNum> rng, ushort sampleSize = 50)
    {
        this.rng = rng;
        this.uniform = new Uniform<TNum>(rng);
        this.sampleSize = TNum.CreateChecked(sampleSize);
        this.probabilities = new TNum[sampleSize];

        var sampleStepSize = TNum.One / TNum.CreateChecked(sampleSize);
        var nextStep = TNum.Zero + sampleStepSize;
        var maxValue = TNum.Zero;
        for (var n = 0; n < sampleSize; n++)
        {
            this.probabilities[n] = shapeFunction(nextStep);
            if (this.probabilities[n] > maxValue)
                maxValue = this.probabilities[n];
                
            nextStep += sampleStepSize;
        }

        this.max = maxValue;
    }

    /// <summary>
    /// Returns a random number regarding the given shape.
    /// </summary>
    /// <param name="token">An optional cancellation token.</param>
    /// <returns>The next value regarding the given shape.</returns>
    public TNum NextNumber(CancellationToken token = default)
    {
        while (!token.IsCancellationRequested)
        {
            var x = this.rng.GetUniform(token);
            if (TNum.IsNaN(x))
                return x;
                
            var nextBucket = int.CreateChecked(TNum.Floor(x * this.sampleSize));
            if (nextBucket >= this.probabilities.Length)
                nextBucket = this.probabilities.Length - 1;
                
            var threshold = this.probabilities[nextBucket];
            var y = this.uniform.NextNumber(TNum.Zero, this.max, token);
            if (TNum.IsNaN(y))
                return y;
                
            if(y > threshold)
                continue;

            return x;
        }

        return TNum.NaN;
    }
}