// File: kernel.cpp
// Reference-scalar moe Q4_K (weight-only, genuine low-bit dot product).
// LLM target: replace this file with an optimised inner_moe (SVE2/NEON
// nibble-unpack + int8 dot-product intrinsics are the intended optimization
// surface here).
//
// Router logits use hidden_states (widened bf16) and router_weight (plain
// fp32) directly, with NO quantization anywhere in the routing path --
// routing decisions are categorical top-k and must not flip from
// quantization rounding. gate/up/down all mirror real GGML's
// ggml_vec_dot_q4_K_q8_K: the relevant fp32 activation vector (hidden state
// for gate/up, the gated silu(gate)*up intermediate for down) is dynamically
// quantized into 256-element blocks (one scale, int8 values), then each
// 32-element sub-block does an integer dot product against the unpacked
// Q4_K nibbles, scale-corrected and accumulated in fp32. This matches
// gemm_q4_k_m.kernel.cpp.tmpl's dot-product core, applied once per expert
// weight matrix.
//
// All per-definition constants live in moe_def:: (moe.h).
#include "moe.h"
#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstring>
#include <vector>
using namespace moe_def;

namespace {

// bf16 shares fp32's exponent field, so widening is an exact bit-shift.
inline float bf16_to_f32(uint16_t bits) {
    uint32_t b = (uint32_t)bits << 16;
    float f;
    std::memcpy(&f, &b, sizeof(f));
    return f;
}

// Proper IEEE-754 half->single conversion (fp16, not bf16).
inline float f16_to_f32(uint16_t h) {
    uint32_t sign = (uint32_t)(h & 0x8000) << 16;
    uint32_t exp  = (h >> 10) & 0x1F;
    uint32_t mant = h & 0x3FF;
    uint32_t bits;
    if (exp == 0) {
        if (mant == 0) {
            bits = sign;
        } else {
            exp = 1;
            while ((mant & 0x400) == 0) { mant <<= 1; --exp; }
            mant &= 0x3FF;
            bits = sign | ((exp + 112u) << 23) | (mant << 13);
        }
    } else if (exp == 0x1F) {
        bits = sign | 0x7F800000u | (mant << 13);
    } else {
        bits = sign | ((exp + 112u) << 23) | (mant << 13);
    }
    float f;
    std::memcpy(&f, &bits, sizeof(f));
    return f;
}

// Dynamically quantize a fp32 vector into 256-element Q8_K-style blocks:
// one scale per block, int8 values in [-127, 127]. len must be a multiple
// of 256 (guaranteed here -- n_embd/n_ff are both chosen divisible by 256).
void quantize_q8k(const float* x, int len, std::vector<float>& d_out, std::vector<int8_t>& q8_out) {
    const int nblk = len / 256;
    d_out.resize(nblk);
    q8_out.resize(len);
    for (int sb = 0; sb < nblk; ++sb) {
        float amax = 0.0f;
        for (int l = 0; l < 256; ++l) {
            const float v = std::fabs(x[sb * 256 + l]);
            if (v > amax) amax = v;
        }
        const float d = amax / 127.0f;
        d_out[sb] = d;
        const float id = d > 0.0f ? 1.0f / d : 0.0f;
        for (int l = 0; l < 256; ++l) {
            int q = (int)std::lround(x[sb * 256 + l] * id);
            if (q > 127) q = 127;
            if (q < -127) q = -127;
            q8_out[sb * 256 + l] = (int8_t)q;
        }
    }
}

// Dot product of one Q4_K-quantized weight row [K elements, K%256==0]
// against a pre-quantized Q8_K-style activation of the same length.
float q4k_dot(const uint8_t* w_q4, const uint16_t* w_scales, const uint16_t* w_mins,
              const std::vector<float>& d_act, const std::vector<int8_t>& q8, int K)
{
    const int K_sub = K / 32;
    float acc = 0.0f;
    for (int sb32 = 0; sb32 < K_sub; ++sb32) {
        const float scale = f16_to_f32(w_scales[sb32]);
        const float minv = f16_to_f32(w_mins[sb32]);
        const float dact = d_act[sb32 / 8];  // 8 sub-blocks per 256-elem superblock

        int32_t sumi = 0;
        int32_t bsum = 0;
        const int base = sb32 * 32;
        for (int l = 0; l < 32; ++l) {
            const int k = base + l;
            const uint8_t byte = w_q4[k / 2];
            const int nib = (k % 2 == 0) ? (byte & 0xF) : (byte >> 4);
            const int8_t qv = q8[k];
            sumi += nib * qv;
            bsum += qv;
        }
        acc += dact * (scale * (float)sumi - minv * (float)bsum);
    }
    return acc;
}

} // namespace

