//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Processed/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/Processed/ReSample.h"
#include "Base/Type/Span.h"
#include "Base/Util/Assert.h"
#include "Resample/Coherence/CoheringSubparticles.h"
#include "Resample/Flux/IFlux.h"
#include "Resample/Option/SimulationOptions.h"
#include "Resample/Particle/IReParticle.h"
#include "Resample/Processed/ReLayout.h"
#include "Resample/Processed/Slicer.h"
#include "Resample/Specular/ComputeFluxMagnetic.h"
#include "Resample/Specular/ComputeFluxScalar.h"
#include "Sample/Aggregate/IInterference.h"
#include "Sample/Aggregate/ParticleLayout.h"
#include "Sample/Interface/Roughness.h"
#include "Sample/Material/MaterialBySLDImpl.h"
#include "Sample/Material/MaterialUtil.h"
#include "Sample/Material/RefractiveMaterialImpl.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Sample.h"
#include "Sample/Particle/IParticle.h"
#include <iomanip>
#include <iostream>
#include <stdexcept>

namespace {

//! Returns for each layer the z span that contains some particles.
//! However, if use_slicing is false, then each span is set to -inf..+inf.
std::vector<ZLimits> particleSpans(const Sample& sample)
{
    size_t N = sample.numberOfLayers();
    ASSERT(N > 0);

    std::vector<ZLimits> result(N); // default span: -inf .. inf

    for (size_t i = 0; i < N; ++i) {
        const Layer* layer = sample.layer(i);
        const double offset = (i == 0) ? 0 : sample.hig(i);
        for (const auto* layout : layer->layouts())
            for (const IParticle* particle : layout->particles()) {
                auto [bottom, top] = particle->zSpan().pair();
                bottom += offset;
                top += offset;
                if (bottom == top) // zero-height particle
                    continue;
                ASSERT(bottom < top);
                ASSERT(isfinite(bottom) && isfinite(top));

                for (size_t ii = 0; ii < N; ++ii) {
                    if ((ii != 0 && bottom >= sample.hig(ii))
                        || (ii != N - 1 && top <= sample.low(ii)))
                        continue;

                    double layer_ref = ii ? sample.hig(ii) : sample.low(ii);
                    double upper = ii ? std::min(top, layer_ref) : top;
                    double lower = (ii == N - 1) ? bottom : std::max(bottom, sample.low(ii));
                    ZLimits bounded_limits(lower - layer_ref, upper - layer_ref);

                    if (result[ii].isFinite())
                        result[ii] = ZLimits::unite(result[ii], bounded_limits);
                    else
                        result[ii] = bounded_limits;
                }
            }
    }
    return result;
}

//! Returns a vector of slices that represent the layer structure of the sample.
//!
//! Each slice is either a layer, or a fraction of a layer.
//! Layers are fractioned into several slices if
//! - option useAvgMaterials is true,
//! - layer->numberOfSlices > 1, and
//! - there are particles in the layer.
//!
//! Used in ReSample constructor.

SliceStack slicify(const Sample& sample, bool useAvgMaterials)
{
    SliceStack result;

    size_t nLayers = sample.numberOfLayers();
    if (nLayers == 0)
        throw std::runtime_error("Sample has no layer");

    const bool use_slicing = useAvgMaterials && nLayers > 1;

    if (!use_slicing) {
        for (size_t i = 0; i < nLayers; ++i) {
            const Layer* const layer = sample.layer(i);
            double tl = layer->thickness();
            const Material* const material = layer->material();

            if (i == 0) {
                ASSERT(tl == 0);
                result.addTopSlice(tl, *material);
            } else {
                if (i == nLayers - 1)
                    ASSERT(tl == 0);
                const Roughness* roughness = sample.layer(i)->roughness();
                const double rms = sample.roughnessRMS(i);
                result.addSlice(tl, *material, roughness, rms);
            }
        }
        return result;
    }

    const std::vector<ZLimits> particle_spans = particleSpans(sample);

    for (size_t i = 0; i < nLayers; ++i) {
        const Layer* const layer = sample.layer(i);
        double tl = layer->thickness();
        const Material* const material = layer->material();
        const ZLimits& particle_span = particle_spans[i];
        const Roughness* roughness = sample.layer(i)->roughness();
        const double rms = sample.roughnessRMS(i);

        // if no slicing is needed, create single slice for the layer
        if (!particle_span.isFinite()) { // also if layer contains no particles
            if (i == nLayers - 1)
                tl = 0.0;
            if (i == 0)
                result.addTopSlice(tl, *material);
            else
                result.addSlice(tl, *material, roughness, rms);
            continue;
        }

        const double top = particle_span.hig();
        const double bottom = particle_span.low();
        const size_t nSlices = layer->numberOfSlices();
        ASSERT(nSlices > 0);

        // top layer
        if (i == 0) {
            ASSERT(top > 0);
            result.addTopSlice(top, *material); // semiinfinite top slice
            result.addNSlices(nSlices, top - bottom, *material);
            if (bottom > 0)
                result.addSlice(bottom, *material);
        }
        // middle or bottom layer
        else {
            ASSERT(top <= 0);
            if (top < 0) {
                result.addSlice(-top, *material, roughness, rms);
                result.addNSlices(nSlices, top - bottom, *material);
            } else { // top == 0
                result.addNSlices(nSlices, top - bottom, *material, roughness, rms);
            }
            // middle layer
            if (i < nLayers - 1 && bottom > -tl)
                result.addSlice(bottom + tl, *material);
            // bottom layer
            if (i == nLayers - 1)
                result.addSlice(0.0, *material); // semiinfinite bottom slice
        }
    }
    return result;
}

SliceStack setAvgMaterials(const SliceStack& stack, const OwningVector<const ReLayout>& layouts)
{
    SliceStack result = stack;

    // loop over inner layers only; there is no admixture in semi-infinite layers
    for (size_t i_slice = 1; i_slice < result.size() - 1; ++i_slice) {
        Admixtures admixtures;
        for (const ReLayout* layout : layouts) {
            for (const CoheringSubparticles* group : layout->subparticles()) {
                for (const IReParticle* sp : group->terms()) {
                    if (sp->i_layer_or_0() == i_slice)
                        admixtures.push_back(sp->admixed());
                }
            }
        }
        const auto slice_mat = result[i_slice].material();
        result[i_slice].setMaterial(MaterialUtil::averagedMaterial(slice_mat, admixtures));
    }
    return result;
}

size_t zToSliceIndex(double z, const SliceStack& slices)
{
    auto n_layers = slices.size();
    if (n_layers < 2 || z > 0.0)
        return 0;

    double z0 = slices[0].low();
    size_t result = 0;
    while (result < n_layers - 1) {
        ++result;
        if (slices[result].low() - z0 < z)
            break;
    }
    return result;
}

std::pair<size_t, size_t> SliceIndexSpan(const IParticle& particle, const SliceStack& slices,
                                         double z_ref)
{
    const auto [zbottom, ztop] = particle.zSpan().pair();
    const double eps =
        (ztop - zbottom) * 1e-6; // allow for relatively small crossing
                                 // due to numerical approximations (like rotation over 180 degrees)
    const double zmax = ztop + z_ref - eps;
    const double zmin = zbottom + z_ref + eps;
    size_t top_index = zToSliceIndex(zmax, slices);
    const size_t bottom_index = zToSliceIndex(zmin, slices);
    if (top_index > bottom_index) // happens for zero size particles
        top_index = bottom_index;
    return {top_index, bottom_index};
}

ReLayout* makeReLayout(const ParticleLayout& layout, const SliceStack& slices, double z_ref,
                       const SimulationOptions& options, bool polarized)
{
    const double layout_abundance = layout.totalAbundance();
    if (layout_abundance <= 0)
        throw std::runtime_error("Particle layout has invalid total abundance <= 0");
    const double surface_density = layout.totalParticleSurfaceDensity() * layout.absoluteWeight();
    if (surface_density <= 0)
        throw std::runtime_error("Particle layout has invalid surface density <= 0");

    OwningVector<const CoheringSubparticles> coherentParticles;

    for (const auto& particle : layout.particles()) {

        OwningVector<IReParticle> terms;

        for (const auto* subparticle : particle->decompose()) {
            double z_sub = z_ref;
            const auto slice_indices = SliceIndexSpan(*subparticle, slices, z_sub);
            bool single_layer = (slice_indices.first == slice_indices.second);
            double z0 = slices.at(0).low();
            std::unique_ptr<IParticle> myparticle(subparticle->clone());

            for (size_t i = slice_indices.first; i < slice_indices.second + 1; ++i) {
                const Slice& slice = slices[i];
                double z_top_i = i == 0 ? 0 : slice.hig() - z0;
                R3 translation(0.0, 0.0, z_sub - z_top_i);
                z_sub = z_top_i; // subparticle now has coordinates relative to z_top_i
                myparticle->translate(translation);

                // if subparticle is contained in this layer, set limits to infinite:
                ZLimits limits;
                if (!single_layer) {
                    double z1 = slice.higOr0();
                    limits = {slice.low() - z1, slice.hig() - z1};
                }
                const Material& ambientMat = slices.at(i).material();
                OwningVector<IReParticle> subparticles = Compute::Slicing::particlesInSlice(
                    myparticle.get(), limits, ambientMat, options.mesoOptions());

                while (IReParticle* sp = subparticles.release_front()) {
                    if (slices.size() > 1)
                        sp->setLayerIndex(i);

                    double factor = particle->abundance() / layout_abundance * surface_density;
                    if (double thickness = slice.thicknessOr0(); thickness > 0.0)
                        factor /= thickness;
                    sp->setAdmixedFraction(factor * sp->admixedFraction());

                    terms.push_back(sp);
                }
            }
        }

        coherentParticles.push_back(
            new CoheringSubparticles(particle->abundance() / layout_abundance, std::move(terms)));
    }

    const auto* iff = layout.interferenceFunction();

    return new ReLayout{surface_density, std::move(coherentParticles), iff ? iff->clone() : nullptr,
                        options, polarized};
}

//! Returns collection of all ReLayout%s, for use in ReSample constructor.

OwningVector<const ReLayout> collectLayouts(const Sample& sample, const SliceStack& slices,
                                            const SimulationOptions& options, bool polarized)
{
    OwningVector<const ReLayout> result;

    ASSERT(!slices.empty());
    double z_ref = -slices.front().low();

    for (size_t i = 0; i < sample.numberOfLayers(); ++i) {
        if (i > 1)
            z_ref -= sample.layer(i - 1)->thickness();
        for (const auto* layout : sample.layer(i)->layouts())
            result.push_back(makeReLayout(*layout, slices, z_ref, options, polarized));
    }
    return result;
}

} // namespace


