//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Correlation/Profiles2D.cpp
//! @brief     Implements class interface IProfile2D and children thereof.
//!
//! @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 "Sample/Correlation/Profiles2D.h"
#include "Base/Math/Bessel.h"
#include "Base/Math/IntegratorGK.h"
#include "Base/Py/PyFmt.h"
#include "Base/Util/Assert.h"
#include <limits>

//  ************************************************************************************************
//  interface IProfile1D
//  ************************************************************************************************

IProfile2D::IProfile2D(const std::vector<double>& PValues)
    : INode(PValues)
    , m_omega_x(m_P[0])
    , m_omega_y(m_P[1])
    , m_gamma(m_P[2])
{
}

std::string IProfile2D::pythonConstructor() const
{
    ASSERT(m_P.size() == 3);
    return Py::Fmt::printFunction(className(),
                                  {{m_omega_x, "nm"}, {m_omega_y, "nm"}, {m_gamma, "rad"}});
}

double IProfile2D::sumsq(double qx, double qy) const
{
    ASSERT(m_validated);
    return qx * qx * m_omega_x * m_omega_x + qy * qy * m_omega_y * m_omega_y;
}

std::string IProfile2D::validate() const
{
    std::vector<std::string> errs;
    requestGt0(errs, m_omega_x, "omega_x");
    requestGt0(errs, m_omega_y, "omega_y");
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}

//  ************************************************************************************************
//  class Profile2DCauchy
//  ************************************************************************************************

Profile2DCauchy::Profile2DCauchy(const std::vector<double> P)
    : IProfile2D(P)
{
    validateOrThrow();
}

Profile2DCauchy::Profile2DCauchy(double omega_x, double omega_y, double gamma)
    : Profile2DCauchy(std::vector<double>{omega_x, omega_y, gamma})
{
}

Profile2DCauchy* Profile2DCauchy::clone() const
{
    return new Profile2DCauchy(m_omega_x, m_omega_y, m_gamma);
}

double Profile2DCauchy::standardizedFT2D(double qx, double qy) const
{
    ASSERT(m_validated);
    return std::pow(1.0 + sumsq(qx, qy), -1.5);
}

double Profile2DCauchy::decayFT2D(double qx, double qy) const
{
    ASSERT(m_validated);
    double sum_sq = qx * qx * m_omega_x * m_omega_x + qy * qy * m_omega_y * m_omega_y;
    return (2 * pi) * m_omega_x * m_omega_y * std::pow(1.0 + sum_sq, -1.5);
}

std::unique_ptr<IDistribution2DSampler> Profile2DCauchy::createSampler() const
{
    return std::make_unique<Distribution2DCauchySampler>(m_omega_x, m_omega_y);
}

//  ************************************************************************************************
//  class Profile2DGauss
//  ************************************************************************************************

Profile2DGauss::Profile2DGauss(const std::vector<double> P)
    : IProfile2D(P)
{
    validateOrThrow();
}

Profile2DGauss::Profile2DGauss(double omega_x, double omega_y, double gamma)
    : Profile2DGauss(std::vector<double>{omega_x, omega_y, gamma})
{
}

Profile2DGauss* Profile2DGauss::clone() const
{
    return new Profile2DGauss(m_omega_x, m_omega_y, m_gamma);
}

double Profile2DGauss::standardizedFT2D(double qx, double qy) const
{
    ASSERT(m_validated);
    return std::exp(-sumsq(qx, qy) / 2);
}

double Profile2DGauss::decayFT2D(double qx, double qy) const
{
    ASSERT(m_validated);
    double sum_sq = qx * qx * m_omega_x * m_omega_x + qy * qy * m_omega_y * m_omega_y;
    return (2 * pi) * m_omega_x * m_omega_y * std::exp(-sum_sq / 2.0);
}

std::unique_ptr<IDistribution2DSampler> Profile2DGauss::createSampler() const
{
    return std::make_unique<Distribution2DGaussSampler>(m_omega_x, m_omega_y);
}

//  ************************************************************************************************
//  class Profile2DGate
//  ************************************************************************************************

Profile2DGate::Profile2DGate(const std::vector<double> P)
    : IProfile2D(P)
{
    validateOrThrow();
}

Profile2DGate::Profile2DGate(double omega_x, double omega_y, double gamma)
    : Profile2DGate(std::vector<double>{omega_x, omega_y, gamma})
{
}

