#include "Resample/Sample/ReSample.h"

#include "Sample/Aggregate/Dilute3D.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"

TEST(AvgMaterials, Random3DIndependentOfLayerThickness)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial matrix_mat("Matrix", 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);

    // 3D random assembly particle (must be thinner than each finite layer)
    const Particle random_particle(particle_mat, Box(1.0, 1.0, 1.0));
    const double density = 1e-3; // 1/nm^3

    Layer fronting(vacuum);

    Layer layer1(matrix_mat, 5.0);
    layer1.fill3D(Dilute3D(density, random_particle));

    Layer layer2(matrix_mat, 10.0);
    layer2.fill3D(Dilute3D(density, random_particle));

    Layer substrate(vacuum);

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

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

    // Slice indices: 0 fronting, 1 layer1, 2 layer2, 3 substrate.
    const auto mat_1 = re_sample.avgeSlice(1).refractingMaterial()->rawData();
    const auto mat_2 = re_sample.avgeSlice(2).refractingMaterial()->rawData();

    EXPECT_NEAR(mat_1.real(), mat_2.real(), 1e-14);
    EXPECT_NEAR(mat_1.imag(), mat_2.imag(), 1e-14);
}

TEST(AvgMaterials, Random3DCorrectnessVolumeWeighting)
{
    // Test that material averaging for 3D structures correctly computes volume-weighted average
    // Formula: averaged = f_particle * particle_mat + (1 - f_particle) * layer_mat
    // where f_particle = density * particle_volume

    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), 1.0e-6, 5.0e-8);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 4.0e-6, 2.0e-7);

    // Create a particle with known volume
    const double box_size = 2.0; // nm
    const Particle random_particle(particle_mat, Box(box_size, box_size, box_size));
    const double particle_volume = box_size * box_size * box_size; // 8.0 nm^3

    // Set density to get a specific volume fraction
    const double density = 0.0125;                            // 1/nm^3
    const double volume_fraction = density * particle_volume; // 0.1 (10%)

    // Create sample with layers
    const double layer_thickness = 10.0; // nm (must be > box_size)
    Layer fronting(vacuum);
    Layer layer(layer_mat, layer_thickness);
    layer.fill3D(Dilute3D(density, random_particle));
    Layer substrate(vacuum);

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

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

    // Get the averaged material from the layer slice (slice 1, not 0)
    const auto averaged_mat = re_sample.avgeSlice(1).refractingMaterial()->rawData();

    // Calculate expected averaged material properties
    // For refractive materials, we average the "averageableData" which is related to n^2-1
    // But the rawData() returns delta + i*beta where n = 1 - delta + i*beta
    // The averaging is done in the "averageableData" space: (1-n)^2 - 1 = -2*delta + delta^2 + ...
    // For small delta, beta, this simplifies to approximately -2*delta + i*2*beta
    // But the actual averaging uses the correct formula, so we need to use averageableData

    // Convert materials to averageableData
    const auto layer_avge_data = layer_mat.averageableData();
    const auto particle_avge_data = particle_mat.averageableData();

    // Compute weighted average in averageableData space
    const auto expected_avge_data =
        volume_fraction * particle_avge_data + (1.0 - volume_fraction) * layer_avge_data;

    // Convert back to delta, beta (rawData)
    // averageableData = (1-n)^2 - 1, so n = 1 - sqrt(1 + averageableData)
    // rawData = conj(n - 1) = conj(-sqrt(1 + averageableData))
    const complex_t expected_mat = std::conj(1.0 - std::sqrt(1.0 + expected_avge_data));

    // Compare with tolerance
    const double rel_tol = 1e-12;
    EXPECT_NEAR(averaged_mat.real(), expected_mat.real(),
                std::abs(expected_mat.real()) * rel_tol + 1e-15);
    EXPECT_NEAR(averaged_mat.imag(), expected_mat.imag(),
                std::abs(expected_mat.imag()) * rel_tol + 1e-15);
}

TEST(AvgMaterials, Random3DCorrectnessMagnetization)
{
    // Test that magnetization is also correctly averaged for 3D structures
    const R3 layer_magnetization(1.0, 0.0, 0.0);
    const R3 particle_magnetization(0.0, 2.0, 0.0);

    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), 1.0e-6, 5.0e-8,
                                       layer_magnetization);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 4.0e-6, 2.0e-7,
                                          particle_magnetization);

    // Create particle with known volume
    const double box_size = 2.0;
    const Particle random_particle(particle_mat, Box(box_size, box_size, box_size));
    const double particle_volume = box_size * box_size * box_size;

    // Volume fraction of 25%
    const double density = 0.03125;                           // 1/nm^3
    const double volume_fraction = density * particle_volume; // 0.25

    // Create sample
    const double layer_thickness = 10.0;
    Layer fronting(vacuum);
    Layer layer(layer_mat, layer_thickness);
    layer.fill3D(Dilute3D(density, random_particle));
    Layer substrate(vacuum);

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

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

    // Get averaged material (slice 1, not 0)
    const auto* averaged_mat_ptr = re_sample.avgeSlice(1).refractingMaterial();

    // Calculate expected magnetization
    const R3 expected_magnetization =
        volume_fraction * particle_magnetization + (1.0 - volume_fraction) * layer_magnetization;

    // Compare
    EXPECT_NEAR(averaged_mat_ptr->magnetization().x(), expected_magnetization.x(), 1e-14);
    EXPECT_NEAR(averaged_mat_ptr->magnetization().y(), expected_magnetization.y(), 1e-14);
    EXPECT_NEAR(averaged_mat_ptr->magnetization().z(), expected_magnetization.z(), 1e-14);
}

TEST(AvgMaterials, Random3DEdgeCaseLowDensity)
{
    // Test with very low density - averaged material should be close to layer material
    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), 1.0e-6, 5.0e-8);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 1.0e-5, 1.0e-6);

    const Particle random_particle(particle_mat, Box(1.0, 1.0, 1.0));
    const double density = 1e-6; // very low density

    Layer fronting(vacuum);
    Layer layer(layer_mat, 10.0);
    layer.fill3D(Dilute3D(density, random_particle));
    Layer substrate(vacuum);

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

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

    const auto averaged_mat = re_sample.avgeSlice(1).refractingMaterial()->rawData();
    const auto layer_rawdata = layer_mat.rawData();

    // Should be very close to layer material (within 1% relative difference)
    const double rel_tol = 1e-2;
    EXPECT_NEAR(averaged_mat.real(), layer_rawdata.real(),
                std::abs(layer_rawdata.real()) * rel_tol);
    EXPECT_NEAR(averaged_mat.imag(), layer_rawdata.imag(),
                std::abs(layer_rawdata.imag()) * rel_tol);
}
