/*
 * M4 — concurrent CPU(NEON) + QPU(V3D) throughput.
 *
 * Phase 1 §"Decision rules" YELLOW-band rule says: at 0.5 ≤ R < 1.0,
 * the question isn't "is QPU faster" but "does QPU offload buy total
 * system throughput when CPU is also working."
 *
 * Modes (selected with --mode):
 *   neon-only     N NEON pthread workers, pinned 0..N-1, no QPU
 *   qpu-only      QPU dispatch loop on main thread, no NEON
 *   mixed         N NEON pthread workers + QPU dispatch on its own thread
 *
 * Time-based loop (--duration seconds). Workers all start at a
 * pthread_barrier release, stop when a shared volatile flag is set
 * by the timer thread. Each worker counts blocks completed; sum is
 * the system aggregate.
 *
 * Decision (from this binary's output, by inspection):
 *   if mixed (--neon 3 + qpu) > neon-only --threads 4   → offload wins
 *   if mixed ≈ neon-only --threads 4                    → offload neutral
 *   if mixed < neon-only --threads 4                    → bandwidth contention hurts
 *
 * License: BSD-2-Clause; links FFmpeg NEON snapshot (LGPL-2.1+).
 */
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <stddef.h>
#include <time.h>
#include <getopt.h>
#include <pthread.h>
#include <sched.h>
#include <vulkan/vulkan.h>

#include "v3d_runner.h"

extern void ff_vp9_idct_idct_8x8_add_neon(
    uint8_t *dst, ptrdiff_t stride, int16_t *block, int eob);

/* --- RNG + block gen (same shape as bench_neon_idct.c) ----------- */

static uint64_t xs_seed_init(uint64_t s) { return s ? s : 0xdeadbeefcafebabeULL; }
static inline uint64_t xs_step(uint64_t *s) {
    uint64_t x = *s; x ^= x << 13; x ^= x >> 7; x ^= x << 17; return *s = x;
}
static int gen_block(int16_t block[64], uint64_t *s) {
    memset(block, 0, 64 * sizeof(*block));
    int eob = 0;
    int n_nonzero = 1 + (int)(xs_step(s) % 16);
    for (int i = 0; i < n_nonzero; i++) {
        int pos = (int)(xs_step(s) % 64);
        int16_t coef = (int16_t)((int)(xs_step(s) % 8192) - 4096);
        block[pos] = coef;
        if (pos + 1 > eob) eob = pos + 1;
    }
    if (eob == 0) eob = 1;
    return eob;
}
static double now_seconds(void) {
    struct timespec ts;
    clock_gettime(CLOCK_MONOTONIC_RAW, &ts);
    return ts.tv_sec + ts.tv_nsec * 1e-9;
}

/* --- Shared between timer thread and workers ---------------------- */

static volatile int g_stop = 0;
static pthread_barrier_t g_start_barrier;

/* --- NEON worker --------------------------------------------------- */

typedef struct {
    int      worker_id;
    int      affinity_core;
    uint64_t blocks_done;     /* output */
    double   elapsed_s;       /* output */
} neon_args;

static const int NEON_BATCH = 8192;   /* blocks held in memory per worker */

static void *neon_worker(void *p)
{
    neon_args *a = p;

    /* Pin to core. Hertz has 4 A76 cores (0..3). */
    cpu_set_t cs; CPU_ZERO(&cs); CPU_SET(a->affinity_core, &cs);
    pthread_setaffinity_np(pthread_self(), sizeof(cs), &cs);

    /* Per-worker random blocks + preds. Pre-generate to keep gen cost
     * out of the timed loop. */
    uint64_t s = xs_seed_init((uint64_t)a->worker_id * 0xc01dbeefULL);
    int16_t *blocks_master = malloc((size_t)NEON_BATCH * 64 * sizeof(int16_t));
    int16_t *blocks_work   = malloc((size_t)NEON_BATCH * 64 * sizeof(int16_t));
    uint8_t *preds         = malloc((size_t)NEON_BATCH * 64);
    uint8_t *dsts          = malloc((size_t)NEON_BATCH * 64);
    int     *eobs          = malloc(NEON_BATCH * sizeof(int));
    for (int i = 0; i < NEON_BATCH; i++) {
        eobs[i] = gen_block(blocks_master + i * 64, &s);
        for (int j = 0; j < 64; j++) preds[i * 64 + j] = (uint8_t)(xs_step(&s) & 0xff);
    }

    /* Barrier: every worker (and the timer thread) waits here.
     * The timer thread starts its clock immediately after release. */
    pthread_barrier_wait(&g_start_barrier);
    double t0 = now_seconds();

    uint64_t done = 0;
    while (!g_stop) {
        memcpy(blocks_work, blocks_master, (size_t)NEON_BATCH * 64 * sizeof(int16_t));
        memcpy(dsts, preds, (size_t)NEON_BATCH * 64);
        for (int i = 0; i < NEON_BATCH; i++)
            ff_vp9_idct_idct_8x8_add_neon(dsts + i * 64, 8,
                                          blocks_work + i * 64, eobs[i]);
        done += NEON_BATCH;
    }
    a->elapsed_s   = now_seconds() - t0;
    a->blocks_done = done;
    free(blocks_master); free(blocks_work); free(preds); free(dsts); free(eobs);
    return NULL;
}

