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

namespace {

// maximum value for qx*Lambdax and qy*lambday
const int nmax = 20;
// minimum number of neighboring reciprocal lattice points to use
const int min_points = 4;

std::pair<double, double> toReciprocal(double qX, double qY, double a, double b, double alpha,
                                       double gamma)
{
    double qa = (a * qX * std::cos(gamma) - a * qY * std::sin(gamma)) / (2 * pi);
    double qb = (b * qX * std::cos(alpha - gamma) + b * qY * std::sin(alpha - gamma)) / (2 * pi);
    return {qa, qb};
}

//! Calculates bounding values of reciprocal lattice coordinates that contain the centered
//! rectangle with a corner defined by qX and qY
std::pair<double, double> boundingQs(double qX, double qY, double a, double b, double alpha,
                                     double gamma)
{
    auto q_bounds_1 = toReciprocal(qX, qY, a, b, alpha, gamma);
    auto q_bounds_2 = toReciprocal(qX, -qY, a, b, alpha, gamma);
    double qa_max = std::max(std::abs(q_bounds_1.first), std::abs(q_bounds_2.first));
    double qb_max = std::max(std::abs(q_bounds_1.second), std::abs(q_bounds_2.second));
    return {qa_max, qb_max};
}

} // namespace

Interference2DLattice::Interference2DLattice(const Lattice2D& lattice)
    : IInterference(0)
    , m_integrate_xi(false)
{
    m_lattice.reset(lattice.clone());

    BasicLattice2D base_lattice(m_lattice->length1(), m_lattice->length2(),
                                m_lattice->latticeAngle(), 0.);
    m_sbase = base_lattice.reciprocalBases();
}

Interference2DLattice::~Interference2DLattice() = default;

Interference2DLattice* Interference2DLattice::clone() const
{
    auto* result = new Interference2DLattice(*m_lattice);
    result->setPositionVariance(m_position_var);
    result->setIntegrationOverXi(integrationOverXi());
    if (m_decay)
        result->setDecayFunction(*m_decay);
    return result;
}

//! Sets two-dimensional decay function.
//! @param decay: two-dimensional decay function in reciprocal space
void Interference2DLattice::setDecayFunction(const IProfile2D& decay)
{
    m_decay.reset(decay.clone());

    // number of reciprocal lattice points to use
    auto q_bounds = boundingQs(nmax / m_decay->decayLengthX(), nmax / m_decay->decayLengthY(),
                               m_lattice->length1(), m_lattice->length2(),
                               m_lattice->latticeAngle(), m_decay->gamma());
    m_na = static_cast<int>(std::lround(q_bounds.first + 0.5));
    m_nb = static_cast<int>(std::lround(q_bounds.second + 0.5));
    m_na = std::max(m_na, min_points);
    m_nb = std::max(m_nb, min_points);
}

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

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

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

std::vector<const INode*> Interference2DLattice::nodeChildren() const
{
    return std::vector<const INode*>() << m_decay << m_lattice;
}

double Interference2DLattice::iff_without_dw(const R3 q) const
{
    if (!m_decay)
        throw std::runtime_error("Interference2DLattice needs a decay function");
    if (!m_integrate_xi)
        return interferenceForXi(m_lattice->rotationAngle(), q.x(), q.y());
    double range = pi;
    return RealIntegrator().integrate(
               [&](double xi) -> double { return interferenceForXi(xi, q.x(), q.y()); }, 0.0, range)
           / range;
}

double Interference2DLattice::interferenceForXi(double xi, double qx, double qy) const
{
    double result = 0.0;
    auto q_frac = calculateReciprocalVectorFraction(qx, qy, xi);

    for (int i = -m_na - 1; i < m_na + 2; ++i) {
        for (int j = -m_nb - 1; j < m_nb + 2; ++j) {
            const double px = q_frac.first + i * m_sbase.m_asx + j * m_sbase.m_bsx;
            const double py = q_frac.second + i * m_sbase.m_asy + j * m_sbase.m_bsy;
            result += interferenceAtOneRecLatticePoint(px, py);
        }
    }
    return particleDensity() * result;
}

double Interference2DLattice::interferenceAtOneRecLatticePoint(double qx, double qy) const
{
    if (!m_decay)
        throw std::runtime_error("Interference2DLattice needs a decay function");
    double gamma = m_decay->gamma();
    auto qXY = rotateOrthonormal(qx, qy, gamma);
    return m_decay->decayFT2D(qXY.first, qXY.second);
}

// Rotate by angle gamma between orthonormal systems
std::pair<double, double> Interference2DLattice::rotateOrthonormal(double qx, double qy,
                                                                   double gamma) const
{
    double q_X = qx * std::cos(gamma) + qy * std::sin(gamma);
    double q_Y = -qx * std::sin(gamma) + qy * std::cos(gamma);
    return {q_X, q_Y};
}

// (qx, qy) are in the global reciprocal reference frame
// the returned values (qx_frac, qy_frac) are in the rotated frame with first lattice basis
// vector aligned with the real-space x-axis (same frame as the one stored in m_sbase)
std::pair<double, double>
Interference2DLattice::calculateReciprocalVectorFraction(double qx, double qy, double xi) const
{
    double a = m_lattice->length1();
    double b = m_lattice->length2();
    double alpha = m_lattice->latticeAngle();
    // first rotate the input to the system of m_sbase:
    double qx_rot = qx * std::cos(xi) + qy * std::sin(xi);
    double qy_rot = -qx * std::sin(xi) + qy * std::cos(xi);

    // find the reciprocal lattice coordinates of (qx_rot, qy_rot):
    int qa_int = static_cast<int>(std::lround(a * qx_rot / (2 * pi)));
    int qb_int = static_cast<int>(
        std::lround(b * (qx_rot * std::cos(alpha) + qy_rot * std::sin(alpha)) / (2 * pi)));
    // take the fractional part only (in m_sbase coordinates)
    double qx_frac = qx_rot - qa_int * m_sbase.m_asx - qb_int * m_sbase.m_bsx;
    double qy_frac = qy_rot - qa_int * m_sbase.m_asy - qb_int * m_sbase.m_bsy;
    return {qx_frac, qy_frac};
}
