#include "Sample/Particle/CoreAndShell.h"

#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/HardParticle/HardParticles.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/Particle.h"
#include "Sample/Particle/Rotations.h"
#include "Sample/Sample/Sample.h"
#include "Tests/GTestWrapper/google_test.h"
#include <stdexcept>

namespace {

const Particle* findParticleByMaterial(const std::vector<std::unique_ptr<IParticle>>& particles,
                                       const complex_t& material_sld)
{
    for (const auto& particle : particles) {
        const auto* p = dynamic_cast<const Particle*>(particle.get());
        if (p && p->material()->rawData() == material_sld)
            return p;
    }
    return nullptr;
}

} // namespace

TEST(CoreAndShell, ComplexCoreShellClone)
{
    const RefractiveMaterial core_material("Ag", MaterialColor(0.5, 0.5, 0.5), 1.245e-5, 5.419e-7);
    const RefractiveMaterial shell_material("AgO2", MaterialColor(0.5, 0.5, 0.5), 8.600e-6,
                                            3.442e-7);

    double shell_length(50);
    double shell_width(20);
    double shell_height(10);
    double core_length = shell_length / 2;
    double core_width = shell_width / 2;
    double core_height = shell_height / 2;

    R3 relative_pos(0, 0, (shell_height - core_height) / 2);

    Particle core(core_material, Box(core_length, core_width, core_height));
    core.translate(relative_pos);
    Particle shell(shell_material, Box(shell_length, shell_width, shell_height));
    CoreAndShell coreshell(core, shell);
    coreshell.rotate(RotationY(90 * Units::deg));
    coreshell.translate(R3(0, 0, -10));

    CoreAndShell* clone = coreshell.clone();
    EXPECT_EQ(coreshell.coreParticle()->particlePosition(), relative_pos);
    EXPECT_EQ(clone->coreParticle()->particlePosition(), relative_pos);
    delete clone;
}

TEST(CoreAndShell, OnionClonePreservesNestedCoreStructure)
{
    const RefractiveMaterial inner_core_material("InnerCore", MaterialColor(0.5, 0.5, 0.5),
                                                 1.245e-5, 5.419e-7);
    const RefractiveMaterial inner_shell_material("InnerShell", MaterialColor(0.5, 0.5, 0.5),
                                                  8.600e-6, 3.442e-7);
    const RefractiveMaterial outer_shell_material("OuterShell", MaterialColor(0.5, 0.5, 0.5),
                                                  6.100e-6, 2.111e-7);

    Particle inner_core(inner_core_material, Box(20, 20, 4));
    inner_core.translate(R3(0, 0, 1));
    Particle inner_shell(inner_shell_material, Box(30, 30, 6));
    CoreAndShell inner_coreshell(inner_core, inner_shell);
    inner_coreshell.translate(R3(0, 0, 1));

    Particle outer_shell(outer_shell_material, Box(40, 40, 8));
    CoreAndShell onion(inner_coreshell, outer_shell);
    onion.rotate(RotationY(90 * Units::deg));
    onion.translate(R3(0, 0, -10));

    std::unique_ptr<CoreAndShell> clone(onion.clone());
    const auto* cloned_inner = dynamic_cast<const CoreAndShell*>(clone->coreParticle());
    ASSERT_TRUE(cloned_inner);
    EXPECT_EQ(cloned_inner->particlePosition(), inner_coreshell.particlePosition());

    const auto* cloned_inner_core = dynamic_cast<const Particle*>(cloned_inner->coreParticle());
    ASSERT_TRUE(cloned_inner_core);
    EXPECT_EQ(cloned_inner_core->particlePosition(), inner_core.particlePosition());
}

TEST(CoreAndShell, OnionUsesOutermostVolume)
{
    const RefractiveMaterial inner_core_material("InnerCore", MaterialColor(0.5, 0.5, 0.5),
                                                 1.245e-5, 5.419e-7);
    const RefractiveMaterial inner_shell_material("InnerShell", MaterialColor(0.5, 0.5, 0.5),
                                                  8.600e-6, 3.442e-7);
    const RefractiveMaterial outer_shell_material("OuterShell", MaterialColor(0.5, 0.5, 0.5),
                                                  6.100e-6, 2.111e-7);

    Particle inner_core(inner_core_material, Box(4, 4, 4));
    inner_core.translate(R3(0, 0, 1));
    Particle inner_shell(inner_shell_material, Box(6, 6, 6));
    inner_shell.setVolumeToAdmix(inner_shell.pFormfactor()->volume()
                                 - inner_core.pFormfactor()->volume());
    CoreAndShell inner_coreshell(inner_core, inner_shell);
    inner_coreshell.translate(R3(0, 0, 1));

    Particle outer_shell(outer_shell_material, Box(8, 8, 8));
    outer_shell.setVolumeToAdmix(outer_shell.pFormfactor()->volume()
                                 - inner_shell.pFormfactor()->volume());
    CoreAndShell onion(inner_coreshell, outer_shell);

    EXPECT_DOUBLE_EQ(onion.volumeToAdmix(), outer_shell.pFormfactor()->volume());
}

