// File: kernel.cpp
// Reference-scalar moe Q8_0 (full quantization: activations + weights, per-32-block
// fp16 scale). LLM target: replace this file with an optimised inner_moe.
// All per-definition constants live in moe_def:: (moe.h).
#include "moe.h"
#include <algorithm>
#include <cmath>
#include <cstring>
#include <vector>
using namespace moe_def;

namespace {

// 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;
}

// Dequantize one row of `cols` int8 values (32-wide blocks, one fp16 scale per
// block) into `out`.
void dequant_row(const int8_t* x, const uint16_t* scales, int cols, float* out) {
    const int nblk = cols / 32;
    for (int blk = 0; blk < nblk; ++blk) {
        const float d = f16_to_f32(scales[blk]);
        for (int l = 0; l < 32; ++l) out[blk * 32 + l] = d * (float)x[blk * 32 + l];
    }
}

} // namespace

extern "C" void inner_moe(
    const int8_t* hidden_states, float* output, const uint16_t* hs_scales,
    const float* router_weight,
    const int8_t* gate_proj, const uint16_t* gate_scales,
    const int8_t* up_proj, const uint16_t* up_scales,
    const int8_t* down_proj, const uint16_t* down_scales,
    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> gate_row(NEmbd), up_row(NEmbd), down_row(NFf), gu(NFf);

    for (int t = 0; t < n_tokens; ++t) {
        const int8_t* hs_row = hidden_states + (long)t * NEmbd;
        const uint16_t* hs_sc_row = hs_scales + (long)t * NEmbdBlk;
        dequant_row(hs_row, hs_sc_row, NEmbd, h.data());

        // 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;

        for (int ki = 0; ki < NExpertUsed; ++ki) {
            const int e = top_idx[ki];
            const int8_t* gate_e  = gate_proj   + (long)e * NFf * NEmbd;
            const uint16_t* gsc_e = gate_scales + (long)e * NFf * NEmbdBlk;
            const int8_t* up_e    = up_proj     + (long)e * NFf * NEmbd;
            const uint16_t* usc_e = up_scales   + (long)e * NFf * NEmbdBlk;
            const int8_t* down_e  = down_proj   + (long)e * NEmbd * NFf;
            const uint16_t* dsc_e = down_scales + (long)e * NEmbd * NFfBlk;

            for (int fi = 0; fi < NFf; ++fi) {
                dequant_row(gate_e + (long)fi * NEmbd, gsc_e + (long)fi * NEmbdBlk, NEmbd, gate_row.data());
                dequant_row(up_e + (long)fi * NEmbd, usc_e + (long)fi * NEmbdBlk, NEmbd, up_row.data());
                float gv = 0.0f, uv = 0.0f;
                for (int d = 0; d < NEmbd; ++d) { gv += gate_row[d] * h[d]; uv += up_row[d] * h[d]; }
                const float silu = gv / (1.0f + std::exp(-gv));
                gu[fi] = silu * uv;
            }

            for (int d = 0; d < NEmbd; ++d) {
                dequant_row(down_e + (long)d * NFf, dsc_e + (long)d * NFfBlk, NFf, down_row.data());
                float acc = 0.0f;
                for (int fi = 0; fi < NFf; ++fi) acc += down_row[fi] * gu[fi];
                out_row[d] += top_w[ki] * acc;
            }
        }
    }
}