//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Tests/Unit/Resample/ShareFormfactorsTest.cpp
//! @brief     Tests for FFMemo and 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/Vector/WavevectorInfo.h"
#include "Resample/Sample/ReSample.h"
#include "Sample/Aggregate/Dilute2D.h"
#include "Sample/HardParticle/Cylinder.h"
#include "Sample/HardParticle/Polyhedra.h"
#include "Sample/HardParticle/Sphere.h"
#include "Sample/Lattice/Lattice3D.h"
#include "Sample/Material/RefractiveMaterial.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Option/SimulationOptions.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/Crystal.h"
#include "Sample/Particle/FFMemo.h"
#include "Sample/Particle/Mesocrystal.h"
#include "Sample/Particle/Mixture.h"
#include "Sample/Particle/Particle.h"
#include "Sample/Particle/Rotations.h"
#include "Sample/Sample/Sample.h"
#include "Sample/SoftParticle/FuzzySphere.h"
#include "Tests/GTestWrapper/google_test.h"

namespace {

const RefractiveMaterial vac("vac", MaterialColor(0.9, 0.93, 0.97), 0.0, 0.0);
const RefractiveMaterial stone("stone", MaterialColor(0.4, 0.4, 0.4), 4e-4, 8e-7);
const RefractiveMaterial gold("gold", MaterialColor(0.8, 0.7, 0.2), 6e-4, 2e-6);

WavevectorInfo makeWavevectors(double qx, double qy, double qz)
{
    const C3 ki(0.1, 0.0, -0.02);
    return {ki, ki - C3(qx, qy, qz), 0.1};
}

} // namespace

//! Checks that FFMemo::theFF reproduces the direct formfactor evaluation exactly,
//! on first use, on reuse, and after eviction by other wavevectors.
TEST(ShareFormfactors, MemoReproducesDirectEvaluation)
{
    const Pyramid4 ff(10.0, 5.0, 1.2);
    const FFMemo memo(2);

    const WavevectorInfo wv1 = ::makeWavevectors(0.3, 0.1, 1.1);
    const complex_t direct1 = ff.theFF(wv1);
    EXPECT_EQ(memo.theFF(wv1, ff), direct1); // first evaluation (miss)
    EXPECT_EQ(memo.theFF(wv1, ff), direct1); // memoized value (hit)

    // Fill the memo with other wavevectors, then recheck the first one.
    for (int i = 1; i <= 100; ++i) {
        const WavevectorInfo wv = ::makeWavevectors(0.01 * i, 0.2, 0.9);
        EXPECT_EQ(memo.theFF(wv, ff), ff.theFF(wv));
    }
    EXPECT_EQ(memo.theFF(wv1, ff), direct1);
}

//! Checks that innermost particles are grouped by formfactor geometry and rotation:
//! position and material may differ within a group, while different geometric
//! parameters, types, or rotations separate groups, and singletons get no memo.
TEST(ShareFormfactors, GroupsByGeometryAndRotation)
{
    const Pyramid4 pyramid(10.0, 5.0, 1.2);
    const Pyramid4 other_pyramid(12.0, 5.0, 1.2);
    const Sphere sphere(5.0);

    Compound compound;
    compound.addComponent(Particle(stone, pyramid), R3(0, 0, 0));
    compound.addComponent(Particle(gold, pyramid), R3(20, 0, 0)); // other material
    Particle rotated(stone, pyramid);
    rotated.rotate(RotationZ(0.5));
    compound.addComponent(rotated, R3(40, 0, 0));
    compound.addComponent(rotated, R3(60, 0, 0)); // equally rotated copy
    compound.addComponent(Particle(stone, other_pyramid), R3(80, 0, 0));
    compound.addComponent(Particle(stone, sphere), R3(100, 0, 0));

    Refine::shareFormfactors(compound);

    std::vector<Particle*> innermost;
    compound.collectInnermostParticles(innermost);
    ASSERT_EQ(innermost.size(), 6u);

    const FFMemo* group_memo = innermost[0]->ffMemo();
    EXPECT_TRUE(group_memo);
    EXPECT_EQ(innermost[1]->ffMemo(), group_memo);
    const FFMemo* rotated_memo = innermost[2]->ffMemo();
    EXPECT_TRUE(rotated_memo);
    EXPECT_NE(rotated_memo, group_memo); // rotation separates groups
    EXPECT_EQ(innermost[3]->ffMemo(), rotated_memo);
    EXPECT_FALSE(innermost[4]->ffMemo()); // different base edge: own group of one
    EXPECT_FALSE(innermost[5]->ffMemo()); // different type: own group of one
}

