#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/Struct2D.h"
#include "Sample/Aggregate/Struct3D.h"
#include "Sample/Element/DiffuseElement.h"
#include "Sample/Flux/Flux.h"
#include "Sample/HardParticle/HardParticles.h"
#include "Sample/HardParticle/Sphere.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 "Sim/Evaluate/EvaluateScattering.h"
#include "Tests/GTestWrapper/google_test.h"

namespace {

struct ScatteringData {
    double total = 0.0;
    double structure = 0.0;
    double weight = 0.0;
    double z_range_width = 0.0;
};

double bornLikeDilute3DScattering(double volume_density, double layer_thickness, bool polarized)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 2e-5, 0.0);

    Particle particle(particle_mat, Sphere(0.5));

    Layer fronting(vacuum);
    Layer layer(vacuum, layer_thickness);
    layer.fill3D(Dilute3D(volume_density, particle));
    Layer substrate(vacuum);

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

    SimulationOptions options;
    const ReSample re_sample = ReSample::processedSample(sample, options);

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

    return Evaluate::scatteredContribution(re_sample, element, options, polarized);
}

ScatteringData fixedHeightEquivalentLayerScattering(bool use_3d, bool polarized,
                                                    bool use_avg_materials)
{
    const RefractiveMaterial vacuum("Vacuum", MaterialColor(0.5, 0.5, 0.5), 0.0, 0.0);
    const RefractiveMaterial cap_mat("Cap", MaterialColor(0.5, 0.5, 0.5), 3e-6, 2e-7);
    const RefractiveMaterial host_mat("Host", MaterialColor(0.5, 0.5, 0.5), 7e-6, 4e-7);
    const RefractiveMaterial buffer_mat("Buffer", MaterialColor(0.5, 0.5, 0.5), 1.4e-5, 1e-7);
    const RefractiveMaterial substrate_mat("Substrate", MaterialColor(0.5, 0.5, 0.5), 2.2e-5, 3e-7);
    const RefractiveMaterial particle_mat("Particle", MaterialColor(0.5, 0.5, 0.5), 2.8e-5, 6e-7,
                                          R3(7e5, -2e5, 4e5));

    constexpr double particle_height = 6.0;
    constexpr double volume_density = 1.25e-3;
    constexpr double area_density = volume_density * particle_height;

    Particle particle(particle_mat, Box(2.4, 1.7, particle_height));

    Layer fronting(vacuum);
    Layer cap(cap_mat, 11.0);
    Layer host(host_mat, particle_height);
    if (use_3d)
        host.fill3D(Dilute3D(volume_density, particle));
    else
        host.deposit2D(Dilute2D(area_density, particle));
    Layer buffer(buffer_mat, 9.0);
    Layer substrate(substrate_mat);

    Sample sample;
    sample.addLayer(fronting);
    sample.addLayer(cap);
    sample.addLayer(host);
    sample.addLayer(buffer);
    sample.addLayer(substrate);

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

    Pixel pixel(Bin1D::At(-0.31), Bin1D::At(0.27));
    const SpinMatrix polarizer = SpinMatrix::FromBlochVector(R3(0.35, -0.25, 0.45));
    const SpinMatrix analyzer = SpinMatrix::FromBlochVector(R3(-0.15, 0.20, 0.40));
    DiffuseElement element(1.54, 0.23, -0.17, &pixel, polarizer, analyzer);

    ScatteringData result;
    result.total = Evaluate::scatteredContribution(re_sample, element, options, polarized);
    const Fluxes fluxes_in = re_sample.fluxesIn(element.getKi(), polarized);
    const Fluxes fluxes_out = re_sample.fluxesOut(element.meanKf(), polarized);
    element.setFluxes(&fluxes_in, &fluxes_out);

    if (re_sample.structs().size() != 1)
        throw std::runtime_error("Expected a single structure");
    const Struct* structure = re_sample.structs().front().get();
    if (use_3d) {
        const auto* structure3d = dynamic_cast<const Struct3D*>(structure);
        if (!structure3d)
            throw std::runtime_error("Expected Struct3D");
        result.structure = polarized ? structure3d->polarizedScattering(element)
                                     : structure3d->scalarScattering(element);
        result.weight =
            structure3d->volumeNumberDensity() * structure3d->thickness() * structure3d->coverage();
        result.z_range_width = structure3d->zRangeHigh() - structure3d->zRangeLow();
    } else {
        const auto* structure2d = dynamic_cast<const Struct2D*>(structure);
        if (!structure2d)
            throw std::runtime_error("Expected Struct2D");
        result.structure = polarized ? structure2d->polarizedScattering(element)
                                     : structure2d->scalarScattering(element);
        result.weight = structure2d->areaNumberDensity() * structure2d->coverage();
    }

    return result;
}

} // namespace

