#include <stdint.h>
#include <string.h>
#include <math.h>
#include <pthread.h>
#include <stdatomic.h>
#if defined(__AVX2__)
#include <immintrin.h>
#endif

static inline uint64_t be64(const uint8_t* p) { uint64_t v; memcpy(&v, p, 8); return __builtin_bswap64(v); }

static inline void row_words(const uint8_t* tp, int step, int r, uint64_t* w, uint32_t* b8) {
    const int TB = 16 * step;
    if (step == 6) {
        if (r < 15) *w = be64(tp + 6 * r);
        else { uint8_t t[8]; memcpy(t, tp + 90, 6); memcpy(t + 6, tp, 2); *w = be64(t); }
        *b8 = 0;
    } else {
        *w = be64(tp + 8 * r);
        *b8 = tp[(8 * r + 8) % TB];
    }
}

static inline uint32_t state_k(uint64_t w, uint32_t b8, int step, int k) {
    if (step == 6) return (uint32_t)(w >> (48 - 6 * k)) & 0xFFFF;
    if (k < 7) return (uint32_t)(w >> (48 - 8 * k)) & 0xFFFF;
    return (uint32_t)(((w & 0xFF) << 8) | b8);
}

void tmk_tlut10(const float* tl, float* t10) {
    for (int e = 0; e < 1024; e++) {
        t10[2 * e] = (e >> 9) ? -tl[2 * (e & 511)] : tl[2 * (e & 511)];
        t10[2 * e + 1] = tl[2 * (e & 511) + 1];
    }
}

#if defined(__AVX512F__)
#define VEC512 1
static inline __m512 row16(uint64_t w, uint32_t b8, int step, const uint64_t* t10) {
    __m512i s;
    if (step == 6) {
        const __m512i sh = _mm512_setr_epi64(48, 42, 36, 30, 24, 18, 12, 6);
        s = _mm512_srlv_epi64(_mm512_set1_epi64((long long)w), sh);
    } else {
        const __m512i sh = _mm512_setr_epi64(48, 40, 32, 24, 16, 8, 0, 0);
        __m512i wv = _mm512_mask_set1_epi64(_mm512_set1_epi64((long long)w), 0x80, (long long)((w << 8) | b8));
        s = _mm512_srlv_epi64(wv, sh);
    }
    s = _mm512_and_si512(s, _mm512_set1_epi64(0xFFFF));
    __m512i p = _mm512_mul_epu32(s, _mm512_add_epi64(s, _mm512_set1_epi64(1)));
    __m512i ix = _mm512_and_si512(_mm512_srli_epi64(p, 6), _mm512_set1_epi64(1023));
    return _mm512_castsi512_ps(_mm512_i64gather_epi64(ix, (const long long*)t10, 8));
}
#elif defined(__AVX2__)
static inline __m256i row_idx8(const uint8_t* tp, int step, int r) {
    if (step == 6) {
        uint64_t raw;
        if (r < 15) memcpy(&raw, tp + 6 * r, 8);
        else { uint8_t t[8]; memcpy(t, tp + 90, 6); memcpy(t + 6, tp, 2); memcpy(&raw, t, 8); }
        const __m256i shuf = _mm256_setr_epi8(1, 0, -1, -1, 2, 1, 0, -1, 3, 2, 1, -1, 4, 3, 2, -1,
                                              4, 3, 2, -1, 5, 4, 3, -1, 6, 5, 4, -1, 7, 6, 5, -1);
        __m256i v = _mm256_shuffle_epi8(_mm256_set1_epi64x((long long)raw), shuf);
        __m256i s = _mm256_and_si256(_mm256_srlv_epi32(v, _mm256_setr_epi32(0, 2, 4, 6, 0, 2, 4, 6)), _mm256_set1_epi32(0xFFFF));
        return _mm256_srli_epi32(_mm256_mullo_epi16(s, _mm256_add_epi32(s, _mm256_set1_epi32(1))), 6);
    } else {
        __m128i raw;
        if (r < 15) raw = _mm_loadu_si128((const __m128i*)(tp + 8 * r));
        else { uint8_t t[16] = {0}; memcpy(t, tp + 120, 8); t[8] = tp[0]; raw = _mm_loadu_si128((const __m128i*)t); }
        const __m256i shuf = _mm256_setr_epi8(1, 0, -1, -1, 2, 1, -1, -1, 3, 2, -1, -1, 4, 3, -1, -1,
                                              5, 4, -1, -1, 6, 5, -1, -1, 7, 6, -1, -1, 8, 7, -1, -1);
        __m256i s = _mm256_shuffle_epi8(_mm256_broadcastsi128_si256(raw), shuf);
        return _mm256_srli_epi32(_mm256_mullo_epi16(s, _mm256_add_epi32(s, _mm256_set1_epi32(1))), 6);
    }
}
static inline void row16(const uint8_t* tp, int step, int r, const uint64_t* t10, __m256* lo, __m256* hi) {
    __m256i ix = row_idx8(tp, step, r);
    *lo = _mm256_castsi256_ps(_mm256_i32gather_epi64((const long long*)t10, _mm256_castsi256_si128(ix), 8));
    *hi = _mm256_castsi256_ps(_mm256_i32gather_epi64((const long long*)t10, _mm256_extracti128_si256(ix, 1), 8));
}
#endif

