#include "Base/Axis/Bin.h"
#include "Base/Axis/Pixel.h"
#include "Base/Const/Units.h"
#include "Base/Spin/SpinMatrix.h"
#include "Resample/Sample/ReSample.h"
#include "Sample/Aggregate/Dilute2D.h"
#include "Sample/Aggregate/Dilute3D.h"
#include "Sample/Aggregate/Struct3D.h"
#include "Sample/Element/DiffuseElement.h"
#include "Sample/Flux/Flux.h"
#include "Sample/HardParticle/HardParticles.h"
#include "Sample/Material/RefractiveMaterial.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Option/SimulationOptions.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 <array>
#include <cmath>
#include <stdexcept>

namespace {

ReSample makeDilute3DFillSample(const R3& offset, double layer_thickness = 10.0,
                                double particle_height = 2.5)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial layer_mat("Layer", MaterialColor(0.5, 0.5, 0.5), 2e-6, 1e-7);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 2e-5, 4e-7);

    Particle particle(particle_mat, Box(1.3, 0.8, particle_height));
    particle.translate(offset);

    Layer fronting(vacuum);
    Layer layer(layer_mat, layer_thickness);
    layer.fill3D(Dilute3D(1e-3, particle));
    Layer substrate(vacuum);

    Sample sample;
    sample.addLayer(fronting);
    sample.addLayer(layer);
    sample.addLayer(substrate);

    return ReSample::processedSample(sample, {});
}

double dilute3DFillScattering(const ReSample& re_sample)
{
    if (re_sample.structs().size() != 1)
        throw std::runtime_error("Expected a single 3D structure");

    const auto* structure = dynamic_cast<const Struct3D*>(re_sample.structs().front().get());
    if (!structure)
        throw std::runtime_error("Expected Struct3D after 3D fill");

    Pixel pixel(Bin1D::At(-0.21), Bin1D::At(0.29));
    DiffuseElement element(1.54, 0.17, 0.11, &pixel, SpinMatrix::One() / 2.0, SpinMatrix::One());
    const Fluxes fluxes_in = re_sample.fluxesIn(element.getKi(), false);
    const Fluxes fluxes_out = re_sample.fluxesOut(element.meanKf(), false);
    element.setFluxes(&fluxes_in, &fluxes_out);

    return structure->volumeNumberDensity() * structure->thickness() * structure->coverage()
           * structure->scalarScattering(element);
}

//! As makeDilute3DFillSample, but the reference particle is wrapped in a one-component Mixture.
ReSample makeDilute3DFillMixtureSample(const R3& offset)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial layer_mat("Layer", MaterialColor(0.5, 0.5, 0.5), 2e-6, 1e-7);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 2e-5, 4e-7);

    Particle particle(particle_mat, Box(1.3, 0.8, 2.5));
    particle.translate(offset);
    Mixture mixture;
    mixture.addParticle(particle, 1.0);

    Layer fronting(vacuum);
    Layer layer(layer_mat, 10.0);
    layer.fill3D(Dilute3D(1e-3, mixture));
    Layer substrate(vacuum);

    Sample sample;
    sample.addLayer(fronting);
    sample.addLayer(layer);
    sample.addLayer(substrate);

    return ReSample::processedSample(sample, {});
}

} // namespace

// A 3D structure must live in exactly one host slice. If material averaging forces the host layer
// to be subdivided (e.g. due to 2D particles crossing interfaces), we must reject the sample.
TEST(Struct3DLayer, RejectsSlicedHostLayer)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial layer_mat("Layer", MaterialColor(0.5, 0.5, 0.5), 2e-6, 1e-7);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 2e-5, 4e-7);

    const Particle bulk_particle(particle_mat, Box(1.0, 1.0, 1.0));
    const Particle interface_particle(particle_mat, Sphere(1.0));

    Layer fronting(vacuum);
    Layer layer(layer_mat, 10.0);
    layer.fill3D(Dilute3D(1e-3, bulk_particle));
    layer.deposit2D(Dilute2D(1e-2, interface_particle));
    Layer substrate(vacuum);

    Sample sample;
    sample.addLayer(fronting);
    sample.addLayer(layer);
    sample.addLayer(substrate);

    SimulationOptions options;
    options.setUseAvgMaterials(true);

    try {
        const auto re_sample = ReSample::processedSample(sample, options);
        (void)re_sample;
        FAIL() << "Expected sliced 3D host layer to be rejected";
    } catch (const std::runtime_error& err) {
        EXPECT_STREQ("3D structures require a host layer represented by a single slice",
                     err.what());
    }
}

