// File: kernel.cpp
// Reference-scalar moe bf16 (fp32 accumulate/output).
// 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 {
// bf16 shares fp32's exponent field, so widening is an exact bit-shift — no
// rounding/subnormal handling needed (unlike fp16->fp32).
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;
}
} // namespace

extern "C" void inner_moe(const uint16_t* hidden_states, float* output,
                           const uint16_t* router_weight, const uint16_t* gate_proj,
                           const uint16_t* up_proj, const uint16_t* down_proj,
                           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> gu(NFf);

    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 logits + softmax over all experts.
        for (int e = 0; e < NExpert; ++e) {
            const uint16_t* rw = router_weight + (long)e * NEmbd;
            float dot = 0.0f;
            for (int d = 0; d < NEmbd; ++d) dot += h[d] * bf16_to_f32(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 uint16_t* gate_e = gate_proj + (long)e * NFf * NEmbd;
            const uint16_t* up_e   = up_proj   + (long)e * NFf * NEmbd;
            const uint16_t* down_e = down_proj + (long)e * NEmbd * NFf;

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

            for (int d = 0; d < NEmbd; ++d) {
                float acc = 0.0f;
                const uint16_t* down_row = down_e + (long)d * NFf;
                for (int f = 0; f < NFf; ++f) {
                    acc += bf16_to_f32(down_row[f]) * gu[f];
                }
                out_row[d] += top_w[ki] * acc;
            }
        }
    }
}