static inline void row16_scalar(uint64_t w, uint32_t b8, int step, const float* t10, float* out) {
    for (int k = 0; k < 8; k++) {
        uint32_t s = state_k(w, b8, step, k), e = ((s * (s + 1)) & 0xFFFF) >> 6;
        out[2 * k] = t10[2 * e]; out[2 * k + 1] = t10[2 * e + 1];
    }
}

#if defined(__AVX512VBMI__)
#define VBMI 1
static __m512i PI6[8], PI8[8], MS6;
static volatile int vb_ready = 0;
__attribute__((constructor)) static void vb_init(void) {
    uint8_t b[64];
    for (int p = 0; p < 8; p++) {
        for (int l = 0; l < 8; l++) { int r = 2 * p + (l >= 4); for (int q = 0; q < 8; q++) b[8 * l + q] = (uint8_t)(6 * r + 7 - q); }
        PI6[p] = _mm512_loadu_si512(b);
        for (int d = 0; d < 16; d++) {
            int r = 2 * p + (d >= 8), k = d % 8;
            b[4 * d] = (uint8_t)((8 * r + k + 1) % 128); b[4 * d + 1] = (uint8_t)(8 * r + k); b[4 * d + 2] = 0; b[4 * d + 3] = 0;
        }
        PI8[p] = _mm512_loadu_si512(b);
    }
    for (int l = 0; l < 8; l++) for (int j = 0; j < 2; j++) {
        int k = 2 * (l % 4) + j, o = 48 - 6 * k;
        b[8 * l + 4 * j] = (uint8_t)o; b[8 * l + 4 * j + 1] = (uint8_t)(o + 8); b[8 * l + 4 * j + 2] = 0; b[8 * l + 4 * j + 3] = 0;
    }
    MS6 = _mm512_loadu_si512(b);
    vb_ready = 1;
}
static inline void tile_regs(const uint8_t* tp, int step, __m512i* lo, __m512i* hi) {
    *lo = _mm512_loadu_si512(tp);
    if (step == 6)
        *hi = _mm512_inserti64x4(_mm512_castsi256_si512(_mm256_loadu_si256((const __m256i*)(tp + 64))), _mm512_castsi512_si256(*lo), 1);
    else
        *hi = _mm512_loadu_si512(tp + 64);
}
static inline __m512i pair_idx(__m512i lo, __m512i hi, int step, int p) {
    __m512i s;
    if (step == 6) s = _mm512_and_si512(_mm512_multishift_epi64_epi8(MS6, _mm512_permutex2var_epi8(lo, PI6[p], hi)), _mm512_set1_epi32(0xFFFF));
    else s = _mm512_maskz_permutex2var_epi8(0x3333333333333333ULL, lo, PI8[p], hi);
    return _mm512_srli_epi32(_mm512_mullo_epi16(s, _mm512_add_epi32(s, _mm512_set1_epi32(1))), 6);
}
#endif

