//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Correlation/IPeakShape.cpp
//! @brief     Implements the interface IPeakShape and subclasses.
//!
//! @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/IPeakShape.h"
#include "Base/Math/Bessel.h"
#include "Base/Math/IntegratorGK.h"
#include <limits>
#include <numbers>

using std::numbers::pi;

namespace {

const double maxkappa = std::log(1.0 / std::numeric_limits<double>::epsilon()) / 2.0;
const double maxkappa2 = std::log(std::numeric_limits<double>::max());

double FisherDistribution(double x, double kappa)
{
    if (kappa <= 0.0)
        return 1.0 / (4.0 * pi);
    double prefactor = kappa / (4.0 * pi);
    if (kappa > maxkappa)
        return 2.0 * prefactor * std::exp(kappa * (x - 1.0));
    return prefactor * std::exp(kappa * x) / std::sinh(kappa);
}

double FisherPrefactor(double kappa)
{
    if (kappa <= 0.0)
        return 1.0 / (4.0 * pi);
    if (kappa > maxkappa)
        return kappa / 2.0 / pi;
    return kappa * std::exp(kappa) / 4.0 / pi / std::sinh(kappa);
}

double MisesPrefactor(double kappa)
{
    if (kappa <= 0.0)
        return 1.0 / ((2 * pi));
    if (kappa > maxkappa2)
        return std::sqrt(kappa / 2.0 / pi) / (1.0 + 1.0 / (8.0 * kappa));
    return std::exp(kappa) / ((2 * pi) * Math::Bessel::I0(kappa));
}

double Gauss3D(double q2, double domainsize)
{
    double norm_factor = std::pow(domainsize / std::sqrt((2 * pi)), 3.0);
    double exponent = -q2 * domainsize * domainsize / 2.0;
    return norm_factor * std::exp(exponent);
}

double Cauchy3D(double q2, double domainsize)
{
    double lorentz1 = domainsize / (1.0 + q2 * domainsize * domainsize) / pi;
    return domainsize * lorentz1 * lorentz1;
}

} // namespace

//  ************************************************************************************************
//  interface IPeakShape
//  ************************************************************************************************

IPeakShape::IPeakShape(const std::vector<double>& PValues)
    : INode(PValues)
{
}

IPeakShape::~IPeakShape() = default;

//  ************************************************************************************************
//  class IsotropicGaussPeakShape
//  ************************************************************************************************

IsotropicGaussPeakShape::IsotropicGaussPeakShape(double max_intensity, double domainsize)
    : m_max_intensity(max_intensity)
    , m_domainsize(domainsize)
{
}

IsotropicGaussPeakShape::~IsotropicGaussPeakShape() = default;

IsotropicGaussPeakShape* IsotropicGaussPeakShape::clone() const
{
    return new IsotropicGaussPeakShape(m_max_intensity, m_domainsize);
}

double IsotropicGaussPeakShape::peakDistribution(const R3& q) const
{
    double q_norm = q.mag2();
    return m_max_intensity * Gauss3D(q_norm, m_domainsize);
}

double IsotropicGaussPeakShape::peakDistribution(const R3& q, const R3& q_lattice_point) const
{
    return peakDistribution(q - q_lattice_point);
}

//  ************************************************************************************************
//  class IsotropicLorentzPeakShape
//  ************************************************************************************************

IsotropicLorentzPeakShape::IsotropicLorentzPeakShape(double max_intensity, double domainsize)
    : m_max_intensity(max_intensity)
    , m_domainsize(domainsize)
{
}

IsotropicLorentzPeakShape::~IsotropicLorentzPeakShape() = default;

IsotropicLorentzPeakShape* IsotropicLorentzPeakShape::clone() const
{
    return new IsotropicLorentzPeakShape(m_max_intensity, m_domainsize);
}

double IsotropicLorentzPeakShape::peakDistribution(const R3& q) const
{
    double q_norm = q.mag2();
    return m_max_intensity * Cauchy3D(q_norm, m_domainsize);
}

double IsotropicLorentzPeakShape::peakDistribution(const R3& q, const R3& q_lattice_point) const
{
    return peakDistribution(q - q_lattice_point);
}

