#include "distributions.hpp"

#include <algorithm>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <random>
#include <string>
#include <vector>
#if defined(_MSC_VER)
    #pragma warning(push)
    #pragma warning(disable : 4189) // Disable some compiler warnings that come from fmt
#endif
#include <fmt/std.h>
#if defined(_MSC_VER)
    #pragma warning(pop)
#endif

/*
    @param bins: std::vector<DistributionPair>&, reference to vector of bins
    @param element: std::uint32_t, the number to insert - fractional values
           should be `floored` as to meet the fractional requirements of the
   program
    @param start: std::uint8_t, the beginning of the current window
    @param end: std::uint8_t, the end of the current window
    @return DistributionPair*: pointer to the bin element fits in
*/
DistributionPair* binarySearchBins(std::vector<DistributionPair>& bins,
                                   std::uint32_t element, std::uint8_t start,
                                   std::uint8_t end);
/*
    @param min: uint32_t, the minvalue value of the first bin
    @param max: uint32_t, the maxvalue of the last bin
    @param howMany: uint32_t, the number of bins to create
    @return std::vector<DistributionPair>: a vector of bins initialized
            with count = 0
*/
std::vector<DistributionPair>
initializeBins(std::uint32_t min, std::uint32_t max, std::uint8_t howMany);

DistributionPair* binarySearchBins(std::vector<DistributionPair>& bins,
                                   std::uint32_t element, std::uint8_t start,
                                   std::uint8_t end)
{
    std::uint8_t mid = (end + start) / 2;
    DistributionPair* dP = &bins.at(mid);

    // Base cases:
    // 1. Our element is less than the minimum in the bins
    // 2. Our element is greater than maximum in the bins
    // 3. Our element belongs in the current bin
    if (end <= start || mid == bins.size() - 1 ||
        (dP->minValue <= element && dP->maxValue >= element))
        return dP;

    // RESCURSION!!!
    if (dP->maxValue < element)
        return binarySearchBins(bins, element, mid, end);
    return binarySearchBins(bins, element, start, mid);
}

std::vector<DistributionPair>
initializeBins(std::uint32_t min, std::uint32_t max, std::uint8_t howMany)
{
    std::vector<DistributionPair> bins;
    // The difference between each minValue and maxValue of each bin can be zero
    // if numBins = num. of elements
    std::uint32_t delta = static_cast<uint32_t>(
        std::ceil((static_cast<float>(max - min) / howMany) - 1));

    for (std::uint32_t i = 0; i < howMany; i++)
    {
        std::uint32_t newMin = min;
        if (i)
            newMin = bins.at(i - 1).maxValue + 1;
        DistributionPair newBin(newMin, newMin + delta);
        bins.push_back(newBin);
    }

    return bins;
}

void plotDistribution(std::string title,
                      const std::vector<DistributionPair>& distribution,
                      const std::uint8_t maxPlotLineSize)
{
    std::cout << fmt::format("{}", title) << std::endl;

    std::uint32_t maxCount =
        std::max_element(
            std::begin(distribution), std::end(distribution),
            [](const DistributionPair bin1, const DistributionPair bin2)
            { return bin1.count < bin2.count; })
            ->count;

    std::uint32_t maxBinElementWidth = static_cast<uint32_t>(ceil(log10(distribution.back().maxValue)) + 1);

    std::uint32_t binThreshold = maxCount / maxPlotLineSize;
    for (const DistributionPair bin : distribution)
    {
        std::cout << fmt::format("[{0:>{1}}, {2:>{3}}] {4:<{5}}", bin.minValue,
                                 maxBinElementWidth, bin.maxValue,
                                 maxBinElementWidth, std::string(bin.count / binThreshold, '*'), maxPlotLineSize)
                  << std::endl;
    }
}

std::vector<DistributionPair>
generateUniformDistribution(std::uint32_t howMany, std::uint32_t min,
                            std::uint32_t max, std::uint8_t numberBins)
{
    std::vector<DistributionPair> bins = initializeBins(min, max, numberBins);

    std::default_random_engine generator;
    std::uniform_int_distribution<int> distribution(min, max);

    for (std::uint32_t i = 0; i < howMany; i++)
    {
        uint32_t randomElement = distribution(generator);
        binarySearchBins(bins, randomElement, 0, numberBins)->count++;
    }

    return bins;
}

std::vector<DistributionPair>
generateNormalDistribution(std::uint32_t howMany, float mean, float stdev,
                           std::uint8_t numberBins)
{
    std::uint32_t min = static_cast<std::uint32_t>(mean - 4 * stdev);
    std::uint32_t max = static_cast<std::uint32_t>(mean + 4 * stdev);

    // Get the right number of bins
    if (((max - min) % numberBins > 0) && (((max - 1) - min) % numberBins == 0))
        max -= 1;

    std::vector<DistributionPair> bins = initializeBins(min, max, numberBins);

    std::default_random_engine generator;
    std::normal_distribution<float> distribution(mean, stdev);

    for (std::uint32_t i = 0; i < howMany; i++)
    {
        uint32_t randomElement = static_cast<uint32_t>(distribution(generator));
        binarySearchBins(bins, randomElement, 0, numberBins)->count++;
    }

    return bins;
}

std::vector<DistributionPair>
generatePoissonDistribution(std::uint32_t howMany, std::uint8_t howOften,
                            std::uint8_t numberBins)
{
    std::vector<DistributionPair> bins =
        initializeBins(0, numberBins - 1, numberBins);

    std::default_random_engine generator;
    std::poisson_distribution<std::uint32_t> distribution(howOften);

    for (std::uint32_t i = 0; i < howMany; i++)
    {
        uint32_t randomElement = distribution(generator);
        binarySearchBins(bins, randomElement, 0, numberBins)->count++;
    }

    return bins;
}