const char* tmk_isa(void) {
#if defined(VBMI)
    return "avx512vbmi";
#elif defined(VEC512)
    return "avx512";
#elif defined(__AVX2__)
    return "avx2";
#else
    return "scalar";
#endif
}

void tmk_qgemv_rows(const uint8_t* codes, int step, int n, const float* tlut, const float* x, float* z, int t0, int t1) {
    const int nb = n / 16, TB = 16 * step;
#if defined(VBMI)
    if (!vb_ready) vb_init();
    for (int t = t0; t < t1; t++) {
        const uint8_t* strip = codes + (size_t)t * nb * TB;
        __m512 acc[16];
        for (int r = 0; r < 16; r++) acc[r] = _mm512_setzero_ps();
        for (int j = 0; j < nb; j++) {
            __m512i lo, hi; tile_regs(strip + (size_t)j * TB, step, &lo, &hi);
            const __m512 xv = _mm512_loadu_ps(x + 16 * j);
            for (int p = 0; p < 8; p++) {
                __m512i ix = pair_idx(lo, hi, step, p);
                acc[2 * p] = _mm512_fmadd_ps(_mm512_castsi512_ps(_mm512_i32gather_epi64(_mm512_castsi512_si256(ix), (const long long*)tlut, 8)), xv, acc[2 * p]);
                acc[2 * p + 1] = _mm512_fmadd_ps(_mm512_castsi512_ps(_mm512_i32gather_epi64(_mm512_extracti64x4_epi64(ix, 1), (const long long*)tlut, 8)), xv, acc[2 * p + 1]);
            }
        }
        for (int r = 0; r < 16; r++) z[16 * t + r] = _mm512_reduce_add_ps(acc[r]);
    }
#else
    for (int t = t0; t < t1; t++) {
        const uint8_t* strip = codes + (size_t)t * nb * TB;
        for (int r = 0; r < 16; r++) {
#if defined(VEC512)
            __m512 a = _mm512_setzero_ps(), b = _mm512_setzero_ps();
            int j = 0;
            for (; j + 1 < nb; j += 2) {
                uint64_t w0, w1; uint32_t c0, c1;
                row_words(strip + j * TB, step, r, &w0, &c0); row_words(strip + (j + 1) * TB, step, r, &w1, &c1);
                a = _mm512_fmadd_ps(row16(w0, c0, step, (const uint64_t*)tlut), _mm512_loadu_ps(x + 16 * j), a);
                b = _mm512_fmadd_ps(row16(w1, c1, step, (const uint64_t*)tlut), _mm512_loadu_ps(x + 16 * j + 16), b);
            }
            for (; j < nb; j++) {
                uint64_t w0; uint32_t c0; row_words(strip + j * TB, step, r, &w0, &c0);
                a = _mm512_fmadd_ps(row16(w0, c0, step, (const uint64_t*)tlut), _mm512_loadu_ps(x + 16 * j), a);
            }
            z[16 * t + r] = _mm512_reduce_add_ps(_mm512_add_ps(a, b));
#elif defined(__AVX2__)
            __m256 a0 = _mm256_setzero_ps(), a1 = _mm256_setzero_ps();
            for (int j = 0; j < nb; j++) {
                __m256 lo, hi;
                row16(strip + (size_t)j * TB, step, r, (const uint64_t*)tlut, &lo, &hi);
                a0 = _mm256_fmadd_ps(lo, _mm256_loadu_ps(x + 16 * j), a0);
                a1 = _mm256_fmadd_ps(hi, _mm256_loadu_ps(x + 16 * j + 8), a1);
            }
            __m256 s = _mm256_add_ps(a0, a1);
            __m128 h = _mm_add_ps(_mm256_castps256_ps128(s), _mm256_extractf128_ps(s, 1));
            h = _mm_hadd_ps(h, h); h = _mm_hadd_ps(h, h);
            z[16 * t + r] = _mm_cvtss_f32(h);
#else
            float acc = 0.f, v[16];
            for (int j = 0; j < nb; j++) {
                uint64_t w; uint32_t c; row_words(strip + j * TB, step, r, &w, &c);
                row16_scalar(w, c, step, tlut, v);
                for (int q = 0; q < 16; q++) acc += v[q] * x[16 * j + q];
            }
            z[16 * t + r] = acc;
#endif
        }
    }
#endif
}

