#include "Resample/Flux/ScalarFlux.h"

#include "Resample/Processed/ReSample.h"
#include "Resample/Specular/ComputeFluxScalar.h"
#include "Sample/Aggregate/ParticleLayout.h"
#include "Sample/Interface/LayerRoughness.h"
#include "Sample/Material/MaterialFactoryFuncs.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Multilayer/Sample.h"
#include "Tests/GTestWrapper/google_test.h"

class RTTest : public ::testing::Test {
protected:
    void printCoeffs(const std::vector<ScalarFlux>& coeffs)
    { // for debug phases
        for (size_t i = 0; i < coeffs.size(); ++i) {
            const ScalarFlux& coeff = coeffs[i];
            std::cout << i << " " << coeff.getScalarT() << " " << coeff.getScalarR() << "\n";
        }
    }
    void compareCoeffs(const ScalarFlux& coeff1, const ScalarFlux& coeff2)
    {
        EXPECT_NEAR(abs(coeff1.getScalarT()), abs(coeff2.getScalarT()), 5e-14);
        EXPECT_NEAR(coeff1.getScalarT().real(), coeff2.getScalarT().real(), 1e-10);
        EXPECT_NEAR(coeff1.getScalarT().imag(), coeff2.getScalarT().imag(), 1e-10);
        EXPECT_NEAR(abs(coeff1.getScalarR()), abs(coeff2.getScalarR()), 5e-14);
        EXPECT_NEAR(coeff1.getScalarR().real(), coeff2.getScalarR().real(), 1e-10);
        EXPECT_NEAR(coeff1.getScalarR().imag(), coeff2.getScalarR().imag(), 1e-10);
    }
    std::vector<ScalarFlux> getCoeffs(Fluxes&& fluxes)
    {
        std::vector<ScalarFlux> result;
        for (auto& flux : fluxes)
            result.push_back(*dynamic_cast<const ScalarFlux*>(flux));
        return result;
    }
    const Material air = RefractiveMaterial("Air", 1e-8, 1e-8);
    const Material amat = RefractiveMaterial("material A", 2e-6, 8e-7);
    const Material bmat = RefractiveMaterial("material B (high absorption)", 3e-5, 2e-4);
    const Material stone = RefractiveMaterial("substrate material", 1e-6, 1e-7);
    const Layer topLayer{air};
    const Layer substrate{stone};
    const R3 k{1, 0, -1e-3};
    Sample sample1, sample2;
    std::vector<ScalarFlux> coeffs1, coeffs2;
};

TEST_F(RTTest, SplitLayer)
{
    const int n = 40;

    sample1.addLayer(topLayer);
    sample1.addLayer(Layer(amat, n * 10));
    sample1.addLayer(substrate);

    sample2.addLayer(topLayer);
    for (size_t i = 0; i < n; ++i)
        sample2.addLayer(Layer(amat, 10));
    sample2.addLayer(substrate);

    const ReSample sample_1 = ReSample::make(sample1);
    const ReSample sample_2 = ReSample::make(sample2);

    coeffs1 = getCoeffs(Compute::scalarFluxes(sample_1.averageSlices(), k));
    coeffs2 = getCoeffs(Compute::scalarFluxes(sample_2.averageSlices(), k));

    // printCoeffs( coeffs1 );
    // printCoeffs( coeffs2 );

    compareCoeffs(coeffs1[0], coeffs2[0]);
    compareCoeffs(coeffs1[1], coeffs2[1]);
    compareCoeffs(coeffs1.back(), coeffs2.back());
}

TEST_F(RTTest, SplitBilayers)
{
    // With exaggerated values of #layers, layer thickness, and absorption
    // so that we also test correct handling of floating-point overflow.
    const int n = 250;

    sample1.addLayer(topLayer);
    for (size_t i = 0; i < n; ++i) {
        sample1.addLayer(Layer(amat, 100));
        sample1.addLayer(Layer(bmat, 200));
    }
    sample1.addLayer(substrate);

    sample2.addLayer(topLayer);
    for (size_t i = 0; i < n; ++i) {
        sample2.addLayer(Layer(amat, 100));
        sample2.addLayer(Layer(bmat, 100));
        sample2.addLayer(Layer(bmat, 100));
    }
    sample2.addLayer(substrate);

    ReSample sample_1 = ReSample::make(sample1);
    ReSample sample_2 = ReSample::make(sample2);

    coeffs1 = getCoeffs(Compute::scalarFluxes(sample_1.averageSlices(), k));
    coeffs2 = getCoeffs(Compute::scalarFluxes(sample_2.averageSlices(), k));

    printCoeffs(coeffs1);
    printCoeffs(coeffs2);

    compareCoeffs(coeffs1[0], coeffs2[0]);
    compareCoeffs(coeffs1[1], coeffs2[1]);

    // Amplitudes at bottom must be strictly zero.
    // The new algorithm handles this without an overflow
    EXPECT_EQ(complex_t(), coeffs1[coeffs1.size() - 2].getScalarT());
    EXPECT_EQ(complex_t(), coeffs1[coeffs1.size() - 2].getScalarR());
    EXPECT_EQ(complex_t(), coeffs2[coeffs2.size() - 2].getScalarT());
    EXPECT_EQ(complex_t(), coeffs2[coeffs2.size() - 2].getScalarR());
}

TEST_F(RTTest, Overflow)
{
    // Text extra thick layers to also provoke an overflow in the new algorithm
    const int n = 5;

    sample1.addLayer(topLayer);
    for (size_t i = 0; i < n; ++i) {
        sample1.addLayer(Layer(amat, 1000));
        sample1.addLayer(Layer(bmat, 200000));
    }
    sample1.addLayer(substrate);

    sample2.addLayer(topLayer);
    for (size_t i = 0; i < n; ++i) {
        sample2.addLayer(Layer(amat, 1000));
        sample2.addLayer(Layer(bmat, 100000));
        sample2.addLayer(Layer(bmat, 100000));
    }
    sample2.addLayer(substrate);

    ReSample sample_1 = ReSample::make(sample1);
    ReSample sample_2 = ReSample::make(sample2);

    coeffs1 = getCoeffs(Compute::scalarFluxes(sample_1.averageSlices(), k));
    coeffs2 = getCoeffs(Compute::scalarFluxes(sample_2.averageSlices(), k));

    printCoeffs(coeffs1);
    printCoeffs(coeffs2);

    compareCoeffs(coeffs1[0], coeffs2[0]);
    compareCoeffs(coeffs1[1], coeffs2[1]);

    // If floating-point overflow is handled correctly, amplitudes at bottom must be strictly zero.
    EXPECT_EQ(complex_t(), coeffs1[coeffs1.size() - 2].getScalarT());
    EXPECT_EQ(complex_t(), coeffs1[coeffs1.size() - 2].getScalarR());
    EXPECT_EQ(complex_t(), coeffs2[coeffs2.size() - 2].getScalarT());
    EXPECT_EQ(complex_t(), coeffs2[coeffs2.size() - 2].getScalarR());
}
