//  ************************************************************************************************
//
//  BornAgain: simulate and fit reflection and scattering
//
//! @file      Tests/Unit/Sample/CompositeAlignmentTest.cpp
//! @brief     Tests composite-particle transform semantics for rigid composites and
//!            mixtures.
//!
//! @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/Type/Span.h"
#include "Resample/Refine/ParticlesInSlice.h"
#include "Resample/Sample/ReSample.h"
#include "Sample/Aggregate/Dilute2D.h"
#include "Sample/Aggregate/Struct.h"
#include "Sample/Aggregate/Struct2D.h"
#include "Sample/HardParticle/Polyhedra.h"
#include "Sample/HardParticle/Sphere.h"
#include "Sample/Material/RefractiveMaterial.h"
#include "Sample/Material/Vacuum.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Slice.h"
#include "Sample/Option/SimulationOptions.h"
#include "Sample/Particle/Compound.h"
#include "Sample/Particle/CoreAndShell.h"
#include "Sample/Particle/Mixture.h"
#include "Sample/Particle/Particle.h"
#include "Sample/Particle/Rotations.h"
#include "Sample/Sample/Sample.h"
#include "Tests/GTestWrapper/google_test.h"
#include <algorithm>
#include <cmath>
#include <limits>
#include <memory>
#include <vector>

namespace {

//! Returns the first sliced Particle whose refractive data real part matches delta.
const Particle* findByDelta(const std::vector<std::unique_ptr<IParticle>>& parts, double delta)
{
    for (const auto& p : parts) {
        const auto* particle = dynamic_cast<const Particle*>(p.get());
        if (particle && std::abs(particle->material()->rawData().real() - delta) < 1e-12)
            return particle;
    }
    return nullptr;
}

} // namespace

