//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Tests/Unit/Sample/ParticleFFAlignmentTest.cpp
//! @brief     Tests that particle alignment reaches the scattering amplitude.
//!
//! @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 "Base/Const/Units.h"
#include "Base/Vector/WavevectorInfo.h"
#include "Sample/HardParticle/HardParticles.h"
#include "Sample/Material/RefractiveMaterial.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/CoreAndShell.h"
#include "Sample/Particle/Formfactor.h"
#include "Sample/Particle/Mixture.h"
#include "Sample/Particle/Particle.h"
#include "Sample/Particle/Rotations.h"
#include "Tests/GTestWrapper/google_test.h"
#include <algorithm>
#include <complex>
#include <memory>
#include <vector>

namespace {

const RefractiveMaterial& materialA()
{
    static const RefractiveMaterial mat("A", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    return mat;
}

const RefractiveMaterial& materialB()
{
    static const RefractiveMaterial mat("B", MaterialColor(0.5, 0.5, 0.5), 2e-5, 0.);
    return mat;
}

std::vector<WavevectorInfo> wavevectors()
{
    return {{R3(0.10, 0.05, -0.20), R3(-0.15, 0.10, 0.30), 1.54},
            {R3(0.31, -0.22, -0.13), R3(0.04, 0.11, 0.27), 1.54},
            {R3(-0.08, 0.14, -0.19), R3(0.21, -0.07, 0.24), 1.54}};
}

void expectComplexNear(const complex_t actual, const complex_t expected, const double rel_eps)
{
    const double scale = std::max({1.0, std::abs(actual), std::abs(expected)});
    EXPECT_LE(std::abs(actual - expected), rel_eps * scale);
}

complex_t phaseFactor(const R3& position, const WavevectorInfo& wv)
{
    return std::exp(complex_t(0.0, 1.0) * position.dot(wv.getQ()));
}

complex_t expectedParticleFF(const Particle& particle, const Formfactor& formfactor,
                             const WavevectorInfo& wv)
{
    return formfactor.theFF(wv.backTransformed(particle.rotMatrix()))
           * phaseFactor(particle.particlePosition(), wv);
}

std::unique_ptr<Particle> transformedParticle(const Formfactor& formfactor, const Rotation* own_rot,
                                              const R3& own_pos, const Rotation* parent_rot,
                                              const R3& parent_pos)
{
    auto particle = std::make_unique<Particle>(materialA(), formfactor);
    particle->rotate(own_rot);
    particle->translate(own_pos);
    particle->applyParentTransform(parent_rot, parent_pos);
    return particle;
}

} // namespace


//! Verifies the affine form-factor rule for a single particle: rotation is represented by
//! back-transforming the wavevectors, while translation contributes the positive phase factor.
TEST(ParticleFFAlignment, ParticleAffineTransformCovariance)
{
    const Box formfactor(2.0, 3.0, 5.0);
    const RotationEuler rotation(0.37, 0.53, 0.19);
    const R3 translation(1.2, -0.7, 2.3);

    Particle particle(materialA(), formfactor);
    particle.rotate(rotation);
    particle.translate(translation);

    for (const WavevectorInfo& wv : wavevectors()) {
        const complex_t actual = particle.theFF(wv);
        const complex_t expected = expectedParticleFF(particle, formfactor, wv);
        expectComplexNear(actual, expected, 1e-12);
    }
}

//! Verifies that the rigid transform of a Compound, once pushed to its leaves as in resampling,
//! gives the coherent sum of the affinely transformed component amplitudes.
TEST(ParticleFFAlignment, CompoundPreparedLeavesMatchCoherentAffineSum)
{
    const Box box(2.0, 3.0, 5.0);
    const Cylinder cylinder(1.4, 4.0);
    const RotationZ child_rot_a(23.0 * Units::deg);
    const RotationX child_rot_b(31.0 * Units::deg);
    const RotationEuler parent_rot(0.27, 0.41, 0.33);
    const R3 parent_pos(0.8, -1.1, 1.7);

    Compound compound;
    compound.addComponent(
        *transformedParticle(box, &child_rot_a, R3(1.0, 0.4, -0.3), nullptr, R3()));
    compound.addComponent(
        *transformedParticle(cylinder, &child_rot_b, R3(-0.6, 1.3, 0.5), nullptr, R3()));
    compound.rotate(parent_rot);
    compound.translate(parent_pos);

    const std::unique_ptr<Particle> expected_a =
        transformedParticle(box, &child_rot_a, R3(1.0, 0.4, -0.3), compound.particleRotation(),
                            compound.particlePosition());
    const std::unique_ptr<Particle> expected_b =
        transformedParticle(cylinder, &child_rot_b, R3(-0.6, 1.3, 0.5), compound.particleRotation(),
                            compound.particlePosition());

    for (const WavevectorInfo& wv : wavevectors()) {
        complex_t actual = 0.0;
        for (const auto& leaf : compound.decompose())
            actual += leaf->theFF(wv);

        const complex_t expected = expectedParticleFF(*expected_a, box, wv)
                                   + expectedParticleFF(*expected_b, cylinder, wv);
        expectComplexNear(actual, expected, 1e-12);
    }
}

//! Verifies the CoreAndShell amplitude in the prepared representation: parent alignment has
//! already been pushed into the shell and core leaves before their amplitudes are summed.
TEST(ParticleFFAlignment, PreparedCoreAndShellMatchesCoherentAffineSum)
{
    const Box core_ff(2.0, 2.0, 2.0);
    const Box shell_ff(4.0, 4.0, 6.0);
    const RotationY core_rot(17.0 * Units::deg);
    const RotationEuler parent_rot(0.23, 0.35, 0.47);
    const R3 parent_pos(-0.9, 0.6, 1.4);

    Particle core(materialA(), core_ff);
    core.rotate(core_rot);
    core.translate(R3(0.0, 0.0, 2.0));
    Particle shell(materialB(), shell_ff);

    CoreAndShell source(core, shell);
    source.rotate(parent_rot);
    source.translate(parent_pos);

    std::unique_ptr<IParticle> prepared_core(core.clone());
    prepared_core->applyParentTransform(source.particleRotation(), source.particlePosition());
    std::unique_ptr<Particle> prepared_shell(shell.clone());
    prepared_shell->applyParentTransform(source.particleRotation(), source.particlePosition());
    CoreAndShell prepared(*prepared_core, *prepared_shell);

    const auto* prepared_core_particle = dynamic_cast<const Particle*>(prepared_core.get());
    ASSERT_TRUE(prepared_core_particle);

    for (const WavevectorInfo& wv : wavevectors()) {
        const complex_t actual = prepared.theFF(wv);
        const complex_t expected = expectedParticleFF(*prepared_shell, shell_ff, wv)
                                   + expectedParticleFF(*prepared_core_particle, core_ff, wv);
        expectComplexNear(actual, expected, 1e-12);
    }
}

//! Verifies that a prepared Mixture amplitude is the normalized weighted sum of its affinely
//! transformed alternatives.
TEST(ParticleFFAlignment, PreparedMixtureMatchesWeightedAffineAlternatives)
{
    const Box box(2.0, 3.0, 5.0);
    const Sphere sphere(1.7);
    const RotationX child_rot_a(28.0 * Units::deg);
    const RotationZ child_rot_b(41.0 * Units::deg);
    const RotationEuler parent_rot(0.34, 0.29, 0.51);
    const R3 parent_pos(1.5, -0.8, 0.4);

    Mixture source;
    source.addParticle(*transformedParticle(box, &child_rot_a, R3(0.7, -0.2, 0.9), nullptr, R3()),
                       2.0);
    source.addParticle(
        *transformedParticle(sphere, &child_rot_b, R3(-1.1, 0.8, -0.4), nullptr, R3()), 3.0);
    source.rotate(parent_rot);
    source.translate(parent_pos);

    std::unique_ptr<Particle> expected_a =
        transformedParticle(box, &child_rot_a, R3(0.7, -0.2, 0.9), source.particleRotation(),
                            source.particlePosition());
    std::unique_ptr<Particle> expected_b =
        transformedParticle(sphere, &child_rot_b, R3(-1.1, 0.8, -0.4), source.particleRotation(),
                            source.particlePosition());

    Mixture prepared;
    prepared.addParticle(*expected_a, 2.0);
    prepared.addParticle(*expected_b, 3.0);
    prepared.validate();

    for (const WavevectorInfo& wv : wavevectors()) {
        const complex_t actual = prepared.theFF(wv);
        const complex_t expected = 0.4 * expectedParticleFF(*expected_a, box, wv)
                                   + 0.6 * expectedParticleFF(*expected_b, sphere, wv);
        expectComplexNear(actual, expected, 1e-12);
    }
}

//! Verifies that one-component containers in prepared form do not alter the particle amplitude.
TEST(ParticleFFAlignment, OneComponentContainersMatchParticleAmplitude)
{
    const Box formfactor(2.0, 3.0, 5.0);
    const RotationEuler rotation(0.37, 0.53, 0.19);
    Particle particle(materialA(), formfactor);
    particle.rotate(rotation);
    particle.translate(R3(1.2, -0.7, 2.3));

    Compound compound;
    compound.addComponent(particle);

    Mixture mixture;
    mixture.addParticle(particle, 1.0);
    mixture.validate();

    for (const WavevectorInfo& wv : wavevectors()) {
        const complex_t expected = particle.theFF(wv);
        expectComplexNear(compound.theFF(wv), expected, 1e-12);
        expectComplexNear(mixture.theFF(wv), expected, 1e-12);
    }
}
