//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Sample/ReSample.cpp
//! @brief     Implements class ReSample.
//!
//! @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/Sample/ReSample.h"
#include "Base/Util/Assert.h"
#include "Resample/Refine/RefineStructs.h"
#include "Resample/Refine/SubdivideLayers.h"
#include "Resample/Specular/ComputeFluxMagnetic.h"
#include "Resample/Specular/ComputeFluxScalar.h"
#include "Sample/Aggregate/Struct2D.h"
#include "Sample/Aggregate/Struct3D.h"
#include "Sample/Flux/Flux.h"
#include "Sample/Material/Admixtures.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/LayerStack.h"
#include "Sample/Option/SimulationOptions.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/IParticle.h"
#include "Sample/Particle/Mixture.h"
#include <stdexcept>

namespace {

//! Recursively collects admixture contributions from the particle tree for a given slice.
//! @param particle The particle node to process.
//! @param weight Accumulated weight factor (density_per_area * mixture abundance fraction).
//! @param slice_index The slice index to match against leaf particles' sliceNumber.
//! @param admixtures Collected (weighted_volume, material) pairs (mutable).
void collectAdmixtures(const IParticle* particle, double weight, size_t slice_index,
                       std::vector<OneAdmixture>& admixtures)
{
    ASSERT(particle);

    if (const auto* mix = dynamic_cast<const Mixture*>(particle)) {
        const auto& abundances = mix->abundances();
        const auto children = mix->particles();
        double total = 0;
        for (double a : abundances)
            total += a;
        if (total > 0)
            for (size_t j = 0; j < children.size(); ++j)
                ::collectAdmixtures(children[j], weight * abundances[j] / total, slice_index,
                                    admixtures);
    } else if (const auto* comp = dynamic_cast<const Compound*>(particle)) {
        for (const IParticle* child : comp->components())
            ::collectAdmixtures(child, weight, slice_index, admixtures);
    } else if (particle->sliceNumber() == slice_index) {
        // Leaf: Particle, Mesocrystal, CoreShell (with sliceNumber set by particlesInSlice)
        admixtures.emplace_back(particle->volumeToAdmix() * weight, *particle->materialToAdmix());
    }
}

//! For inner slices in `stack`, overwrite material to account for particles from `structs`.
void setAvgMaterials(std::vector<std::unique_ptr<Slice>>& slices,
                     const std::vector<std::unique_ptr<const Struct>>& structs)
{
    // loop over inner slices only; there is no admixture in semi-infinite slices
    for (size_t i = 1; i < slices.size() - 1; ++i) {

        // Get list of materials and coverages of particles in slice
        std::vector<OneAdmixture> admixtures;
        for (const auto& s : structs) {
            if (const auto* s2 = dynamic_cast<const Struct2D*>(s.get())) {
                // Struct2D: walk particle tree
                const double density_per_area = s2->areaNumberDensity() * s->coverage();
                ::collectAdmixtures(s->particle(), density_per_area, i, admixtures);
            } else if (const auto* s3 = dynamic_cast<const Struct3D*>(s.get())) {
                // Struct3D: prepared particle tree is already assigned to one host slice.
                if (s3->sliceNumber() != i)
                    continue;
                const double density_per_area =
                    s3->volumeNumberDensity() * s3->thickness() * s->coverage();
                ::collectAdmixtures(s->particle(), density_per_area, i, admixtures);
            } else
                ASSERT_NEVER;
        }

        // Compute and store effective slice material
        const double slice_volume = slices[i]->thickness(); // per unit area
        const std::unique_ptr<const Material> averaged_material =
            MixUtil::basePlusAdmixtures(*slices[i]->layerMaterial(), slice_volume, admixtures);
        slices[i]->setRefractingMaterial(*averaged_material);
    }
}

void setBField(std::vector<std::unique_ptr<Slice>>& slices, const R3& externalField)
{
    if (slices.empty())
        return;

    // Temporary forbid non-zero magnetization in fronting medium (see issue #654)
    if (slices.at(0)->refractingMaterial()->magnetization() != R3())
        throw std::runtime_error("Processing fronting magnetization is not implemented yet.");

    // Temporary forbid non-zero external field (see issue #654)
    if (externalField != R3())
        throw std::runtime_error("Processing external field is not implemented yet.");
    const double M_z0 = slices.at(0)->refractingMaterial()->magnetization().z();
    const double H_z = externalField.z() + M_z0;
    for (auto& slice : slices)
        slice->initBField(externalField, H_z);
}

} // namespace


//! Static factory function, the sole way to create a ReSample.
//!
//! Cannot easily be merged into the c'tor (to keep members const, they must
//! be initialized in the initializer block).

ReSample ReSample::processedSample(const Sample& sample, const SimulationOptions& options)
{
    // Reprocessing step R1:
    // Flatten the user-defined sample model by expanding periodic structures.
    const auto& unwrapped = sample.unwrapped();
    std::vector<const Layer*> flat;
    flat.reserve(unwrapped.size());
    for (const auto& layer : unwrapped)
        flat.push_back(layer.get());

    // Reprocessing step R2:
    // Divide layers into slices.
    // With material averaging, analyze particle positions and subdivide layers
    // according to particle spans. Without averaging, wrap each layer in a slice.
    std::vector<std::unique_ptr<Slice>> slices =
        options.useAvgMaterials() ? ::Refine::subdivideLayers(flat) : ::Refine::wrapLayers(flat);

    // Reprocessing step R3:
    // Process particles: decompose, geometrically slice formfactors across layer boundaries
    // (with averaging), and assemble into coherent piles assigned to appropriate slices.
    std::vector<std::unique_ptr<const Struct>> structs =
        ::Refine::refineStructs(flat, slices, options);

    // Reprocessing step R4:
    // Compute the average SLD of each slice.
    if (options.useAvgMaterials())
        ::setAvgMaterials(slices, structs);

    // Reprocessing step R5:
    // Recalculate full magnetic induction.
    ::setBField(slices, sample.externalField());

    // Convert to const slices (modifications are complete)
    std::vector<std::unique_ptr<const Slice>> const_slices;
    const_slices.reserve(slices.size());
    for (auto& s : slices)
        const_slices.push_back(std::move(s));

    return {std::move(structs), std::move(const_slices)};
}

ReSample::~ReSample() = default;

const Slice& ReSample::avgeSlice(size_t i) const
{
    return *m_slices.at(i);
}

bool ReSample::hasRoughness() const
{
    for (const auto& slice : m_slices)
        if (slice->rms() > 0)
            return true;
    return false;
}

Fluxes ReSample::fluxesIn(const R3& k, bool polarized) const
{
    if (polarized)
        return Compute::polarizedFluxes(m_slices, k, true);
    return Compute::scalarFluxes(m_slices, k);
}

Fluxes ReSample::fluxesOut(const R3& k, bool polarized) const
{
    if (polarized)
        return Compute::polarizedFluxes(m_slices, -k, false);
    return Compute::scalarFluxes(m_slices, -k);
}

//! Private c'tor, always called through ReSample::processSample.
ReSample::ReSample(std::vector<std::unique_ptr<const Struct>>&& structs,
                   std::vector<std::unique_ptr<const Slice>>&& refined_slices)
    : m_structs(std::move(structs))
    , m_slices(std::move(refined_slices))
{
    // Mark as validated since we're setting slices directly
    m_validated = true;
}
