#include "Resample/Sample/ReSample.h"

#include "Sample/Aggregate/Crystal2D.h"
#include "Sample/Aggregate/Dilute2D.h"
#include "Sample/HardParticle/Cylinder.h"
#include "Sample/HardParticle/Polyhedra.h"
#include "Sample/Material/Admixtures.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"

TEST(MultilayerAveraging, AverageMultilayer)
{
    const auto mat0 = RefractiveMaterial("vac", MaterialColor(0.9, 0.93, 0.97), 0.0, 0.0);
    const auto mat1 = RefractiveMaterial("stone", MaterialColor(0.4, 0.4, 0.4), 4e-4, 8e-7);

    // particles
    Cylinder cylinder_ff(1.0, 3.0);
    Particle particle(mat1, cylinder_ff);

    // interferences
    Crystal2D interf_1(particle, BasicLattice2D(10.0, 10.0, 120.0, 0.0));
    Crystal2D interf_2(particle, BasicLattice2D(10.0, 10.0, 120.0, 0.0));

    Layer layer1(mat0);
    Layer layer2(mat1);

    layer1.deposit2D(interf_1);

    Sample sample;
    sample.addLayer(layer1);
    sample.addLayer(layer2);

    SimulationOptions options;
    options.setUseAvgMaterials(false);

    const auto re_sample1 = ReSample::processedSample(sample, options);

    Layer layer21(mat0);
    Layer layer22(mat1);

    layer21.deposit2D(interf_1);
    layer21.deposit2D(interf_2);

    Sample sample2;
    sample2.addLayer(layer21);
    sample2.addLayer(layer22);

    const auto re_sample2 = ReSample::processedSample(sample2, options);

    EXPECT_EQ(re_sample1.numberOfSlices(), re_sample2.numberOfSlices());

    for (size_t i = 0; i < re_sample1.numberOfSlices(); ++i) {
        auto mat_1 = re_sample1.avgeSlice(i).refractingMaterial()->rawData();
        auto mat_2 = re_sample2.avgeSlice(i).refractingMaterial()->rawData();
        EXPECT_DOUBLE_EQ(mat_1.real(), mat_2.real());
        EXPECT_DOUBLE_EQ(mat_1.imag(), mat_2.imag());
    }
}

// Sliced material averaging must collect several layout contributions
// before computing the effective material of each slice.
TEST(MultilayerAveraging, SlicedSeveralLayoutsAreAveragedTogether)
{
    const auto vacuum = RefractiveMaterial("vac", MaterialColor(0.9, 0.93, 0.97), 0.0, 0.0);
    const auto matrix_mat = RefractiveMaterial("matrix", MaterialColor(0.4, 0.4, 0.4), 1e-6, 2e-8);
    const auto particle_mat_1 =
        RefractiveMaterial("particle1", MaterialColor(0.86, 0.24, 0.18), 4e-6, 8e-8);
    const auto particle_mat_2 =
        RefractiveMaterial("particle2", MaterialColor(0.05, 0.62, 0.55), 9e-6, 1e-7);

    const double footprint_area = 1.0;
    const double slice_thickness = 2.0;
    const double density_1 = 0.1;
    const double density_2 = 0.2;

    const Particle tall_particle(particle_mat_1, Box(1.0, 1.0, 4.0));
    const Particle short_particle(particle_mat_2, Box(1.0, 1.0, 2.0));

    Layer fronting(vacuum);
    Layer particle_layer(matrix_mat, 10.0);
    particle_layer.setNumberOfSlices(2);
    particle_layer.deposit2D(Dilute2D(density_1, tall_particle));
    particle_layer.deposit2D(Dilute2D(density_2, short_particle));
    Layer substrate(vacuum);

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

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

    ASSERT_EQ(re_sample.structs().size(), 2u);
    ASSERT_EQ(re_sample.numberOfSlices(), 5u);
    EXPECT_DOUBLE_EQ(re_sample.avgeSlice(1).thickness(), 6.0);
    EXPECT_DOUBLE_EQ(re_sample.avgeSlice(2).thickness(), slice_thickness);
    EXPECT_DOUBLE_EQ(re_sample.avgeSlice(3).thickness(), slice_thickness);

    const auto top_expected = MixUtil::basePlusAdmixtures(matrix_mat, 6.0, {});
    const auto upper_expected = MixUtil::basePlusAdmixtures(
        matrix_mat, slice_thickness,
        {OneAdmixture(density_1 * footprint_area * slice_thickness, particle_mat_1)});
    const auto lower_expected = MixUtil::basePlusAdmixtures(
        matrix_mat, slice_thickness,
        {OneAdmixture(density_1 * footprint_area * slice_thickness, particle_mat_1),
         OneAdmixture(density_2 * footprint_area * slice_thickness, particle_mat_2)});

    const auto top_data = re_sample.avgeSlice(1).refractingMaterial()->rawData();
    const auto upper_data = re_sample.avgeSlice(2).refractingMaterial()->rawData();
    const auto lower_data = re_sample.avgeSlice(3).refractingMaterial()->rawData();
    const auto top_expected_data = top_expected->rawData();
    const auto upper_expected_data = upper_expected->rawData();
    const auto lower_expected_data = lower_expected->rawData();

    EXPECT_NEAR(top_data.real(), top_expected_data.real(), 1e-15);
    EXPECT_NEAR(top_data.imag(), top_expected_data.imag(), 1e-15);
    EXPECT_NEAR(upper_data.real(), upper_expected_data.real(), 1e-15);
    EXPECT_NEAR(upper_data.imag(), upper_expected_data.imag(), 1e-15);
    EXPECT_NEAR(lower_data.real(), lower_expected_data.real(), 1e-15);
    EXPECT_NEAR(lower_data.imag(), lower_expected_data.imag(), 1e-15);
}