//! Checks that an orientation distribution (Mixture of differently rotated copies of one
//! particle) gets no memos: each orientation has its own cache key, so a shared memo
//! could never produce a hit and would only cost overhead.
TEST(ShareFormfactors, OrientationDistributionGetsNoMemos)
{
    const Pyramid4 pyramid(10.0, 5.0, 1.2);

    Mixture mixture;
    for (int i = 0; i < 5; ++i) {
        Particle particle(stone, pyramid);
        particle.rotate(RotationZ(0.1 * (i + 1)));
        mixture.addParticle(particle, 1.0);
    }

    Refine::shareFormfactors(mixture);

    std::vector<Particle*> innermost;
    mixture.collectInnermostParticles(innermost);
    ASSERT_EQ(innermost.size(), 5u);
    for (Particle* particle : innermost)
        EXPECT_FALSE(particle->ffMemo());
}

//! Checks that soft particles are grouped like hard ones: equal radius and fuzziness
//! group together, while a different fuzziness separates.
TEST(ShareFormfactors, GroupsSoftParticles)
{
    const FuzzySphere fuzzy(5.0, 1.0);
    const FuzzySphere other_fuzzy(5.0, 2.0);

    Compound compound;
    compound.addComponent(Particle(stone, fuzzy), R3(0, 0, 0));
    compound.addComponent(Particle(stone, fuzzy), R3(20, 0, 0));
    compound.addComponent(Particle(stone, other_fuzzy), R3(40, 0, 0));

    Refine::shareFormfactors(compound);

    std::vector<Particle*> innermost;
    compound.collectInnermostParticles(innermost);
    ASSERT_EQ(innermost.size(), 3u);
    EXPECT_TRUE(innermost[0]->ffMemo());
    EXPECT_EQ(innermost[0]->ffMemo(), innermost[1]->ffMemo());
    EXPECT_FALSE(innermost[2]->ffMemo()); // different fuzziness: own group of one
}

//! Checks that innermost particles are collected across nested Compound|s, so that identical
//! particles in different subtrees end up in one group.
TEST(ShareFormfactors, GroupsAcrossNestedCompounds)
{
    const Cylinder cylinder(2.0, 4.0);

    Compound inner;
    inner.addComponent(Particle(stone, cylinder), R3(0, 0, 0));

    Compound outer;
    outer.addComponent(inner, R3(0, 10, 0));
    outer.addComponent(Particle(gold, cylinder), R3(0, 20, 0));

    Refine::shareFormfactors(outer);

    std::vector<Particle*> innermost;
    outer.collectInnermostParticles(innermost);
    ASSERT_EQ(innermost.size(), 2u);
    EXPECT_TRUE(innermost[0]->ffMemo());
    EXPECT_EQ(innermost[0]->ffMemo(), innermost[1]->ffMemo());
}

//! Checks that in reciprocal-space mode a mesocrystal attaches a memo even to a single
//! basis particle (its form factor recurs at reciprocal lattice vectors across pixels),
//! while in real-space mode a singleton basis gets none.
TEST(ShareFormfactors, MesocrystalBasisSingletonMemo)
{
    const Lattice3D lattice(R3(8, 0, 0), R3(0, 8, 0), R3(0, 0, 8));
    const Crystal crystal(Particle(stone, Sphere(2.5)), lattice);
    const Cylinder outer_shape(20, 50);

    Mesocrystal meso_recip(crystal, outer_shape);
    meso_recip.prepareForComputation(vac, {true, 2.5});
    std::vector<Particle*> innermost;
    meso_recip.collectInnermostParticles(innermost);
    ASSERT_EQ(innermost.size(), 1u);
    EXPECT_TRUE(innermost[0]->ffMemo());

    Mesocrystal meso_real(crystal, outer_shape);
    meso_real.prepareForComputation(vac, {false, 2.5});
    innermost.clear();
    meso_real.collectInnermostParticles(innermost);
    ASSERT_EQ(innermost.size(), 1u);
    EXPECT_FALSE(innermost[0]->ffMemo());
}

//! Checks that resampling (ReSample::processedSample) attaches shared memos to
//! identical innermost particles, i.e. that the grouping pass is hooked into Refine.
TEST(ShareFormfactors, AppliedDuringResampling)
{
    const Pyramid4 pyramid(10.0, 5.0, 1.2);
    const Sphere sphere(5.0);

    Compound compound;
    compound.addComponent(Particle(stone, pyramid), R3(0, 0, 0));
    compound.addComponent(Particle(stone, pyramid), R3(20, 0, 0));
    compound.addComponent(Particle(stone, sphere), R3(40, 0, 0));

    Dilute2D layout(1e-4, compound);
    Layer top_layer(vac);
    top_layer.deposit2D(layout);

    Sample sample;
    sample.addLayer(top_layer);
    sample.addLayer(Layer(stone));

    SimulationOptions options;
    const ReSample re_sample = ReSample::processedSample(sample, options);

    ASSERT_EQ(re_sample.structs().size(), 1u);
    auto* root = const_cast<IParticle*>(re_sample.structs()[0]->particle());
    std::vector<Particle*> innermost;
    root->collectInnermostParticles(innermost);
    ASSERT_EQ(innermost.size(), 3u);
    EXPECT_TRUE(innermost[0]->ffMemo());
    EXPECT_EQ(innermost[0]->ffMemo(), innermost[1]->ffMemo());
    EXPECT_FALSE(innermost[2]->ffMemo());
}
