// Bit-level unit tests for the DT3 dual-plane ternary format
//
// DT3 packs 128 weights as two ternary planes (w = d1*t1 + d2*t2), each plane
// laid out exactly like tq1_0 with all constants halved. A block with
// misplaced bits still loads and generates, so the layout is checked here
// bit by bit against an independent packer that implements the format
// specification directly.
//
// Extra mode for cross-implementation parity checks (see
// tests/test-dt3-rust-parity.py):
//   test-dt3 --dequant IN.bin OUT.f32
// dequantizes raw DT3 blocks from IN.bin into float32 little-endian OUT.f32.

#include "ggml.h"
#include "ggml-cpu.h"

#undef NDEBUG
#include <assert.h>
#include <inttypes.h>
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <vector>

constexpr int    QK_DT3         = 128;
constexpr size_t DT3_QS_BYTES   = 24; // per plane
constexpr size_t DT3_QH_BYTES   = 2;  // per plane
constexpr size_t DT3_BLOCK_SIZE = 2*DT3_QS_BYTES + 2*DT3_QH_BYTES + 2*sizeof(uint16_t);

// byte offsets inside a block (spec: qs[2][24] | qh[2][2] | d[2])
constexpr size_t OFF_QS = 0;
constexpr size_t OFF_QH = 2*DT3_QS_BYTES;
constexpr size_t OFF_D  = 2*DT3_QS_BYTES + 2*DT3_QH_BYTES;

// Independent packer, written from the format specification (not from
// ggml-quants.c): trits in {-1, 0, 1}, element i of the plane goes to
//   region A: qs[m],    m in [0,16), digit n: elements m + n*16   (0..79)
//   region B: qs[16+m], m in [0,8),  digit n: elements 80 + m + n*8 (80..119)
//   region C: qh[j],    j in [0,2),  digit n: elements 120 + j + n*2 (120..127)
// with the first element in the most significant trit, an extra *3 shift in
// region C, and ceiling division by 243 to fit 5 trits per byte.
static void ref_pack_plane(const int8_t * t, uint8_t * qs, uint8_t * qh) {
    for (int m = 0; m < 16; ++m) {
        uint32_t q = 0;
        for (int n = 0; n < 5; ++n) {
            q = q*3 + (uint32_t)(t[m + n*16] + 1);
        }
        qs[m] = (uint8_t)((q*256 + 242)/243);
    }
    for (int m = 0; m < 8; ++m) {
        uint32_t q = 0;
        for (int n = 0; n < 5; ++n) {
            q = q*3 + (uint32_t)(t[80 + m + n*8] + 1);
        }
        qs[16 + m] = (uint8_t)((q*256 + 242)/243);
    }
    for (int j = 0; j < 2; ++j) {
        uint32_t q = 0;
        for (int n = 0; n < 4; ++n) {
            q = q*3 + (uint32_t)(t[120 + j + n*2] + 1);
        }
        q *= 3; // shift the first value to the most significant trit
        qh[j] = (uint8_t)((q*256 + 242)/243);
    }
}

static void ref_pack_block(const int8_t * t1, float d1, const int8_t * t2, float d2, uint8_t * block) {
    ref_pack_plane(t1, block + OFF_QS,                block + OFF_QH);
    ref_pack_plane(t2, block + OFF_QS + DT3_QS_BYTES, block + OFF_QH + DT3_QH_BYTES);
    const uint16_t h1 = ggml_fp32_to_fp16(d1);
    const uint16_t h2 = ggml_fp32_to_fp16(d2);
    memcpy(block + OFF_D,     &h1, sizeof(h1));
    memcpy(block + OFF_D + 2, &h2, sizeof(h2));
}

// deterministic PRNG so failures are reproducible
static uint32_t rng_state = 0x12345678;
static uint32_t rng_next(void) {
    rng_state ^= rng_state << 13;
    rng_state ^= rng_state >> 17;
    rng_state ^= rng_state << 5;
    return rng_state;
}
static int8_t rng_trit(void) {
    return (int8_t)(rng_next() % 3) - 1;
}

