//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Sim/Export/SceneParticleCollector.cpp
//! @brief     Converts IParticle trees into ParticleData templates for the 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/SceneParticleCollector.h"
#include "Base/Util/Assert.h"
#include "Sample/HardParticle/Bar.h"
#include "Sample/HardParticle/Cone.h"
#include "Sample/HardParticle/CosineRipple.h"
#include "Sample/HardParticle/Cylinder.h"
#include "Sample/HardParticle/Ellipsoid.h"
#include "Sample/HardParticle/EllipsoidalCylinder.h"
#include "Sample/HardParticle/FormfactorPolyhedron.h"
#include "Sample/HardParticle/HorizontalCylinder.h"
#include "Sample/HardParticle/LongBoxGauss.h"
#include "Sample/HardParticle/LongBoxLorentz.h"
#include "Sample/HardParticle/SawtoothRipple.h"
#include "Sample/HardParticle/Spheroid.h"
#include "Sample/HardParticle/SpheroidalSegment.h"
#include "Sample/Material/Material.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/CoreAndShell.h"
#include "Sample/Particle/Crystal.h"
#include "Sample/Particle/Formfactor.h"
#include "Sample/Particle/IParticle.h"
#include "Sample/Particle/Mesocrystal.h"
#include "Sample/Particle/Mixture.h"
#include "Sample/Particle/Particle.h"
#include "Sample/Particle/Rotations.h"
#include "Sample/SoftParticle/FuzzySphere.h"
#include "Sample/SoftParticle/GaussSphere.h"
#include "Sim/Export/SceneGeometry.h"
#include <ff/Polyhedron.h>
#include <ff/Topology.h>
#include <utility>

namespace {

// ---- Particle data collection ----

void fillPolyhedronData(const Formfactor* ff, ParticleData& pd)
{
    const auto* poly_ff = dynamic_cast<const FormfactorPolyhedron*>(ff);
    if (!poly_ff)
        return;
    auto poly = poly_ff->getPolyhedron();
    if (!poly)
        return;

    for (const auto& v : poly->vertices()) {
        pd.vert_x.push_back(v.x());
        pd.vert_y.push_back(v.y());
        pd.vert_z.push_back(v.z());
    }
    for (const auto& face : poly->topology().faces) {
        pd.face_vertex_counts.push_back(static_cast<int>(face.vertexIndices.size()));
        for (int idx : face.vertexIndices)
            pd.face_vertex_indices.push_back(idx);
    }
}

//! Carries state while flattening a particle tree into ParticleData records.
//! Needed so Mixture leaves keep their activation path and share per-struct weights.
struct CollectContext {
    explicit CollectContext(std::vector<std::vector<double>>& weights)
        : mixture_weights(weights)
    {
    }

    int next_mixture_group = 0;
    std::vector<std::pair<int, int>> mixture_path;
    std::vector<std::vector<double>>& mixture_weights;
};

//! RAII guard for one Mixture choice in the current activation path.
//! Ensures recursive collection pops the path even if a child collector throws.
struct PathFrame {
    PathFrame(std::vector<std::pair<int, int>>& path_, int group, int variant)
        : path(path_)
    {
        path.emplace_back(group, variant);
    }

    ~PathFrame() { path.pop_back(); }

    PathFrame(const PathFrame&) = delete;
    PathFrame& operator=(const PathFrame&) = delete;

    std::vector<std::pair<int, int>>& path;
};

//! Pushes a (name, value) parameter pair onto a ParticleData.
//! Parameter names must match the keys read by ba_view3d.py's _make_smooth_polydata;
//! polyhedra take a separate path via fillPolyhedronData and don't need these.
void pushParam(ParticleData& pd, const std::string& name, double value)
{
    pd.param_names.push_back(name);
    pd.param_values.push_back(value);
}

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)};
}