//! Verifies that a CoreAndShell applies its own rotation and translation to both core and
//! shell as a single rigid whole. After slicing, the shell (no intrinsic offset) sits at the
//! translated origin, while the core keeps its intra-particle z offset rotated along with the
//! shell. A full-span slice is used so no z-clamping perturbs the positions.
TEST(CompositeAlignment, CoreAndShellPropagatesTransformRigidly)
{
    const RefractiveMaterial core_mat("core", MaterialColor(0.5, 0.5, 0.5), 2e-5, 0.);
    const RefractiveMaterial shell_mat("shell", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const Vacuum vacuum;

    // shell occupies z in [0, 8]; core in [2, 6], i.e. offset by +2 in z within the shell
    Particle core(core_mat, Box(4., 4., 4.));
    core.translate(R3(0., 0., 2.));
    const Particle shell(shell_mat, Box(8., 8., 8.));
    CoreAndShell coreshell(core, shell);
    coreshell.rotate(RotationX(90 * Units::deg));
    coreshell.translate(R3(1., 2., 3.));

    const Layer layer(vacuum);
    const Slice slice(0, layer, Span(-100., 100.), R3());
    const Span limits(-Span::inf, Span::inf); // full span: the slice does not clamp positions

    const auto parts = Refine::particlesInSlice(&coreshell, limits, &slice, vacuum, MesoOptions{});
    ASSERT_EQ(parts.size(), 2u);

    const Particle* sliced_shell = ::findByDelta(parts, 1e-5);
    const Particle* sliced_core = ::findByDelta(parts, 2e-5);
    ASSERT_TRUE(sliced_shell);
    ASSERT_TRUE(sliced_core);

    // shell: Rx(90)*(0,0,0) + (1,2,3) = (1,2,3)
    const R3 shell_pos = sliced_shell->particlePosition();
    EXPECT_NEAR(shell_pos.x(), 1., 1e-12);
    EXPECT_NEAR(shell_pos.y(), 2., 1e-12);
    EXPECT_NEAR(shell_pos.z(), 3., 1e-12);

    // core: Rx(90)*(0,0,2) + (1,2,3) = (0,-2,0) + (1,2,3) = (1,0,3)
    const R3 core_pos = sliced_core->particlePosition();
    EXPECT_NEAR(core_pos.x(), 1., 1e-12);
    EXPECT_NEAR(core_pos.y(), 0., 1e-12);
    EXPECT_NEAR(core_pos.z(), 3., 1e-12);

    // the intra-particle offset is rotated as a rigid whole: core - shell = Rx(90)*(0,0,2)
    EXPECT_NEAR(core_pos.y() - shell_pos.y(), -2., 1e-12);

    // both core and shell carry the same orientation, Rx(90): (0,0,1) -> (0,-1,0)
    for (const Particle* p : {sliced_shell, sliced_core}) {
        ASSERT_TRUE(p->particleRotation());
        const R3 probe = p->particleRotation()->transformed(R3(0., 0., 1.));
        EXPECT_NEAR(probe.x(), 0., 1e-12);
        EXPECT_NEAR(probe.y(), -1., 1e-12);
        EXPECT_NEAR(probe.z(), 0., 1e-12);
    }
}

//! Verifies that a CoreAndShell composes a child's own rotation with its parent rotation: a
//! core that is itself rotated arrives with orientation parent-after-child, while the shell
//! carries only the parent rotation. This exercises the rotation-composition path that the
//! plain rigid-whole test (translated-only children) does not reach.
TEST(CompositeAlignment, CoreAndShellComposesChildRotation)
{
    const RefractiveMaterial core_mat("core", MaterialColor(0.5, 0.5, 0.5), 2e-5, 0.);
    const RefractiveMaterial shell_mat("shell", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const Vacuum vacuum;

    Particle core(core_mat, Box(4., 4., 4.));
    core.rotate(RotationZ(90 * Units::deg)); // child's own rotation (about z keeps z-extent)
    core.translate(R3(0., 0., 2.));          // core in z [2,6], inside shell z [0,8]
    const Particle shell(shell_mat, Box(8., 8., 8.));
    CoreAndShell coreshell(core, shell);
    coreshell.rotate(RotationX(90 * Units::deg));
    coreshell.translate(R3(1., 2., 3.));

    const Layer layer(vacuum);
    const Slice slice(0, layer, Span(-100., 100.), R3());
    const Span limits(-Span::inf, Span::inf);

    const auto parts = Refine::particlesInSlice(&coreshell, limits, &slice, vacuum, MesoOptions{});
    ASSERT_EQ(parts.size(), 2u);
    const Particle* sliced_shell = ::findByDelta(parts, 1e-5);
    const Particle* sliced_core = ::findByDelta(parts, 2e-5);
    ASSERT_TRUE(sliced_shell);
    ASSERT_TRUE(sliced_core);

    // shell: only the parent rotation Rx(90): (0,0,1) -> (0,-1,0)
    ASSERT_TRUE(sliced_shell->particleRotation());
    const R3 shell_probe = sliced_shell->particleRotation()->transformed(R3(0., 0., 1.));
    EXPECT_NEAR(shell_probe.x(), 0., 1e-12);
    EXPECT_NEAR(shell_probe.y(), -1., 1e-12);
    EXPECT_NEAR(shell_probe.z(), 0., 1e-12);

    // core: parent after child, Rx(90)*Rz(90): (1,0,0) -Rz-> (0,1,0) -Rx-> (0,0,1)
    ASSERT_TRUE(sliced_core->particleRotation());
    const R3 core_probe = sliced_core->particleRotation()->transformed(R3(1., 0., 0.));
    EXPECT_NEAR(core_probe.x(), 0., 1e-12);
    EXPECT_NEAR(core_probe.y(), 0., 1e-12);
    EXPECT_NEAR(core_probe.z(), 1., 1e-12);

    // position is unaffected by the core's own (reference-point) rotation:
    // Rx(90)*(0,0,2) + (1,2,3) = (1,0,3)
    const R3 core_pos = sliced_core->particlePosition();
    EXPECT_NEAR(core_pos.x(), 1., 1e-12);
    EXPECT_NEAR(core_pos.y(), 0., 1e-12);
    EXPECT_NEAR(core_pos.z(), 3., 1e-12);
}

//! Verifies that a Mixture transform defines the local frame for each alternative.
TEST(CompositeAlignment, MixtureUsesLocalFrameForAlternativeOffsets)
{
    const RefractiveMaterial mat_a("A", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const RefractiveMaterial mat_b("B", MaterialColor(0.5, 0.5, 0.5), 2e-5, 0.);
    const Vacuum vacuum;

    Particle a(mat_a, Sphere(2.));
    a.translate(R3(10., 0., 0.));
    Particle b(mat_b, Sphere(2.));
    b.translate(R3(0., 10., 0.));

    Mixture mixture;
    mixture.addParticle(a, 0.5);
    mixture.addParticle(b, 0.5);
    mixture.rotate(RotationZ(90 * Units::deg));
    mixture.translate(R3(1., 2., 0.));

    Layer layer(vacuum);
    layer.deposit2D(Dilute2D(1e-3, mixture));
    Sample sample;
    sample.addLayer(layer);
    sample.addLayer(Layer(vacuum));

    SimulationOptions options;
    options.setUseAvgMaterials(false);
    const auto re_sample = ReSample::processedSample(sample, options);

    const Mixture* processed = nullptr;
    for (const auto& s : re_sample.structs())
        if (const auto* m = dynamic_cast<const Mixture*>(s->particle()))
            processed = m;
    ASSERT_TRUE(processed);

    const auto components = processed->particles();
    ASSERT_EQ(components.size(), 2u);

    bool seen_a = false;
    bool seen_b = false;
    for (const IParticle* component : components) {
        // a component may be a single Particle or, if sliced, a Compound: decompose to leaves
        for (const auto& leaf : component->decompose()) {
            const auto* p = dynamic_cast<const Particle*>(leaf.get());
            ASSERT_TRUE(p);
            const R3 pos = p->particlePosition();
            const double delta = p->material()->rawData().real();

            if (std::abs(delta - 1e-5) < 1e-12) {
                seen_a = true;
                EXPECT_NEAR(pos.x(), 1., 1e-12);
                EXPECT_NEAR(pos.y(), 12., 1e-12);
            } else if (std::abs(delta - 2e-5) < 1e-12) {
                seen_b = true;
                EXPECT_NEAR(pos.x(), -9., 1e-12);
                EXPECT_NEAR(pos.y(), 2., 1e-12);
            } else {
                FAIL() << "unexpected component material, delta=" << delta;
            }

            // Each alternative carries the Mixture's z rotation: (1,0,0) -> (0,1,0).
            ASSERT_TRUE(p->particleRotation());
            const R3 probe = p->particleRotation()->transformed(R3(1., 0., 0.));
            EXPECT_NEAR(probe.x(), 0., 1e-12);
            EXPECT_NEAR(probe.y(), 1., 1e-12);
            EXPECT_NEAR(probe.z(), 0., 1e-12);
        }
    }
    EXPECT_TRUE(seen_a);
    EXPECT_TRUE(seen_b);
}

//! Verifies that a rigid parent transform composes with a Mixture local-frame transform
//! before reaching the alternative offset.
TEST(CompositeAlignment, ParentAndMixtureRotationsComposeForAlternativeOffsets)
{
    const RefractiveMaterial mat_a("A", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const Vacuum vacuum;

    Particle a(mat_a, Sphere(2.));
    a.translate(R3(10., 0., 0.));

    Mixture mixture;
    mixture.addParticle(a, 1.);
    mixture.rotate(RotationZ(90 * Units::deg));
    mixture.translate(R3(1., 0., 0.));

    Compound compound;
    compound.addComponent(mixture);
    compound.rotate(RotationZ(90 * Units::deg));

    Layer layer(vacuum);
    layer.deposit2D(Dilute2D(1e-3, compound));
    Sample sample;
    sample.addLayer(layer);
    sample.addLayer(Layer(vacuum));

    SimulationOptions options;
    options.setUseAvgMaterials(false);
    const auto re_sample = ReSample::processedSample(sample, options);

    const Compound* processed_compound = nullptr;
    for (const auto& s : re_sample.structs())
        if (const auto* c = dynamic_cast<const Compound*>(s->particle()))
            processed_compound = c;
    ASSERT_TRUE(processed_compound);

    const auto compound_components = processed_compound->components();
    ASSERT_EQ(compound_components.size(), 1u);
    const auto* processed_mixture = dynamic_cast<const Mixture*>(compound_components[0]);
    ASSERT_TRUE(processed_mixture);

    const auto mixture_components = processed_mixture->particles();
    ASSERT_EQ(mixture_components.size(), 1u);
    const auto leaves = mixture_components[0]->decompose();
    ASSERT_EQ(leaves.size(), 1u);
    const auto* p = dynamic_cast<const Particle*>(leaves[0].get());
    ASSERT_TRUE(p);

    // Parent and Mixture Rz rotations compose:
    // Rz(90) * (Rz(90) * (10,0,0) + (1,0,0)) = (-10,1,0).
    const R3 pos = p->particlePosition();
    EXPECT_NEAR(pos.x(), -10., 1e-12);
    EXPECT_NEAR(pos.y(), 1., 1e-12);
    EXPECT_NEAR(pos.z(), 0., 1e-12);

    ASSERT_TRUE(p->particleRotation());
    const R3 probe = p->particleRotation()->transformed(R3(1., 0., 0.));
    EXPECT_NEAR(probe.x(), -1., 1e-12);
    EXPECT_NEAR(probe.y(), 0., 1e-12);
    EXPECT_NEAR(probe.z(), 0., 1e-12);
}

//! Regression test for #1415: a composite's own rotation must enter its reported z-extent. A
//! Compound holding a single Box(10,2,2) has intrinsic height 2 (the box H); rotating the
//! Compound by 90 deg about y tips the 10 nm length vertical, so the reported height must
//! become 10. Before the fix the Compound ignored its own rotation and kept reporting 2.
TEST(CompositeAlignment, RotatedCompoundReportsRotatedExtent)
{
    const RefractiveMaterial mat("m", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const Particle box(mat, Box(10., 2., 2.));

    Compound upright;
    upright.addComponent(box);
    EXPECT_NEAR(upright.extensionZ(), 2., 1e-12);

    Compound tipped;
    tipped.addComponent(box);
    tipped.rotate(RotationY(90 * Units::deg));
    EXPECT_NEAR(tipped.extensionZ(), 10., 1e-12);
}

//! Regression test for #1415: same as above for a Mixture of one alternative.
TEST(CompositeAlignment, RotatedMixtureReportsRotatedExtent)
{
    const RefractiveMaterial mat("m", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const Particle box(mat, Box(10., 2., 2.));

    Mixture upright;
    upright.addParticle(box, 1.);
    EXPECT_NEAR(upright.extensionZ(), 2., 1e-12);

    Mixture tipped;
    tipped.addParticle(box, 1.);
    tipped.rotate(RotationY(90 * Units::deg));
    EXPECT_NEAR(tipped.extensionZ(), 10., 1e-12);
}

//! Regression test for #1415: same as above for a CoreAndShell, whose extent is the shell's.
TEST(CompositeAlignment, RotatedCoreShellReportsRotatedExtent)
{
    const RefractiveMaterial core_mat("core", MaterialColor(0.5, 0.5, 0.5), 2e-5, 0.);
    const RefractiveMaterial shell_mat("shell", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const Particle core(core_mat, Box(4., 1., 1.));
    const Particle shell(shell_mat, Box(10., 2., 2.));

    CoreAndShell upright(core, shell);
    EXPECT_NEAR(upright.extensionZ(), 2., 1e-12);

    CoreAndShell tipped(core, shell);
    tipped.rotate(RotationY(90 * Units::deg));
    EXPECT_NEAR(tipped.extensionZ(), 10., 1e-12);
}

//! Regression test for #1415: for a Compound, bottomZ()/topZ() must equal the min/max over the
//! z-extents of its decompose() leaves, which carry the Compound's own transform. This pins the
//! reported extent to the actually-placed (rotated) geometry for an arbitrary Euler rotation.
TEST(CompositeAlignment, CompoundExtentMatchesDecomposedLeaves)
{
    const RefractiveMaterial mat("m", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    Particle a(mat, Box(10., 2., 2.));
    a.translate(R3(0., 0., 3.));
    Particle b(mat, Box(4., 4., 8.));
    b.translate(R3(0., 5., -2.));

    Compound compound;
    compound.addComponent(a);
    compound.addComponent(b);
    compound.rotate(RotationEuler(0.3, 0.4, 0.5));
    compound.translate(R3(1., 2., 3.));

    double leaf_bottom = std::numeric_limits<double>::max();
    double leaf_top = std::numeric_limits<double>::lowest();
    for (const auto& leaf : compound.decompose()) {
        leaf_bottom = std::min(leaf_bottom, leaf->bottomZ());
        leaf_top = std::max(leaf_top, leaf->topZ());
    }
    EXPECT_NEAR(compound.bottomZ(), leaf_bottom, 1e-12);
    EXPECT_NEAR(compound.topZ(), leaf_top, 1e-12);
}

//! Regression test for #1415 at the level where the bug actually bit: deposit2D anchors a
//! particle's bottom at the layer's lower interface using intrinsicBottomZ(). A Compound holding
//! a Box(10,2,2) tipped 90 deg about y is a 10 nm tall column; once deposited it must rest with
//! its bottom on the interface (z=0) and its top at z=10. The check uses the decomposed leaf,
//! whose extent comes from the always-correct single-particle path, as an oracle independent of
//! Compound::bottomZ(). With the old rotation-blind extent the anchor shift was zero, so the
//! column straddled the interface (z in [-5,5]) instead of resting on it.
TEST(CompositeAlignment, RotatedCompoundDeposit2DAnchorsAtInterface)
{
    const RefractiveMaterial mat("m", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const Vacuum vacuum;
    const Particle box(mat, Box(10., 2., 2.));

    Compound compound;
    compound.addComponent(box);
    compound.rotate(RotationY(90 * Units::deg));

    Layer layer(vacuum);
    layer.deposit2D(Dilute2D(1e-3, compound));

    ASSERT_EQ(layer.structsAtBottom().size(), 1u);
    const IParticle* deposited = layer.structsAtBottom()[0]->particle();
    double leaf_bottom = std::numeric_limits<double>::max();
    double leaf_top = std::numeric_limits<double>::lowest();
    for (const auto& leaf : deposited->decompose()) {
        leaf_bottom = std::min(leaf_bottom, leaf->bottomZ());
        leaf_top = std::max(leaf_top, leaf->topZ());
    }
    EXPECT_NEAR(leaf_bottom, 0., 1e-12);
    EXPECT_NEAR(leaf_top, 10., 1e-12);
}

//! Counterpart of the deposit2D test for the top-anchored path: suspend2D anchors a particle's
//! top at the layer's upper interface using intrinsicTopZ(). The same 10 nm column must hang
//! with its top on the interface (z=0) and its bottom at z=-10. Before the fix the rotation-blind
//! topZ() gave the wrong anchor shift, so the column was placed too high.
TEST(CompositeAlignment, RotatedCompoundSuspend2DAnchorsAtInterface)
{
    const RefractiveMaterial mat("m", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.);
    const Vacuum vacuum;
    const Particle box(mat, Box(10., 2., 2.));

    Compound compound;
    compound.addComponent(box);
    compound.rotate(RotationY(90 * Units::deg));

    Layer layer(vacuum);
    layer.suspend2D(Dilute2D(1e-3, compound));

    ASSERT_EQ(layer.structsAtTop().size(), 1u);
    const IParticle* suspended = layer.structsAtTop()[0]->particle();
    double leaf_bottom = std::numeric_limits<double>::max();
    double leaf_top = std::numeric_limits<double>::lowest();
    for (const auto& leaf : suspended->decompose()) {
        leaf_bottom = std::min(leaf_bottom, leaf->bottomZ());
        leaf_top = std::max(leaf_top, leaf->topZ());
    }
    EXPECT_NEAR(leaf_top, 0., 1e-12);
    EXPECT_NEAR(leaf_bottom, -10., 1e-12);
}