//  ************************************************************************************************
//  class GaussFisherPeakShape
//  ************************************************************************************************

GaussFisherPeakShape::GaussFisherPeakShape(double max_intensity, double radial_size, double kappa)
    : m_max_intensity(max_intensity)
    , m_radial_size(radial_size)
    , m_kappa(kappa)
{
}

GaussFisherPeakShape::~GaussFisherPeakShape() = default;

GaussFisherPeakShape* GaussFisherPeakShape::clone() const
{
    return new GaussFisherPeakShape(m_max_intensity, m_radial_size, m_kappa);
}

double GaussFisherPeakShape::peakDistribution(const R3& q, const R3& q_lattice_point) const
{
    const double q_r = q.mag();
    const double q_lat_r = q_lattice_point.mag();
    const double dq2 = (q_r - q_lat_r) * (q_r - q_lat_r);
    if (q_lat_r == 0.0)
        return m_max_intensity * Gauss3D(dq2, m_radial_size);
    const double norm_factor = m_radial_size / std::sqrt((2 * pi));
    const double radial_part = norm_factor * std::exp(-dq2 * m_radial_size * m_radial_size / 2.0);
    double angular_part = 1.0;
    if (q_r * q_lat_r > 0.0) {
        const double dot_norm = q.dot(q_lattice_point) / q_r / q_lat_r;
        angular_part = FisherDistribution(dot_norm, m_kappa) / (q_r * q_r);
    }
    return m_max_intensity * radial_part * angular_part;
}

//  ************************************************************************************************
//  class LorentzFisherPeakShape
//  ************************************************************************************************

LorentzFisherPeakShape::LorentzFisherPeakShape(double max_intensity, double radial_size,
                                               double kappa)
    : m_max_intensity(max_intensity)
    , m_radial_size(radial_size)
    , m_kappa(kappa)
{
}

LorentzFisherPeakShape::~LorentzFisherPeakShape() = default;

LorentzFisherPeakShape* LorentzFisherPeakShape::clone() const
{
    return new LorentzFisherPeakShape(m_max_intensity, m_radial_size, m_kappa);
}

double LorentzFisherPeakShape::peakDistribution(const R3& q, const R3& q_lattice_point) const
{
    const double q_r = q.mag();
    const double q_lat_r = q_lattice_point.mag();
    const double dq2 = (q_r - q_lat_r) * (q_r - q_lat_r);
    if (q_lat_r == 0.0)
        return m_max_intensity * Cauchy3D(dq2, m_radial_size);
    const double radial_part = m_radial_size / (1.0 + dq2 * m_radial_size * m_radial_size) / pi;
    double angular_part = 1.0;
    if (q_r * q_lat_r > 0.0) {
        const double dot_norm = q.dot(q_lattice_point) / q_r / q_lat_r;
        angular_part = FisherDistribution(dot_norm, m_kappa) / (q_r * q_r);
    }
    return m_max_intensity * radial_part * angular_part;
}

//  ************************************************************************************************
//  class MisesFisherGaussPeakShape
//  ************************************************************************************************

MisesFisherGaussPeakShape::MisesFisherGaussPeakShape(double max_intensity, double radial_size,
                                                     const R3& zenith, double kappa_1,
                                                     double kappa_2)
    : m_max_intensity(max_intensity)
    , m_radial_size(radial_size)
    , m_zenith(zenith.unit_or_throw())
    , m_kappa_1(kappa_1)
    , m_kappa_2(kappa_2)
{
}

MisesFisherGaussPeakShape::~MisesFisherGaussPeakShape() = default;

MisesFisherGaussPeakShape* MisesFisherGaussPeakShape::clone() const
{
    return new MisesFisherGaussPeakShape(m_max_intensity, m_radial_size, m_zenith, m_kappa_1,
                                         m_kappa_2);
}