static int test_layout_constants(void) {
    int num_failed = 0;

    if (ggml_blck_size(GGML_TYPE_DT3) != QK_DT3) {
        printf("FAILED: blck_size is %" PRId64 ", expected %d\n", ggml_blck_size(GGML_TYPE_DT3), QK_DT3);
        num_failed++;
    }
    if (ggml_type_size(GGML_TYPE_DT3) != DT3_BLOCK_SIZE) {
        printf("FAILED: type_size is %zu, expected %zu\n", ggml_type_size(GGML_TYPE_DT3), DT3_BLOCK_SIZE);
        num_failed++;
    }
    // 56 bytes / 128 weights = 3.5 bpw exactly
    if (DT3_BLOCK_SIZE*8 != (size_t)QK_DT3*7/2) {
        printf("FAILED: not 3.5 bpw\n");
        num_failed++;
    }

    return num_failed;
}

// a single trit set to -1 or +1 at each position of each plane must come back
// at the same position, scaled by that plane's scale only
static int test_single_trits(const ggml_type_traits * qfns) {
    int num_failed = 0;

    const float d1 = 1.0f;   // exact in fp16
    const float d2 = 0.25f;  // exact in fp16

    for (int plane = 0; plane < 2; ++plane) {
        for (int pos = 0; pos < QK_DT3; ++pos) {
            for (int val = -1; val <= 1; val += 2) {
                int8_t t1[QK_DT3] = {0};
                int8_t t2[QK_DT3] = {0};
                (plane == 0 ? t1 : t2)[pos] = (int8_t)val;

                uint8_t block[DT3_BLOCK_SIZE];
                ref_pack_block(t1, d1, t2, d2, block);

                float out[QK_DT3];
                qfns->to_float(block, out, QK_DT3);

                for (int j = 0; j < QK_DT3; ++j) {
                    const float expected = j == pos ? (plane == 0 ? d1 : d2)*val : 0.0f;
                    if (out[j] != expected) {
                        printf("FAILED: plane %d pos %d val %d: out[%d] = %f, expected %f\n",
                               plane, pos, val, j, out[j], expected);
                        num_failed++;
                    }
                }
            }
        }
    }

    return num_failed;
}

// spot-check the byte positions of the region boundaries against values
// computed by hand from the packing formula
static int test_byte_positions(void) {
    int num_failed = 0;

    // verify the element -> byte mapping structurally: flipping element `pos`
    // must change EXACTLY the one byte the formula says, and no other
    struct pos_case {
        int    pos;
        size_t off; // expected changed byte, offset inside plane data
    };
    const pos_case pcases[] = {
        {   0, 0  }, // region A, byte 0
        {  15, 15 }, // region A, byte 15
        {  16, 0  }, // region A, digit 1 of byte 0
        {  64, 0  }, // region A, digit 4 of byte 0
        {  79, 15 }, // region A boundary: last element, byte 15
        {  80, 16 }, // region B boundary: first element, byte 16
        {  87, 23 }, // region B, byte 23
        { 119, 23 }, // region B boundary: last element, byte 23
        { 120, 24 }, // region C boundary: first element, qh[0]
        { 121, 25 }, // region C, qh[1]
        { 126, 24 }, // region C, digit 3 of qh[0]
        { 127, 25 }, // region C, last element, qh[1]
    };

    uint8_t zero_plane[DT3_QS_BYTES + DT3_QH_BYTES];
    {
        const int8_t t0[QK_DT3] = {0};
        ref_pack_plane(t0, zero_plane, zero_plane + DT3_QS_BYTES);
    }

    for (size_t c = 0; c < sizeof(pcases)/sizeof(pcases[0]); ++c) {
        int8_t t[QK_DT3] = {0};
        t[pcases[c].pos] = -1;

        uint8_t plane[DT3_QS_BYTES + DT3_QH_BYTES];
        ref_pack_plane(t, plane, plane + DT3_QS_BYTES);

        for (size_t b = 0; b < sizeof(plane); ++b) {
            const bool should_differ = b == pcases[c].off;
            const bool differs       = plane[b] != zero_plane[b];
            if (differs != should_differ) {
                printf("FAILED: pos %d: byte %zu %s, expected %s\n",
                       pcases[c].pos, b,
                       differs ? "changed" : "unchanged",
                       should_differ ? "changed" : "unchanged");
                num_failed++;
            }
        }
    }

    return num_failed;
}