void fillSmoothShapeParams(const Formfactor* ff, ParticleData& pd)
{
    if (const auto* cyl = dynamic_cast<const Cylinder*>(ff)) {
        pushParam(pd, "Radius", cyl->radius());
        pushParam(pd, "Height", cyl->height());
    } else if (const auto* ecyl = dynamic_cast<const EllipsoidalCylinder*>(ff)) {
        pushParam(pd, "RadiusX", ecyl->radiusX());
        pushParam(pd, "RadiusY", ecyl->radiusY());
        pushParam(pd, "Height", ecyl->height());
    } else if (const auto* cone = dynamic_cast<const Cone*>(ff)) {
        pushParam(pd, "Radius", cone->radius());
        pushParam(pd, "Height", cone->height());
    } else if (const auto* ell = dynamic_cast<const Ellipsoid*>(ff)) {
        pushParam(pd, "Radius X", ell->radiusX());
        pushParam(pd, "Radius Y", ell->radiusY());
        pushParam(pd, "Radius Z", ell->radiusZ());
    } else if (const auto* sph = dynamic_cast<const Spheroid*>(ff)) {
        pushParam(pd, "Radius XY", sph->radiusXY());
        pushParam(pd, "Radius Z", sph->radiusZ());
    } else if (const auto* seg = dynamic_cast<const SpheroidalSegment*>(ff)) {
        pushParam(pd, "Radius XY", seg->radiusXY());
        pushParam(pd, "Radius Z", seg->radiusZ());
        pushParam(pd, "Top cut", seg->cutFromTop());
        pushParam(pd, "Bottom cut", seg->cutFromBottom());
    } else if (const auto* hc = dynamic_cast<const HorizontalCylinder*>(ff)) {
        pushParam(pd, "Radius", hc->radius());
        pushParam(pd, "Length", hc->length());
    } else if (const auto* fs = dynamic_cast<const FuzzySphere*>(ff)) {
        pushParam(pd, "Radius", fs->radialExtension());
        pushParam(pd, "Fuzziness", fs->fuzziness());
    } else if (const auto* gs = dynamic_cast<const GaussSphere*>(ff)) {
        pushParam(pd, "MeanRadius", gs->meanRadius());
    } else if (const auto* sr = dynamic_cast<const SawtoothRipple*>(ff)) {
        pushParam(pd, "Length", sr->length());
        pushParam(pd, "Width", sr->width());
        pushParam(pd, "Height", sr->height());
        pushParam(pd, "Asymmetry", sr->asymmetry());
    } else if (const auto* cr = dynamic_cast<const CosineRipple*>(ff)) {
        // Catches CosineRippleBox/Gauss/Lorentz; profile envelope is scattering-only
        // and irrelevant to the 3D mesh (rendered as a single cosine-shape ridge).
        pushParam(pd, "Length", cr->length());
        pushParam(pd, "Width", cr->width());
        pushParam(pd, "Height", cr->height());
    } else if (const auto* bg = dynamic_cast<const BarGauss*>(ff)) {
        pushParam(pd, "Length", bg->length());
        pushParam(pd, "Width", bg->width());
        pushParam(pd, "Height", bg->height());
    } else if (const auto* bl = dynamic_cast<const BarLorentz*>(ff)) {
        pushParam(pd, "Length", bl->length());
        pushParam(pd, "Width", bl->width());
        pushParam(pd, "Height", bl->height());
    } else if (const auto* lbg = dynamic_cast<const LongBoxGauss*>(ff)) {
        pushParam(pd, "Length", lbg->length());
        pushParam(pd, "Width", lbg->width());
        pushParam(pd, "Height", lbg->height());
    } else if (const auto* lbl = dynamic_cast<const LongBoxLorentz*>(ff)) {
        pushParam(pd, "Length", lbl->length());
        pushParam(pd, "Width", lbl->width());
        pushParam(pd, "Height", lbl->height());
    } else {
        // Fallback for unrecognised smooth shapes: only the radial extension.
        // Sphere falls through to here as well, since "Radius" is all it needs.
        pushParam(pd, "Radius", ff->radialExtension());
    }
}

void fillParticleData(const Particle* p, const R3& base_pos, const RotMatrix& base_rot,
                      double opacity, ParticleData& pd)
{
    const Formfactor* ff = p->pFormfactor();
    ASSERT(ff);

    pd.shape = ff->className();
    fillSmoothShapeParams(ff, pd);

    // Rotate the particle's local offset by accumulated parent rotation, then translate.
    const R3 lpos = base_rot.transformed(p->particlePosition());
    pd.pos_x = lpos.x() + base_pos.x();
    pd.pos_y = lpos.y() + base_pos.y();
    pd.pos_z = lpos.z() + base_pos.z();

    // Combined rotation = base_rot * particle's own rotation.
    RotMatrix total = base_rot;
    if (const Rotation* rot = p->particleRotation())
        total = total * rot->rotMatrix();
    if (!total.isIdentity()) {
        const auto angles = total.zxzEulerAngles();
        pd.euler_alpha = angles[0];
        pd.euler_beta = angles[1];
        pd.euler_gamma = angles[2];
    }

    pd.material = p->material()->materialName();
    pd.color = ::colorRgba(p->material(), opacity);

    fillPolyhedronData(ff, pd);
}

void collectParticles(const IParticle* ipart, const R3& base_pos, const RotMatrix& base_rot,
                      double opacity, CollectContext& ctx, std::vector<ParticleData>& out);

void collectFromParticle(const Particle* p, const R3& base_pos, const RotMatrix& base_rot,
                         double opacity, CollectContext& ctx, std::vector<ParticleData>& out)
{
    out.emplace_back();
    ParticleData& pd = out.back();
    fillParticleData(p, base_pos, base_rot, opacity, pd);
    pd.mixture_groups.reserve(ctx.mixture_path.size());
    pd.mixture_variants.reserve(ctx.mixture_path.size());
    for (const auto& choice : ctx.mixture_path) {
        pd.mixture_groups.push_back(choice.first);
        pd.mixture_variants.push_back(choice.second);
    }
}

//! Combines a child node's own rotation with the accumulated parent rotation,
//! and rotates the child's local-frame position into the parent frame.
void accumulate(const IParticle* node, const R3& base_pos, const RotMatrix& base_rot, R3& out_pos,
                RotMatrix& out_rot)
{
    out_pos = base_rot.transformed(node->particlePosition()) + base_pos;
    out_rot = base_rot;
    if (const Rotation* rot = node->particleRotation())
        out_rot = out_rot * rot->rotMatrix();
}

