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

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

Crystal2D::Crystal2D(const IParticle& particle, const Lattice2D& lattice)
    : Ordered2D(particle)
    , LongRangeOrdered2D(particle, 0)
    , ICrystal2D(particle, lattice)
{
    BasicLattice2D base_lattice(m_lattice->length1(), m_lattice->length2(),
                                m_lattice->latticeAngle(), 0.);
    m_sbase = base_lattice.reciprocalBases();
}

Crystal2D::~Crystal2D() = default;

void Crystal2D::setDecayFunction(const Profile2D& 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);
}

double Crystal2D::decayLengthX() const
{
    return m_decay ? m_decay->decayLengthX() : std::numeric_limits<double>::infinity();
}

double Crystal2D::decayLengthY() const
{
    return m_decay ? m_decay->decayLengthY() : std::numeric_limits<double>::infinity();
}

double Crystal2D::decayGamma() const
{
    return m_decay ? m_decay->gamma() : 0.0;
}

Crystal2D* Crystal2D::clone() const
{
    auto* result = new Crystal2D(*m_particle, *m_lattice);
    instructLRClone(result);
    instructCrystalClone(result);
    if (m_decay)
        result->setDecayFunction(*m_decay);
    return result;
}

double Crystal2D::iff_without_dw(const R3& q) const
{
    if (!m_decay)
        throw std::runtime_error(className() + " needs a decay function");
    return ICrystal2D::iff_without_dw(q);
}

double Crystal2D::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);
        }
    }
    // Normalization: S(q) = (1/A_uc) * sum_i H(q - q_i).
    return result / m_lattice->unitCellArea();
}

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

std::pair<double, double> Crystal2D::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};
}

std::pair<double, double> Crystal2D::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};
}