TEST(CoreAndShell, RejectsUnsupportedCompoundCore)
{
    const Vacuum material;

    Compound compound_core{};
    compound_core.addComponent(Particle(material, Box(4, 4, 4)));
    Particle shell(material, Box(8, 8, 8));
    CoreAndShell invalid(compound_core, shell);

    EXPECT_THROW(invalid.validate(), std::runtime_error);
}

TEST(CoreAndShell, RejectsNestedCoreLargerThanOuterShell)
{
    const Vacuum material;

    Particle inner_core(material, Box(4, 4, 4));
    Particle inner_shell(material, Box(10, 10, 10));
    CoreAndShell inner_coreshell(inner_core, inner_shell);

    Particle outer_shell(material, Box(8, 8, 8));
    CoreAndShell invalid(inner_coreshell, outer_shell);

    EXPECT_THROW(invalid.validate(), std::runtime_error);
}

TEST(CoreAndShell, RejectsCoreOutsideShellAlongZ)
{
    const Vacuum material;

    Particle core(material, Box(4, 4, 4));
    core.translate(R3(0, 0, 6));
    Particle shell(material, Box(8, 8, 8));
    CoreAndShell invalid(core, shell);

    EXPECT_THROW(invalid.validate(), std::runtime_error);
}

TEST(CoreAndShell, OnionSlicePreservesRecursiveShellVolumesAndAmbients)
{
    const complex_t vacuum_sld(0.0, 0.0);
    const complex_t outer_shell_sld(6.0e-5, 2.0e-8);
    const complex_t inner_shell_sld(9.0e-5, 3.0e-8);
    const complex_t inner_core_sld(1.2e-4, 4.0e-8);

    const Vacuum vacuum;
    const RefractiveMaterial outer_shell_mat("OuterShell", MaterialColor(0.5, 0.5, 0.5),
                                             outer_shell_sld.real(), outer_shell_sld.imag());
    const RefractiveMaterial inner_shell_mat("InnerShell", MaterialColor(0.5, 0.5, 0.5),
                                             inner_shell_sld.real(), inner_shell_sld.imag());
    const RefractiveMaterial inner_core_mat("InnerCore", MaterialColor(0.5, 0.5, 0.5),
                                            inner_core_sld.real(), inner_core_sld.imag());

    Particle inner_core(inner_core_mat, Box(2, 2, 4));
    inner_core.translate(R3(0, 0, 1));
    Particle inner_shell(inner_shell_mat, Box(4, 4, 6));
    CoreAndShell inner_coreshell(inner_core, inner_shell);

    Particle outer_shell(outer_shell_mat, Box(6, 6, 8));
    CoreAndShell onion(inner_coreshell, outer_shell);

    Layer layer(vacuum);
    Slice slice(3, layer, Span(2.0, 5.0), R3());
    const auto slices =
        Refine::particlesInSlice(&onion, Span(2.0, 5.0), &slice, vacuum, MesoOptions{});

    ASSERT_EQ(slices.size(), 3u);

    const Particle* sliced_inner_core = findParticleByMaterial(slices, inner_core_sld);
    const Particle* sliced_inner_shell = findParticleByMaterial(slices, inner_shell_sld);
    const Particle* sliced_outer_shell = findParticleByMaterial(slices, outer_shell_sld);

    ASSERT_TRUE(sliced_inner_core);
    ASSERT_TRUE(sliced_inner_shell);
    ASSERT_TRUE(sliced_outer_shell);

    EXPECT_DOUBLE_EQ(sliced_inner_core->volumeToAdmix(), 12.0);
    EXPECT_DOUBLE_EQ(sliced_inner_shell->volumeToAdmix(), 36.0);
    EXPECT_DOUBLE_EQ(sliced_outer_shell->volumeToAdmix(), 60.0);

    EXPECT_EQ(sliced_inner_core->ambientMaterial()->rawData(), inner_shell_sld);
    EXPECT_EQ(sliced_inner_shell->ambientMaterial()->rawData(), outer_shell_sld);
    EXPECT_EQ(sliced_outer_shell->ambientMaterial()->rawData(), vacuum_sld);
}

