//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      GUI/View/Realspace/RealspaceBuilder.cpp
//! @brief     Implements class RealspaceBuilder.
//!
//! @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 "GUI/View/Realspace/RealspaceBuilder.h"
#include "Base/Util/Assert.h"
#include "GUI/Model/Sample/CompoundItem.h"
#include "GUI/Model/Sample/CoreAndShellItem.h"
#include "GUI/Model/Sample/InterferenceItems.h"
#include "GUI/Model/Sample/LayerItem.h"
#include "GUI/Model/Sample/MesocrystalItem.h"
#include "GUI/Model/Sample/ParticleItem.h"
#include "GUI/Model/Sample/ParticleLayoutItem.h"
#include "GUI/Model/Sample/SampleItem.h"
#include "Img3D/Build/BuilderUtil.h"
#include "Img3D/Build/Particle3DContainer.h"
#include "Img3D/Build/PositionBuilders.h"
#include "Img3D/Model/Layer.h"
#include "Img3D/Model/Model.h"
#include "Img3D/Model/ParticleFromFF.h"
#include "Img3D/Type/SceneGeometry.h"
#include "Sample/Aggregate/Interferences.h"
#include "Sample/Particle/CoreAndShell.h"
#include "Sample/Particle/Mesocrystal.h"
#include "Sample/Particle/Particle.h"

using Img3D::F3;
using Img3D::Model;
using Img3D::Particle3DContainer;

namespace {

std::vector<std::vector<double>> generatePositions(IInterference* const iff, double layerSize,
                                                   double density)
{
    if (!iff)
        return RandomPositionBuilder().generatePositions(layerSize, density);

    if (auto* p = dynamic_cast<Interference1DLattice*>(iff))
        return Lattice1DPositionBuilder(p).generatePositions(layerSize, density);

    if (auto* p = dynamic_cast<Interference2DLattice*>(iff))
        return Lattice2DPositionBuilder(p).generatePositions(layerSize, density);

    if (auto* p = dynamic_cast<Interference2DParacrystal*>(iff))
        return Paracrystal2DPositionBuilder(p).generatePositions(layerSize, density);

    if (auto* p = dynamic_cast<InterferenceFinite2DLattice*>(iff))
        return Finite2DLatticePositionBuilder(p).generatePositions(layerSize, density);

    if (auto* p = dynamic_cast<InterferenceRadialParacrystal*>(iff))
        return RadialParacrystalPositionBuilder(p).generatePositions(layerSize, density);

    // TODO https://jugit.fz-juelich.de/mlz/bornagain/-/issues/538
    if (dynamic_cast<InterferenceHardDisk*>(iff))
        throw std::runtime_error("Percus-Yevick model 3D vizualization is not implemented yet");

    ASSERT_NEVER;
}

double visualLayerThickness(const LayerItem& layerItem, const SceneGeometry& sceneGeometry)
{
    double thickness(0.0);
    if (layerItem.isTopLayer() || layerItem.isBottomLayer())
        thickness = sceneGeometry.topOrBottomLayerThickness;
    else
        thickness = layerItem.thickness();

    return thickness == 0.0 ? sceneGeometry.layerMinimumThickness : thickness;
}

std::unique_ptr<Img3D::Layer> createLayer(const LayerItem& layerItem,
                                          const SceneGeometry& sceneGeometry, const F3& origin)
{
    double thickness = visualLayerThickness(layerItem, sceneGeometry);
    double s2 = sceneGeometry.layerSize;

    auto ztop = static_cast<double>(origin.z());
    double zbottom = static_cast<double>(origin.z()) - thickness;

    std::unique_ptr<Img3D::Layer> result = std::make_unique<Img3D::Layer>(
        Img3D::F3Range(Img3D::F3fromR3({-s2, -s2, ztop}), Img3D::F3fromR3({s2, s2, zbottom})));

    QColor color = layerItem.materialColor();
    color.setAlphaF(.3);

    result->setColor(color);

    return result;
}

} // namespace


RealspaceBuilder::RealspaceBuilder(std::function<QColor(const QString&)> fnColorFromMaterialName)
{
    m_builder_utils = std::make_unique<Img3D::BuilderUtils>(fnColorFromMaterialName);
}

RealspaceBuilder::~RealspaceBuilder() = default;

void RealspaceBuilder::populate(Model* model, Item3D* item, const SceneGeometry& sceneGeometry,
                                unsigned& numParticles) const
{
    ASSERT(item);

    // depending on item type, visualize the full sample model, or only parts of it
    if (const auto* p = dynamic_cast<SampleItem*>(item))
        populateSample(model, *p, sceneGeometry, numParticles);

    else if (const auto* p = dynamic_cast<LayerItem*>(item))
        populateLayer(model, *p, sceneGeometry, numParticles);

    else if (const auto* p = dynamic_cast<ParticleLayoutItem*>(item))
        populateLayout(model, *p, sceneGeometry, numParticles);

    else if (const auto* p = dynamic_cast<ItemWithParticles*>(item))
        // visualize one generalized particle (simple particle or core/shell or compound or meso..)
        translateContainer(model, particlesFromItem(*p), numParticles);

    else
        ASSERT_NEVER;
}

void RealspaceBuilder::populateSample(Model* model, const SampleItem& mlayerItem,
                                      const SceneGeometry& sceneGeometry, unsigned& numParticles,
                                      const F3&) const
{
    double total_height(0.0);
    int index(0);
    for (LayerItem* layer : mlayerItem.layerItems()) {
        bool isTopLayer = index == 0;
        populateLayer(model, *layer, sceneGeometry, numParticles,
                      F3(0, 0, static_cast<float>(-total_height)), isTopLayer);

        if (index != 0)
            total_height += visualLayerThickness(*layer, sceneGeometry);

        ++index;
    }
}

