//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sample/Particle/Mesocrystal.cpp
//! @brief     Implements class Mesocrystal.
//!
//! @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/Particle/Mesocrystal.h"
#include "Base/Type/Span.h"
#include "Base/Util/Assert.h"
#include "Sample/Lattice/ISelectionRule.h"
#include "Sample/Lattice/Lattice3D.h"
#include "Sample/Particle/Crystal.h"
#include "Sample/Particle/IFormFactor.h"
#include "Sample/Particle/Particle.h"
#include "Sample/Scattering/Rotations.h"
#include <algorithm>

Mesocrystal::Mesocrystal(Crystal* crystal, IFormFactor* formfactor)
    : m_crystal(crystal)
    , m_meso_formfactor(formfactor)
{
}

Mesocrystal::Mesocrystal(const Crystal& crystal, const IFormFactor& formfactor)
    : Mesocrystal(crystal.clone(), formfactor.clone())
{
}

Mesocrystal::~Mesocrystal() = default;

Mesocrystal* Mesocrystal::clone() const
{
    auto* result = new Mesocrystal(m_crystal->clone(), m_meso_formfactor->clone());
    result->setAbundance(m_abundance);
    if (rotation())
        result->rotate(*rotation());
    result->translate(particlePosition());
    return result;
}

std::vector<const INode*> Mesocrystal::nodeChildren() const
{
    return std::vector<const INode*>()
           << IParticle::nodeChildren() << m_crystal << m_meso_formfactor;
}

const Crystal& Mesocrystal::particleStructure() const
{
    ASSERT(m_crystal);
    return *m_crystal;
}

Span Mesocrystal::zSpan() const
{
    return m_meso_formfactor->spanZ(rotation()) + particlePosition().z();
}

std::vector<R3> Mesocrystal::calcBasisPositions() const
{
    std::vector<I3> basisIndexes = calcBasisIndexes().basisIndexes;
    std::vector<R3> result(basisIndexes.size());
    size_t counter = 0;
    for (const I3& index : basisIndexes)
        result[counter++] = index.x() * m_crystal->lattice()->basisVectorA()
                            + index.y() * m_crystal->lattice()->basisVectorB()
                            + index.z() * m_crystal->lattice()->basisVectorC();
    return result;
}

ShapeIndexes Mesocrystal::calcBasisIndexes() const
{
    const Lattice3D& lattice = *m_crystal->lattice();
    R3 shift = m_crystal->basis()->particlePosition();

    auto contains = [&](const I3& index) -> bool {
        R3 position = index.x() * lattice.basisVectorA() + index.y() * lattice.basisVectorB()
                      + index.z() * lattice.basisVectorC();

        return m_meso_formfactor->contains(position + shift);
    };

    // positive and negative boundaries are changeable
    struct Limits {
        int i_min = 0, i_max = 1;
        int j_min = 0, j_max = 1;
        int k_min = 0, k_max = 1;
    } limits;

    // On the first pass we find min and max indices along each axis WITHOUT selection rule
    auto check_and_add_particle = [&](int i, int j, int k) {
        if (contains(I3(i, j, k))) {
            // if boundary particle belongs to the outer shape, extend the boundaries
            if (i == limits.i_min)
                limits.i_min--;
            if (i == limits.i_max)
                limits.i_max++;

            if (j == limits.j_min)
                limits.j_min--;
            if (j == limits.j_max)
                limits.j_max++;

            if (k == limits.k_min)
                limits.k_min--;
            if (k == limits.k_max)
                limits.k_max++;
        }
    };

    auto iterate_i = [&](int j, int k) {
        for (int i = 1; i <= limits.i_max; i++)
            check_and_add_particle(i, j, k);
        for (int i = 0; i >= limits.i_min; i--)
            check_and_add_particle(i, j, k);
    };
    auto iterate_j = [&](int k) {
        for (int j = 1; j <= limits.j_max; j++)
            iterate_i(j, k);
        for (int j = 0; j >= limits.j_min; j--)
            iterate_i(j, k);
    };
    auto iterate_k = [&]() {
        for (int k = 1; k <= limits.k_max; k++)
            iterate_j(k);
        for (int k = 0; k >= limits.k_min; k--)
            iterate_j(k);
    };
    iterate_k();

    // min and max indexes belonging to the outer shape
    I3 min_index(limits.i_min + 1, limits.j_min + 1, limits.k_min + 1);
    I3 max_index(limits.i_max - 1, limits.j_max - 1, limits.k_max - 1);
    I3 range = max_index - min_index;

    // On the second pass we find indices WITH selection rule
    auto contains_shifted = [&](int i, int j, int k) -> bool {
        const I3 shifted_index = I3(i, j, k) + min_index;
        if (lattice.selectionRule() && !lattice.selectionRule()->coordinateSelected(shifted_index))
            return false;
        return contains(shifted_index);
    };

    // reuse to find updated index range WITH selection rule
    limits = {range.x(), 0, range.y(), 0, range.z(), 0};

    // For each [i,j] point we collect continuous ranges of k indexes and store them as of pairs
    // (k_begin, k_end). These ranges will be used to efficiently sum the geometric progression.
    std::vector<std::vector<std::vector<std::pair<int, int>>>> k_pairs(
        range.x() + 1, std::vector<std::vector<std::pair<int, int>>>(range.y() + 1));
    std::vector<I3> basisIndexes;
    for (int i = 0; i <= range.x(); i++) {
        for (int j = 0; j <= range.y(); j++) {
            for (int k = 0; k <= range.z(); k++) {
                if (!contains_shifted(i, j, k))
                    continue;

                basisIndexes.push_back(I3(i, j, k) + min_index);

                limits.i_min = std::min(limits.i_min, i);
                limits.j_min = std::min(limits.j_min, j);
                limits.k_min = std::min(limits.k_min, k);

                limits.i_max = std::max(limits.i_max, i);
                limits.j_max = std::max(limits.j_max, j);
                limits.k_max = std::max(limits.k_max, k);

                // open pair
                if (k == 0 || !contains_shifted(i, j, k - 1))
                    k_pairs[i][j].push_back(std::make_pair(k, k));

                // close pair
                if (k == range.z() || !contains_shifted(i, j, k + 1))
                    k_pairs[i][j].back().second = k;
            }
        }
    }

    // min and max indexes belonging to the outer shape with selection rule
    const I3 upd_min_index = I3(limits.i_min, limits.j_min, limits.k_min) + min_index;
    const I3 upd_max_index = I3(limits.i_max, limits.j_max, limits.k_max) + min_index;

    return {upd_min_index, upd_max_index, basisIndexes, k_pairs};
}
