//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Particle/ReMesocrystal.cpp
//! @brief     Implements class ReMesocrystal.
//!
//! @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 "Resample/Particle/ReMesocrystal.h"
#include "Base/Math/Numeric.h"
#include "Base/Spin/SpinMatrix.h"
#include "Base/Type/Span.h"
#include "Base/Util/Assert.h"
#include "Base/Vector/WavevectorInfo.h"
#include "Resample/Particle/ReParticle.h"
#include <numbers>

using std::numbers::pi;

ReMesocrystal::ReMesocrystal(const std::optional<size_t>& i_layer, const Lattice3D& lattice,
                             const IReParticle& basis, const ReParticle& outer_shape,
                             const MesoOptions& meso_options, double position_variance)
    : IReParticle(i_layer)
    , m_lattice(lattice)
    , m_basis(basis.clone())
    , m_outer_shape(outer_shape.clone())
    , m_position_variance(position_variance)
    , m_meso_options(meso_options)
{
    if (m_meso_options.use_reciprocal_sum) {
        m_compute_FF = [this](const WavevectorInfo& wv) { return reciprocalSpaceSum(wv); };
        m_compute_FF_pol = [this](const WavevectorInfo& wv) { return reciprocalSpaceSumPol(wv); };
        calculateLargestReciprocalDistance();
    } else {
        m_compute_FF = [this](const WavevectorInfo& wv) { return realSpaceSum(wv); };
        m_compute_FF_pol = [this](const WavevectorInfo& wv) { return realSpaceSumPol(wv); };
    }
}

ReMesocrystal::~ReMesocrystal() = default;

ReMesocrystal* ReMesocrystal::clone() const
{
    return new ReMesocrystal(i_layer(), m_lattice, *m_basis, *m_outer_shape, m_meso_options,
                             m_position_variance);
}

double ReMesocrystal::radialExtension() const
{
    return m_outer_shape->radialExtension();
}

Span ReMesocrystal::zSpan() const
{
    return m_outer_shape->zSpan();
}

complex_t ReMesocrystal::structureFactor(const WavevectorInfo& wavevectors) const
{
    // The underlying implementation computes formfactor exactly, summing over all particles
    // positions inside the mesocrystal shape. Its computational complexity is ~O(Nx * Ny), where
    // 'Nx' and 'Ny' are numbers of particles along the 'x' and 'y' axes. For mesocrystals larger
    // than ~30x30xNz particles, the approximate algorithm may be preferred.

    C3 q = wavevectors.getQ();
    const auto& k_pairs = m_shape_indexes.k_pairs;
    const I3& min = m_shape_indexes.min;
    const I3& max = m_shape_indexes.max;
    I3 range = max - min;

    // The code below is the algorithm to compute the following sum (pseudocode):
    // ***
    // for (R3 position : basis_positions_inside_mesocrystal)
    //     field_factor += exp_I(position.dot(q));
    // ***

    // precompute and cache values

    complex_t exp_a = exp_I(m_lattice.basisVectorA().dot(q));
    complex_t exp_b = exp_I(m_lattice.basisVectorB().dot(q));
    complex_t exp_c = exp_I(m_lattice.basisVectorC().dot(q));

    std::vector<complex_t> exponents_b(range.y() + 1);

    for (int j = 0; j <= range.y(); j++)
        exponents_b[j] = pow(exp_b, j + min.y());

    // main summation

    complex_t field_factor(0.0, 0.0);

    if (Numeric::almostEqual(exp_c, 1., 1)) {
        // if exp_c == 1, the sum is even simpler

        for (int i = 0; i <= range.x(); i++) {
            const complex_t e_a = pow(exp_a, i + min.x());

            for (int j = 0; j <= range.y(); j++) {
                const complex_t e_ab = e_a * exponents_b[j];

                for (const auto& p : k_pairs[i][j])
                    field_factor += e_ab * (p.second - p.first + 1.);
            }
        }
    } else {
        std::vector<complex_t> exponents_c(range.z() + 1);
        std::vector<complex_t> exp_c_sum_factor(range.z() + 1);
        for (int k = 0; k <= range.z(); k++) {
            exponents_c[k] = pow(exp_c, k + min.z());
            exp_c_sum_factor[k] = (pow(exp_c, k + 1) - 1.) / (exp_c - 1.);
        }

        for (int i = 0; i <= range.x(); i++) {
            const complex_t e_a = pow(exp_a, i + min.x());

            for (int j = 0; j <= range.y(); j++) {
                const complex_t e_ab = e_a * exponents_b[j];

                for (const auto& p : k_pairs[i][j]) {
                    int k_begin = p.first;
                    int k_end = p.second;
                    int N = k_end - k_begin;
                    field_factor += e_ab * exponents_c[k_begin - min.z()] * exp_c_sum_factor[N];
                }
            }
        }
    }
    return field_factor * debyeWallerFactor(q.real());
}

complex_t ReMesocrystal::realSpaceSum(const WavevectorInfo& wavevectors) const
{
    return structureFactor(wavevectors) * m_basis->theFF(wavevectors);
}

SpinMatrix ReMesocrystal::realSpaceSumPol(const WavevectorInfo& wavevectors) const
{
    return structureFactor(wavevectors) * m_basis->thePolFF(wavevectors);
}

