/* Randomized property tests of sift_bitmap against sorted-array reference implementations. */
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "../../src/bm25/bitmap.h"
#include "../../src/core/cmp.h"

#define MAXN 200000

static uint64_t rng = 0x9e3779b97f4a7c15ull;
static uint64_t rnd(void)
{
	rng ^= rng << 13, rng ^= rng >> 7, rng ^= rng << 17;
	return rng;
}
static uint32_t rndn(uint32_t n) { return (uint32_t)(rnd() % n); }

static long fails, checks;
#define CHECK(c, ...) do { checks++; if (!(c)) { if (fails++ < 20) { printf("FAIL %s:%d: ", __FILE__, __LINE__); printf(__VA_ARGS__); printf("\n"); } } } while (0)

typedef struct { sift_tid *t; uint8_t *tf; uint32_t n; } set;

/* random set with a shape: base block, page span, offsets 1..maxoff, density */
static void gen(set *s, uint32_t n, uint64_t base, uint32_t span, uint32_t maxoff)
{
	for (uint32_t i = 0; i < n; i++) {
		uint64_t blk = base + rndn(span);
		if (blk > 0xffffffffull) blk = 0xffffffffull;
		uint32_t off = rnd() % 4 == 0 ? maxoff : 1 + rndn(maxoff);
		s->t[i] = SIFT_TID(blk, off);
	}
	qsort(s->t, n, 8, sift_cmp_u64);
	uint32_t k = 0;
	for (uint32_t i = 0; i < n; i++)
		if (!k || s->t[i] != s->t[k - 1]) s->t[k++] = s->t[i];
	s->n = k;
	for (uint32_t i = 0; i < k; i++) s->tf[i] = (uint8_t)(rnd() % 8 ? 1 + rndn(3) : rndn(256));
}

static void gen_shape(set *s)
{
	static const uint32_t offs[] = { 1, 7, 8, 9, 16, 17, 32, 33, 64, 65, 100, 180, 255, 256, 257, 291, 400, 450, 512 };
	static const uint64_t bases[] = { 0, 12345, 0xffffffffull - 3000, 0xfffffe00ull, 1u << 24 };
	uint32_t maxoff = offs[rndn(19)], span = 1 + rndn(rnd() % 2 ? 600 : 100000);
	uint32_t n = rnd() % 10 == 0 ? rndn(3) : rndn(rnd() % 3 ? 3000 : 60000);
	if (rnd() % 8 == 0) { /* full pages */
		uint64_t b0 = bases[rndn(5)];
		uint32_t k = 0, pages = 1 + rndn(300);
		for (uint32_t p = 0; p < pages && k < MAXN; p++)
			for (uint32_t o = 1; o <= maxoff && k < MAXN; o++) s->t[k] = SIFT_TID(b0 + p > 0xffffffffull ? 0xffffffffull : b0 + p, o), k++;
		qsort(s->t, k, 8, sift_cmp_u64);
		uint32_t u = 0;
		for (uint32_t i = 0; i < k; i++)
			if (!u || s->t[i] != s->t[u - 1]) s->t[u++] = s->t[i];
		s->n = u;
		for (uint32_t i = 0; i < u; i++) s->tf[i] = (uint8_t)(1 + rndn(20));
		return;
	}
	gen(s, n, bases[rndn(5)] + rndn(1000), span, maxoff);
}

static uint32_t ref_and(const set *a, const set *b, sift_tid *o)
{
	uint32_t i = 0, j = 0, n = 0;
	while (i < a->n && j < b->n)
		if (a->t[i] < b->t[j]) i++;
		else if (a->t[i] > b->t[j]) j++;
		else o[n++] = a->t[i++], j++;
	return n;
}

static uint32_t ref_or(const set *a, const set *b, sift_tid *o)
{
	uint32_t i = 0, j = 0, n = 0;
	while (i < a->n || j < b->n)
		if (j == b->n || (i < a->n && a->t[i] < b->t[j])) o[n++] = a->t[i++];
		else if (i == a->n || b->t[j] < a->t[i]) o[n++] = b->t[j++];
		else o[n++] = a->t[i++], j++;
	return n;
}

