//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Aggregate/Interference2DParacrystal.cpp
//! @brief     Implements class Interference2DParacrystal.
//!
//! @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/Aggregate/Interference2DParacrystal.h"
#include "Base/Math/IntegratorGK.h"
#include "Base/Util/Assert.h"
#include <limits>

Interference2DParacrystal::Interference2DParacrystal(const Lattice2D& lattice,
                                                     double damping_length, double domain_size_1,
                                                     double domain_size_2)
    : IInterference(0)
    , m_integrate_xi(false)
    , m_damping_length(damping_length)
    , m_domain_sizes({domain_size_1, domain_size_2})
{
    m_lattice.reset(lattice.clone());
    validateOrThrow();
}

Interference2DParacrystal::~Interference2DParacrystal() = default;

Interference2DParacrystal* Interference2DParacrystal::clone() const
{
    auto* result = new Interference2DParacrystal(*m_lattice, m_damping_length, m_domain_sizes[0],
                                                 m_domain_sizes[1]);
    result->setPositionVariance(m_position_var);
    if (m_pdf1 && m_pdf2)
        result->setProbabilityDistributions(*m_pdf1, *m_pdf2);
    result->setIntegrationOverXi(m_integrate_xi);
    return result;
}

//! Sets the probability distributions (Fourier transformed) for the two lattice directions.
//! @param pdf_1: probability distribution in first lattice direction
//! @param pdf_2: probability distribution in second lattice direction

void Interference2DParacrystal::setProbabilityDistributions(const IProfile2D& pdf_1,
                                                            const IProfile2D& pdf_2)
{
    m_pdf1.reset(pdf_1.clone());
    m_pdf2.reset(pdf_2.clone());
}

//! Sets the damping length.
//! @param damping_length: the damping (coherence) length of the paracrystal in nanometers

void Interference2DParacrystal::setDampingLength(double damping_length)
{
    m_damping_length = damping_length;
}

double Interference2DParacrystal::particleDensity() const
{
    double area = m_lattice->unitCellArea();
    return area == 0.0 ? 0.0 : 1.0 / area;
}

std::vector<const INode*> Interference2DParacrystal::nodeChildren() const
{
    return std::vector<const INode*>() << m_pdf1 << m_pdf2 << m_lattice;
}

double Interference2DParacrystal::iff_without_dw(const R3& q) const
{
    double range = pi;
    if (!m_integrate_xi)
        return interferenceForXi(m_lattice->rotationAngle(), q.x(), q.y());
    return RealIntegrator().integrate(
               [&](double xi) -> double { return interferenceForXi(xi, q.x(), q.y()); }, 0.0, range)
           / range;
}

//! Sets the sizes of coherence domains.
//! @param size_1: coherence domain size along the first basis vector in nanometers
//! @param size_2: coherence domain size along the second basis vector in nanometers

void Interference2DParacrystal::setDomainSizes(double size_1, double size_2)
{
    m_domain_sizes = {size_1, size_2};
}

void Interference2DParacrystal::transformToPrincipalAxes(double qx, double qy, double gamma,
                                                         double delta, double& q_pa_1,
                                                         double& q_pa_2) const
{
    q_pa_1 = qx * std::cos(gamma) + qy * std::sin(gamma);
    q_pa_2 = qx * std::cos(gamma + delta) + qy * std::sin(gamma + delta);
}

//! Returns interference function for fixed angle xi.
double Interference2DParacrystal::interferenceForXi(double xi, double qx, double qy) const
{
    // don't touch order of computation; problems under Windows
    double rx = interference1D(qx, qy, xi, 0);
    double ry = interference1D(qx, qy, xi + m_lattice->latticeAngle(), 1);
    return rx * ry;
}

//! Returns interference function for fixed xi in the dimension determined by the given index.
double Interference2DParacrystal::interference1D(double qx, double qy, double xi,
                                                 size_t index) const
{
    ASSERT(m_validated);
    ASSERT(index <= 1);
    ASSERT(m_pdf1 && m_pdf2);

    double length = index ? m_lattice->length2() : m_lattice->length1();
    int n = static_cast<int>(std::abs(m_domain_sizes[index] / length));
    auto nd = static_cast<double>(n);
    complex_t fp = FTPDF(qx, qy, xi, index);
    if (n < 1)
        return ((1.0 + fp) / (1.0 - fp)).real();
    if (std::norm(1.0 - fp) < std::numeric_limits<double>::epsilon())
        return nd;
    // for (1-fp)*nd small, take the series expansion to second order in nd*(1-fp)
    if (std::abs(1.0 - fp) * nd < 2e-4) {
        complex_t intermediate =
            (nd - 1.0) / 2.0 + (nd * nd - 1.0) * (fp - 1.0) / 6.0
            + (nd * nd * nd - 2.0 * nd * nd - nd + 2.0) * (fp - 1.0) * (fp - 1.0) / 24.0;
        return 1.0 + 2.0 * intermediate.real();
    }
    complex_t tmp;
    if (std::abs(fp) == 0.0
        || std::log(std::abs(fp)) * nd < std::log(std::numeric_limits<double>::min()))
        tmp = 0.0;
    else
        tmp = std::pow(fp, n);
    complex_t intermediate = fp / (1.0 - fp) - fp * (1.0 - tmp) / nd / (1.0 - fp) / (1.0 - fp);
    return 1.0 + 2.0 * intermediate.real();
}

complex_t Interference2DParacrystal::FTPDF(double qx, double qy, double xi, size_t index) const
{
    ASSERT(m_validated);
    double length = (index ? m_lattice->length2() : m_lattice->length1());

    const IProfile2D* pdf = (index ? m_pdf2.get() : m_pdf1.get());
    double qa = qx * length * std::cos(xi) + qy * length * std::sin(xi);
    complex_t phase = exp_I(qa);
    // transform q to principal axes:
    double qp1, qp2;
    double gamma = xi + pdf->gamma();
    double delta = pdf->delta();
    transformToPrincipalAxes(qx, qy, gamma, delta, qp1, qp2);
    double amplitude = pdf->standardizedFT2D(qp1, qp2);
    complex_t result = phase * amplitude;
    if (m_damping_length != 0.0)
        result *= std::exp(-length / m_damping_length);
    return result;
}

std::vector<double> Interference2DParacrystal::domainSizes() const
{
    return {m_domain_sizes[0], m_domain_sizes[1]};
}

//! Enables/disables averaging over the lattice rotation angle.
//! @param integrate_xi: integration flag

void Interference2DParacrystal::setIntegrationOverXi(bool integrate_xi)
{
    m_integrate_xi = integrate_xi;
    m_lattice->setRotationEnabled(!m_integrate_xi); // deregister Xi in the case of integration
}

const Lattice2D& Interference2DParacrystal::lattice() const
{
    ASSERT(m_lattice);
    return *m_lattice;
}


std::string Interference2DParacrystal::validate() const
{
    std::vector<std::string> errs;
    requestGe0(errs, m_damping_length, "DampingLength");
    requestGe0(errs, m_domain_sizes[0], "DomainSize1");
    requestGe0(errs, m_domain_sizes[1], "DomainSize2");
    if (!errs.empty())
        return jointError(errs);
    m_validated = true;
    return "";
}
