//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Img3D/Build/PositionBuilders.cpp
//! @brief     Implements subclasses of IPositionBuilder.
//!
//! @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 "Img3D/Build/PositionBuilders.h"
#include "Base/Math/Functions.h"
#include "Img3D/Build/ParacrystalLatticePositions.h"
#include "Sample/Aggregate/Interferences.h"
#include <random>

namespace {

double2d_t Generate2DLatticePoints(double l1, double l2, double alpha, double xi, unsigned n1,
                                   unsigned n2)
{
    double2d_t lattice_positions;
    std::vector<double> position;

    const unsigned nn1 = std::max(1u, n1);
    const unsigned nn2 = std::max(1u, n2);
    const int n1m = -static_cast<int>((nn1 - 1) / 2);
    const int n1M = static_cast<int>(nn1 / 2);
    const int n2m = -static_cast<int>((nn2 - 1) / 2);
    const int n2M = static_cast<int>(nn2 / 2);

    for (int i = n1m; i <= n1M; ++i) {
        for (int j = n2m; j <= n2M; ++j) {
            // For calculating lattice position vector v, we use: v = i*l1 + j*l2
            // where l1 and l2 are the lattice vectors,
            position.push_back(i * l1 * std::cos(xi)
                               + j * l2 * std::cos(alpha + xi)); // x coordinate
            position.push_back(i * l1 * std::sin(xi)
                               + j * l2 * std::sin(alpha + xi)); // y coordinate

            lattice_positions.push_back(position);
            position.clear();
        }
    }
    return lattice_positions;
}

} // namespace


//  ************************************************************************************************
//  class IPositionBuilder
//  ************************************************************************************************

IPositionBuilder::~IPositionBuilder() = default;

double2d_t IPositionBuilder::generatePositions(double layer_size, double density, int seed) const
{
    double2d_t positions = generatePositionsImpl(layer_size, density, Math::GenerateNextSeed(seed));
    const double pos_var = positionVariance();
    if (pos_var > 0.0) {
        // random generator and distribution
        std::random_device rd; // Will be used to obtain a seed for the random number engine
        std::mt19937 gen(seed < 0 ? rd()
                                  : seed); // Standard mersenne_twister_engine seeded with rd()
        std::normal_distribution<double> dis(0.0, std::sqrt(pos_var));
        for (auto& position : positions) {
            for (double& coordinate : position)
                coordinate += dis(gen);
        }
    }
    return positions;
}


//  ************************************************************************************************
//  class RandomPositionBuilder
//  ************************************************************************************************

RandomPositionBuilder::RandomPositionBuilder() = default;

RandomPositionBuilder::~RandomPositionBuilder() = default;

double2d_t RandomPositionBuilder::generatePositionsImpl(double layer_size, double density,
                                                        int seed) const
{
    double2d_t lattice_positions;
    std::vector<double> position;

    // to compute total number of particles we use the total particle density
    // and multiply by the area of the layer
    int n_particles = static_cast<int>(density * (2 * layer_size) * (2 * layer_size));

    // random generator and distribution
    std::random_device rd; // Will be used to obtain a seed for the random number engine
    std::mt19937 gen(seed < 0 ? rd() : seed); // Standard mersenne_twister_engine seeded with rd()
    std::uniform_real_distribution<double> dis(0.0, 1.0);

    for (int i = 1; i <= n_particles; ++i) {
        // generate random x and y coordinates
        position.push_back(dis(gen) * 2 * layer_size - layer_size); // x
        position.push_back(dis(gen) * 2 * layer_size - layer_size); // y

        lattice_positions.push_back(position);
        position.clear();
    }
    return lattice_positions;
}

double RandomPositionBuilder::positionVariance() const
{
    return 0.0; // no need for extra randomness here
}


//  ************************************************************************************************
//  class Lattice1DPositionBuilder
//  ************************************************************************************************

Lattice1DPositionBuilder::Lattice1DPositionBuilder(const Interference1DLattice* p_iff)
    : m_iff(p_iff->clone())
{
}

Lattice1DPositionBuilder::~Lattice1DPositionBuilder() = default;

double2d_t Lattice1DPositionBuilder::generatePositionsImpl(double layer_size, double, int) const
{
    const double length = m_iff->length();
    const double xi = m_iff->xi();

    // Take the maximum possible integer multiple of the lattice vector required
    // for populating particles correctly within the 3D model's boundaries
    unsigned n1 =
        length == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / length);

    return Generate2DLatticePoints(length, 0.0, 0.0, xi, n1, 1u);
}

double Lattice1DPositionBuilder::positionVariance() const
{
    return m_iff->positionVariance();
}


//  ************************************************************************************************
//  class Lattice2DPositionBuilder
//  ************************************************************************************************

Lattice2DPositionBuilder::Lattice2DPositionBuilder(const Interference2DLattice* p_iff)
    : m_iff(p_iff->clone())
{
}

Lattice2DPositionBuilder::~Lattice2DPositionBuilder() = default;