static uint32_t ref_andnot(const set *a, const set *b, sift_tid *o)
{
	uint32_t i = 0, j = 0, n = 0;
	while (i < a->n)
		if (j == b->n || a->t[i] < b->t[j]) o[n++] = a->t[i++];
		else if (a->t[i] > b->t[j]) j++;
		else i++, j++;
	return n;
}

static sift_tid *outbuf;
static uint8_t *blob[4];
static uint8_t cbuf[8][SIFT_CUR_BUF];

static int same(const sift_tid *x, size_t nx, const sift_tid *y, size_t ny)
{
	return nx == ny && !memcmp(x, y, nx * 8);
}

static void open_set(sift_cur *c, const set *s, int mode, int slot)
{
	static const sift_enc_opts opts[] = { { 0, 0, 0, 1 }, { SIFT_LIST, 0, 0, 1 }, { SIFT_GROUPED, SIFT_CT_BITMAP, 0, 1 },
					      { SIFT_GROUPED, SIFT_CT_ARRAY, 0, 0 }, { SIFT_GROUPED, SIFT_CT_PACKED, 0, 1 } };
	if (mode == 5) {
		sift_cur_tids(c, s->t, s->n, cbuf[slot]);
		return;
	}
	size_t len = sift_encode(s->t, s->tf, s->tf, s->n, &opts[mode], blob[slot]);
	CHECK(len <= sift_encode_bound(s->n), "bound %zu > %zu", len, sift_encode_bound(s->n));
	CHECK(len == sift_encode(s->t, s->tf, s->tf, s->n, &opts[mode], NULL), "size query mismatch mode %d: %zu vs %zu n %u", mode, len, sift_encode(s->t, s->tf, s->tf, s->n, &opts[mode], NULL), s->n);
	CHECK(sift_blob_n(blob[slot]) == s->n, "n");
	sift_cur_blob(c, blob[slot], len, cbuf[slot]);
}

static void check_tf(const set *s, int mode, int slot, sift_cur *src, const sift_grp *res, const sift_tid *rt)
{
	static uint32_t rank[SIFT_GRP_MAXTIDS];
	if (mode == 5 || mode == 3 || !src->tf) return;
	uint32_t n = sift_grp_ranks(&src->g, res, rank);
	for (uint32_t i = 0; i < n; i++) {
		CHECK(rank[i] < s->n && s->t[rank[i]] == rt[i], "rank %u -> %u", i, rank[i]);
		if (rank[i] < s->n) CHECK(sift_tf_at(src, rank[i]) == sift_tf_quant(s->tf[rank[i]]), "tf");
	}
	(void)slot;
}