#define MAXT 16
typedef struct { const uint8_t* codes; const float* tlut; const float* x; float* z; int step, n, t0, t1; } gjob;
static gjob jobs[MAXT];
static pthread_t thr[MAXT];
static int nthr = 1;
static atomic_int gen = 0, pending = 0, quit = 0;
static pthread_mutex_t mu = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t cv = PTHREAD_COND_INITIALIZER;

static inline void cpu_relax(void) {
#if defined(__x86_64__) || defined(__i386__)
    __builtin_ia32_pause();
#endif
}

static void* worker(void* arg) {
    const int id = (int)(intptr_t)arg;
    int seen = 0;
    for (;;) {
        int g, spins = 0;
        while ((g = atomic_load_explicit(&gen, memory_order_acquire)) == seen && spins < 20000) { cpu_relax(); spins++; }
        if (g == seen) {
            pthread_mutex_lock(&mu);
            while ((g = atomic_load_explicit(&gen, memory_order_acquire)) == seen) pthread_cond_wait(&cv, &mu);
            pthread_mutex_unlock(&mu);
        }
        seen = g;
        if (atomic_load(&quit)) return NULL;
        const gjob* j = &jobs[id];
        if (j->t1 > j->t0) tmk_qgemv_rows(j->codes, j->step, j->n, j->tlut, j->x, j->z, j->t0, j->t1);
        atomic_fetch_sub_explicit(&pending, 1, memory_order_acq_rel);
    }
}

int tmk_threads(int n) {
    if (n < 1) n = 1;
    if (n > MAXT) n = MAXT;
    if (n == nthr) return nthr;
    if (nthr > 1) {
        atomic_store(&quit, 1);
        pthread_mutex_lock(&mu); atomic_fetch_add(&gen, 1); pthread_cond_broadcast(&cv); pthread_mutex_unlock(&mu);
        for (int i = 1; i < nthr; i++) pthread_join(thr[i], NULL);
        atomic_store(&quit, 0);
    }
    atomic_store(&gen, 0);
    nthr = n;
    for (int i = 1; i < nthr; i++) pthread_create(&thr[i], NULL, worker, (void*)(intptr_t)i);
    return nthr;
}

void tmk_qgemv(const uint8_t* codes, int step, int m, int n, const float* tlut, const float* x, float* z) {
    const int tiles = m / 16, T = nthr;
    if (T == 1 || tiles < 2 * T || (size_t)m * n < (1u << 15)) { tmk_qgemv_rows(codes, step, n, tlut, x, z, 0, tiles); return; }
    const int per = (tiles + T - 1) / T;
    for (int w = 1; w < T; w++) {
        int t0 = w * per, t1 = (w + 1) * per;
        if (t0 > tiles) t0 = tiles;
        if (t1 > tiles) t1 = tiles;
        jobs[w] = (gjob){codes, tlut, x, z, step, n, t0, t1};
    }
    atomic_store_explicit(&pending, T - 1, memory_order_release);
    pthread_mutex_lock(&mu);
    atomic_fetch_add_explicit(&gen, 1, memory_order_acq_rel);
    pthread_cond_broadcast(&cv);
    pthread_mutex_unlock(&mu);
    tmk_qgemv_rows(codes, step, n, tlut, x, z, 0, per < tiles ? per : tiles);
    while (atomic_load_explicit(&pending, memory_order_acquire) > 0) cpu_relax();
}

