//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Base/Distrib/Distributions.cpp
//! @brief     Implements classes representing one-dimensional distributions.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2018
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Base/Distrib/Distributions.h"
#include "Base/Distrib/ParameterSample.h"
#include "Base/Py/PyFmt.h"
#include "Base/Util/Assert.h"
#include <cmath>
#include <limits>
#include <numbers>
#include <sstream>

using std::numbers::pi;

namespace {

bool DoubleEqual(double a, double b)
{
    double eps = 10.0
                 * std::max(std::abs(a) * std::numeric_limits<double>::epsilon(),
                            std::numeric_limits<double>::min());
    return std::abs(a - b) < eps;
}

std::vector<double> equidistantPointsInRange(size_t n_samples, double xmin, double xmax)
{
    if (n_samples < 2 || DoubleEqual(xmin, xmax))
        return {(xmin + xmax) / 2};
    std::vector<double> result(n_samples);
    for (size_t i = 0; i < n_samples; ++i)
        result[i] = xmin + i * (xmax - xmin) / (n_samples - 1.0);
    return result;
}

} // namespace


//  ************************************************************************************************
//  class Distribution1D
//  ************************************************************************************************

Distribution1D::Distribution1D(size_t n_samples, double rel_sampling_width)
    : m_n_samples(n_samples)
    , m_relative_sampling_width(rel_sampling_width)
{
}

std::vector<std::pair<double, double>> Distribution1D::plotGraph() const
{
    size_t N = 400;
    std::vector<std::pair<double, double>> result(N);
    const auto [xmin, xmax] = plotRange();
    for (size_t i = 0; i < N; ++i) {
        const double x = xmin + i * (xmax - xmin) / (N - 1);
        result[i] = {x, probabilityDensity(x)};
    }
    return result;
}

size_t Distribution1D::nSamples() const
{
    return isDelta() ? 1L : m_n_samples;
}

std::vector<ParameterSample> Distribution1D::samplesInRange(double xmin, double xmax) const
{
    const size_t N = nSamples();
    ASSERT(N > 0);
    if (N == 1)
        return {{(xmin + xmax) / 2, 1.}};

    std::vector<ParameterSample> result;
    result.reserve(N);
    double wgt = 0;
    for (size_t i = 0; i < N; ++i) {
        const double x = xmin + (xmax - xmin) * (i) / (N - 1);
        result.push_back({x, probabilityDensity(x)});
        wgt += probabilityDensity(x);
    }
    for (size_t i = 0; i < N; ++i)
        result[i].weight /= wgt;
    return result;
}

std::pair<double, double> Distribution1D::plotRange() const
{
    throw std::runtime_error(
        "This distribution is not defined by a finite range. Therefore cannot be displayed");
}

//  ************************************************************************************************
//  class DistributionGate
//  ************************************************************************************************

DistributionGate::DistributionGate(double min, double max, size_t n_samples)
    : Distribution1D(n_samples, 1.)
    , m_min(min)
    , m_max(max)
{
    DistributionGate::validate();
}

std::vector<ParameterSample> DistributionGate::distributionSamples() const
{
    std::vector<double> xx = equidistantPointsInRange(nSamples(), m_min, m_max);
    std::vector<ParameterSample> result;
    result.reserve(xx.size());
    for (double x : xx)
        result.push_back({x, 1. / xx.size()});
    return result;
}

void DistributionGate::setMean(double val)
{
    double shift = val - mean();
    m_min += shift;
    m_max += shift;
    validate();
}

DistributionGate* DistributionGate::clone() const
{
    return new DistributionGate(m_min, m_max, m_n_samples);
}

void DistributionGate::validate() const
{
    if (m_max < m_min)
        throw std::runtime_error("parameters violate condition min<=max");
    m_validated = true;
}

double DistributionGate::probabilityDensity(double x) const
{
    if (x < m_min || x > m_max)
        return 0.0;
    ASSERT(!isDelta());
    return 1.0 / (m_max - m_min);
}

bool DistributionGate::isDelta() const
{
    return m_min == m_max;
}

std::pair<double, double> DistributionGate::plotRange() const
{
    const double margin = std::abs(m_max - m_min) / 10;
    return {m_min - margin, m_max + margin};
}


//  ************************************************************************************************
//  class DistributionLorentz
//  ************************************************************************************************

DistributionLorentz::DistributionLorentz(double mean, double hwhm, size_t n_samples,
                                         double rel_sampling_width)
    : Distribution1D(n_samples, rel_sampling_width)
    , m_mean(mean)
    , m_hwhm(hwhm)
{
    DistributionLorentz::validate();
}

std::vector<ParameterSample> DistributionLorentz::distributionSamples() const
{
    ASSERT(m_hwhm >= 0);
    const double range = m_hwhm * relSamplingWidth();
    return samplesInRange(m_mean - range, m_mean + range);
}

void DistributionLorentz::setMean(double val)
{
    m_mean = val;
    validate();
}

DistributionLorentz* DistributionLorentz::clone() const
{
    return new DistributionLorentz(m_mean, m_hwhm, m_n_samples, relSamplingWidth());
}

void DistributionLorentz::validate() const
{
    if (m_hwhm < 0)
        throw std::runtime_error("hwhm < 0");
    m_validated = true;
}

double DistributionLorentz::probabilityDensity(double x) const
{
    ASSERT(m_validated);
    ASSERT(!isDelta());
    return m_hwhm / (m_hwhm * m_hwhm + (x - m_mean) * (x - m_mean)) / pi;
}