/* --- QPU worker (runs on its own pthread for fair pacing) --------- */

typedef struct {
    int      affinity_core;       /* core to pin the host thread to */
    int      frame_blocks_x;      /* blocks_per_row */
    int      frame_blocks_y;      /* rows_of_blocks */
    int      blocks_per_wg;
    uint64_t blocks_done;
    double   elapsed_s;
} qpu_args;

typedef struct {
    uint32_t n_blocks;
    uint32_t blocks_per_row;
    uint32_t dst_stride_u8;
    uint32_t _pad;
} push_consts;

static void *qpu_worker(void *p)
{
    qpu_args *a = p;

    cpu_set_t cs; CPU_ZERO(&cs); CPU_SET(a->affinity_core, &cs);
    pthread_setaffinity_np(pthread_self(), sizeof(cs), &cs);

    v3d_runner *r = v3d_runner_create();
    if (!r) { fprintf(stderr, "qpu worker: v3d_runner_create failed\n"); return NULL; }

    int dst_width  = a->frame_blocks_x * 8;
    int dst_height = a->frame_blocks_y * 8;
    int dst_stride = dst_width;
    size_t n_blocks = (size_t) a->frame_blocks_x * a->frame_blocks_y;
    size_t dst_bytes = (size_t) dst_height * dst_stride;

    v3d_buffer buf_coeffs = {0}, buf_dst = {0}, buf_meta = {0};
    v3d_runner_create_buffer(r, n_blocks * 64 * sizeof(int16_t), &buf_coeffs);
    v3d_runner_create_buffer(r, dst_bytes,                       &buf_dst);
    v3d_runner_create_buffer(r, n_blocks * 2 * sizeof(uint32_t), &buf_meta);

    /* Fill with deterministic content; we don't check correctness in
     * this bench (Phase 6 already verified M1' = 100%). */
    uint64_t s = 0xfeedfacecafebabeULL;
    int16_t  *m_coeffs = malloc(n_blocks * 64 * sizeof(int16_t));
    uint8_t  *m_pred   = malloc(dst_bytes);
    for (size_t b = 0; b < n_blocks; b++) gen_block(m_coeffs + b * 64, &s);
    for (size_t i = 0; i < dst_bytes; i++) m_pred[i] = (uint8_t)(xs_step(&s) & 0xff);
    memcpy(buf_coeffs.mapped, m_coeffs, buf_coeffs.size);
    uint32_t *meta = buf_meta.mapped;
    for (size_t b = 0; b < n_blocks; b++) {
        meta[2*b+0] = (uint32_t)(b % a->frame_blocks_x);
        meta[2*b+1] = (uint32_t)(b / a->frame_blocks_x);
    }

    v3d_pipeline pipe = {0};
    v3d_runner_create_pipeline(r, "v3d_idct8.spv", 3, sizeof(push_consts), &pipe);
    v3d_buffer bind_bufs[3] = { buf_coeffs, buf_dst, buf_meta };
    v3d_runner_bind_buffers(r, &pipe, bind_bufs, 3);

    uint32_t group_count_x = (uint32_t)((n_blocks + a->blocks_per_wg - 1)
                                        / a->blocks_per_wg);
    push_consts pc = {
        .n_blocks       = (uint32_t)n_blocks,
        .blocks_per_row = (uint32_t)a->frame_blocks_x,
        .dst_stride_u8  = (uint32_t)dst_stride,
        ._pad           = 0,
    };
    VkCommandBuffer cb = v3d_runner_alloc_cmdbuf(r);
    VkCommandBufferBeginInfo cbbi = { .sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO };
    vkBeginCommandBuffer(cb, &cbbi);
    vkCmdBindPipeline(cb, VK_PIPELINE_BIND_POINT_COMPUTE, pipe.pipeline);
    vkCmdBindDescriptorSets(cb, VK_PIPELINE_BIND_POINT_COMPUTE,
                            pipe.layout, 0, 1, &pipe.desc_set, 0, NULL);
    vkCmdPushConstants(cb, pipe.layout, VK_SHADER_STAGE_COMPUTE_BIT,
                       0, sizeof(pc), &pc);
    vkCmdDispatch(cb, group_count_x, 1, 1);
    vkEndCommandBuffer(cb);

    /* Warm-up */
    for (int i = 0; i < 5; i++) v3d_runner_submit_wait(r, cb);

    pthread_barrier_wait(&g_start_barrier);
    double t0 = now_seconds();

    uint64_t done = 0;
    while (!g_stop) {
        memcpy(buf_dst.mapped, m_pred, dst_bytes);
        v3d_runner_submit_wait(r, cb);
        done += n_blocks;
    }
    a->elapsed_s   = now_seconds() - t0;
    a->blocks_done = done;

    free(m_coeffs); free(m_pred);
    v3d_runner_destroy_pipeline(r, &pipe);
    v3d_runner_destroy_buffer(r, &buf_meta);
    v3d_runner_destroy_buffer(r, &buf_dst);
    v3d_runner_destroy_buffer(r, &buf_coeffs);
    v3d_runner_destroy(r);
    return NULL;
}