void tmk_qdense(const uint8_t* codes, int step, int m, int n, const float* tlut, float scale, float* out) {
    const int nb = n / 16, TB = 16 * step;
#if defined(VBMI)
    if (!vb_ready) vb_init();
    const __m512 sc = _mm512_set1_ps(scale);
    for (int t = 0; t < m / 16; t++)
        for (int j = 0; j < nb; j++) {
            __m512i lo, hi; tile_regs(codes + ((size_t)t * nb + j) * TB, step, &lo, &hi);
            float* o = out + (size_t)(16 * t) * n + 16 * j;
            for (int p = 0; p < 8; p++) {
                __m512i ix = pair_idx(lo, hi, step, p);
                _mm512_storeu_ps(o + (size_t)(2 * p) * n, _mm512_mul_ps(_mm512_castsi512_ps(_mm512_i32gather_epi64(_mm512_castsi512_si256(ix), (const long long*)tlut, 8)), sc));
                _mm512_storeu_ps(o + (size_t)(2 * p + 1) * n, _mm512_mul_ps(_mm512_castsi512_ps(_mm512_i32gather_epi64(_mm512_extracti64x4_epi64(ix, 1), (const long long*)tlut, 8)), sc));
            }
        }
#else
    for (int t = 0; t < m / 16; t++)
        for (int j = 0; j < nb; j++) {
            const uint8_t* tp = codes + ((size_t)t * nb + j) * TB;
            for (int r = 0; r < 16; r++) {
                uint64_t w; uint32_t c;
                row_words(tp, step, r, &w, &c);
                float* o = out + (size_t)(16 * t + r) * n + 16 * j;
#if defined(VEC512)
                _mm512_storeu_ps(o, _mm512_mul_ps(row16(w, c, step, (const uint64_t*)tlut), _mm512_set1_ps(scale)));
#elif defined(__AVX2__)
                __m256 lo, hi; row16(tp, step, r, (const uint64_t*)tlut, &lo, &hi); (void)w; (void)c;
                _mm256_storeu_ps(o, _mm256_mul_ps(lo, _mm256_set1_ps(scale)));
                _mm256_storeu_ps(o + 8, _mm256_mul_ps(hi, _mm256_set1_ps(scale)));
#else
                float v[16]; row16_scalar(w, c, step, tlut, v);
                for (int q = 0; q < 16; q++) o[q] = v[q] * scale;
#endif
            }
        }
#endif
}

#if defined(VEC512)
static inline __m512 fwht16(__m512 v) {
    __m512 b;
    b = _mm512_permute_ps(v, 0xB1);       v = _mm512_mask_sub_ps(_mm512_add_ps(v, b), 0xAAAA, b, v);
    b = _mm512_permute_ps(v, 0x4E);       v = _mm512_mask_sub_ps(_mm512_add_ps(v, b), 0xCCCC, b, v);
    b = _mm512_shuffle_f32x4(v, v, 0xB1); v = _mm512_mask_sub_ps(_mm512_add_ps(v, b), 0xF0F0, b, v);
    b = _mm512_shuffle_f32x4(v, v, 0x4E); v = _mm512_mask_sub_ps(_mm512_add_ps(v, b), 0xFF00, b, v);
    return v;
}
static inline __m512 signs16(__m512 v, const uint8_t* bits, int i) {
    __mmask16 k = (__mmask16)~(bits[i >> 3] | (bits[(i >> 3) + 1] << 8));
    __m512i vi = _mm512_castps_si512(v);
    return _mm512_castsi512_ps(_mm512_mask_xor_epi32(vi, k, vi, _mm512_set1_epi32((int)0x80000000)));
}
#define VW 16
#elif defined(__AVX2__)
static inline __m256 fwht8(__m256 v) {
    __m256 b;
    b = _mm256_permute_ps(v, 0xB1);            v = _mm256_blend_ps(_mm256_add_ps(v, b), _mm256_sub_ps(b, v), 0xAA);
    b = _mm256_permute_ps(v, 0x4E);            v = _mm256_blend_ps(_mm256_add_ps(v, b), _mm256_sub_ps(b, v), 0xCC);
    b = _mm256_permute2f128_ps(v, v, 0x01);    v = _mm256_blend_ps(_mm256_add_ps(v, b), _mm256_sub_ps(b, v), 0xF0);
    return v;
}
static inline __m256 signs8(__m256 v, const uint8_t* bits, int i) {
    const __m256i bitv = _mm256_setr_epi32(1, 2, 4, 8, 16, 32, 64, 128);
    __m256i on = _mm256_cmpeq_epi32(_mm256_and_si256(_mm256_set1_epi32(bits[i >> 3]), bitv), _mm256_setzero_si256());
    return _mm256_xor_ps(v, _mm256_castsi256_ps(_mm256_and_si256(on, _mm256_set1_epi32((int)0x80000000))));
}
#define VW 8
#endif