void RealspaceBuilder::populateLayer(Model* model, const LayerItem& layerItem,
                                     const SceneGeometry& sceneGeometry, unsigned& numParticles,
                                     const F3& origin, const bool isTopLayer) const
{
    auto layer = createLayer(layerItem, sceneGeometry, origin);
    if (layer && !isTopLayer)
        model->emplaceTransparentBody(layer.release());

    for (ParticleLayoutItem* layout : layerItem.layoutItems())
        populateLayout(model, *layout, sceneGeometry, numParticles, origin);
}

void RealspaceBuilder::populateLayout(Model* model, const ParticleLayoutItem& layoutItem,
                                      const SceneGeometry& sceneGeometry, unsigned& numParticles,
                                      const F3& origin) const
{
    if (layoutItem.itemsWithParticles().empty())
        return;

    const double layer_size = sceneGeometry.layerSize;
    const double total_density = layoutItem.totalDensityValue();

    auto particle3DContainer_vector = particle3DContainerVector(layoutItem, origin);
    auto* interferenceItem = layoutItem.interferenceSelection().certainItem();
    std::unique_ptr<IInterference> iff;
    if (interferenceItem)
        iff = interferenceItem->createInterference();

    const auto latticePositions = generatePositions(iff.get(), layer_size, total_density);
    populateParticlesInLattice(latticePositions, particle3DContainer_vector, model, sceneGeometry,
                               numParticles);
}

std::vector<Img3D::Particle3DContainer>
RealspaceBuilder::particle3DContainerVector(const ParticleLayoutItem& layoutItem,
                                            const F3& origin) const
{
    double total_abundance = 0;
    for (const auto* particle : layoutItem.itemsWithParticles())
        total_abundance += particle->abundance();

    double cumulative_abundance = 0;
    std::vector<Particle3DContainer> result;

    for (auto* particleItem : layoutItem.itemsWithParticles()) {
        Particle3DContainer out = particlesFromItem(*particleItem, total_abundance, origin);

        cumulative_abundance += out.cumulativeAbundance();
        out.setCumulativeAbundance(cumulative_abundance);

        result.emplace_back(std::move(out));
    }

    return result;
}

Img3D::Particle3DContainer
RealspaceBuilder::particlesFromItem(const ItemWithParticles& particleItem, double total_abundance,
                                    const Img3D::F3& origin) const
{
    if (const auto* item = dynamic_cast<const ParticleItem*>(&particleItem)) {
        auto particle = item->createParticle();
        return m_builder_utils->singleParticle3DContainer(*particle, total_abundance, origin);

    } else if (const auto* item = dynamic_cast<const CoreAndShellItem*>(&particleItem)) {
        if (!item->coreItem() || !item->shellItem())
            return {};
        auto particle = item->createCoreAndShell();
        return m_builder_utils->particleCoreShell3DContainer(*particle, total_abundance, origin);

    } else if (const auto* item = dynamic_cast<const CompoundItem*>(&particleItem)) {
        if (item->itemsWithParticles().empty())
            return {};
        auto particle = item->createCompound();
        return m_builder_utils->particleComposition3DContainer(*particle, total_abundance, origin);

    } else if (const auto* item = dynamic_cast<const MesocrystalItem*>(&particleItem)) {
        if (!item->basisItem())
            return {};
        return m_builder_utils->mesocrystal3DContainer(item->createMesocrystal().get(),
                                                       total_abundance, origin);

    } else
        ASSERT_NEVER;
}

void RealspaceBuilder::translateContainer(Model* model,
                                          const Particle3DContainer& particle3DContainer,
                                          unsigned int& numParticles,
                                          const F3& lattice_position) const
{
    numParticles += particle3DContainer.containerSize();

    for (size_t i = 0; i < particle3DContainer.containerSize(); ++i) {
        auto particle3D = particle3DContainer.createParticle(i);
        ASSERT(particle3D);
        particle3D->addTranslation(lattice_position);
        if (particle3D->isTransparent())
            model->emplaceTransparentBody(particle3D.release());
        else
            model->emplaceSolidBody(particle3D.release());
    }
}

void RealspaceBuilder::populateParticlesInLattice(
    const std::vector<std::vector<double>>& lattice_positions,
    const std::vector<Particle3DContainer>& particle3DContainer_vector, Model* model,
    const SceneGeometry& sceneGeometry, unsigned& numParticles) const
{
    const double layer_border_width = sceneGeometry.layerBorderWidth;
    const double layer_size = sceneGeometry.layerSize;
    const double layer_thickness = sceneGeometry.topOrBottomLayerThickness;

    for (std::vector<double> position : lattice_positions) {
        // for random selection of particles based on their abundances
        double rand_num = (rand() / static_cast<double>(RAND_MAX)); // (between 0 and 1)
        for (const auto& particle3DContainer : particle3DContainer_vector) {
            if (rand_num <= particle3DContainer.cumulativeAbundance()) {
                // lattice position + location (TO BE ADDED)
                double pos_x = position[0];
                double pos_y = position[1];
                double pos_z = 0;

                if (std::abs(pos_x) <= layer_size - 2 * layer_border_width
                    && std::abs(pos_y) <= layer_size - 2 * layer_border_width
                    && std::abs(pos_z) <= layer_thickness) {
                    translateContainer(model, particle3DContainer, numParticles,
                                       Img3D::F3fromR3({position[0], position[1], 0}));
                }
                break;
            }
        }
    }
}