bool DistributionLorentz::isDelta() const
{
    return m_hwhm == 0.0;
}

std::pair<double, double> DistributionLorentz::plotRange() const
{
    return {m_mean - 4 * m_hwhm, m_mean + 4 * m_hwhm};
}


//  ************************************************************************************************
//  class DistributionGaussian
//  ************************************************************************************************

DistributionGaussian::DistributionGaussian(double mean, double std_dev, size_t n_samples,
                                           double rel_sampling_width)
    : Distribution1D(n_samples, rel_sampling_width)
    , m_mean(mean)
    , m_std_dev(std_dev)
{
    DistributionGaussian::validate();
}

std::vector<ParameterSample> DistributionGaussian::distributionSamples() const
{
    ASSERT(m_std_dev >= 0);
    const double range = m_std_dev * relSamplingWidth();
    return samplesInRange(m_mean - range, m_mean + range);
}

void DistributionGaussian::setMean(double val)
{
    m_mean = val;
    validate();
}

DistributionGaussian* DistributionGaussian::clone() const
{
    return new DistributionGaussian(m_mean, m_std_dev, m_n_samples, relSamplingWidth());
}

void DistributionGaussian::validate() const
{
    if (m_std_dev < 0)
        throw std::runtime_error("std_dev < 0");
    m_validated = true;
}

double DistributionGaussian::probabilityDensity(double x) const
{
    ASSERT(m_validated);
    ASSERT(!isDelta());
    double exponential = std::exp(-(x - m_mean) * (x - m_mean) / (2.0 * m_std_dev * m_std_dev));
    return exponential / m_std_dev / std::sqrt((2 * pi));
}

bool DistributionGaussian::isDelta() const
{
    return m_std_dev == 0.0;
}

std::pair<double, double> DistributionGaussian::plotRange() const
{
    return {m_mean - 3 * m_std_dev, m_mean + 3 * m_std_dev};
}


//  ************************************************************************************************
//  class DistributionLogNormal
//  ************************************************************************************************

DistributionLogNormal::DistributionLogNormal(double median, double scale_param, size_t n_samples,
                                             double rel_sampling_width)
    : Distribution1D(n_samples, rel_sampling_width)
    , m_median(median)
    , m_scale_param(scale_param)
{
    DistributionLogNormal::validate();
}

std::vector<ParameterSample> DistributionLogNormal::distributionSamples() const
{
    ASSERT(m_scale_param >= 0);
    const double range = m_scale_param * relSamplingWidth();
    double xmin = m_median * std::exp(-range);
    double xmax = m_median * std::exp(range);
    return samplesInRange(xmin, xmax);
}

void DistributionLogNormal::setMean(double val)
{
    m_median = val;
    validate();
}

DistributionLogNormal* DistributionLogNormal::clone() const
{
    return new DistributionLogNormal(m_median, m_scale_param, m_n_samples, relSamplingWidth());
}

void DistributionLogNormal::validate() const
{
    if (m_scale_param < 0)
        throw std::runtime_error("scale_param < 0");
    if (m_median <= 0)
        throw std::runtime_error("median <= 0");
    m_validated = true;
}

double DistributionLogNormal::probabilityDensity(double x) const
{
    ASSERT(m_validated);
    ASSERT(!isDelta());
    double t = std::log(x / m_median) / m_scale_param;
    return std::exp(-t * t / 2.0) / (x * m_scale_param * std::sqrt((2 * pi)));
}

double DistributionLogNormal::mean() const
{
    ASSERT(m_validated);
    double exponent = m_scale_param * m_scale_param / 2.0;
    return m_median * std::exp(exponent);
}

bool DistributionLogNormal::isDelta() const
{
    return m_scale_param == 0.0;
}

std::pair<double, double> DistributionLogNormal::plotRange() const
{
    double xmin = m_median * std::exp(-3 * m_scale_param);
    double xmax = m_median * std::exp(3 * m_scale_param);
    return {xmin, xmax};
}

//  ************************************************************************************************
//  class DistributionCosine
//  ************************************************************************************************

DistributionCosine::DistributionCosine(double mean, double sigma, size_t n_samples)
    : Distribution1D(n_samples)
    , m_mean(mean)
    , m_hwhm(sigma)
{
    DistributionCosine::validate();
}

std::vector<ParameterSample> DistributionCosine::distributionSamples() const
{
    auto [xmin, xmax] = plotRange();
    return samplesInRange(xmin, xmax);
}

void DistributionCosine::setMean(double val)
{
    m_mean = val;
    validate();
}

DistributionCosine* DistributionCosine::clone() const
{
    return new DistributionCosine(m_mean, m_hwhm, m_n_samples);
}

void DistributionCosine::validate() const
{
    if (m_hwhm < 0)
        throw std::runtime_error("hwhm < 0");
    m_validated = true;
}

double DistributionCosine::probabilityDensity(double x) const
{
    ASSERT(m_validated);
    ASSERT(!isDelta());
    if (std::abs(x - m_mean) > pi * m_hwhm)
        return 0.0;
    return (1.0 + std::cos(((x - m_mean) / m_hwhm) * (pi / 2))) / (4 * m_hwhm);
}

bool DistributionCosine::isDelta() const
{
    return m_hwhm == 0.0;
}

std::pair<double, double> DistributionCosine::plotRange() const
{
    double xmin = m_mean - 2 * m_hwhm;
    double xmax = m_mean + 2 * m_hwhm;
    return {xmin, xmax};
}