static void rot_row(float* v, int n, const uint8_t* pre, const uint8_t* post, float scale) {
#if defined(VW)
    if (n >= VW) {
#if defined(VEC512)
        const __m512 sc = _mm512_set1_ps(scale);
        for (int i = 0; i < n; i += 16) {
            __m512 a = _mm512_loadu_ps(v + i);
            if (pre) a = signs16(a, pre, i);
            _mm512_storeu_ps(v + i, fwht16(_mm512_mul_ps(a, sc)));
        }
        for (int h = 16; h < n; h <<= 1)
            for (int i = 0; i < n; i += 2 * h)
                for (int j = i; j < i + h; j += 16) {
                    __m512 a = _mm512_loadu_ps(v + j), b = _mm512_loadu_ps(v + j + h);
                    _mm512_storeu_ps(v + j, _mm512_add_ps(a, b)); _mm512_storeu_ps(v + j + h, _mm512_sub_ps(a, b));
                }
        if (post) for (int i = 0; i < n; i += 16) _mm512_storeu_ps(v + i, signs16(_mm512_loadu_ps(v + i), post, i));
#else
        const __m256 sc = _mm256_set1_ps(scale);
        for (int i = 0; i < n; i += 8) {
            __m256 a = _mm256_loadu_ps(v + i);
            if (pre) a = signs8(a, pre, i);
            _mm256_storeu_ps(v + i, fwht8(_mm256_mul_ps(a, sc)));
        }
        for (int h = 8; h < n; h <<= 1)
            for (int i = 0; i < n; i += 2 * h)
                for (int j = i; j < i + h; j += 8) {
                    __m256 a = _mm256_loadu_ps(v + j), b = _mm256_loadu_ps(v + j + h);
                    _mm256_storeu_ps(v + j, _mm256_add_ps(a, b)); _mm256_storeu_ps(v + j + h, _mm256_sub_ps(a, b));
                }
        if (post) for (int i = 0; i < n; i += 8) _mm256_storeu_ps(v + i, signs8(_mm256_loadu_ps(v + i), post, i));
#endif
        return;
    }
#endif
    uint32_t* u = (uint32_t*)v;
    if (pre) for (int i = 0; i < n; i++) u[i] ^= (uint32_t)(~(pre[i >> 3] >> (i & 7)) & 1) << 31;
    for (int i = 0; i < n; i++) v[i] *= scale;
    for (int h = 1; h < n; h <<= 1)
        for (int i = 0; i < n; i += 2 * h)
            for (int j = i; j < i + h; j++) { float a = v[j], b = v[j + h]; v[j] = a + b; v[j + h] = a - b; }
    if (post) for (int i = 0; i < n; i++) u[i] ^= (uint32_t)(~(post[i >> 3] >> (i & 7)) & 1) << 31;
}

void tmk_rot(float* x, int T, int n, const uint8_t* pre, const uint8_t* post, float scale) {
    for (int t = 0; t < T; t++) rot_row(x + (size_t)t * n, n, pre, post, scale);
}

void tmk_fwht(float* x, int T, int n) { tmk_rot(x, T, n, NULL, NULL, 1.0f / sqrtf((float)n)); }

void tmk_signs(float* x, int T, int n, const uint8_t* bits) {
    for (int t = 0; t < T; t++) {
        uint32_t* v = (uint32_t*)(x + (size_t)t * n);
        for (int i = 0; i < n; i++) v[i] ^= (uint32_t)(~(bits[i >> 3] >> (i & 7)) & 1) << 31;
    }
}

void tmk_linear_vec(const uint8_t* codes, int step, int m, int n, const float* tlut, const uint8_t* su,
                    const uint8_t* sv, float wscale, const float* x, float* y, float* tmp) {
    memcpy(tmp, x, sizeof(float) * n);
    rot_row(tmp, n, su, NULL, 1.0f / sqrtf((float)n));
    tmk_qgemv(codes, step, m, n, tlut, tmp, y);
    rot_row(y, m, NULL, sv, wscale / sqrtf((float)m));
}