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

#include "grotto.hpp"

// Haar and bior(5,3) lookup tables (Reis, Ugurbil, Wagh, Henry, de Vega,
// PoPETs 2025, ePrint 2025/013). The grid is sigmoid on [0, 4), stored as
// Q4.4. Depth 2 keeps the top 4 bits of a 6-bit index.
//
//   c++ -std=c++17 -march=native -I include -I thirdparty examples/grotto/dwt_lut.cpp

int main()
{
    //! [dwt-lut]
    constexpr unsigned domain_bits = 6;
    constexpr unsigned fractional_bits = 4;
    constexpr unsigned depth = 2;
    auto samples = grotto::sample_dwt_signal(domain_bits, fractional_bits,
        [](double x) {
            return 1.0 / (1.0 + std::exp(-(x - 2.0)));
        });
    auto haar = grotto::make_haar_dwt_lut(samples, fractional_bits, depth);
    auto bior = grotto::make_bior53_dwt_lut(samples, fractional_bits, depth);

    // Haar is the mean of each block of 2^depth samples, then quantized.
    const std::uint64_t raw = 32;
    double block = 0;
    for (unsigned k = 0; k < 4; ++k)
        block += samples[(raw & ~std::uint64_t{3}) + k];
    const auto haar_expect = static_cast<std::int64_t>(
        std::floor(block / 4.0 * 16.0));

    // bior(5,3), lsb = 0: only the first tap, at index msb+2, divided by 2^j.
    const std::uint64_t msb = raw >> depth;
    const auto c0 = bior.coeff[(msb + 2) % bior.coeff.size()];
    const auto bior_at_32 = c0 / 4;

    // lsb = 1: both taps, weights (2^j - lsb) and lsb, then divide by 2^{2j}.
    const auto c1 = bior.coeff[(msb + 3) % bior.coeff.size()];
    const auto bior_at_33 = (c0 * 3 + c1) / 16;
    //! [dwt-lut]

    if (haar(raw) != haar_expect || haar(raw) != 8)
    {
        std::cerr << "haar lut\n";
        return 1;
    }
    if (bior(raw) != bior_at_32 || bior(33) != bior_at_33 || bior(raw) != 8)
    {
        std::cerr << "bior lut\n";
        return 1;
    }

    std::cout << haar(raw) << " " << bior(raw) << " " << bior(33) << "\n";
    return 0;
}