// random trits and scales (negative scales included) must round-trip exactly
static int test_roundtrip(const ggml_type_traits * qfns) {
    int num_failed = 0;

    const float scales[][2] = {
        { 1.0f,  0.25f  },
        { 0.5f, -0.125f }, // negative second plane
        {-2.0f,  0.75f  }, // negative first plane
        { 0.0f,  0.0f   }, // all-zero scales
    };

    for (size_t sc = 0; sc < sizeof(scales)/sizeof(scales[0]); ++sc) {
        for (int rep = 0; rep < 64; ++rep) {
            int8_t t1[QK_DT3];
            int8_t t2[QK_DT3];
            for (int j = 0; j < QK_DT3; ++j) {
                t1[j] = rng_trit();
                t2[j] = rng_trit();
            }

            const float d1 = ggml_fp16_to_fp32(ggml_fp32_to_fp16(scales[sc][0]));
            const float d2 = ggml_fp16_to_fp32(ggml_fp32_to_fp16(scales[sc][1]));

            uint8_t block[DT3_BLOCK_SIZE];
            ref_pack_block(t1, d1, t2, d2, block);

            float out[QK_DT3];
            qfns->to_float(block, out, QK_DT3);

            for (int j = 0; j < QK_DT3; ++j) {
                const float expected = d1*t1[j] + d2*t2[j];
                if (out[j] != expected) {
                    printf("FAILED: roundtrip scales (%f, %f) rep %d: out[%d] = %f, expected %f\n",
                           d1, d2, rep, j, out[j], expected);
                    num_failed++;
                }
            }
        }
    }

    return num_failed;
}

// the in-tree quantizer must produce the same bytes as the independent packer
// when the input is already exactly ternary (plane 1 = input, plane 2 = 0)
static int test_quantize_pack_parity(const ggml_type_traits_cpu * qfns_cpu) {
    int num_failed = 0;

    for (int rep = 0; rep < 64; ++rep) {
        int8_t t1[QK_DT3];
        const int8_t t2[QK_DT3] = {0};
        float  x[QK_DT3];
        for (int j = 0; j < QK_DT3; ++j) {
            t1[j] = rng_trit();
            x[j]  = (float) t1[j];
        }
        // make sure the block is not all zeros so that d1 == 1.0
        t1[0] = 1;
        x[0]  = 1.0f;

        uint8_t expected[DT3_BLOCK_SIZE];
        ref_pack_block(t1, 1.0f, t2, 0.0f, expected);

        uint8_t block[DT3_BLOCK_SIZE];
        qfns_cpu->from_float(x, block, QK_DT3);

        if (memcmp(block, expected, DT3_BLOCK_SIZE) != 0) {
            for (size_t b = 0; b < DT3_BLOCK_SIZE; ++b) {
                if (block[b] != expected[b]) {
                    printf("FAILED: quantize pack parity rep %d: byte %zu is 0x%02x, expected 0x%02x\n",
                           rep, b, block[b], expected[b]);
                }
            }
            num_failed++;
        }
    }

    return num_failed;
}