extern "C" void inner_moe(
    const uint16_t* hidden_states, float* output, const float* router_weight,
    const uint8_t* gate_proj_q4, const uint16_t* gate_proj_scales, const uint16_t* gate_proj_mins,
    const uint8_t* up_proj_q4, const uint16_t* up_proj_scales, const uint16_t* up_proj_mins,
    const uint8_t* down_proj_q4, const uint16_t* down_proj_scales, const uint16_t* down_proj_mins,
    int n_tokens)
{
    std::vector<float> h(NEmbd);
    std::vector<float> logits(NExpert);
    std::vector<float> probs(NExpert);
    std::vector<bool> used(NExpert);
    std::vector<int> top_idx(NExpertUsed);
    std::vector<float> top_w(NExpertUsed);

    std::vector<float> d_act_h;
    std::vector<int8_t> q8_h;
    std::vector<float> gu(NFf);
    std::vector<float> d_act_gu;
    std::vector<int8_t> q8_gu;

    for (int t = 0; t < n_tokens; ++t) {
        const uint16_t* hs_row = hidden_states + (long)t * NEmbd;
        for (int d = 0; d < NEmbd; ++d) h[d] = bf16_to_f32(hs_row[d]);

        // Router: plain fp32 dot product, no quantization anywhere.
        for (int e = 0; e < NExpert; ++e) {
            const float* rw = router_weight + (long)e * NEmbd;
            float dot = 0.0f;
            for (int d = 0; d < NEmbd; ++d) dot += h[d] * rw[d];
            logits[e] = dot;
        }
        float max_logit = logits[0];
        for (int e = 1; e < NExpert; ++e) if (logits[e] > max_logit) max_logit = logits[e];
        float sum_exp = 0.0f;
        for (int e = 0; e < NExpert; ++e) {
            probs[e] = std::exp(logits[e] - max_logit);
            sum_exp += probs[e];
        }
        for (int e = 0; e < NExpert; ++e) probs[e] /= sum_exp;

        // Top-k selection (repeated argmax over unselected experts) + renormalize.
        std::fill(used.begin(), used.end(), false);
        for (int ki = 0; ki < NExpertUsed; ++ki) {
            int best = -1;
            for (int e = 0; e < NExpert; ++e) {
                if (!used[e] && (best < 0 || probs[e] > probs[best])) best = e;
            }
            used[best] = true;
            top_idx[ki] = best;
            top_w[ki] = probs[best];
        }
        float wsum = 0.0f;
        for (int ki = 0; ki < NExpertUsed; ++ki) wsum += top_w[ki];
        for (int ki = 0; ki < NExpertUsed; ++ki) top_w[ki] /= wsum;

        float* out_row = output + (long)t * NEmbd;
        for (int d = 0; d < NEmbd; ++d) out_row[d] = 0.0f;

        // Quantize the hidden state once per token -- shared by gate & up
        // for every selected expert (same activation vector each time).
        quantize_q8k(h.data(), NEmbd, d_act_h, q8_h);

        for (int ki = 0; ki < NExpertUsed; ++ki) {
            const int e = top_idx[ki];
            const uint8_t* gate_q4_e  = gate_proj_q4     + (long)e * NFf * NEmbdHalf;
            const uint16_t* gate_sc_e = gate_proj_scales + (long)e * NFf * NEmbdSub;
            const uint16_t* gate_mn_e = gate_proj_mins   + (long)e * NFf * NEmbdSub;
            const uint8_t* up_q4_e    = up_proj_q4       + (long)e * NFf * NEmbdHalf;
            const uint16_t* up_sc_e   = up_proj_scales   + (long)e * NFf * NEmbdSub;
            const uint16_t* up_mn_e   = up_proj_mins     + (long)e * NFf * NEmbdSub;
            const uint8_t* down_q4_e  = down_proj_q4     + (long)e * NEmbd * NFfHalf;
            const uint16_t* down_sc_e = down_proj_scales + (long)e * NEmbd * NFfSub;
            const uint16_t* down_mn_e = down_proj_mins   + (long)e * NEmbd * NFfSub;

            for (int fi = 0; fi < NFf; ++fi) {
                const uint8_t* gate_row_q4  = gate_q4_e  + (long)fi * NEmbdHalf;
                const uint16_t* gate_row_sc = gate_sc_e  + (long)fi * NEmbdSub;
                const uint16_t* gate_row_mn = gate_mn_e  + (long)fi * NEmbdSub;
                const uint8_t* up_row_q4    = up_q4_e    + (long)fi * NEmbdHalf;
                const uint16_t* up_row_sc   = up_sc_e    + (long)fi * NEmbdSub;
                const uint16_t* up_row_mn   = up_mn_e    + (long)fi * NEmbdSub;

                const float gv = q4k_dot(gate_row_q4, gate_row_sc, gate_row_mn, d_act_h, q8_h, NEmbd);
                const float uv = q4k_dot(up_row_q4, up_row_sc, up_row_mn, d_act_h, q8_h, NEmbd);
                const float silu = gv / (1.0f + std::exp(-gv));
                gu[fi] = silu * uv;
            }

            quantize_q8k(gu.data(), NFf, d_act_gu, q8_gu);

            for (int d = 0; d < NEmbd; ++d) {
                const uint8_t* down_row_q4  = down_q4_e + (long)d * NFfHalf;
                const uint16_t* down_row_sc = down_sc_e + (long)d * NFfSub;
                const uint16_t* down_row_mn = down_mn_e + (long)d * NFfSub;
                const float acc = q4k_dot(down_row_q4, down_row_sc, down_row_mn, d_act_gu, q8_gu, NFf);
                out_row[d] += top_w[ki] * acc;
            }
        }
    }
}