// Pre-deposit translation of the particle must not affect fill3D scattering.
TEST(Struct3DLayer, FillIsInvariantUnderReferenceParticleTranslation)
{
    const std::array<R3, 3> offsets{R3(0.0, 0.0, 0.0), R3(1.7, -0.9, 3.4), R3(-4.2, 2.3, -5.1)};

    const double baseline = dilute3DFillScattering(makeDilute3DFillSample(offsets[0]));
    const double shifted_1 = dilute3DFillScattering(makeDilute3DFillSample(offsets[1]));
    const double shifted_2 = dilute3DFillScattering(makeDilute3DFillSample(offsets[2]));
    const double tol = std::abs(baseline) * 1e-12 + 1e-14;

    EXPECT_NEAR(shifted_1, baseline, tol);
    EXPECT_NEAR(shifted_2, baseline, tol);
}

// Wrapping the reference particle in a one-component Mixture must not disturb fill3D's reference
// normalization: translating the alternative's reference still leaves the scattering invariant,
// as for a bare particle. Guards the 3D path against the 2D anchoring rule leaking in.
TEST(Struct3DLayer, FillMixtureIsInvariantUnderReferenceParticleTranslation)
{
    const std::array<R3, 3> offsets{R3(0.0, 0.0, 0.0), R3(1.7, -0.9, 3.4), R3(-4.2, 2.3, -5.1)};

    const double baseline = dilute3DFillScattering(makeDilute3DFillMixtureSample(offsets[0]));
    const double shifted_1 = dilute3DFillScattering(makeDilute3DFillMixtureSample(offsets[1]));
    const double shifted_2 = dilute3DFillScattering(makeDilute3DFillMixtureSample(offsets[2]));
    const double tol = std::abs(baseline) * 1e-12 + 1e-14;

    EXPECT_NEAR(shifted_1, baseline, tol);
    EXPECT_NEAR(shifted_2, baseline, tol);
}

//! Verifies that 3D-structure preparation uses the Mixture transform as the local frame
//! for each alternative.
TEST(Struct3DLayer, Fill3DRotatedMixtureUsesLocalFrameForAlternatives)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial layer_mat("Layer", MaterialColor(0.5, 0.5, 0.5), 2e-6, 1e-7);
    const RefractiveMaterial mat_a("A", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.0);
    const RefractiveMaterial mat_b("B", MaterialColor(0.5, 0.5, 0.5), 2e-5, 0.0);

    Particle particle_a(mat_a, Box(1.0, 1.0, 1.0));
    particle_a.translate(R3(10.0, 0.0, 0.0));
    Particle particle_b(mat_b, Box(1.0, 1.0, 1.0));
    particle_b.translate(R3(0.0, 10.0, 0.0));

    Mixture mixture;
    mixture.addParticle(particle_a, 1.0);
    mixture.addParticle(particle_b, 1.0);
    mixture.rotate(RotationZ(90.0 * Units::deg));
    mixture.translate(R3(1.0, 2.0, 0.0));

    Layer fronting(vacuum);
    Layer layer(layer_mat, 30.0);
    layer.fill3D(Dilute3D(1e-3, mixture));
    Layer substrate(vacuum);

    Sample sample;
    sample.addLayer(fronting);
    sample.addLayer(layer);
    sample.addLayer(substrate);

    const ReSample re_sample = ReSample::processedSample(sample, {});
    ASSERT_EQ(re_sample.structs().size(), 1u);
    const auto* structure = dynamic_cast<const Struct3D*>(re_sample.structs().front().get());
    ASSERT_TRUE(structure);
    const auto* prepared = dynamic_cast<const Mixture*>(structure->particle());
    ASSERT_TRUE(prepared);

    bool seen_a = false;
    bool seen_b = false;
    for (const IParticle* component : prepared->particles()) {
        const auto* p = dynamic_cast<const Particle*>(component);
        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.0, 1e-12);
            EXPECT_NEAR(pos.y(), 12.0, 1e-12);
        } else if (std::abs(delta - 2e-5) < 1e-12) {
            seen_b = true;
            EXPECT_NEAR(pos.x(), -9.0, 1e-12);
            EXPECT_NEAR(pos.y(), 2.0, 1e-12);
        } else {
            FAIL() << "unexpected component material, delta=" << delta;
        }

        ASSERT_TRUE(p->particleRotation());
        const R3 probe = p->particleRotation()->transformed(R3(1.0, 0.0, 0.0));
        EXPECT_NEAR(probe.x(), 0.0, 1e-12);
        EXPECT_NEAR(probe.y(), 1.0, 1e-12);
        EXPECT_NEAR(probe.z(), 0.0, 1e-12);
    }
    EXPECT_TRUE(seen_a);
    EXPECT_TRUE(seen_b);
}

