#include "Sample/Flux/ScalarFlux.h"

#include "Resample/Sample/ReSample.h"
#include "Resample/Specular/ComputeFluxScalar.h"
#include "Sample/Interface/Roughness.h"
#include "Sample/Material/RefractiveMaterial.h"
#include "Sample/Multilayer/Layer.h"
#include "Sample/Option/SimulationOptions.h"
#include "Sample/Sample/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(const Fluxes& fluxes)
    {
        std::vector<ScalarFlux> result;
        for (const auto& flux : fluxes)
            result.push_back(*dynamic_cast<const ScalarFlux*>(flux.get()));
        return result;
    }
    const RefractiveMaterial air = {"Air", MaterialColor(0.5, 0.5, 0.5), 1e-8, 1e-8};
    const RefractiveMaterial amat = {"material A", MaterialColor(0.5, 0.5, 0.5), 2e-6, 8e-7};
    const RefractiveMaterial bmat = {"material B = {high absorption}", MaterialColor(0.5, 0.5, 0.5),
                                     3e-5, 2e-4};
    const RefractiveMaterial stone = {"substrate material", MaterialColor(0.5, 0.5, 0.5), 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::processedSample(sample1, {});
    const ReSample sample_2 = ReSample::processedSample(sample2, {});

    coeffs1 = getCoeffs(Compute::scalarFluxes(sample_1.slices(), k));
    coeffs2 = getCoeffs(Compute::scalarFluxes(sample_2.slices(), 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::processedSample(sample1, {});
    ReSample sample_2 = ReSample::processedSample(sample2, {});

    coeffs1 = getCoeffs(Compute::scalarFluxes(sample_1.slices(), k));
    coeffs2 = getCoeffs(Compute::scalarFluxes(sample_2.slices(), 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::processedSample(sample1, {});
    ReSample sample_2 = ReSample::processedSample(sample2, {});

    coeffs1 = getCoeffs(Compute::scalarFluxes(sample_1.slices(), k));
    coeffs2 = getCoeffs(Compute::scalarFluxes(sample_2.slices(), 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());
}
