#include "Sample/Aggregate/Dilute3D.h"
#include "Base/Axis/Bin.h"
#include "Base/Axis/Pixel.h"
#include "Base/Math/Functions.h"
#include "Base/Spin/SpinMatrix.h"
#include "Base/Spin/Spinor.h"
#include "Base/Vector/WavevectorInfo.h"
#include "Sample/Element/DiffuseElement.h"
#include "Sample/Flux/Flux.h"
#include "Sample/Flux/ScalarFlux.h"
#include "Sample/HardParticle/Sphere.h"
#include "Sample/Material/SLDMaterial.h"
#include "Sample/Material/Vacuum.h"
#include "Sample/Particle/DwbaChannels.h"
#include "Sample/Particle/Particle.h"
#include "Tests/GTestWrapper/google_test.h"
#include <stdexcept>

namespace {

constexpr size_t TEST_SLICE = 0;

struct ScalarDwbaElement {
    Pixel pixel;
    Fluxes fluxes_in;
    Fluxes fluxes_out;
    DiffuseElement element;

    ScalarDwbaElement()
        : pixel(Bin1D::At(-0.21), Bin1D::At(0.29))
        , fluxes_in()
        , fluxes_out()
        , element(1.54, 0.17, 0.11, &pixel, SpinMatrix::One() / 2.0, SpinMatrix::One(), &fluxes_in,
                  &fluxes_out)
    {
        fluxes_in.emplace_back(std::make_unique<ScalarFlux>(
            complex_t(1.13, 0.07), Spinor(complex_t(0.82, 0.04), complex_t(-0.21, 0.16))));
        fluxes_out.emplace_back(std::make_unique<ScalarFlux>(
            complex_t(0.91, -0.05), Spinor(complex_t(0.76, -0.03), complex_t(0.19, -0.12))));
    }
};

void expectComplexNear(complex_t actual, complex_t expected, double tol)
{
    EXPECT_NEAR(actual.real(), expected.real(), tol);
    EXPECT_NEAR(actual.imag(), expected.imag(), tol);
}

complex_t expectedPositionAverage(complex_t q_z, double shift_low, double shift_high)
{
    if (shift_low > shift_high)
        throw std::runtime_error("Invalid shift range");
    const double width = shift_high - shift_low;
    const double center = 0.5 * (shift_low + shift_high);
    return std::exp(I * q_z * center) * Math::sinc(q_z * (width / 2.0));
}

std::unique_ptr<Dilute3D> makePreparedReferenceStructure()
{
    const Vacuum ambient;
    const SLDMaterial particle_mat("particle", MaterialColor(0.86, 0.24, 0.18), 6e-6, 0.0);

    auto particle = std::make_unique<Particle>(particle_mat, Sphere(2.0));
    particle->setAmbientMaterial(ambient);
    particle->setSliceNumber(TEST_SLICE);

    auto result = std::make_unique<Dilute3D>(1e-3, *particle);
    result->setParticle(std::move(particle));
    result->setSliceNumber(TEST_SLICE);
    result->setThickness(10.0);
    result->setZRange(2.0, 8.0);
    return result;
}

} // namespace

// For Dilute3D (no pair correlations), structureCrossFactor reduces to the self-position average.
TEST(Dilute3D, StructuralFactorsMatchSelfPositionAverage)
{
    const auto structure = makePreparedReferenceStructure();
    const WavevectorInfo q_m(C3(0.17, -0.23, complex_t(0.35, -0.12)), C3(), 1.54);
    const WavevectorInfo q_n(C3(0.11, -0.19, complex_t(-0.08, 0.21)), C3(), 1.54);
    const complex_t delta_qz = q_m.getQ().z() - std::conj(q_n.getQ().z());
    const complex_t expected = expectedPositionAverage(delta_qz, 2.0, 8.0);

    expectComplexNear(structure->selfPositionFactor(q_m, q_n), expected, 1e-12);
    expectComplexNear(structure->structureCrossFactor(q_m, q_n), expected, 1e-12);
    EXPECT_GT(std::abs(expected - complex_t(1.0)), 1e-3);
}

// Struct3D::scalarScattering must use structureCrossFactor for coherent means and
// selfPositionFactor for the covariance term, matching the analytic hook formula.
TEST(Dilute3D, ScalarDwbaScatteringUsesStructuralFactors)
{
    const auto structure = makePreparedReferenceStructure();
    const ScalarDwbaElement dwba_element;

    const double actual = structure->scalarScattering(dwba_element.element);

    const DwbaChannels dwba = extractDwbaChannels(dwba_element.element, TEST_SLICE);
    const ChannelMoments moments = structure->particle()->channelMoments(dwba);

    complex_t intensity_from_hooks = 0.0;
    for (size_t m = 0; m < N_DWBA; ++m) {
        for (size_t n = 0; n < N_DWBA; ++n) {
            const complex_t weight = dwba.w[m] * std::conj(dwba.w[n]);
            intensity_from_hooks +=
                weight
                * (structure->structureCrossFactor(dwba.q[m], dwba.q[n]) * moments.mu[m]
                       * std::conj(moments.mu[n])
                   + structure->selfPositionFactor(dwba.q[m], dwba.q[n]) * moments.C[m][n]);
        }
    }

    EXPECT_NEAR(actual, intensity_from_hooks.real(), 1e-11);
}