void collectFromCompound(const Compound* compound, const R3& base_pos, const RotMatrix& base_rot,
                         double opacity, CollectContext& ctx, std::vector<ParticleData>& out)
{
    R3 cpos;
    RotMatrix crot;
    ::accumulate(compound, base_pos, base_rot, cpos, crot);
    for (const auto* child : compound->components())
        collectParticles(child, cpos, crot, opacity, ctx, out);
}

void collectFromCoreAndShell(const CoreAndShell* cas, const R3& base_pos, const RotMatrix& base_rot,
                             double opacity, CollectContext& ctx, std::vector<ParticleData>& out)
{
    R3 cpos;
    RotMatrix crot;
    ::accumulate(cas, base_pos, base_rot, cpos, crot);
    if (const Particle* shell = cas->shellParticle())
        collectFromParticle(shell, cpos, crot, 0.5 * opacity, ctx, out);
    if (const IParticle* core = cas->coreParticle())
        collectParticles(core, cpos, crot, opacity, ctx, out);
}

void collectFromMixture(const Mixture* mix, const R3& base_pos, const RotMatrix& base_rot,
                        double opacity, CollectContext& ctx, std::vector<ParticleData>& out)
{
    R3 mpos;
    RotMatrix mrot;
    ::accumulate(mix, base_pos, base_rot, mpos, mrot);

    const std::vector<const IParticle*> children = mix->particles();
    const int gid = ctx.next_mixture_group++;
    ASSERT(static_cast<size_t>(gid) == ctx.mixture_weights.size());
    ctx.mixture_weights.push_back(mix->normalizedWeights());
    ASSERT(ctx.mixture_weights.back().size() == children.size());

    for (size_t v = 0; v < children.size(); ++v) {
        const PathFrame frame(ctx.mixture_path, gid, static_cast<int>(v));
        collectParticles(children[v], mpos, mrot, opacity, ctx, out);
    }
}

void collectFromMesocrystal(const Mesocrystal* meso, const R3& base_pos, const RotMatrix& base_rot,
                            double opacity, CollectContext& ctx, std::vector<ParticleData>& out)
{
    // Basis particles placed at each lattice point inside the outer shape.
    // calcBasisPositions() returns positions in meso-local coordinates where
    // z=0 is the bottom of the outer shape (matches Cylinder::contains convention).
    const Crystal& crystal = meso->particleStructure();
    const IParticle* basis = crystal.basis();
    if (!basis)
        return;

    // Compose the meso's own rotation with the accumulated parent rotation; lattice
    // points inside the meso are then rotated by this combined rotation.
    RotMatrix meso_rot = base_rot;
    if (const Rotation* rot = meso->particleRotation())
        meso_rot = meso_rot * rot->rotMatrix();

    // Meso reference position in world coords (rotated lateral offset, bottomZ in z).
    const R3 reg_lat = base_rot.transformed(meso->particlePosition());
    const double bottom_z = base_pos.z() + meso->bottomZ();
    const R3 meso_base(reg_lat.x() + base_pos.x(), reg_lat.y() + base_pos.y(), bottom_z);

    // A Mixture basis intentionally gets an independent group for each mesocrystal node.
    // Large lattices may therefore create many ParticleData entries in the 3D viewer.
    const std::vector<R3> positions = meso->calcBasisPositions();
    for (const R3& lpos : positions) {
        const R3 wpos = meso_rot.transformed(lpos) + meso_base;
        collectParticles(basis, wpos, meso_rot, opacity, ctx, out);
    }
}

void collectParticles(const IParticle* ipart, const R3& base_pos, const RotMatrix& base_rot,
                      double opacity, CollectContext& ctx, std::vector<ParticleData>& out)
{
    if (const auto* p = dynamic_cast<const Particle*>(ipart))
        collectFromParticle(p, base_pos, base_rot, opacity, ctx, out);
    else if (const auto* c = dynamic_cast<const Compound*>(ipart))
        collectFromCompound(c, base_pos, base_rot, opacity, ctx, out);
    else if (const auto* cs = dynamic_cast<const CoreAndShell*>(ipart))
        collectFromCoreAndShell(cs, base_pos, base_rot, opacity, ctx, out);
    else if (const auto* m = dynamic_cast<const Mixture*>(ipart))
        collectFromMixture(m, base_pos, base_rot, opacity, ctx, out);
    else if (const auto* mc = dynamic_cast<const Mesocrystal*>(ipart))
        collectFromMesocrystal(mc, base_pos, base_rot, opacity, ctx, out);
}


} // namespace

void SceneGeometryDetail::collectParticleTemplates(const IParticle* ipart, const R3& base_pos,
                                                   const RotMatrix& base_rot, double opacity,
                                                   StructData& sd)
{
    CollectContext ctx(sd.mixture_weights);
    ::collectParticles(ipart, base_pos, base_rot, opacity, ctx, sd.particles);
}
