// Exactness test for the bitshuffle block decoders: classic (malloc), S1 (caller scratch),
// hperf-NEON port, hperf (x86 SIMD or scalar), and the shipped JFJochBitUnshuffleBlock selector.
// Prints a hash over every encoded input and decoded output so x86 and aarch64 runs can be compared.
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <vector>
#include <string>
#include <fstream>

#include <BitShuffleBlock.h>
#include <lz4/lz4.h>
extern "C" uint64_t bshuf_read_uint64_BE(const void* buf);
#if defined(__aarch64__)
#include <bitshuffle_hperf/bitshuffle_neon.h>
#endif

static uint64_t rng_state = 0x1234567887654321ULL;
static uint64_t next_rand() {
    uint64_t z = (rng_state += 0x9e3779b97f4a7c15ULL);
    z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
    z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
    return z ^ (z >> 31);
}

static uint64_t hash = 1469598103934665603ULL;
static void hash_bytes(const char *p, size_t n) {
    for (size_t i = 0; i < n; i++) { hash ^= (uint8_t) p[i]; hash *= 1099511628211ULL; }
}

static long failures = 0, checks = 0;
static void expect_same(const std::vector<char> &ref, const std::vector<char> &got, const char *what,
                        size_t size, size_t elem, int pattern) {
    checks++;
    if (memcmp(ref.data(), got.data(), ref.size()) != 0) {
        if (failures < 20)
            printf("MISMATCH %s size=%zu elem=%zu pattern=%d\n", what, size, elem, pattern);
        failures++;
    }
}

static void fill(std::vector<char> &v, size_t elem, int pattern) {
    const size_t n = v.size() / elem;
    for (size_t i = 0; i < v.size(); i++) {
        switch (pattern) {
            case 0: v[i] = (char) next_rand(); break;                               // random
            case 1: v[i] = 0; break;                                                // zeros
            case 2: v[i] = (char) 0xff; break;                                      // ones
            case 3: v[i] = (char) ((i % elem == 0) ? (next_rand() % 7 == 0 ? next_rand() % 5 : 0) : 0); break; // sparse counts
            case 4: v[i] = (char) (i & 1 ? 0xaa : 0x55); break;                     // alternating
            case 5: v[i] = (char) (((i / elem) % (8 * elem)) / 8 == i % elem ? 1u << ((i / elem) % 8) : 0); break; // walking bit
            case 6: v[i] = (char) ((i % elem == elem - 1) ? 0x80 : (char) next_rand()); break; // sign bit set + noise
        }
    }
    (void) n;
}

static void run_block(size_t size, size_t elem, int pattern) {
    std::vector<char> orig(size * elem), enc(size * elem), enc_h(size * elem), scratch(size * elem);
    fill(orig, elem, pattern);

    if (bshuf_trans_bit_elem(orig.data(), enc.data(), size, elem) < 0) { printf("encode failed\n"); failures++; return; }
    if (bitshuf_encode_block(enc_h.data(), orig.data(), scratch.data(), size, elem) != 0) { printf("hperf encode failed\n"); failures++; return; }
    expect_same(enc, enc_h, "encode classic vs hperf", size, elem, pattern);
    hash_bytes(enc.data(), enc.size());

    std::vector<char> out(size * elem);

    std::fill(out.begin(), out.end(), 0x5a);
    bshuf_untrans_bit_elem(enc.data(), out.data(), size, elem);
    expect_same(orig, out, "classic bshuf_untrans_bit_elem", size, elem, pattern);

    std::fill(out.begin(), out.end(), 0x5a);
    bitshuf_decode_block(out.data(), enc.data(), scratch.data(), size, elem);
    expect_same(orig, out, "hperf bitshuf_decode_block", size, elem, pattern);

    std::fill(out.begin(), out.end(), 0x5a);
    std::fill(scratch.begin(), scratch.end(), 0x33);
    if (JFJochBitUnshuffleBlock(out.data(), enc.data(), scratch.data(), size, elem) < 0) failures++;
    expect_same(orig, out, "JFJochBitUnshuffleBlock", size, elem, pattern);

#if defined(__aarch64__)
    std::fill(out.begin(), out.end(), 0x5a);
    bshuf_trans_byte_bitrow_NEON(enc.data(), scratch.data(), size, elem);
    bshuf_shuffle_bit_eightelem_NEON(scratch.data(), out.data(), size, elem);
    expect_same(orig, out, "S1 two-stage NEON with scratch", size, elem, pattern);

    if (elem == 1 || elem == 2 || elem == 4) {
        std::fill(out.begin(), out.end(), 0x5a);
        if (bitshuf_decode_block_neon(out.data(), enc.data(), scratch.data(), size, elem) != 0) failures++;
        expect_same(orig, out, "hperf NEON port", size, elem, pattern);
    } else if (bitshuf_decode_block_neon(out.data(), enc.data(), scratch.data(), size, elem) != -1)
        failures++;
#endif
    hash_bytes(out.data(), out.size());
}

