//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Refine/RefineStructs.cpp
//! @brief     Implements function Refine::refineStructs.
//!
//! @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/RefineStructs.h"
#include "Base/Util/Assert.h"
#include "Resample/Refine/ParticlesInSlice.h"
#include "Sample/Aggregate/Struct2D.h"
#include "Sample/Aggregate/Struct3D.h"
#include "Sample/Material/Material.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Slice.h"
#include "Sample/Option/SimulationOptions.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/IParticle.h"
#include "Sample/Particle/Mesocrystal.h"
#include "Sample/Particle/Mixture.h"
#include <algorithm>
#include <cmath>
#include <stdexcept>

namespace {

//! Returns the index of the slice containing z-coordinate z (absolute coords).
size_t zToSliceIndex(double z /*absolute*/, const std::vector<std::unique_ptr<Slice>>& slices)
{
    const size_t n_slices = slices.size();
    if (n_slices < 2 || z >= slices[0]->low())
        return 0;
    for (size_t i = 1; i < n_slices; ++i)
        if (slices[i]->low() < z)
            return i;
    return n_slices - 1;
}

//! Returns the range of slice indices [top, bottom] that a particle spans.
//! @param z0_particle Z-offset to convert particle z-extents to absolute coordinates.
std::pair<size_t, size_t> sliceIndexSpan(const IParticle& particle,
                                         const std::vector<std::unique_ptr<Slice>>& slices,
                                         double z0_particle /*absolute*/)
{
    double particle_bot = particle.bottomZ(); // particle-local, then absolute after adjustment
    double particle_top = particle.topZ();    // particle-local, then absolute after adjustment
    const double eps = (particle_top - particle_bot)
                       * 1e-6; // allow for relatively small crossing
                               // due to numerical approximations (like rotation over 180 degrees)
    particle_top += z0_particle - eps;
    particle_bot += z0_particle + eps;
    size_t top_index = ::zToSliceIndex(particle_top, slices);
    const size_t bottom_index = ::zToSliceIndex(particle_bot, slices);
    if (top_index > bottom_index) // happens for zero size particles
        top_index = bottom_index;
    return {top_index, bottom_index};
}

//! Preserves the original object's radius in a rebuilt Compound.
void copyRadius(const IParticle& original, Compound& compound)
{
    // If the source has no explicit radius, leave the rebuilt Compound without one;
    // Dense2D validates that later.
    if (original.hasRadius())
        compound.setRadius(original.radius());
}

//! Returns the slicing result as a single particle:
//! one part as-is, multiple parts wrapped in a Compound.
std::unique_ptr<IParticle> wrapParticles(std::vector<std::unique_ptr<IParticle>>&& particles,
                                         const IParticle& prototype)
{
    ASSERT(!particles.empty());
    if (particles.size() == 1)
        return std::move(particles[0]);

    auto result = std::make_unique<Compound>();
    for (auto& particle : particles)
        result->addComponent(std::move(particle));
    ::copyRadius(prototype, *result);
    return result;
}

//! Slices a leaf particle (Particle, CoreAndShell, or Mesocrystal) across slice
//! boundaries and wraps the results in a Compound.
//! @param particle Leaf particle with position in Struct-frame (parent transforms applied).
//! @param z0_particle Z-offset to convert particle z-extents to absolute coordinates.
std::unique_ptr<IParticle> processParticle(const IParticle& particle,
                                           const std::vector<std::unique_ptr<Slice>>& slices,
                                           double z0_particle /*absolute*/,
                                           const SimulationOptions& options)
{
    const auto slice_reference_z = [](const Slice& slice) {
        return std::isfinite(slice.hig()) ? slice.hig() : slice.low();
    };
    std::vector<std::unique_ptr<IParticle>> subparticles;

    if (options.useAvgMaterials()) {
        const auto slice_indices = ::sliceIndexSpan(particle, slices, z0_particle);
        ASSERT(slice_indices.first > 0); // slice[0] is above all particles
        const bool single_slice = (slice_indices.first == slice_indices.second);
        for (size_t i = slice_indices.first; i <= slice_indices.second; ++i) {
            const Slice& slice = *slices[i];
            // Clone particle and translate to slice-local coordinates.
            std::unique_ptr<IParticle> p(particle.clone());
            p->translate(R3(0, 0, z0_particle - slice_reference_z(slice)));

            const Span limits = // slice-local
                single_slice ? Span{-Span::inf, Span::inf} : Span{-slice.thickness(), 0};
            const Material* ambientMat = slice.layerMaterial();
            auto sliced = Refine::particlesInSlice(p.get(), limits, &slice, *ambientMat,
                                                   options.mesoReciprocalSum());
            for (auto& sp : sliced)
                subparticles.push_back(std::move(sp));
        }
    } else {
        // Without averaging: place particle in its containing slice
        const double particle_center = particle.centerZ() + z0_particle; // absolute
        const size_t slice_index = ::zToSliceIndex(particle_center, slices);
        const Slice& slice = *slices[slice_index];
        // Clone particle and translate to slice-local coordinates.
        std::unique_ptr<IParticle> p(particle.clone());
        p->translate(R3(0, 0, z0_particle - slice_reference_z(slice)));

        const Span limits(-Span::inf, Span::inf); // slice-local
        const Material* ambientMat = slice.layerMaterial();
        subparticles = Refine::particlesInSlice(p.get(), limits, &slice, *ambientMat,
                                                options.mesoReciprocalSum());
    }

    return ::wrapParticles(std::move(subparticles), particle);
}

//! Recursively walks the particle tree, slicing leaf particles.
//! Container nodes (Mixture, Compound) are rebuilt with recursively processed children.
//! Parent transforms (rotation/translation) are pushed down to children at each level,
//! so that leaves arrive with positions in the Struct-frame.
//! @param z0_particle Z-offset to convert particle z-extents to absolute coordinates.
//!                    For Top alignment with Mixture, recomputed per child.
//! @param layer Needed for TopOfLayer offset calculation.
static std::unique_ptr<IParticle> sliceRecursive(const IParticle& particle,
                                                 const std::vector<std::unique_ptr<Slice>>& slices,
                                                 double z0_particle /*absolute*/,
                                                 const SimulationOptions& options,
                                                 const Layer* layer)
{
    ASSERT(layer);

    if (const auto* mix = dynamic_cast<const Mixture*>(&particle)) {
        auto result = std::make_unique<Mixture>();
        const auto& abundances = mix->abundances();
        const auto children = mix->particles();
        for (size_t i = 0; i < children.size(); ++i) {
            std::unique_ptr<IParticle> child(children[i]->clone());
            child->rotate(mix->particleRotation());
            child->translate(mix->particlePosition());
            auto sliced = ::sliceRecursive(*child, slices, z0_particle, options, layer);
            result->addParticle(std::move(sliced), abundances[i]);
        }
        result->validate();
        return result;
    }

    if (const auto* comp = dynamic_cast<const Compound*>(&particle)) {
        auto result = std::make_unique<Compound>();
        for (const IParticle* child_ptr : comp->components()) {
            std::unique_ptr<IParticle> child_clone(child_ptr->clone());
            child_clone->rotate(comp->particleRotation());
            child_clone->translate(comp->particlePosition());

            auto sliced = ::sliceRecursive(*child_clone, slices, z0_particle, options, layer);
            result->addComponent(std::move(sliced));
        }
        ::copyRadius(particle, *result);
        return result;
    }

    // Leaf: Particle, CoreAndShell, Mesocrystal
    return ::processParticle(particle, slices, z0_particle, options);
}

//! Recursively prepares a particle tree for 3D scattering within a single slice.
//! Container nodes (Mixture, Compound) are rebuilt with recursively prepared children;
//! leaf particles are resolved via particlesInSlice and assigned the host slice number.
//! @param particle  Node of the particle tree, with position in slice-local coordinates.
//! @param slice     The single host slice that contains the 3D assembly.
//! @param options   Simulation options (e.g. mesocrystal reciprocal-sum flag).
//! @return          A new particle tree ready for scattering evaluation.
std::unique_ptr<IParticle> prepare3DRecursive(const IParticle& particle, const Slice& slice,
                                              const SimulationOptions& options)
{
    if (const auto* mix = dynamic_cast<const Mixture*>(&particle)) {
        auto result = std::make_unique<Mixture>();
        const auto& abundances = mix->abundances();
        const auto children = mix->particles();
        for (size_t i = 0; i < children.size(); ++i) {
            std::unique_ptr<IParticle> child(children[i]->clone());
            child->rotate(mix->particleRotation());
            child->translate(mix->particlePosition());
            result->addParticle(::prepare3DRecursive(*child, slice, options), abundances[i]);
        }
        result->validate();
        result->setSliceNumber(slice.number());
        return result;
    }

    if (const auto* comp = dynamic_cast<const Compound*>(&particle)) {
        auto result = std::make_unique<Compound>();
        for (const IParticle* child_ptr : comp->components()) {
            std::unique_ptr<IParticle> child_clone(child_ptr->clone());
            child_clone->rotate(comp->particleRotation());
            child_clone->translate(comp->particlePosition());
            result->addComponent(::prepare3DRecursive(*child_clone, slice, options));
        }
        ::copyRadius(particle, *result);
        result->setSliceNumber(slice.number());
        return result;
    }

    // Leaf: Particle, CoreAndShell, Mesocrystal
    const Span limits(-Span::inf, Span::inf); // slice-local
    const Material* ambientMat = slice.layerMaterial();
    auto prepared = Refine::particlesInSlice(&particle, limits, &slice, *ambientMat,
                                             options.mesoReciprocalSum());
    if (prepared.empty())
        throw std::runtime_error("Failed to prepare 3D particle in host slice");
    auto result = ::wrapParticles(std::move(prepared), particle);
    result->setSliceNumber(slice.number());
    return result;
}

//! Creates a resampled 2D structure using recursive tree slicing.
//! The particle tree is preserved: Mixture/Compound containers are rebuilt with
//! recursively processed children, and leaf particles are sliced into Compounds.
//! @param ref_z Reference z-coordinate in absolute coordinates (layer bottom or top interface).
Struct* remake2D(const Struct2D& structure, const std::vector<std::unique_ptr<Slice>>& slices,
                 double ref_z /*absolute*/, const SimulationOptions& options, const Layer* layer)
{
    if (structure.areaNumberDensity() <= 0)
        throw std::runtime_error("Particle structure has invalid surface density <= 0");

    const IParticle* root = structure.particle();

    // After translate-on-deposit, the root's alignment reference is at z=0 in the struct
    // frame and the struct origin is at absolute ref_z, so z0_particle = ref_z.
    const double z0_particle = ref_z;

    auto sliced = ::sliceRecursive(*root, slices, z0_particle, options, layer);

    auto* result = structure.clone();
    result->setParticle(std::move(sliced));
    return result;
}

//! Creates a resampled 3D structure (random assembly) confined to a single slice.
Struct* remake3D(const Struct3D& structure, const Slice& slice, const SimulationOptions& options)
{
    if (structure.volumeNumberDensity() <= 0)
        throw std::runtime_error("Particle structure has invalid volume density <= 0");
    const double thickness = slice.thickness();
    if (!std::isfinite(thickness))
        throw std::runtime_error("3D assembly not allowed in infinite layer");

    const IParticle* root = structure.particle();
    const double layer_bottom = -thickness; // slice-local
    std::unique_ptr<IParticle> anchored(root->clone());
    // Re-anchor at the slice bottom.
    anchored->translate(R3(0.0, 0.0, layer_bottom - root->intrinsicBottomZ()));

    // Then average over the allowed z-shift in the dilute case.
    std::unique_ptr<IParticle> prepared = ::prepare3DRecursive(*anchored, slice, options);

    const double shift_low = layer_bottom - prepared->bottomZ();
    const double shift_high = -prepared->topZ();
    if (shift_low > shift_high)
        throw std::runtime_error("Particle in 3D assembly does not fit within the layer thickness");

    auto* result = structure.clone();
    result->setSliceNumber(slice.number());
    result->setThickness(thickness);
    result->setZRange(shift_low, shift_high);
    result->setParticle(std::move(prepared));
    return result;
}

std::vector<const Slice*> slicesInLayer(const Layer& layer,
                                        const std::vector<std::unique_ptr<Slice>>& slices)
{
    std::vector<const Slice*> result;
    for (const auto& slice : slices)
        if (!(slice->hig() <= layer.low() || slice->low() >= layer.hig()))
            result.push_back(slice.get());
    return result;
}

} // namespace


std::vector<std::unique_ptr<const Struct>>
Refine::refineStructs(const std::vector<const Layer*>& flat,
                      const std::vector<std::unique_ptr<Slice>>& slices,
                      const SimulationOptions& options)
{
    std::vector<std::unique_ptr<const Struct>> result;

    ASSERT(!slices.empty());

    for (const Layer* layer : flat) {
        // Process 2D structures at bottom interface
        for (const auto& s : layer->structsAtBottom())
            result.emplace_back(::remake2D(*s, slices, layer->low(), options, layer));

        // Process 2D structures at top interface
        for (const auto& s : layer->structsAtTop())
            result.emplace_back(::remake2D(*s, slices, layer->hig(), options, layer));

        // Process 3D structures
        if (!layer->structs3D().empty()) {
            const auto layer_slices = ::slicesInLayer(*layer, slices);
            if (layer_slices.size() != 1)
                throw std::runtime_error(
                    "3D structures require a host layer represented by a single slice");
            for (const auto& s : layer->structs3D())
                result.emplace_back(::remake3D(*s, *layer_slices.front(), options));
        }
    }
    return result;
}