double2d_t Lattice2DPositionBuilder::generatePositionsImpl(double layer_size, double, int) const
{
    const Lattice2D& lattice = m_iff->lattice();
    const double l1 = lattice.length1();
    const double l2 = lattice.length2();
    const double alpha = lattice.latticeAngle();
    const double xi = lattice.rotationAngle();

    // Estimate the limits n1 and n2 of the maximum integer multiples of the lattice vectors
    // required for populating particles correctly within the 3D model's boundaries
    unsigned n1, n2;
    const double sina = std::abs(std::sin(alpha));
    if (sina <= 1e-4) {
        n1 = l1 == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / l1);
        n2 = l2 == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / l2);
    } else {
        n1 = l1 == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / l1 / sina);
        n2 = l2 == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / l2 / sina);
    }
    return Generate2DLatticePoints(l1, l2, alpha, xi, n1, n2);
}

double Lattice2DPositionBuilder::positionVariance() const
{
    return m_iff->positionVariance();
}


//  ************************************************************************************************
//  class Paracrystal2DPositionBuilder
//  ************************************************************************************************

Paracrystal2DPositionBuilder::Paracrystal2DPositionBuilder(const Interference2DParacrystal* p_iff)
    : m_iff(p_iff->clone())
{
}

Paracrystal2DPositionBuilder::~Paracrystal2DPositionBuilder() = default;

double2d_t Paracrystal2DPositionBuilder::generatePositionsImpl(double layer_size, double,
                                                               int seed) const
{
    return Paracrystal::latticePositions(m_iff.get(), layer_size, seed);
}

double Paracrystal2DPositionBuilder::positionVariance() const
{
    return m_iff->positionVariance();
}


//  ************************************************************************************************
//  class Finite2DLatticePositionBuilder
//  ************************************************************************************************

Finite2DLatticePositionBuilder::Finite2DLatticePositionBuilder(
    const InterferenceFinite2DLattice* p_iff)
    : m_iff(p_iff->clone())
{
}

Finite2DLatticePositionBuilder::~Finite2DLatticePositionBuilder() = default;

double2d_t Finite2DLatticePositionBuilder::generatePositionsImpl(double layer_size, double,
                                                                 int) const
{
    const auto& lattice = m_iff->lattice();
    const double l1 = lattice.length1();
    const double l2 = lattice.length2();
    const double alpha = lattice.latticeAngle();
    const double xi = lattice.rotationAngle();

    unsigned n1, n2;
    const double sina = std::abs(std::sin(alpha));
    if (sina <= 1e-4) {
        n1 = l1 == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / l1);
        n2 = l2 == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / l2);
    } else {
        n1 = l1 == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / l1 / sina);
        n2 = l2 == 0.0 ? 2 : static_cast<unsigned>(2.0 * layer_size * std::sqrt(2.0) / l2 / sina);
    }
    n1 = std::min(n1, m_iff->numberUnitCells1());
    n2 = std::min(n2, m_iff->numberUnitCells2());

    return Generate2DLatticePoints(l1, l2, alpha, xi, n1, n2);
}

double Finite2DLatticePositionBuilder::positionVariance() const
{
    return m_iff->positionVariance();
}


//  ************************************************************************************************
//  class RadialParacrystalPositionBuilder
//  ************************************************************************************************

RadialParacrystalPositionBuilder::RadialParacrystalPositionBuilder(
    const InterferenceRadialParacrystal* p_iff)
    : m_iff(p_iff->clone())
{
}

RadialParacrystalPositionBuilder::~RadialParacrystalPositionBuilder() = default;

double2d_t RadialParacrystalPositionBuilder::generatePositionsImpl(double layer_size, double,
                                                                   int seed) const
{
    double2d_t lattice_positions;

    const double distance = m_iff->peakDistance();

    // Estimate the limit n of the integer multiple i of the peakDistance required
    // for populating particles correctly within the 3D model's boundaries
    const int n = distance <= 0.0 ? 1 : static_cast<int>(layer_size * std::sqrt(2.0) / distance);

    lattice_positions.resize(2 * n + 1);
    for (auto& it : lattice_positions)
        it.resize(2);

    lattice_positions[0][0] = 0.0; // x coordinate of reference particle - at the origin
    lattice_positions[0][1] = 0.0; // y coordinate of reference particle - at the origin

    int upd_seed = Math::GenerateNextSeed(seed);
    for (int i = 1; i <= n; ++i) {
        // update seeds for each position
        upd_seed = Math::GenerateNextSeed(upd_seed);
        int upd_seed_2 = Math::GenerateNextSeed(upd_seed);

        // positions of particles located along +x (store at odd index)
        const unsigned i_left = static_cast<unsigned>(std::max(0, 2 * i - 3));

        double offset = m_iff->randomSample(upd_seed);
        lattice_positions[2 * i - 1][0] = lattice_positions[i_left][0] + distance + offset;
        lattice_positions[2 * i - 1][1] = 0.0;

        // positions of particles located along -x (store at even index)
        offset = m_iff->randomSample(upd_seed_2);
        lattice_positions[2 * i][0] = lattice_positions[2 * (i - 1)][0] - distance + offset;
        lattice_positions[2 * i][1] = 0.0;
    }
    return lattice_positions;
}

double RadialParacrystalPositionBuilder::positionVariance() const
{
    return m_iff->positionVariance();
}