Profile2DGate* Profile2DGate::clone() const
{
    return new Profile2DGate(m_omega_x, m_omega_y, m_gamma);
}

double Profile2DGate::standardizedFT2D(double qx, double qy) const
{
    ASSERT(m_validated);
    double scaled_q = std::sqrt(sumsq(qx, qy));
    return Math::Bessel::J1c(scaled_q) * 2.0;
}

double Profile2DGate::decayFT2D(double, double) const
{
    ASSERT_NEVER; // not yet needed, not yet implemented
}

std::unique_ptr<IDistribution2DSampler> Profile2DGate::createSampler() const
{
    return std::make_unique<Distribution2DGateSampler>(m_omega_x, m_omega_y);
}

//  ************************************************************************************************
//  class Profile2DCone
//  ************************************************************************************************

Profile2DCone::Profile2DCone(const std::vector<double> P)
    : IProfile2D(P)
{
    validateOrThrow();
}

Profile2DCone::Profile2DCone(double omega_x, double omega_y, double gamma)
    : Profile2DCone(std::vector<double>{omega_x, omega_y, gamma})
{
}

Profile2DCone* Profile2DCone::clone() const
{
    return new Profile2DCone(m_omega_x, m_omega_y, m_gamma);
}

double Profile2DCone::standardizedFT2D(double qx, double qy) const
{
    ASSERT(m_validated);
    double scaled_q = std::sqrt(sumsq(qx, qy));
    if (scaled_q < std::numeric_limits<double>::epsilon())
        return 1.0 - 3.0 * scaled_q * scaled_q / 40.0;
    // second part of the integrand: \f$u^2\cdot J_0(u)\f$
    double integral = RealIntegrator().integrate(
        [](double x) -> double { return x * x * Math::Bessel::J0(x); }, 0.0, scaled_q);
    return 6.0 * (Math::Bessel::J1c(scaled_q) - integral / scaled_q / scaled_q / scaled_q);
}

double Profile2DCone::decayFT2D(double, double) const
{
    ASSERT_NEVER; // not yet needed, not yet implemented
}

std::unique_ptr<IDistribution2DSampler> Profile2DCone::createSampler() const
{
    return std::make_unique<Distribution2DConeSampler>(m_omega_x, m_omega_y);
}

//  ************************************************************************************************
//  class Profile2DVoigt
//  ************************************************************************************************

Profile2DVoigt::Profile2DVoigt(const std::vector<double> P)
    : IProfile2D(P)
    , m_eta(m_P[3])
{
    validateOrThrow();
}

Profile2DVoigt::Profile2DVoigt(double omega_x, double omega_y, double gamma, double eta)
    : Profile2DVoigt(std::vector<double>{omega_x, omega_y, gamma, eta})
{
}

Profile2DVoigt* Profile2DVoigt::clone() const
{
    return new Profile2DVoigt(m_omega_x, m_omega_y, m_gamma, m_eta);
}

double Profile2DVoigt::standardizedFT2D(double qx, double qy) const
{
    ASSERT(m_validated);
    double sum_sq = sumsq(qx, qy);
    return m_eta * std::exp(-sum_sq / 2) + (1.0 - m_eta) * std::pow(1.0 + sum_sq, -1.5);
}

double Profile2DVoigt::decayFT2D(double qx, double qy) const
{
    ASSERT(m_validated);
    double sum_sq = qx * qx * m_omega_x * m_omega_x + qy * qy * m_omega_y * m_omega_y;
    return (2 * pi) * m_omega_x * m_omega_y
           * (m_eta * std::exp(-sum_sq / 2.0) + (1.0 - m_eta) * std::pow(1.0 + sum_sq, -1.5));
}

std::unique_ptr<IDistribution2DSampler> Profile2DVoigt::createSampler() const
{
    throw std::runtime_error("Profile2DVoigt::createSampler not yet implemented."
                             " Please contact the maintainers");
}

std::string Profile2DVoigt::pythonConstructor() const
{
    ASSERT(m_P.size() == 4);
    return Py::Fmt::printFunction(
        className(), {{m_omega_x, "nm"}, {m_omega_y, "nm"}, {m_gamma, "rad"}, {m_eta, ""}});
}

std::string Profile2DVoigt::validate() const
{
    std::vector<std::string> errs;
    requestGt0(errs, m_omega_x, "omega_x");
    requestGt0(errs, m_omega_y, "omega_y");
    requestIn(errs, m_eta, "eta", 0, 1);
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}