#include "Base/Axis/Bin.h"
#include "Base/Axis/Pixel.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/Particle.h"
#include "Sample/Sample/Sample.h"
#include "Tests/GTestWrapper/google_test.h"
#include <array>
#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);
}

} // 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);
}

// 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());
    }
}