//! fill3D centers each alternative for the volume distribution, which must use the 3D centering,
//! not the 2D interface anchoring -- even under a non-z rotation. A one-component Mixture rotated
//! by RotationY(90) must resolve to the same z-range and centered placement as the bare rotated
//! particle filled directly (the RotationZ case above leaves the z-extent unchanged, so it cannot
//! catch the 3D path borrowing the 2D anchoring rule).
TEST(Struct3DLayer, TiltedOneComponentMixtureFillsLikeBareParticle)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial layer_mat("Layer", MaterialColor(0.5, 0.5, 0.5), 2e-6, 1e-7);
    const RefractiveMaterial mat("A", MaterialColor(0.5, 0.5, 0.5), 1e-5, 0.0);

    auto fill = [&](const IParticle& particle) {
        Layer fronting(vacuum);
        Layer layer(layer_mat, 30.0);
        layer.fill3D(Dilute3D(1e-3, particle));
        Sample sample;
        sample.addLayer(fronting);
        sample.addLayer(layer);
        sample.addLayer(Layer(vacuum));
        return ReSample::processedSample(sample, {});
    };

    Particle bare(mat, Box(2.0, 2.0, 6.0));
    bare.rotate(RotationY(90.0 * Units::deg));
    const ReSample rb = fill(bare);
    const auto* sb = dynamic_cast<const Struct3D*>(rb.structs().front().get());
    ASSERT_TRUE(sb);

    Particle inner(mat, Box(2.0, 2.0, 6.0));
    Mixture mix;
    mix.addParticle(inner, 1.0);
    mix.rotate(RotationY(90.0 * Units::deg));
    const ReSample rm = fill(mix);
    const auto* sm = dynamic_cast<const Struct3D*>(rm.structs().front().get());
    ASSERT_TRUE(sm);

    EXPECT_NEAR(sm->zRangeLow(), sb->zRangeLow(), 1e-12);
    EXPECT_NEAR(sm->zRangeHigh(), sb->zRangeHigh(), 1e-12);
    EXPECT_NEAR(sm->particle()->centerZ(), sb->particle()->centerZ(), 1e-12);
}

// Reject particles taller than the layer.
TEST(Struct3DLayer, RejectsParticleThatDoesNotFitIntoHostLayer)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial layer_mat("Layer", MaterialColor(0.5, 0.5, 0.5), 2e-6, 1e-7);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 2e-5, 4e-7);

    const Particle oversized_particle(particle_mat, Box(1.0, 1.0, 12.0));

    Layer fronting(vacuum);
    Layer layer(layer_mat, 10.0);
    layer.fill3D(Dilute3D(1e-3, oversized_particle));
    Layer substrate(vacuum);

    Sample sample;
    sample.addLayer(fronting);
    sample.addLayer(layer);
    sample.addLayer(substrate);

    try {
        const auto re_sample = ReSample::processedSample(sample, {});
        (void)re_sample;
        FAIL() << "Expected oversized 3D particle to be rejected";
    } catch (const std::runtime_error& err) {
        EXPECT_STREQ("Particle in 3D assembly does not fit within the layer thickness", err.what());
    }
}