static void test_ops(set *a, set *b, set *c, sift_tid *ref, sift_tid *tmp)
{
	int ma = (int)rndn(6), mb = (int)rndn(6), mc = (int)rndn(6);
	sift_cur ca, cb, cc, op, op2, *kids[3];
	size_t n, rn;

	/* decode */
	open_set(&ca, a, ma, 0);
	n = sift_tids(&ca, outbuf);
	CHECK(same(outbuf, n, a->t, a->n), "decode mode %d n %u got %zu", ma, a->n, n);
	open_set(&ca, a, ma, 0);
	CHECK(sift_count(&ca) == a->n, "count");

	/* and / or / andnot, materialized and counted */
	for (int op_kind = 0; op_kind < 3; op_kind++) {
		for (int counted = 0; counted < 2; counted++) {
			open_set(&ca, a, ma, 0), open_set(&cb, b, mb, 1);
			kids[0] = &ca, kids[1] = &cb;
			if (op_kind == 0) sift_cur_and(&op, kids, 2, cbuf[4]), rn = ref_and(a, b, ref);
			else if (op_kind == 1) sift_cur_or(&op, kids, 2, cbuf[4]), rn = ref_or(a, b, ref);
			else sift_cur_andnot(&op, kids, cbuf[4]), rn = ref_andnot(a, b, ref);
			if (counted) {
				uint64_t cnt = sift_count(&op);
				CHECK(cnt == rn, "count op %d modes %d/%d: %lu vs %zu", op_kind, ma, mb, cnt, rn);
			} else {
				n = sift_tids(&op, outbuf);
				CHECK(same(outbuf, n, ref, rn), "op %d modes %d/%d: %zu vs %zu", op_kind, ma, mb, n, rn);
			}
		}
	}

	/* tf ranks of AND results in both sources, group by group */
	open_set(&ca, a, ma, 0), open_set(&cb, b, mb, 1);
	kids[0] = &ca, kids[1] = &cb;
	sift_cur_and(&op, kids, 2, cbuf[4]);
	for (const sift_grp *g; (g = sift_next(&op));) {
		uint32_t k = sift_grp_tids(g, NULL, tmp);
		check_tf(a, ma, 0, &ca, g, tmp);
		check_tf(b, mb, 1, &cb, g, tmp);
		if (ca.ub) {
			uint32_t m = 0;
			for (uint32_t i = ca.g.base; i < ca.g.base + ca.g.n; i++) m = m > a->tf[i] ? m : a->tf[i];
			CHECK(sift_grp_ub(&ca) == m, "group ub");
		}
		(void)k;
	}

	/* 3-way and nested trees */
	open_set(&ca, a, ma, 0), open_set(&cb, b, mb, 1), open_set(&cc, c, mc, 2);
	kids[0] = &ca, kids[1] = &cb, kids[2] = &cc;
	sift_cur_and(&op, kids, 3, cbuf[4]);
	set t1 = { tmp, NULL, ref_and(a, b, tmp) };
	rn = ref_and(&t1, c, ref);
	n = sift_tids(&op, outbuf);
	CHECK(same(outbuf, n, ref, rn), "and3 %zu vs %zu", n, rn);

	open_set(&ca, a, ma, 0), open_set(&cb, b, mb, 1), open_set(&cc, c, mc, 2);
	sift_cur_or(&op, kids, 3, cbuf[4]);
	t1.n = ref_or(a, b, tmp);
	rn = ref_or(&t1, c, ref);
	n = sift_tids(&op, outbuf);
	CHECK(same(outbuf, n, ref, rn), "or3 %zu vs %zu", n, rn);

	/* a AND (b OR c), then (that) AND NOT c */
	open_set(&ca, a, ma, 0), open_set(&cb, b, mb, 1), open_set(&cc, c, mc, 2);
	sift_cur *ok[2] = { &cb, &cc }, *ak[2] = { &ca, &op };
	sift_cur_or(&op, ok, 2, cbuf[4]);
	sift_cur_and(&op2, ak, 2, cbuf[5]);
	t1.n = ref_or(b, c, tmp);
	rn = ref_and(a, &t1, ref);
	n = sift_tids(&op2, outbuf);
	CHECK(same(outbuf, n, ref, rn), "nested %zu vs %zu", n, rn);

	/* seeks with random ascending targets */
	open_set(&ca, a, ma, 0);
	uint32_t target = 0;
	for (int s = 0; s < 20 && a->n; s++) {
		target += rndn(1 + (SIFT_GID(a->t[a->n - 1]) - target) / 4 + 1);
		const sift_grp *g = sift_seek(&ca, target);
		uint32_t i = 0;
		while (i < a->n && SIFT_GID(a->t[i]) < target) i++;
		if (i == a->n) {
			CHECK(!g, "seek past end");
			break;
		}
		CHECK(g && g->gid == SIFT_GID(a->t[i]), "seek gid");
		if (!g) break;
		uint32_t j = i;
		while (j < a->n && SIFT_GID(a->t[j]) == g->gid) j++;
		uint32_t k = sift_grp_tids(g, NULL, outbuf);
		CHECK(same(outbuf, k, a->t + i, j - i), "seek content");
		CHECK(sift_grp_count(g) == j - i, "grp count");
		/* page masks: count and tids restricted to random pages */
		uint64_t mask[4] = { rnd() & rnd(), rnd(), rnd() | rnd(), rnd() % 2 ? ~0ull : 0 };
		uint32_t rm = 0;
		for (uint32_t q = i; q < j; q++) {
			uint32_t pg = SIFT_BLK(a->t[q]) & 255;
			if (mask[pg >> 6] >> (pg & 63) & 1) tmp[rm++] = a->t[q];
		}
		CHECK(sift_grp_count_mask(g, mask) == rm, "count_mask");
		k = sift_grp_tids(g, mask, outbuf);
		CHECK(same(outbuf, k, tmp, rm), "tids mask");
		target = g->gid;
	}
}