//  ************************************************************************************************
//  class implementation
//  ************************************************************************************************

//... Static factory functions:

ReSample ReSample::make(const Sample& sample, const SimulationOptions& options, bool forcePolarized)
{
    sample.checkAndProcess();

    const bool polarized = forcePolarized || sample.isMagnetic() || sample.externalField() != R3();

    // if requested, slice layers that contain particles
    const SliceStack stack1 = slicify(sample, options.useAvgMaterials());

    OwningVector<const ReLayout> layouts = collectLayouts(sample, stack1, options, polarized);

    SliceStack stack2 = options.useAvgMaterials() ? setAvgMaterials(stack1, layouts) : stack1;

    // Recalculate full magnetic induction inside new slices with new materials
    stack2.setBField(sample.externalField());

    return {sample, polarized, std::move(layouts), stack2};
}

ReSample ReSample::make(const Sample& sample)
{
    return make(sample, {}, false);
}

//...

ReSample::ReSample(const Sample& sample, bool polarized, OwningVector<const ReLayout>&& relayouts,
                   const SliceStack& refined_stack)
    : m_sample(sample)
    , m_polarized(polarized)
    , m_relayouts(std::move(relayouts))
    , m_stack(refined_stack)
{
}

ReSample::~ReSample() = default;

bool ReSample::polarizing() const
{
    return m_polarized || m_sample.externalField() != R3{};
}

bool ReSample::hasRoughness() const
{
    for (const auto& slice : m_stack)
        if (slice.topRMS() > 0)
            return true;
    return false;
}

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

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