//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Param/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 "Param/Distrib/Distributions.h"
#include "Base/Py/PyFmt.h"
#include "Base/Util/Assert.h"
#include "Param/Distrib/ParameterSample.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 IDistribution1D
//  ************************************************************************************************

IDistribution1D::IDistribution1D(const std::vector<double>& PValues, size_t n_samples,
                                 double rel_sampling_width)
    : INode(PValues)
    , m_n_samples(n_samples)
    , m_relative_sampling_width(rel_sampling_width)
{
}

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

std::vector<ParameterSample> IDistribution1D::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::vector<std::pair<double, double>> IDistribution1D::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;
}

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

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

DistributionGate::DistributionGate(const std::vector<double> P, size_t n_samples)
    : IDistribution1D(P, n_samples, 1.)
    , m_min(m_P[0])
    , m_max(m_P[1])
{
    validateOrThrow();
}

DistributionGate::DistributionGate(double min, double max, size_t n_samples)
    : DistributionGate(std::vector<double>{min, max}, n_samples)
{
}

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

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);
}

void DistributionGate::setMean(double val)
{
    double shift = val - mean();
    m_P[0] += shift; // min
    m_P[1] += shift; // max
    validateOrThrow();
}

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

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;
}

std::string DistributionGate::validate() const
{
    if (m_max < m_min)
        return jointError({"parameters violate condition min<=max"});
    m_validated = true;
    return "";
}

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(const std::vector<double> P, size_t n_samples,
                                         double rel_sampling_width)
    : IDistribution1D(P, n_samples, rel_sampling_width)
    , m_mean(m_P[0])
    , m_hwhm(m_P[1])
{
    validateOrThrow();
}

DistributionLorentz::DistributionLorentz(double mean, double hwhm, size_t n_samples,
                                         double rel_sampling_width)
    : DistributionLorentz(std::vector<double>{mean, hwhm}, n_samples, rel_sampling_width)
{
}

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

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;
}

void DistributionLorentz::setMean(double val)
{
    m_P[0] = val; // mean
    validateOrThrow();
}

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

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

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

std::string DistributionLorentz::validate() const
{
    std::vector<std::string> errs;
    requestGe0(errs, m_hwhm, "hwhm");
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}


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

DistributionGaussian::DistributionGaussian(const std::vector<double> P, size_t n_samples,
                                           double rel_sampling_width)
    : IDistribution1D(P, n_samples, rel_sampling_width)
    , m_mean(m_P[0])
    , m_std_dev(m_P[1])
{
    validateOrThrow();
    if (m_std_dev < 0.0)
        throw std::runtime_error("DistributionGaussian: std_dev < 0");
}

DistributionGaussian::DistributionGaussian(double mean, double std_dev, size_t n_samples,
                                           double rel_sampling_width)
    : DistributionGaussian(std::vector<double>{mean, std_dev}, n_samples, rel_sampling_width)
{
}

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

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));
}

void DistributionGaussian::setMean(double val)
{
    m_P[0] = val; // mean
    validateOrThrow();
}

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

std::string DistributionGaussian::validate() const
{
    std::vector<std::string> errs;
    requestGe0(errs, m_std_dev, "stdv");
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}

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

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);
}


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

DistributionLogNormal::DistributionLogNormal(const std::vector<double> P, size_t n_samples,
                                             double rel_sampling_width)
    : IDistribution1D(P, n_samples, rel_sampling_width)
    , m_median(m_P[0])
    , m_scale_param(m_P[1])
{
    validateOrThrow();
}

DistributionLogNormal::DistributionLogNormal(double median, double scale_param, size_t n_samples,
                                             double rel_sampling_width)
    : DistributionLogNormal(std::vector<double>{median, scale_param}, n_samples, rel_sampling_width)
{
}

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

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);
}

void DistributionLogNormal::setMean(double val)
{
    m_P[0] = val; // median
    validateOrThrow();
}

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

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);
}

std::string DistributionLogNormal::validate() const
{
    std::vector<std::string> errs;
    requestGe0(errs, m_scale_param, "scale_param");
    requestGt0(errs, m_median, "median");
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}

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(const std::vector<double> P, size_t n_samples)
    : IDistribution1D(P, n_samples)
    , m_mean(m_P[0])
    , m_hwhm(m_P[1])
{
    validateOrThrow();
}

DistributionCosine::DistributionCosine(double mean, double sigma, size_t n_samples)
    : DistributionCosine(std::vector<double>{mean, sigma}, n_samples)
{
}

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

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);
}

void DistributionCosine::setMean(double val)
{
    m_P[0] = val; // mean
    validateOrThrow();
}

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

std::string DistributionCosine::validate() const
{
    std::vector<std::string> errs;
    requestGe0(errs, m_hwhm, "hwhm");
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}

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

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