//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Export/SceneGeometry.cpp
//! @brief     Assembles layer data, layout positions, and particle templates into a 3D scene.
//!
//! @homepage  http://www.bornagainproject.org
//! @license   GNU General Public License v3 or higher (see COPYING)
//! @copyright Forschungszentrum Jülich GmbH 2026
//! @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
//
//  ************************************************************************************************

#include "Sim/Export/SceneGeometry.h"
#include "Base/Type/Field2D.h"
#include "Base/Util/Assert.h"
#include "Sample/Aggregate/Struct.h"
#include "Sample/Aggregate/Struct2D.h"
#include "Sample/Aggregate/Struct3D.h"
#include "Sample/Material/Material.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Particle/IParticle.h"
#include "Sample/Sample/RoughnessMap.h"
#include "Sample/Sample/Sample.h"
#include "Sim/Export/SceneParticleCollector.h"
#include "Sim/Export/ScenePositionBuilder.h"
#include <cmath>
#include <stdexcept>

namespace {

constexpr unsigned MIN_ROUGHNESS_RESOLUTION = 8;
constexpr unsigned MAX_ROUGHNESS_RESOLUTION = 2048;

// ---- Scene assembly ----

ColorRgba colorRgba(const Material* material, double opacity)
{
    const MaterialColor color = material->color();
    return {static_cast<float>(color.r), static_cast<float>(color.g), static_cast<float>(color.b),
            static_cast<float>(opacity)};
}

double displayThickness(const Layer* layer)
{
    const double t = layer->thickness();
    return std::isfinite(t) ? t : 50.0;
}

unsigned normalizedRoughnessResolution(unsigned points)
{
    if (points == 0)
        return 0;
    if (points < MIN_ROUGHNESS_RESOLUTION)
        throw std::runtime_error("buildScene3D: roughness_points must be 0 or at least "
                                 + std::to_string(MIN_ROUGHNESS_RESOLUTION));
    if (points > MAX_ROUGHNESS_RESOLUTION)
        throw std::runtime_error("buildScene3D: roughness_points must be at most "
                                 + std::to_string(MAX_ROUGHNESS_RESOLUTION));
    return points;
}

std::vector<double2d_t> generateRoughnessMaps(const Sample& sample, unsigned roughness_points,
                                              double lx, double ly, unsigned seed)
{
    const unsigned points = ::normalizedRoughnessResolution(roughness_points);
    if (points == 0)
        return std::vector<double2d_t>(sample.unwrapped().size());

    // Generate the rough interfaces as one stack. This lets the 3D view use the same adjacent
    // cross-spectrum models as scattering, so correlated roughness is visible as replicated
    // low-frequency features instead of independent per-interface maps.
    return RoughnessMap::generateInterfaceMaps(points, points, lx, ly, sample, seed);
}

StructData buildStructData(const Struct* s, const std::string& placement, double ref_z,
                           double lateral_half_extent, double layer_thickness,
                           double layer_z_bottom, double layer_z_top, unsigned seed)
{
    const IParticle* particle = s->particle();
    ASSERT(particle);
    particle->validate();

    StructData sd;
    sd.type = s->className();
    sd.placement = placement;
    sd.ref_z = ref_z;

    // Position builder creates the per-instance layout points; particle collector creates the
    // reusable particle templates placed at each of those points by the Python viewer.
    const SceneGeometryDetail::PositionContext position_ctx{lateral_half_extent, layer_thickness,
                                                            layer_z_bottom, layer_z_top, seed};
    SceneGeometryDetail::fillStructPositions(s, position_ctx, sd);
    const double particle_ref_z = dynamic_cast<const Struct3D*>(s) ? 0.0 : ref_z;
    SceneGeometryDetail::collectParticleTemplates(particle, R3(0, 0, particle_ref_z), RotMatrix(),
                                                  1.0, sd);
    return sd;
}

} // namespace


SceneDescription buildScene3D(const Sample& sample, double lateral_half_extent, unsigned seed,
                              unsigned roughness_points)
{
    if (lateral_half_extent <= 0)
        throw std::runtime_error("buildScene3D: lateral_half_extent must be positive");

    SceneDescription scene;

    const auto& layers = sample.unwrapped();

    scene.lateral_half_extent = lateral_half_extent;
    const std::vector<double2d_t> roughness_maps = ::generateRoughnessMaps(
        sample, roughness_points, 2.0 * lateral_half_extent, 2.0 * lateral_half_extent, seed);

    // Layers
    double z_top = 0;
    for (size_t i = 0; i < layers.size(); ++i) {
        const Layer* layer = layers[i].get();
        const double t = layer->thickness();
        const double z_bottom = z_top - ::displayThickness(layer);

        LayerData ld;
        ld.material = layer->material()->materialName();
        ld.color = ::colorRgba(layer->material(), 1.0);
        ld.z_top = z_top;
        ld.z_bottom = z_bottom;
        ld.semi_infinite = !std::isfinite(t);
        if (i < roughness_maps.size() && !roughness_maps[i].empty()) {
            ld.roughness_ny = static_cast<unsigned>(roughness_maps[i].size());
            ld.roughness_nx = static_cast<unsigned>(roughness_maps[i][0].size());
            ld.top_roughness = FieldUtil::flatten(roughness_maps[i]);
        }
        scene.layers.push_back(ld);

        z_top = z_bottom;
    }

    // Structs
    z_top = 0;
    unsigned struct_seed = seed;
    for (size_t i = 0; i < layers.size(); ++i) {
        const Layer* layer = layers[i].get();
        const double thickness = ::displayThickness(layer);
        const double z_bottom = z_top - thickness;
        const double z_mid = (z_top + z_bottom) / 2.0;

        for (const auto& s : layer->structsAtBottom()) {
            if (!s->particle())
                continue;
            scene.structs.push_back(::buildStructData(s.get(), "deposit", z_bottom,
                                                      lateral_half_extent, thickness, z_bottom,
                                                      z_top, struct_seed++));
        }
        for (const auto& s : layer->structsAtTop()) {
            if (!s->particle())
                continue;
            scene.structs.push_back(::buildStructData(s.get(), "suspend", z_top,
                                                      lateral_half_extent, thickness, z_bottom,
                                                      z_top, struct_seed++));
        }
        for (const auto& s : layer->structs3D()) {
            if (!s->particle())
                continue;
            scene.structs.push_back(::buildStructData(s.get(), "fill", z_mid, lateral_half_extent,
                                                      thickness, z_bottom, z_top, struct_seed++));
        }

        z_top = z_bottom;
    }

    return scene;
}
