//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Refine/SubdivideLayers.cpp
//! @brief     Implements functions Refine::wrapLayers and subdivideLayers
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2025
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Resample/Refine/SubdivideLayers.h"
#include "Base/Util/Assert.h"
#include "Sample/Aggregate/Disordered2D.h"
#include "Sample/Aggregate/Struct2D.h"
#include "Sample/Aggregate/Struct3D.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Slice.h"
#include "Sample/Particle/IParticle.h"
#include <stdexcept>

namespace {

void addTopSlice(std::vector<std::unique_ptr<Slice>>& ss, const Layer* layer, double zbottom)
{
    ASSERT(ss.empty());
    ss.emplace_back(std::make_unique<Slice>(0, *layer, Span(zbottom, Span::inf), R3{}));
}

void addSlice(std::vector<std::unique_ptr<Slice>>& ss, const Layer* layer, double thickness,
              size_t count = 1)
{
    ASSERT(!ss.empty());
    ASSERT(thickness > 0);
    ASSERT(count > 0);
    const double slice_thickness = thickness / count;
    for (size_t i = 0; i < count; ++i) {
        double top = ss.back()->low();
        ss.emplace_back(
            std::make_unique<Slice>(ss.size(), *layer, Span{top - slice_thickness, top}, R3{}));
    }
}

} // namespace


std::vector<std::unique_ptr<Slice>> Refine::wrapLayers(const std::vector<const Layer*>& flat)
{
    const size_t N = flat.size();
    ASSERT(N >= 2); // Single-layer samples are now forbidden

    std::vector<std::unique_ptr<Slice>> result;
    ::addTopSlice(result, flat[0], 0);
    for (size_t i = 1; i < N; ++i)
        ::addSlice(result, flat[i], flat[i]->thickness());
    return result;
}

std::vector<std::unique_ptr<Slice>> Refine::subdivideLayers(const std::vector<const Layer*>& flat)
{
    const size_t N = flat.size();
    ASSERT(N >= 2); // Single-layer samples are now forbidden

    // For each layer, compute the z span that contains particles.
    std::vector<Span> particle_spans(N);

    using ZFn = double (IParticle::*)() const;

    // Helper to process structures given reference z and intrinsic-z function.
    // zFn: &IParticle::intrinsicBottomZ for deposit2D, &IParticle::intrinsicTopZ for suspend2D.
    // useTopOffset: true for suspend2D (particle top sits at ref_z, so offset = ref_z - H).
    auto processStructures = [&](const std::vector<std::unique_ptr<Struct2D>>& structs,
                                 double ref_z, ZFn zFn, bool useTopOffset) {
        for (const auto& structure : structs) {
            const IParticle* root = structure->particle();
            for (const auto* particle : root->particles()) {
                // Recover pre-placeRef z extents by reversing the translate-on-deposit shift.
                const double shift = (particle->*zFn)();
                const double rel_bottom = particle->bottomZ() + shift;
                const double rel_top = particle->topZ() + shift;
                ASSERT(rel_bottom < rel_top);
                ASSERT(std::isfinite(rel_bottom) && std::isfinite(rel_top));

                const double offset = useTopOffset ? ref_z - (rel_top - rel_bottom) : ref_z;
                const double bottom = rel_bottom + offset;
                const double top = rel_top + offset;
                for (size_t ii = 0; ii < N; ++ii) {
                    Span s = Span::intersect({bottom, top}, flat[ii]->zSpan());
                    particle_spans[ii] = Span::unite(particle_spans[ii], s);
                }
            }
        }
    };

    for (size_t i = 0; i < N; ++i) {
        const Layer* layer = flat.at(i);

        // Validate and process structures at bottom (from deposit2D)
        if (!layer->structsAtBottom().empty()) {
            if (i == N - 1)
                throw std::runtime_error(
                    "Cannot deposit particles in semi-infinite substrate layer. "
                    "Use suspend2D() to position particles below substrate surface.");
            processStructures(layer->structsAtBottom(), layer->zSpan().low(),
                              &IParticle::intrinsicBottomZ, false);
        }

        // Validate and process structures at top (from suspend2D)
        if (!layer->structsAtTop().empty()) {
            if (i == 0)
                throw std::runtime_error("Cannot suspend particles in semi-infinite top layer. "
                                         "Use deposit2D() instead.");
            processStructures(layer->structsAtTop(), layer->zSpan().hig(),
                              &IParticle::intrinsicTopZ, true);
        }
    }

    // Construct stack of slices
    std::vector<std::unique_ptr<Slice>> result;
    for (size_t i = 0; i < N; ++i) {
        const Layer* const layer = flat.at(i);
        const double tl = layer->thickness();

        // if no slicing is needed, create single slice for the layer
        if (particle_spans[i].empty()) {
            if (i == 0)
                ::addTopSlice(result, layer, 0);
            else
                ::addSlice(result, layer, tl);
            continue;
        }

        const double top = particle_spans[i].hig();
        const double bottom = particle_spans[i].low();
        const size_t nSlices = layer->numberOfSlices();
        ASSERT(nSlices > 0);

        // Add slice above particle zone
        if (i == 0) {
            ::addTopSlice(result, layer, top); // semiinfinite top slice
        } else {
            ASSERT(top <= tl);
            const double dt = flat[i]->zSpan().hig() - top;
            if (dt > 0)
                ::addSlice(result, layer, dt);
        }
        // Add slice(s) in particle zone
        ::addSlice(result, layer, top - bottom, nSlices);
        // Add slice below particle zone
        const double db = bottom - flat[i]->zSpan().low();
        if (i < N - 1 && db > 0)
            ::addSlice(result, layer, db);
        // Add semiinfinite bottom slice
        // Called every time when there are particles fully or partially located in the substrate
        if (i == N - 1)
            ::addSlice(result, layer, Span::inf);
    }
    return result;
}
