//  ************************************************************************************************
//
//  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 "Base/Util/Vec.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/LayerStackItem.h"
#include "GUI/Model/Sample/MesocrystalItem.h"
#include "GUI/Model/Sample/ParticleItem.h"
#include "GUI/Model/Sample/ParticleLayoutItem.h"
#include "GUI/Model/Sample/RoughnessItems.h"
#include "GUI/Model/Sample/SampleItem.h"
#include "GUI/Model/ToCore/SampleToCore.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/Interface/Roughness.h"
#include "Sample/Interface/RoughnessMap.h"
#include "Sample/Multilayer/LayerStack.h"
#include "Sample/Multilayer/Sample.h"
#include "Sample/Particle/Compound.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 {

double2d_t generatePositions(IInterference* const iff, double layerSize, double density, int seed)
{
    if (!iff)
        return RandomPositionBuilder().generatePositions(layerSize, density, seed);

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

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

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

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

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

    // 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.isAmbient() || layerItem.isSubstrate())
        thickness = sceneGeometry.topOrBottomLayerThickness;
    else
        thickness = layerItem.thickness().dVal();

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

std::unique_ptr<double2d_t> scaledArray(const double2d_t& src, double factor)
{
    ASSERT(src.size());
    return std::make_unique<double2d_t>(
        FieldUtil::make<double>(src.size(), src[0].size(),
                                [&src, factor](size_t i, size_t j) { return src[i][j] * factor; }));
}

std::unique_ptr<const double2d_t> roughnessMap(const Sample& sample, size_t i_layer,
                                               const SceneGeometry& sceneGeometry, int seed)
{
    if (i_layer >= sample.numberOfLayers() || sample.roughnessRMS(i_layer) == 0)
        return nullptr;

    int n = sceneGeometry.numRoughnessPointsAlongAxis;
    double L = 2 * sceneGeometry.layerSize;
    auto rmap = RoughnessMap(n, n, L, L, sample, i_layer, seed); // seed < 0 ==> random every time
    return std::make_unique<const double2d_t>(rmap.generateMap());
}

std::unique_ptr<Img3D::Layer> createLayer(const LayerItem& layerItem,
                                          const SceneGeometry& sceneGeometry, const F3& origin,
                                          const double2d_t* topRoughMap,
                                          const double2d_t* bottomRoughMap, bool drawBottom)
{
    double s2 = sceneGeometry.layerSize;
    double thickness = ::visualLayerThickness(layerItem, sceneGeometry);
    ASSERT(thickness > 0);
    auto ztop = static_cast<double>(origin.z());
    double zbottom = static_cast<double>(origin.z()) - thickness;

    // visual mesh will later be scaled, so we divide roughness by thickness factor in advance
    auto top = topRoughMap ? ::scaledArray(*topRoughMap, 1. / thickness) : nullptr;
    auto bottom = bottomRoughMap ? ::scaledArray(*bottomRoughMap, 1. / thickness) : nullptr;

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

    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, const Item3D* item, const SampleItem* sampleItem,
                                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<const SampleItem*>(item))
        populateSample(model, *p, sceneGeometry, numParticles);

    else if (const auto* p = dynamic_cast<const LayerStackItem*>(item))
        populateStack(model, *p, *sampleItem, sceneGeometry, numParticles);

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

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

    else if (const auto* p = dynamic_cast<const 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& sampleItem,
                                      const SceneGeometry& sceneGeometry,
                                      unsigned& numParticles) const
{
    populateStack(model, sampleItem.outerStackItem(), sampleItem, sceneGeometry, numParticles);
}

void RealspaceBuilder::populateStack(Img3D::Model* model, const LayerStackItem& stackItem,
                                     const SampleItem& sampleItem,
                                     const SceneGeometry& sceneGeometry,
                                     unsigned int& numParticles) const
{
    std::vector<LayerItem*> stack_layers = stackItem.unwrappedLayerItems();
    if (stack_layers.empty())
        return;

    // find part of the sample, occupied by the first instance of the stack counting from top
    std::vector<LayerItem*> all_layers = sampleItem.outerStackItem().unwrappedLayerItems();
    const int i_begin = Vec::indexOfPtr(stack_layers.front(), all_layers);
    ASSERT(i_begin >= 0);
    const size_t i_end = i_begin + stack_layers.size();
    ASSERT(i_end <= all_layers.size());

    // generate roughness maps corresponding to the given stack
    sampleItem.adjustLayerSeeds(false);
    std::unique_ptr<Sample> sample = GUI::ToCore::itemToSample(sampleItem);
    OwningVector<const double2d_t> rough_maps;
    for (size_t i = i_begin; i <= i_end; i++) {
        if ((i < i_end && all_layers[i]->isAmbient()) || i == all_layers.size()
            || !sampleItem.showRoughness)
            rough_maps.push_back(nullptr); // top interface of fronting or bottom of substrate
        else
            rough_maps.push_back(
                ::roughnessMap(*sample, i, sceneGeometry, sampleItem.seeds[i]).release());
    }

    double total_height = 0;
    for (size_t i = 0; i < stack_layers.size(); i++) {
        const LayerItem* layer = stack_layers[i];
        const bool drawBottom = (i + 1 == stack_layers.size());
        populateLayer(model, *layer, sampleItem, sceneGeometry, numParticles,
                      F3(0, 0, static_cast<float>(-total_height)), rough_maps[i], rough_maps[i + 1],
                      drawBottom, false);
        if (!layer->isAmbient())
            total_height += ::visualLayerThickness(*layer, sceneGeometry);
    }
}

void RealspaceBuilder::populateLayer(Model* model, const LayerItem& layerItem,
                                     const SampleItem& sampleItem,
                                     const SceneGeometry& sceneGeometry, unsigned& numParticles,
                                     const F3& origin, const double2d_t* topRoughMap,
                                     const double2d_t* bottomRoughMap, bool drawBottom,
                                     bool independentLayer) const
{
    F3 shift(0, 0, 30); // absolute vertical position of sample

    std::unique_ptr<Img3D::Layer> layer;

    if (independentLayer && sampleItem.showRoughness) { // individual layer visualization
        sampleItem.adjustLayerSeeds(false);
        std::vector<LayerItem*> all_layers = sampleItem.outerStackItem().unwrappedLayerItems();
        const int i_layer = Vec::indexOfPtr(&layerItem, all_layers);
        ASSERT(i_layer >= 0);

        std::unique_ptr<Sample> sample = GUI::ToCore::itemToSample(sampleItem);
        auto new_topRoughMap =
            ::roughnessMap(*sample, i_layer, sceneGeometry, sampleItem.seeds[i_layer]);

        std::unique_ptr<const double2d_t> new_bottomRoughMap;
        const size_t i_layer_below = i_layer + 1;
        if (i_layer_below < all_layers.size())
            new_bottomRoughMap = ::roughnessMap(*sample, i_layer_below, sceneGeometry,
                                                sampleItem.seeds[i_layer_below]);

        layer = ::createLayer(layerItem, sceneGeometry, origin + shift, new_topRoughMap.get(),
                              new_bottomRoughMap.get(), drawBottom);
    } else {
        layer = ::createLayer(layerItem, sceneGeometry, origin + shift, topRoughMap, bottomRoughMap,
                              drawBottom);
    }

    if (layer && !layerItem.isAmbient())
        model->emplaceTransparentBody(layer.release());

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

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, layoutItem.seed);
    populateParticlesInLattice(latticePositions, particle3DContainer_vector, model, sceneGeometry,
                               numParticles, layoutItem.seed);
}

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().dVal();

    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 double2d_t& lattice_positions,
    const std::vector<Particle3DContainer>& particle3DContainer_vector, Model* model,
    const SceneGeometry& sceneGeometry, unsigned& numParticles, int seed) const
{
    const double layer_border_width = sceneGeometry.layerBorderWidth;
    const double layer_size = sceneGeometry.layerSize;
    const double layer_thickness = sceneGeometry.topOrBottomLayerThickness;

    std::random_device rd;
    std::mt19937 gen(seed < 0 ? rd() : seed);
    std::uniform_real_distribution<double> uniformDist(0, 1);

    for (std::vector<double> position : lattice_positions) {
        // for random selection of particles based on their abundances
        double rand_num = uniformDist(gen);
        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;
            }
        }
    }
}