/* chunk API, keylist and canonical mask kernels on the groups of a */
static void test_chunks(const set *a, sift_tid *tmp)
{
	static uint64_t om[SIFT_GROUP_PAGES * 8], om2[SIFT_GROUP_PAGES * 8], om3[SIFT_GROUP_PAGES * 8];
	static uint8_t chunk[4 + SIFT_GRP_MAXTIDS * 4], gb[SIFT_GRP_BYTES];
	size_t len = sift_keylist_encode(a->t, a->n, blob[3]);
	CHECK(len > 0 && sift_keylist_decode(blob[3], len, tmp) == a->n && !memcmp(tmp, a->t, a->n * 8ull), "keylist");
	for (uint32_t i = 0; i < a->n;) {
		uint32_t j = i, gid = SIFT_GID(a->t[i]), np = 0, last = UINT32_MAX;
		uint64_t pm[4] = { 0 };
		memset(om, 0, sizeof om);
		for (; j < a->n && SIFT_GID(a->t[j]) == gid; j++) {
			uint32_t pg = SIFT_BLK(a->t[j]) & 255, o = SIFT_OFF(a->t[j]) - 1;
			if (pg != last) np++, last = pg, pm[pg >> 6] |= 1ull << (pg & 63);
			om[(np - 1) * 8 + o / 64] |= 1ull << (o & 63);
		}
		int ct = (int)rndn(3);
		size_t cl = sift_chunk_encode(pm, om, ct, chunk);
		CHECK(cl == sift_chunk_encode(pm, om, ct, NULL) && (ct == SIFT_CT_ARRAY || cl <= sift_chunk_bound(np)), "chunk size");
		CHECK(sift_chunk_check(chunk, (uint32_t)cl) == NULL, "chunk check rejects the chunk of the writer (type %d)", ct);
		CHECK(sift_chunk_count(chunk) == j - i, "chunk count");
		uint64_t pm2[4], sel[4] = { rnd(), rnd() | rnd(), ~0ull, rnd() & rnd() };
		sift_chunk_pagemask(chunk, pm2);
		CHECK(!memcmp(pm, pm2, 32), "chunk pagemask");
		uint32_t ns = sift_chunk_offsets(chunk, sel, om2), k = 0, r = 0;
		int ok = 1;
		for (int w = 0; w < 4; w++)
			for (uint64_t m = pm[w]; m; m &= m - 1, r++)
				if (sel[w] & m & -m) ok &= !memcmp(om2 + k++ * 8, om + r * 8, 64);
		CHECK(ok && ns == k, "chunk offsets");
		sift_grp g;
		sift_chunk_grp(chunk, gid, &g, gb);
		uint32_t nt = sift_grp_tids(&g, NULL, tmp);
		CHECK(same(tmp, nt, a->t + i, j - i), "chunk grp");
		/* canonical mask kernels against plain loops */
		for (uint32_t q = 0; q < np * 8; q++) om2[q] = rnd() & rnd();
		uint64_t pc = 0;
		for (uint32_t q = 0; q < np * 8; q++) pc += (uint64_t)__builtin_popcountll(om[q] & om2[q]);
		sift_offmask_and(om3, om, om2, np);
		CHECK(sift_offmask_popcount(om3, np) == pc, "offmask and");
		sift_offmask_or(om3, om, om2, np);
		sift_offmask_andnot(om3, om3, om2, np);
		CHECK(sift_offmask_popcount(om3, np) == j - i - pc, "offmask or/andnot");
		i = j;
	}
}