complex_t ReMesocrystal::reciprocalSpaceSum(const WavevectorInfo& wavevectors) const
{
    // The underlying implementation computes formfactor approximately.
    // The higher the 'm_meso_options.radius_factor', the higher the accuracy.
    // The calculation time is ~O(m_meso_options.radius_factor^3) but it does not depend on the
    // number of particles inside the mesocrystal shape. Therefore, it is reasonable to use this
    // approach for large mesocrystals with thousands of particles inside. The parameter
    // 'm_meso_options.radius_factor' should then be empirically adjusted to some reasonable value.
    // For mesocrystals smaller than ~30x30xNz particles, the exact algorithm may be preferred.

    // retrieve reciprocal lattice vectors within reasonable radius
    C3 q = wavevectors.getQ();
    double radius = m_meso_options.radius_factor * m_max_rec_length;
    std::vector<R3> rec_vectors = m_lattice.reciprocalLatticeVectorsWithinRadius(q.real(), radius);

    // perform convolution on these lattice vectors
    complex_t result(0.0, 0.0);
    for (const auto& rec : rec_vectors) {
        auto dw_factor = debyeWallerFactor(rec);
        WavevectorInfo basis_wavevectors(R3(), -rec, wavevectors.vacuumLambda());
        complex_t basis_factor = m_basis->theFF(basis_wavevectors);
        WavevectorInfo meso_wavevectors(C3(), rec.complex() - q, wavevectors.vacuumLambda());
        complex_t meso_factor = m_outer_shape->theFF(meso_wavevectors);
        result += dw_factor * basis_factor * meso_factor;
    }
    // the transformed delta train gets a factor of (2pi)^3/V, but the (2pi)^3
    // is canceled by the convolution of Fourier transforms :
    return result / m_lattice.unitCellVolume();
}

SpinMatrix ReMesocrystal::reciprocalSpaceSumPol(const WavevectorInfo& wavevectors) const
{
    // See the comment to the non-polarized implementation.

    // retrieve reciprocal lattice vectors within reasonable radius
    C3 q = wavevectors.getQ();
    double radius = m_meso_options.radius_factor * m_max_rec_length;
    std::vector<R3> rec_vectors = m_lattice.reciprocalLatticeVectorsWithinRadius(q.real(), radius);

    // perform convolution on these lattice vectors
    SpinMatrix result;
    for (const auto& rec : rec_vectors) {
        auto dw_factor = debyeWallerFactor(rec);
        WavevectorInfo basis_wavevectors(R3(), -rec, wavevectors.vacuumLambda());
        SpinMatrix basis_factor = m_basis->thePolFF(basis_wavevectors);
        WavevectorInfo meso_wavevectors(C3(), rec.complex() - q, wavevectors.vacuumLambda());
        complex_t meso_factor = m_outer_shape->theFF(meso_wavevectors);
        result += dw_factor * basis_factor * meso_factor;
    }
    // the transformed delta train gets a factor of (2pi)^3/V, but the (2pi)^3
    // is canceled by the convolution of Fourier transforms :
    return result / m_lattice.unitCellVolume();
}

complex_t ReMesocrystal::theFF(const WavevectorInfo& wavevectors) const
{
    return m_compute_FF(wavevectors);
}

SpinMatrix ReMesocrystal::thePolFF(const WavevectorInfo& wavevectors) const
{
    return m_compute_FF_pol(wavevectors);
}

void ReMesocrystal::setBasisIndexes(const ShapeIndexes& shape_indexes)
{
    m_shape_indexes = shape_indexes;
}

bool ReMesocrystal::consideredEqualTo(const IReParticle& ire) const
{
    if (const auto* re = dynamic_cast<const ReMesocrystal*>(&ire)) {
        bool same_lattice = m_lattice == re->lattice();
        bool same_variance = m_position_variance == re->positionVariance();

        ASSERT(m_basis);
        ASSERT(re->basis());
        bool same_basis = m_basis->consideredEqualTo(*re->basis());

        ASSERT(m_outer_shape);
        ASSERT(re->outerShape());
        bool same_shape = m_outer_shape->consideredEqualTo(*re->outerShape());

        R3 own_shift = posDiff(m_basis->position(), m_outer_shape->position());
        R3 other_shift = posDiff(re->basis()->position(), re->outerShape()->position());
        bool same_shift = own_shift == other_shift;

        return IReParticle::consideredEqualTo(ire) && same_lattice && same_variance && same_basis
               && same_shape && same_shift;
    }
    return false;
}

const R3* ReMesocrystal::position() const
{
    return m_outer_shape->position();
}

void ReMesocrystal::calculateLargestReciprocalDistance()
{
    R3 a1 = m_lattice.basisVectorA();
    R3 a2 = m_lattice.basisVectorB();
    R3 a3 = m_lattice.basisVectorC();

    m_max_rec_length = std::max(pi / a1.mag(), pi / a2.mag());
    m_max_rec_length = std::max(m_max_rec_length, pi / a3.mag());
}

complex_t ReMesocrystal::debyeWallerFactor(const R3& q_i) const
{
    auto q2 = q_i.mag2();
    return std::exp(-q2 * m_position_variance / 2.0);
}