//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Resample/Refine/ShareFormfactors.cpp
//! @brief     Implements function Refine::shareFormfactors.
//!
//! @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 "Resample/Refine/ShareFormfactors.h"
#include "Base/Util/Assert.h"
#include "Base/Vector/RotMatrix.h"
#include "Sample/HardParticle/HardParticles.h"
#include "Sample/Particle/FFMemo.h"
#include "Sample/Particle/IParticle.h"
#include "Sample/Particle/Particle.h"
#include "Sample/SoftParticle/SoftParticles.h"
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

namespace {

//! Returns the geometry-defining parameters of a supported formfactor type, or nullopt for
//! unsupported types, which are then never grouped (no optimization, no risk).
//! Uses dynamic_cast to dispatch, like doSlice in ParticlesInSlice.cpp, so that the per-type
//! knowledge stays in one place. Where types are related by inheritance, the more derived
//! type must be tested first.
std::optional<std::vector<double>> geometryParameters(const Formfactor* ff)
{
    // Hard formfactors:
    if (const auto* f = dynamic_cast<const Box*>(ff))
        return {{f->length(), f->width(), f->height()}};
    if (const auto* f = dynamic_cast<const Prism3*>(ff))
        return {{f->baseEdge(), f->height()}};
    if (const auto* f = dynamic_cast<const Prism6*>(ff))
        return {{f->baseEdge(), f->height()}};
    if (const auto* f = dynamic_cast<const PlatonicTetrahedron*>(ff))
        return {{f->edge()}};
    if (const auto* f = dynamic_cast<const PlatonicOctahedron*>(ff))
        return {{f->edge()}};
    if (const auto* f = dynamic_cast<const Dodecahedron*>(ff))
        return {{f->edge()}};
    if (const auto* f = dynamic_cast<const Icosahedron*>(ff))
        return {{f->edge()}};
    if (const auto* f = dynamic_cast<const Pyramid2*>(ff))
        return {{f->length(), f->width(), f->height(), f->alpha()}};
    if (const auto* f = dynamic_cast<const Pyramid3*>(ff))
        return {{f->baseEdge(), f->height(), f->alpha()}};
    if (const auto* f = dynamic_cast<const Pyramid4*>(ff))
        return {{f->baseEdge(), f->height(), f->alpha()}};
    if (const auto* f = dynamic_cast<const Pyramid6*>(ff))
        return {{f->baseEdge(), f->height(), f->alpha()}};
    if (const auto* f = dynamic_cast<const Bipyramid4*>(ff))
        return {{f->length(), f->base_height(), f->heightRatio(), f->alpha()}};
    if (const auto* f = dynamic_cast<const CantellatedCube*>(ff))
        return {{f->length(), f->removedLength()}};
    if (const auto* f = dynamic_cast<const TruncatedCube*>(ff))
        return {{f->length(), f->removedLength()}};
    if (const auto* f = dynamic_cast<const Cone*>(ff))
        return {{f->radius(), f->height(), f->alpha()}};
    if (const auto* f = dynamic_cast<const Cylinder*>(ff))
        return {{f->radius(), f->height()}};
    if (const auto* f = dynamic_cast<const EllipsoidalCylinder*>(ff))
        return {{f->radiusX(), f->radiusY(), f->height()}};
    if (const auto* f = dynamic_cast<const Ellipsoid*>(ff))
        return {{f->radiusX(), f->radiusY(), f->radiusZ()}};
    if (const auto* f = dynamic_cast<const EllipsoidalSegment*>(ff))
        return {{f->radiusX(), f->radiusY(), f->radiusZ(), f->cutFromTop(), f->cutFromBottom()}};
    if (const auto* f = dynamic_cast<const HorizontalCylinder*>(ff))
        return {{f->radius(), f->length(), f->cutFromTop(), f->cutFromBottom()}};
    if (const auto* f = dynamic_cast<const LongBoxGauss*>(ff))
        return {{f->length(), f->width(), f->height()}};
    if (const auto* f = dynamic_cast<const LongBoxLorentz*>(ff))
        return {{f->length(), f->width(), f->height()}};
    if (const auto* f = dynamic_cast<const Sphere*>(ff))
        return {{f->radius()}};
    if (const auto* f = dynamic_cast<const SphericalSegment*>(ff))
        return {{f->radius(), f->cutFromTop(), f->cutFromBottom()}};
    if (const auto* f = dynamic_cast<const Spheroid*>(ff))
        return {{f->radiusXY(), f->radiusZ()}};
    if (const auto* f = dynamic_cast<const SpheroidalSegment*>(ff))
        return {{f->radiusXY(), f->radiusZ(), f->cutFromTop(), f->cutFromBottom()}};
    if (const auto* f = dynamic_cast<const SawtoothRipple*>(ff)) // before base ProfileRipple
        return {{f->length(), f->width(), f->height(), f->asymmetry()}};
    if (const auto* f = dynamic_cast<const ProfileRipple*>(ff)) // covers CosineRipple*, Bar*
        return {{f->length(), f->width(), f->height()}};

    // Soft formfactors:
    if (const auto* f = dynamic_cast<const FuzzySphere*>(ff))
        return {{f->radius(), f->fuzziness()}};
    if (const auto* f = dynamic_cast<const GaussSphere*>(ff))
        return {{f->meanRadius()}};

    return std::nullopt;
}

} // namespace


void Refine::shareFormfactors(IParticle& root)
{
    std::vector<Particle*> innermost;
    root.collectInnermostParticles(innermost);

    // Group by exact formfactor identity (same concrete type and parameter values, no
    // tolerance) and by rotation: the cache key is the back-rotated scattering vector, so
    // only equally rotated particles can reuse each other's entries. This also keeps
    // orientation distributions (Mixture of rotated copies) free of useless memos.
    // Material and position may differ within a group; they act outside the memo.
    std::map<std::pair<std::string, std::vector<double>>, std::vector<Particle*>> groups;
    for (Particle* particle : innermost) {
        ASSERT(particle->pFormfactor());
        if (auto key_values = ::geometryParameters(particle->pFormfactor())) {
            const RotMatrix* rot = particle->rotMatrix();
            if (rot && !rot->isIdentity()) // no rotation counts as identity
                // The images of the basis vectors identify the rotation exactly; unlike
                // Euler-angle extraction, they cannot produce NaN, which would break the
                // ordering of the map key.
                for (const R3 axis : {R3(1, 0, 0), R3(0, 1, 0), R3(0, 0, 1)}) {
                    const R3 image = rot->transformed(axis);
                    key_values->insert(key_values->end(), {image.x(), image.y(), image.z()});
                }
            groups[{particle->pFormfactor()->className(), *key_values}].push_back(particle);
        }
    }

    for (const auto& [key, group] : groups) {
        if (group.size() < 2)
            continue;
        // Working set per pixel: up to 16 DWBA channels per distinct slice among the
        // members (the rotation is shared), bounded by group.size().
        const auto memo = std::make_shared<const FFMemo>(16 * group.size());
        for (Particle* particle : group)
            particle->setFFMemo(memo);
    }
}