//! Regression test: core must retain shell as ambient (not layer) after averaging.
TEST(CoreAndShell, CoreAmbientMaterialAfterAveraging)
{
    const complex_t vacuum_sld(0.0, 0.0);
    const complex_t shell_sld(6.0e-5, 2.0e-8);
    const complex_t core_sld(1.2e-4, 4.0e-8);

    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), vacuum_sld.real(),
                                    vacuum_sld.imag());
    const RefractiveMaterial shell_mat("Shell", MaterialColor(0.5, 0.5, 0.5), shell_sld.real(),
                                       shell_sld.imag());
    const RefractiveMaterial core_mat("Core", MaterialColor(0.5, 0.5, 0.5), core_sld.real(),
                                      core_sld.imag());

    const Particle core(core_mat, Box(10, 10, 5));
    const Particle shell(shell_mat, Box(16, 16, 8));
    CoreAndShell particle(core, shell);

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

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

    bool found_core = false, found_shell = false;
    for (const auto& s : re_sample.structs()) {
        ASSERT_TRUE(s->particle());
        for (const auto& leaf : s->particle()->decompose()) {
            auto* p = dynamic_cast<const Particle*>(leaf.get());
            if (!p)
                continue;
            auto mat_sld = p->material()->rawData();
            auto amb_sld = p->ambientMaterial()->rawData();
            if (mat_sld == core_sld) {
                found_core = true;
                EXPECT_EQ(amb_sld, shell_sld) << "Core ambient must be shell";
            } else if (mat_sld == shell_sld) {
                found_shell = true;
                EXPECT_EQ(amb_sld, vacuum_sld) << "Shell ambient must be vacuum";
            }
        }
    }

    EXPECT_TRUE(found_core && found_shell);
}

TEST(CoreAndShell, OnionAmbientMaterialsAfterAveraging)
{
    const complex_t vacuum_sld(0.0, 0.0);
    const complex_t outer_shell_sld(6.0e-5, 2.0e-8);
    const complex_t inner_shell_sld(9.0e-5, 3.0e-8);
    const complex_t inner_core_sld(1.2e-4, 4.0e-8);

    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), vacuum_sld.real(),
                                    vacuum_sld.imag());
    const RefractiveMaterial outer_shell_mat("OuterShell", MaterialColor(0.5, 0.5, 0.5),
                                             outer_shell_sld.real(), outer_shell_sld.imag());
    const RefractiveMaterial inner_shell_mat("InnerShell", MaterialColor(0.5, 0.5, 0.5),
                                             inner_shell_sld.real(), inner_shell_sld.imag());
    const RefractiveMaterial inner_core_mat("InnerCore", MaterialColor(0.5, 0.5, 0.5),
                                            inner_core_sld.real(), inner_core_sld.imag());

    Particle inner_core(inner_core_mat, Box(10, 10, 4));
    inner_core.translate(R3(0, 0, 1));
    Particle inner_shell(inner_shell_mat, Box(12, 12, 6));
    CoreAndShell inner_coreshell(inner_core, inner_shell);
    inner_coreshell.translate(R3(0, 0, 1));

    Particle outer_shell(outer_shell_mat, Box(16, 16, 8));
    CoreAndShell onion(inner_coreshell, outer_shell);

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

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

    bool found_inner_core = false;
    bool found_inner_shell = false;
    bool found_outer_shell = false;
    for (const auto& s : re_sample.structs()) {
        ASSERT_TRUE(s->particle());
        for (const auto& leaf : s->particle()->decompose()) {
            auto* p = dynamic_cast<const Particle*>(leaf.get());
            if (!p)
                continue;
            auto mat_sld = p->material()->rawData();
            auto amb_sld = p->ambientMaterial()->rawData();
            if (mat_sld == inner_core_sld) {
                found_inner_core = true;
                EXPECT_EQ(amb_sld, inner_shell_sld) << "Inner core ambient must be inner shell";
            } else if (mat_sld == inner_shell_sld) {
                found_inner_shell = true;
                EXPECT_EQ(amb_sld, outer_shell_sld) << "Inner shell ambient must be outer shell";
            } else if (mat_sld == outer_shell_sld) {
                found_outer_shell = true;
                EXPECT_EQ(amb_sld, vacuum_sld) << "Outer shell ambient must be vacuum";
            }
        }
    }

    EXPECT_TRUE(found_inner_core && found_inner_shell && found_outer_shell);
}