/* The chunk check must accept what the writer makes at the edges of a chunk and reject it one byte short: the same offsets on each given page, every type */
static void edge_chunk(const char *what, const int *pages, int np, const int *offs, int noff)
{
	static uint64_t om[256 * SIFT_OFFMASK_WORDS];
	static uint8_t chunk[4 + 256 * 512 * 4 + 64];
	uint64_t pm[4] = { 0 };
	memset(om, 0, sizeof om);
	for (int k = 0; k < np; k++) {
		pm[pages[k] >> 6] |= 1ull << (pages[k] & 63);
		for (int q = 0; q < noff; q++) om[k * SIFT_OFFMASK_WORDS + offs[q] / 64] |= 1ull << (offs[q] & 63);
	}
	for (int ct = 0; ct < 4; ct++) {
		size_t cl = sift_chunk_encode(pm, om, ct, chunk);
		CHECK(sift_chunk_check(chunk, (uint32_t)cl) == NULL, "chunk check rejects the writer's %s chunk (type %d): %s", what, ct, sift_chunk_check(chunk, (uint32_t)cl));
		CHECK(sift_chunk_check(chunk, (uint32_t)cl - 1) != NULL, "chunk check accepts a %s chunk (type %d) cut by a byte", what, ct);
		CHECK(sift_chunk_count(chunk) == (uint32_t)(np * noff), "%s chunk count", what);
	}
}

static void test_chunk_edges(void)
{
	int all_pages[256], all_offs[512];
	const int none[1] = { 0 };
	for (int i = 0; i < 256; i++) all_pages[i] = i;
	for (int i = 0; i < 512; i++) all_offs[i] = i;
	edge_chunk("empty", none, 0, none, 0);
	edge_chunk("first entry", all_pages, 1, all_offs, 1);
	edge_chunk("last entry", all_pages + 255, 1, all_offs + 511, 1);
	edge_chunk("corner entries", (const int[]){ 0, 255 }, 2, (const int[]){ 0, 511 }, 2);
	edge_chunk("every page and offset", all_pages, 256, all_offs, 512);
	/* the widest offset of a page decides the width of its row: the bytes needed, then 1, 2, 4 or a multiple of 8 */
	for (int i = 0, m[] = { 7, 8, 15, 16, 31, 32, 63, 64, 255, 256, 257, 511 }; i < (int)(sizeof m / sizeof *m); i++) {
		int offs[2] = { 0, m[i] };
		edge_chunk("width", (const int[]){ 17 }, 1, offs, 2);
	}
}

int main(int argc, char **argv)
{
	int iters = argc > 1 ? atoi(argv[1]) : 3000;
	set s[3];
	for (int i = 0; i < 3; i++) s[i].t = malloc(MAXN * 8), s[i].tf = malloc(MAXN);
	for (int i = 0; i < 4; i++) blob[i] = malloc(sift_encode_bound(MAXN));
	outbuf = malloc(4 * MAXN * 8);
	sift_tid *ref = malloc(4 * MAXN * 8), *tmp = malloc(4 * MAXN * 8);
	sift_init();
	for (int isa = 0; isa < SIFT_NISA; isa++) {
		if (!sift_set_isa(isa)) {
			printf("%s: unsupported, skipped\n", sift_isa_name(isa));
			continue;
		}
		long f0 = fails, c0 = checks;
		for (int it = 0; it < iters; it++) {
			for (int i = 0; i < 3; i++) gen_shape(&s[i]);
			if (rnd() % 3 == 0) memcpy(s[1].t, s[0].t, s[0].n * 8), s[1].n = s[0].n; /* identical sets */
			test_ops(&s[0], &s[1], &s[2], ref, tmp);
			test_chunks(&s[0], tmp);
			if (it == 0) test_chunk_edges();
		}
		printf("%s: %ld checks, %ld failures\n", sift_isa_name(isa), checks - c0, fails - f0);
	}
	for (int i = 0; i < 3; i++) free(s[i].t), free(s[i].tf);
	free(ref), free(tmp);
	return fails != 0;
}