// In homogeneous non-magnetic medium, DWBA reduces to Born.
// Then intensity depends only on rho_V * thickness.
// Scalar: doubling thickness at fixed rho_V doubles intensity.
TEST(ScatteredContribution3D, BornLimitScalesLinearlyWithThickness)
{
    const double baseline = bornLikeDilute3DScattering(1e-3, 10.0, false);
    const double doubled = bornLikeDilute3DScattering(1e-3, 20.0, false);

    EXPECT_GT(baseline, 0.0);
    EXPECT_NEAR(doubled / baseline, 2.0, 1e-5);
}

// Scalar: fixed rho_V * thickness keeps intensity unchanged.
TEST(ScatteredContribution3D, BornLimitDependsOnEffectiveSurfaceDensity)
{
    const double baseline = bornLikeDilute3DScattering(1e-3, 10.0, false);
    const double same_effective_density = bornLikeDilute3DScattering(5e-4, 20.0, false);

    EXPECT_GT(baseline, 0.0);
    EXPECT_NEAR(same_effective_density / baseline, 1.0, 1e-5);
}

// Polarized: doubling thickness at fixed rho_V doubles intensity.
TEST(ScatteredContribution3D, PolarizedBornLimitScalesLinearlyWithThickness)
{
    const double baseline = bornLikeDilute3DScattering(1e-3, 10.0, true);
    const double doubled = bornLikeDilute3DScattering(1e-3, 20.0, true);

    EXPECT_GT(baseline, 0.0);
    EXPECT_NEAR(doubled / baseline, 2.0, 1e-5);
}

// Polarized: fixed rho_V * thickness keeps intensity unchanged.
TEST(ScatteredContribution3D, PolarizedBornLimitDependsOnEffectiveSurfaceDensity)
{
    const double baseline = bornLikeDilute3DScattering(1e-3, 10.0, true);
    const double same_effective_density = bornLikeDilute3DScattering(5e-4, 20.0, true);

    EXPECT_GT(baseline, 0.0);
    EXPECT_NEAR(same_effective_density / baseline, 1.0, 1e-5);
}

// Fixed-height 3D filling is equivalent to 2D deposition at the same absolute z.
TEST(ScatteredContribution3D, FixedHeightDilute3DMatchesEquivalent2D)
{
    for (bool use_avg_materials : {false, true}) {
        SCOPED_TRACE(testing::Message() << "use_avg_materials=" << use_avg_materials);

        const ScatteringData intensity_3d =
            fixedHeightEquivalentLayerScattering(true, false, use_avg_materials);
        const ScatteringData intensity_2d =
            fixedHeightEquivalentLayerScattering(false, false, use_avg_materials);
        const double tol = std::abs(intensity_2d.total) * 1e-10 + 1e-14;

        EXPECT_NEAR(intensity_3d.z_range_width, 0.0, 1e-14);
        EXPECT_NEAR(intensity_3d.weight, intensity_2d.weight, 1e-14);
        EXPECT_NEAR(intensity_3d.structure, intensity_2d.structure,
                    std::abs(intensity_2d.structure) * 1e-10 + 1e-14);
        EXPECT_NEAR(intensity_3d.total, intensity_2d.total, tol);
    }
}

// The same 2D/3D equivalence must hold for the polarized scattering path.
TEST(ScatteredContribution3D, FixedHeightDilute3DMatchesEquivalent2DPolarized)
{
    for (bool use_avg_materials : {false, true}) {
        SCOPED_TRACE(testing::Message() << "use_avg_materials=" << use_avg_materials);

        const ScatteringData intensity_3d =
            fixedHeightEquivalentLayerScattering(true, true, use_avg_materials);
        const ScatteringData intensity_2d =
            fixedHeightEquivalentLayerScattering(false, true, use_avg_materials);
        const double tol = std::abs(intensity_2d.total) * 1e-10 + 1e-14;

        EXPECT_NEAR(intensity_3d.z_range_width, 0.0, 1e-14);
        EXPECT_NEAR(intensity_3d.weight, intensity_2d.weight, 1e-14);
        EXPECT_NEAR(intensity_3d.structure, intensity_2d.structure,
                    std::abs(intensity_2d.structure) * 1e-10 + 1e-14);
        EXPECT_NEAR(intensity_3d.total, intensity_2d.total, tol);
    }
}
