#include <gtest/gtest.h>

#include "grotto/dwt_lut.hpp"

#include <cmath>
#include <cstdint>
#include <vector>

namespace
{

grotto::dwt_lut make(grotto::dwt_family family, const std::vector<double> & samples,
    unsigned fractional_bits, unsigned depth)
{
    if (family == grotto::dwt_family::haar)
        return grotto::make_haar_dwt_lut(samples, fractional_bits, depth);
    return grotto::make_bior53_dwt_lut(samples, fractional_bits, depth);
}

}  // namespace

TEST(DwtLut, HaarIsQuantizedBlockMean)
{
    const std::vector<double> samples{
        0.0, 0.5, -1.25, 3.0, 4.0, 0.25, -0.5, 2.0};
    const auto lut = grotto::make_haar_dwt_lut(samples, 8, 1);
    ASSERT_EQ(lut.coeff.size(), 4u);
    EXPECT_EQ(lut.coeff[0], 64);
    EXPECT_EQ(lut.coeff[1], 224);
    EXPECT_EQ(lut.coeff[2], 544);
    EXPECT_EQ(lut.coeff[3], 192);
    EXPECT_EQ(lut(0), 64);
    EXPECT_EQ(lut(1), 64);
    EXPECT_EQ(lut(7), 192);
}

TEST(DwtLut, HaarNegativeFloor)
{
    const std::vector<double> samples{-1.1, 0.2, -0.3, -4.0};
    const auto lut = grotto::make_haar_dwt_lut(samples, 10, 1);
    ASSERT_EQ(lut.coeff.size(), 2u);
    EXPECT_EQ(lut.coeff[0], -461);
    EXPECT_EQ(lut.coeff[1], -2202);
    EXPECT_EQ(lut(0), -461);
    EXPECT_EQ(lut(3), -2202);
}

TEST(DwtLut, Bior53MatchesSmoothExtension)
{
    const std::vector<double> samples{
        0.0, 0.5, -1.25, 3.0, 4.0, 0.25, -0.5, 2.0};
    const auto one = make(grotto::dwt_family::bior53, samples, 8, 1);
    const std::vector<std::int64_t> expect1{-513, 144, -288, 2064, -449, 2304};
    EXPECT_EQ(one.coeff, expect1);

    const auto two = make(grotto::dwt_family::bior53, samples, 8, 2);
    const std::vector<std::int64_t> expect2{-3649, -753, 912, 319, 10112};
    EXPECT_EQ(two.coeff, expect2);
    const std::int64_t at[8] = {228, 190, 153, 116, 79, 691, 1303, 1915};
    for (std::uint64_t raw = 0; raw < 8; ++raw)
        EXPECT_EQ(two(raw), at[raw]) << raw;
}

TEST(DwtLut, Bior53FloorDivOnNegatives)
{
    const std::vector<double> samples{-1.1, 0.2, -0.3, -4.0};
    const auto lut = grotto::make_bior53_dwt_lut(samples, 10, 1);
    const std::vector<std::int64_t> expect{-7578, -1793, -154, -15770};
    EXPECT_EQ(lut.coeff, expect);
    EXPECT_EQ(lut(0), -77);
    EXPECT_EQ(lut(1), -3981);
    EXPECT_EQ(lut(2), -7885);
    EXPECT_EQ(lut(3), -5837);
}

TEST(DwtLut, SigmoidGridAgreesWithHaarMean)
{
    auto samples = grotto::sample_dwt_signal(6, 4, [](double x) {
        return 1.0 / (1.0 + std::exp(-(x - 2.0)));
    });
    const auto haar = grotto::make_haar_dwt_lut(samples, 4, 2);
    ASSERT_EQ(haar.coeff.size(), 16u);
    for (std::uint64_t bin = 0; bin < 16; ++bin)
    {
        double sum = 0;
        for (unsigned k = 0; k < 4; ++k)
            sum += samples[bin * 4 + k];
        const auto q = static_cast<std::int64_t>(std::floor(sum / 4.0 * 16.0));
        EXPECT_EQ(haar.coeff[bin], q) << bin;
        EXPECT_EQ(haar(bin * 4), q);
    }
    const auto bior = grotto::make_bior53_dwt_lut(samples, 4, 2);
    EXPECT_EQ(bior(32), 8);
    EXPECT_EQ(bior(48), 12);
    EXPECT_EQ(haar(32), 8);
}

TEST(DwtLut, RejectsBadShape)
{
    const std::vector<double> samples{0.0, 1.0, 2.0};
    EXPECT_THROW(grotto::make_haar_dwt_lut(samples, 4, 1), std::invalid_argument);
    const std::vector<double> four{0.0, 1.0, 2.0, 3.0};
    EXPECT_THROW(grotto::make_haar_dwt_lut(four, 4, 0), std::invalid_argument);
    EXPECT_THROW(grotto::make_bior53_dwt_lut(four, 4, 3), std::invalid_argument);
    auto lut = grotto::make_haar_dwt_lut(four, 4, 1);
    EXPECT_THROW(lut(4), std::out_of_range);
}