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

Paracrystal2D::Paracrystal2D(const IParticle& particle, const Lattice2D& lattice,
                             double damping_length, double domain_size_1, double domain_size_2)
    : Ordered2D(particle)
    , ICrystal2D(particle, lattice)
    , m_damping_length(damping_length)
    , m_domain_sizes({domain_size_1, domain_size_2})
{
    Paracrystal2D::validate();
}

Paracrystal2D::~Paracrystal2D() = default;

//! 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 Paracrystal2D::setDomainSizes(double size_1, double size_2)
{
    m_domain_sizes = {size_1, size_2};
}

//! 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 Paracrystal2D::setProbabilityDistributions(const Profile2D& pdf_1, const Profile2D& 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 Paracrystal2D::setDampingLength(double damping_length)
{
    m_damping_length = damping_length;
}

Paracrystal2D* Paracrystal2D::clone() const
{
    auto* result = new Paracrystal2D(*m_particle, *m_lattice, m_damping_length, m_domain_sizes[0],
                                     m_domain_sizes[1]);
    instructClone(result);
    instructCrystalClone(result);
    if (m_pdf1 && m_pdf2)
        result->setProbabilityDistributions(*m_pdf1, *m_pdf2);
    return result;
}

void Paracrystal2D::validate() const
{
    if (m_damping_length < 0)
        throw std::runtime_error("damping_length < 0");
    if (m_domain_sizes[0] < 0)
        throw std::runtime_error("domain_sizes[0] < 0");
    if (m_domain_sizes[1] < 0)
        throw std::runtime_error("domain_sizes[1] < 0");

    Struct::validate();
    m_validated = true;
}

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

//! Returns interference function for fixed angle xi.
double Paracrystal2D::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 Paracrystal2D::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 Paracrystal2D::FTPDF(double qx, double qy, double xi, size_t index) const
{
    ASSERT(m_validated);
    double length = (index ? m_lattice->length2() : m_lattice->length1());

    const Profile2D* 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;
}

void Paracrystal2D::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);
}