double MisesFisherGaussPeakShape::peakDistribution(const R3& q, const R3& q_lattice_point) const
{
    // radial part
    const double q_r = q.mag();
    const double q_lat_r = q_lattice_point.mag();
    const double dq2 = (q_r - q_lat_r) * (q_r - q_lat_r);
    if (q_lat_r == 0.0 || q_r == 0.0)
        return m_max_intensity * Gauss3D(dq2, m_radial_size);
    const double norm_factor = m_radial_size / std::sqrt((2 * pi));
    const double radial_part = norm_factor * std::exp(-dq2 * m_radial_size * m_radial_size / 2.0);
    // angular part
    const R3 vy = m_zenith.cross(q_lattice_point);
    const R3 zxq = m_zenith.cross(q);
    const R3 up = q_lattice_point.unit_or_throw();
    if (vy.mag2() <= 0.0 || zxq.mag2() <= 0.0) {
        const double x = q.unit_or_throw().dot(up);
        const double angular_part = FisherDistribution(x, m_kappa_1);
        return m_max_intensity * radial_part * angular_part;
    }
    const R3 uy = vy.unit_or_throw();
    const R3 ux = uy.cross(m_zenith);
    const R3 q_ortho = q - q.dot(m_zenith) * m_zenith;
    const double phi0 = std::acos(q_ortho.unit_or_throw().dot(ux));
    const double theta = std::acos(q.unit_or_throw().dot(m_zenith));
    const double pre_1 = FisherPrefactor(m_kappa_1);
    const double pre_2 = MisesPrefactor(m_kappa_2);
    const double integral = RealIntegrator().integrate(
        [&](double phi) -> double {
            const R3 u_q = std::sin(theta) * std::cos(phi) * ux
                           + std::sin(theta) * std::sin(phi) * uy + std::cos(theta) * m_zenith;
            const double fisher = std::exp(m_kappa_1 * (u_q.dot(up) - 1.0));
            const double mises = std::exp(m_kappa_2 * (std::cos(phi0 - phi) - 1.0));
            return fisher * mises;
        },
        0.0, (2 * pi));
    return m_max_intensity * radial_part * pre_1 * pre_2 * integral;
}

//  ************************************************************************************************
//  class MisesGaussPeakShape
//  ************************************************************************************************

MisesGaussPeakShape::MisesGaussPeakShape(double max_intensity, double radial_size, const R3& zenith,
                                         double kappa)
    : m_max_intensity(max_intensity)
    , m_radial_size(radial_size)
    , m_zenith(zenith.unit_or_throw())
    , m_kappa(kappa)
{
}

MisesGaussPeakShape::~MisesGaussPeakShape() = default;

MisesGaussPeakShape* MisesGaussPeakShape::clone() const
{
    return new MisesGaussPeakShape(m_max_intensity, m_radial_size, m_zenith, m_kappa);
}

double MisesGaussPeakShape::peakDistribution(const R3& q, const R3& q_lattice_point) const
{
    const R3 vy = m_zenith.cross(q_lattice_point);
    const R3 zxq = m_zenith.cross(q);
    if (vy.mag2() <= 0.0 || zxq.mag2() <= 0.0) {
        const double dq2 = (q - q_lattice_point).mag2();
        return m_max_intensity * Gauss3D(dq2, m_radial_size);
    }
    const double m_qr = q.mag();
    const R3 m_p = q_lattice_point;
    const R3 uy = vy.unit_or_throw();
    const R3 ux = uy.cross(m_zenith);
    const R3 q_ortho = q - q.dot(m_zenith) * m_zenith;
    const double phi0 = std::acos(q_ortho.unit_or_throw().dot(ux));
    const double theta = std::acos(q.unit_or_throw().dot(m_zenith));
    const double pre = MisesPrefactor(m_kappa);
    const double integral = RealIntegrator().integrate(
        [&](double phi) -> double {
            R3 q_rot = m_qr
                       * (std::sin(theta) * std::cos(phi) * ux
                          + std::sin(theta) * std::sin(phi) * uy + std::cos(theta) * m_zenith);
            const double dq2 = (q_rot - m_p).mag2();
            const double gauss = Gauss3D(dq2, m_radial_size);
            const double mises = std::exp(m_kappa * (std::cos(phi0 - phi) - 1.0));
            return gauss * mises;
        },
        0.0, (2 * pi));
    return m_max_intensity * pre * integral;
}