// test-moe-perf: microbench for CPU MoE expert GEMM (IQ4_XS, AVX2 path).
// Exercises the exact production path (GGML_OP_MUL_MAT_ID / fused up-gate)
// with qwen4exp-like shapes across batch sizes and thread counts.
// Reports ms/call + effective GB/s of ACTIVE expert weights.
//
// Usage: test-moe-perf [--quick]   (quick trims shapes/threads for CI)

#include "ggml.h"

#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <string>
#include <time.h>
#include <vector>

static uint64_t rng_state = 0x9e3779b97f4a7c15ULL;
static float frand(float lo, float hi) {
    rng_state ^= rng_state << 13;
    rng_state ^= rng_state >> 7;
    rng_state ^= rng_state << 17;
    const double u = (double)(rng_state >> 11) * (1.0 / 9007199254740992.0);
    return (float)(lo + u * (hi - lo));
}

static double now_ms(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_sec * 1000.0 + ts.tv_nsec / 1e6;
}

struct Shape { int K; int N; int E; int used; };

static void bench_one(const Shape & sh, int Ny, int nth, bool fused, bool quick) {
    const int64_t K = sh.K, N = sh.N, E = sh.E, U = sh.used;
    const size_t wrow = ggml_row_size(GGML_TYPE_IQ4_XS, K);
    const size_t mem = (wrow * N * E) * (fused ? 2 : 1)
        + (size_t)K * U * Ny * 4 + (size_t)N * U * Ny * 4 + (size_t)U * Ny * 4
        + 256 * 1024 * 1024;
    struct ggml_init_params ip = { mem, nullptr, false };
    struct ggml_context * ctx = ggml_init(ip);
    if (!ctx) { printf("ggml_init failed\n"); return; }

    // experts: [K, N, E] IQ4_XS from random f32 (values irrelevant for perf)
    std::vector<float> f32((size_t)K * N);
    for (auto & x : f32) x = frand(-1.0f, 1.0f);
    struct ggml_tensor * w1 = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ4_XS, K, N, E);
    ggml_set_name(w1, "w1");
    {
        std::vector<float> imatrix(K, 1.0f);
        char * dst = (char *)w1->data;
        for (int e = 0; e < E; ++e) {
            for (int64_t r = 0; r < N; ++r) {
                ggml_quantize_chunk(GGML_TYPE_IQ4_XS, f32.data() + r * K,
                    dst + ((size_t)e * N + r) * wrow, 0, 1, K, imatrix.data(), nullptr);
            }
        }
    }
    struct ggml_tensor * w2 = nullptr;
    if (fused) {
        w2 = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ4_XS, K, N, E);
        ggml_set_name(w2, "w2");
        memcpy(w2->data, w1->data, ggml_nbytes(w2));
    }
    // activations F32 [K, U, Ny] (3D like production: one row per expert slot)
    struct ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, K, U, Ny);
    ggml_set_name(b, "b");
    for (int64_t i = 0; i < (int64_t)K * U * Ny; ++i) ((float *)b->data)[i] = frand(-1.0f, 1.0f);
    // ids [used, Ny]: random expert per slot
    struct ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, U, Ny);
    ggml_set_name(ids, "ids");
    for (int64_t i = 0; i < (int64_t)U * Ny; ++i) ((int32_t *)ids->data)[i] = (int32_t)(frand(0, 0.999f) * E);

    struct ggml_tensor * out;
    if (fused) {
        out = ggml_moe_up_gate(ctx, w1, w2, b, ids, GGML_UNARY_OP_SILU);
    } else {
        out = ggml_mul_mat_id(ctx, w1, b, ids);
    }
    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16, false);
    ggml_build_forward_expand(gf, out);

    // warmup (also ramps CPU frequency)
    for (int i = 0; i < 8; ++i) ggml_graph_compute_with_ctx(ctx, gf, nth);
    // timed: min over reps (robust to frequency/noise spikes)
    double t0 = now_ms();
    ggml_graph_compute_with_ctx(ctx, gf, nth);
    double dt = now_ms() - t0;
    int reps = dt > 0 ? (int)(1000.0 / dt) : 5;
    if (reps < 8) reps = 8;
    if (reps > 40) reps = 40;
    if (quick && reps > 10) reps = 10;
    double best = 1e30;
    for (int i = 0; i < reps; ++i) {
        t0 = now_ms();
        ggml_graph_compute_with_ctx(ctx, gf, nth);
        dt = now_ms() - t0;
        if (dt < best) best = dt;
    }
    dt = best;

    // active expert bytes per call (up+gate for fused)
    const double bpp = 4.25 / 8.0;
    double bytes = (double)Ny * U * K * N * bpp * (fused ? 2 : 1);
    printf("%s K=%-5d N=%-5d E=%-3d U=%-2d Ny=%-3d T=%-2d  %8.3f ms/call  %7.1f GB/s\n",
           fused ? "fused" : "plain", sh.K, sh.N, sh.E, sh.used, Ny, nth, dt, bytes / dt / 1e6);
    ggml_free(ctx);
}

int main(int argc, char ** argv) {
    bool quick = false;
    int force_threads = 0;
    for (int i = 1; i < argc; ++i) {
        if (!strcmp(argv[i], "--quick")) quick = true;
        else if ((!strcmp(argv[i], "-t") || !strcmp(argv[i], "--threads")) && i + 1 < argc) force_threads = atoi(argv[++i]);
        else { printf("usage: %s [--quick] [-t N]\n", argv[0]); return 1; }
    }
    ggml_quantize_init(GGML_TYPE_IQ4_XS);
    // E=512 mirrors qwen4exp (expert weights >> L3, no false cache hits)
    const std::vector<Shape> shapes = quick ? std::vector<Shape>{{2048, 2048, 128, 10}}
        : std::vector<Shape>{{2048, 1024, 512, 10}, {2048, 2048, 512, 10}, {4096, 2048, 128, 10}};
    const std::vector<int> nys = quick ? std::vector<int>{1, 8, 32} : std::vector<int>{1, 2, 4, 8, 16, 32, 64};
    std::vector<int> nths = quick ? std::vector<int>{16} : std::vector<int>{4, 8, 12, 16};
    if (force_threads > 0) nths = std::vector<int>{force_threads};
    for (auto sh : shapes) {
        for (int Ny : nys) {
            for (int nth : nths) {
                bench_one(sh, Ny, nth, false, quick);
            }
        }
        if (!quick) {
            for (int Ny : nys) bench_one(sh, Ny, 16, true, quick);
        } else {
            for (int Ny : nys) bench_one(sh, Ny, nths[0], true, quick);
        }
    }
    ggml_quantize_free();
    return 0;
}