// A bitshuffle/LZ4 chunk as HDF5 filter 32008 stores it, decoded block by block through every path.
static void run_chunk(const std::string &path, size_t elem) {
    std::ifstream f(path, std::ios::binary);
    std::vector<char> src((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
    const uint8_t *p = (const uint8_t *) src.data();
    const size_t total = bshuf_read_uint64_BE(p);
    const size_t block = bshuf_read_uint32_BE(p + 8) / elem;
    const size_t nelem = total / elem;
    size_t pos = 12, done = 0;
    std::vector<char> image(total);
    std::vector<char> dec(block * elem), scratch(block * elem), ref(block * elem), out(block * elem);
    long nblocks = 0, npartial = 0;
    while (nelem - done >= 8) {
        size_t cur = std::min(block, nelem - done);
        cur -= cur % 8;
        const size_t csize = bshuf_read_uint32_BE(p + pos);
        pos += 4;
        const int r = LZ4_decompress_safe(src.data() + pos, dec.data(), (int) csize, (int) (cur * elem));
        if (r != (int) (cur * elem)) { printf("LZ4 error %s\n", path.c_str()); failures++; return; }
        pos += csize;
        if (cur != block) npartial++;
        nblocks++;

        bshuf_untrans_bit_elem(dec.data(), ref.data(), cur, elem);
        ref.resize(cur * elem); out.resize(cur * elem);

        bitshuf_decode_block(out.data(), dec.data(), scratch.data(), cur, elem);
        expect_same(ref, out, "chunk hperf", cur, elem, -1);
        JFJochBitUnshuffleBlock(out.data(), dec.data(), scratch.data(), cur, elem);
        expect_same(ref, out, "chunk JFJochBitUnshuffleBlock", cur, elem, -1);
#if defined(__aarch64__)
        bshuf_trans_byte_bitrow_NEON(dec.data(), scratch.data(), cur, elem);
        bshuf_shuffle_bit_eightelem_NEON(scratch.data(), out.data(), cur, elem);
        expect_same(ref, out, "chunk S1", cur, elem, -1);
        bitshuf_decode_block_neon(out.data(), dec.data(), scratch.data(), cur, elem);
        expect_same(ref, out, "chunk NEON port", cur, elem, -1);
#endif
        memcpy(image.data() + done * elem, ref.data(), cur * elem);
        ref.resize(block * elem); out.resize(block * elem);
        done += cur;
    }
    const size_t left = (nelem - done) * elem;
    memcpy(image.data() + done * elem, src.data() + pos, left);
    pos += left;
    if (pos != src.size()) { printf("chunk not fully consumed %s\n", path.c_str()); failures++; }
    uint64_t h = 1469598103934665603ULL;
    for (char c : image) { h ^= (uint8_t) c; h *= 1099511628211ULL; }
    printf("chunk %s: %zu elements, %ld blocks (%ld partial), %zu leftover bytes, image hash %016llx\n",
           path.c_str(), nelem, nblocks, npartial, left, (unsigned long long) h);
    hash_bytes(image.data(), image.size());
}

int main(int argc, char **argv) {
    const size_t elems[] = {1, 2, 3, 4, 8};
    std::vector<size_t> sizes;
    for (size_t s = 8; s <= 1040; s += 8) sizes.push_back(s);       // every row-length tail 1..130 bytes
    for (size_t s : {2048, 2056, 4096, 8192, 8200, 16384, 1544, 1536, 2040})
        sizes.push_back(s);
    for (size_t elem : elems)
        for (size_t size : sizes)
            for (int pattern = 0; pattern < 7; pattern++)
                run_block(size, elem, pattern);
    for (int i = 1; i < argc; i++)
        run_chunk(argv[i], 4);
    printf("%ld checks, %ld failures, total hash %016llx\n", checks, failures, (unsigned long long) hash);
    return failures ? 1 : 0;
}