/* --- Timer thread --------------------------------------------------- */

typedef struct { double duration_s; } timer_args;

static void *timer_thread(void *p)
{
    timer_args *a = p;
    pthread_barrier_wait(&g_start_barrier);
    /* Spin-and-check rather than usleep, for tighter end. Doesn't matter
     * much over 10s but reduces noise. */
    double end = now_seconds() + a->duration_s;
    while (now_seconds() < end) {
        struct timespec ts = {0, 1000000};  /* 1 ms */
        nanosleep(&ts, NULL);
    }
    g_stop = 1;
    return NULL;
}

/* --- Main ---------------------------------------------------------- */

enum mode { MODE_NEON, MODE_QPU, MODE_MIXED };

int main(int argc, char **argv)
{
    enum mode mode = MODE_NEON;
    int n_neon = 4;
    int qpu_core = 3;
    double duration = 10.0;
    int blocks_per_wg = 32;     /* matches v4 production kernel */
    int frame_w = 1920, frame_h = 1088;

    static struct option opts[] = {
        {"mode",        required_argument, 0, 'm'},
        {"neon-threads",required_argument, 0, 'n'},
        {"qpu-core",    required_argument, 0, 'c'},
        {"duration",    required_argument, 0, 'd'},
        {"blocks-per-wg",required_argument,0, 'b'},
        {"width",       required_argument, 0, 'w'},
        {"height",      required_argument, 0, 'h'},
        {0,0,0,0}
    };
    for (int c; (c = getopt_long(argc, argv, "m:n:c:d:b:w:h:", opts, 0)) != -1;) {
        switch (c) {
        case 'm':
            if      (!strcmp(optarg, "neon-only")) mode = MODE_NEON;
            else if (!strcmp(optarg, "qpu-only"))  mode = MODE_QPU;
            else if (!strcmp(optarg, "mixed"))     mode = MODE_MIXED;
            else { fprintf(stderr, "bad mode\n"); return 2; }
            break;
        case 'n': n_neon = atoi(optarg); break;
        case 'c': qpu_core = atoi(optarg); break;
        case 'd': duration = atof(optarg); break;
        case 'b': blocks_per_wg = atoi(optarg); break;
        case 'w': frame_w = atoi(optarg); break;
        case 'h': frame_h = atoi(optarg); break;
        default: return 2;
        }
    }

    int has_qpu  = (mode == MODE_QPU || mode == MODE_MIXED);
    int has_neon = (mode == MODE_NEON || mode == MODE_MIXED);
    int n_workers = (has_neon ? n_neon : 0) + (has_qpu ? 1 : 0);
    /* Barrier participants: every worker + timer + main (which releases). */
    int barrier_count = n_workers + 1 /* timer */ + 1 /* main */;

    printf("=== M4 concurrent bench ===\n");
    printf("  mode:          %s\n",
           mode == MODE_NEON ? "neon-only" :
           mode == MODE_QPU  ? "qpu-only"  : "mixed");
    printf("  neon threads:  %d (cores 0..%d)\n", has_neon ? n_neon : 0,
           has_neon ? n_neon - 1 : -1);
    printf("  qpu host core: %d (driver thread)\n", has_qpu ? qpu_core : -1);
    printf("  duration:      %.1f s\n", duration);
    printf("  qpu frame:     %dx%d (%d blocks/dispatch, %d blocks/WG)\n",
           frame_w, frame_h,
           (frame_w/8) * (frame_h/8), blocks_per_wg);
    printf("  NEON_BATCH per worker: %d blocks\n", NEON_BATCH);
    printf("\n");

    pthread_barrier_init(&g_start_barrier, NULL, barrier_count);

    pthread_t   timer_tid;
    timer_args  t_args = { .duration_s = duration };
    pthread_create(&timer_tid, NULL, timer_thread, &t_args);

    pthread_t   neon_tids[16] = {0};
    neon_args   n_args[16]    = {0};
    if (has_neon) {
        for (int i = 0; i < n_neon; i++) {
            n_args[i] = (neon_args){ .worker_id = i, .affinity_core = i };
            pthread_create(&neon_tids[i], NULL, neon_worker, &n_args[i]);
        }
    }

    pthread_t qpu_tid = 0;
    qpu_args  q_args  = {0};
    if (has_qpu) {
        q_args = (qpu_args){
            .affinity_core  = qpu_core,
            .frame_blocks_x = frame_w / 8,
            .frame_blocks_y = frame_h / 8,
            .blocks_per_wg  = blocks_per_wg,
        };
        pthread_create(&qpu_tid, NULL, qpu_worker, &q_args);
    }

    /* Main thread releases via the barrier. */
    pthread_barrier_wait(&g_start_barrier);

    /* Join everyone. */
    pthread_join(timer_tid, NULL);
    if (has_neon) for (int i = 0; i < n_neon; i++) pthread_join(neon_tids[i], NULL);
    if (has_qpu)  pthread_join(qpu_tid, NULL);

    /* Report. */
    uint64_t total_blocks = 0;
    double max_elapsed = 0.0;

    if (has_neon) {
        printf("NEON per-thread:\n");
        for (int i = 0; i < n_neon; i++) {
            double mbps = n_args[i].blocks_done / n_args[i].elapsed_s / 1e6;
            printf("  core %d: %.3f Mblock/s  (%llu blocks / %.3f s)\n",
                   n_args[i].affinity_core, mbps,
                   (unsigned long long) n_args[i].blocks_done,
                   n_args[i].elapsed_s);
            total_blocks += n_args[i].blocks_done;
            if (n_args[i].elapsed_s > max_elapsed) max_elapsed = n_args[i].elapsed_s;
        }
    }
    if (has_qpu) {
        double mbps = q_args.blocks_done / q_args.elapsed_s / 1e6;
        printf("QPU (host on core %d): %.3f Mblock/s  (%llu blocks / %.3f s)\n",
               q_args.affinity_core, mbps,
               (unsigned long long) q_args.blocks_done,
               q_args.elapsed_s);
        total_blocks += q_args.blocks_done;
        if (q_args.elapsed_s > max_elapsed) max_elapsed = q_args.elapsed_s;
    }

    double total_mbps = total_blocks / max_elapsed / 1e6;
    printf("\n=== AGGREGATE ===\n");
    printf("  total blocks  : %llu\n", (unsigned long long) total_blocks);
    printf("  wall-clock    : %.3f s\n", max_elapsed);
    printf("  Mblock/s      : %.3f\n", total_mbps);
    printf("  equiv 1080p FPS: %.1f  (32400 blocks/frame)\n",
           total_mbps * 1e6 / 32400.0);

    pthread_barrier_destroy(&g_start_barrier);
    return 0;
}