// vec_dot against a hand-made sum over the KNOWN trits (not against our own
// dequantization): sum_i y_i * (d1*t1_i + d2*t2_i) with y from q8_0's own
// to_float. Random trits make the qh bytes non-trivial, which would expose a
// vectorization that reads the padding 5th trit of the qh bytes.
static int test_vec_dot(const ggml_type_traits_cpu * qfns_cpu) {
    int num_failed = 0;

    const auto * vdot_traits = ggml_get_type_traits_cpu(qfns_cpu->vec_dot_type);
    const auto * vdot_qfns   = ggml_get_type_traits(qfns_cpu->vec_dot_type);

    if (qfns_cpu->vec_dot_type != GGML_TYPE_Q8_0) {
        printf("FAILED: vec_dot_type is %s, expected q8_0\n", ggml_type_name(qfns_cpu->vec_dot_type));
        return 1;
    }

    const int nblocks = 4;
    const int n       = nblocks*QK_DT3;

    for (int rep = 0; rep < 64; ++rep) {
        std::vector<int8_t> t1(n);
        std::vector<int8_t> t2(n);
        std::vector<float>  d1(nblocks);
        std::vector<float>  d2(nblocks);
        std::vector<uint8_t> xq(nblocks*DT3_BLOCK_SIZE);

        for (int i = 0; i < nblocks; ++i) {
            for (int j = 0; j < QK_DT3; ++j) {
                t1[i*QK_DT3 + j] = rng_trit();
                t2[i*QK_DT3 + j] = rng_trit();
            }
            // fp16-exact scales of both signs
            d1[i] = (float)((int)(rng_next() % 9) - 4) * 0.25f;
            d2[i] = (float)((int)(rng_next() % 9) - 4) * 0.0625f;
            ref_pack_block(&t1[i*QK_DT3], d1[i], &t2[i*QK_DT3], d2[i], &xq[i*DT3_BLOCK_SIZE]);
        }

        std::vector<float> y(n);
        for (int j = 0; j < n; ++j) {
            y[j] = 0.1f + 2.0f*cosf((float)(j + rep));
        }

        std::vector<uint8_t> yq(ggml_row_size(qfns_cpu->vec_dot_type, n));
        vdot_traits->from_float(y.data(), yq.data(), n);

        // exact values the integer path sees
        std::vector<float> ydq(n);
        vdot_qfns->to_float(yq.data(), ydq.data(), n);

        double ref = 0.0;
        for (int i = 0; i < nblocks; ++i) {
            for (int j = 0; j < QK_DT3; ++j) {
                const int ij = i*QK_DT3 + j;
                ref += (double)ydq[ij] * ((double)d1[i]*t1[ij] + (double)d2[i]*t2[ij]);
            }
        }

        float result = INFINITY;
        qfns_cpu->vec_dot(n, &result, 0, xq.data(), 0, yq.data(), 0, 1);

        const float err = fabsf(result - (float)ref);
        const float tol = 1e-4f * (float)n;
        if (!(err <= tol)) {
            printf("FAILED: vec_dot rep %d: got %f, expected %f (err %f)\n", rep, result, (float)ref, err);
            num_failed++;
        }
    }

    return num_failed;
}

// --dequant IN.bin OUT.f32 : dequantize raw DT3 blocks, for parity checks
// against external packers (ternaria's Rust pack_dt3)
static int run_dequant_file(const char * in_path, const char * out_path) {
    FILE * fin = fopen(in_path, "rb");
    if (!fin) {
        fprintf(stderr, "error: cannot open %s\n", in_path);
        return 1;
    }
    fseek(fin, 0, SEEK_END);
    const long size = ftell(fin);
    fseek(fin, 0, SEEK_SET);
    if (size <= 0 || size % DT3_BLOCK_SIZE != 0) {
        fprintf(stderr, "error: %s size %ld is not a multiple of %zu\n", in_path, size, DT3_BLOCK_SIZE);
        fclose(fin);
        return 1;
    }
    std::vector<uint8_t> data(size);
    if (fread(data.data(), 1, size, fin) != (size_t)size) {
        fprintf(stderr, "error: short read on %s\n", in_path);
        fclose(fin);
        return 1;
    }
    fclose(fin);

    const int64_t nel = (int64_t)(size/DT3_BLOCK_SIZE)*QK_DT3;
    std::vector<float> out(nel);
    ggml_get_type_traits(GGML_TYPE_DT3)->to_float(data.data(), out.data(), nel);

    FILE * fout = fopen(out_path, "wb");
    if (!fout) {
        fprintf(stderr, "error: cannot open %s\n", out_path);
        return 1;
    }
    fwrite(out.data(), sizeof(float), nel, fout);
    fclose(fout);

    return 0;
}

int main(int argc, char * argv[]) {
    if (argc == 4 && strcmp(argv[1], "--dequant") == 0) {
        return run_dequant_file(argv[2], argv[3]);
    }
    if (argc != 1) {
        fprintf(stderr, "usage: %s [--dequant IN.bin OUT.f32]\n", argv[0]);
        return 1;
    }

    ggml_cpu_init();

    const auto * qfns     = ggml_get_type_traits(GGML_TYPE_DT3);
    const auto * qfns_cpu = ggml_get_type_traits_cpu(GGML_TYPE_DT3);

    int num_failed = 0;

    num_failed += test_layout_constants();
    num_failed += test_single_trits(qfns);
    num_failed += test_byte_positions();
    num_failed += test_roundtrip(qfns);
    num_failed += test_quantize_pack_parity(qfns_cpu);
    num_failed += test_vec_dot(qfns_cpu);

    printf("%d tests failed\n", num_failed);

    return num_failed > 0;
}