// Device-side dspark vanilla Markov resample. See dspark-markov.h for the
// contract and the exact math. Self-contained: depends only on the CUDA
// runtime, not on ggml/llama internals.

#include "dspark-markov.h"

#include <cuda_runtime.h>
#include <math_constants.h>

#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <limits>

#define DSPARK_WARP            32
#define DSPARK_WARPS_PER_BLOCK 8
#define DSPARK_BLK_THREADS     (DSPARK_WARP * DSPARK_WARPS_PER_BLOCK)  // 256
#define DSPARK_NBLOCKS         1024                                    // partial-argmax grid width (grid-stride)

#define DSPARK_CUDA_CHECK(call)                                                                      \
    do {                                                                                             \
        cudaError_t err_ = (call);                                                                   \
        if (err_ != cudaSuccess) {                                                                   \
            fprintf(stderr, "dspark-markov cuda error %s at %s:%d: %s\n", #call, __FILE__, __LINE__, \
                    cudaGetErrorString(err_));                                                       \
            return false;                                                                            \
        }                                                                                            \
    } while (0)

struct dspark_markov_cuda {
    int64_t n_vocab = 0;
    int64_t rank    = 0;
    int32_t mask_token_id = -1;

    float * d_w1 = nullptr;  // [n_vocab * rank]
    float * d_w2 = nullptr;  // [n_vocab * rank]

    // per-round scratch (grown lazily to fit n_use)
    int64_t   base_cap_rows = 0;        // rows currently allocated in d_base / h_base
    float *   d_base        = nullptr;  // [base_cap_rows * n_vocab]
    float *   h_base        = nullptr;  // pinned staging for the H2D of base logits
    int32_t * d_out         = nullptr;  // [base_cap_rows]
    int32_t * h_out         = nullptr;  // pinned staging for the D2H of the result ids

    int32_t * d_prev   = nullptr;       // device-resident chained prev token (1 int)
    float *   d_scores = nullptr;
    float *   d_prob   = nullptr;

    float *   d_part_val = nullptr;  // [DSPARK_NBLOCKS]
    int32_t * d_part_idx = nullptr;  // [DSPARK_NBLOCKS]

    cudaStream_t stream = nullptr;
};

// Fused GEMV + add-base + per-block partial argmax for one position k.
// One warp reduces one vocab row's rank-length dot product (coalesced reads
// of w2), lane 0 adds the base logit and tracks the warp's running best over
// its grid-stride rows; the block then reduces its 8 warp bests to one
// (value, index) partial. Ties resolve to the lowest vocab index, matching
// the host path's strict-greater argmax.
__global__ void dspark_gemv_argmax_partial(const float * __restrict__ w1,
                                           const float * __restrict__ w2,
                                           const float * __restrict__ base,
                                           const int32_t * __restrict__ prev,
                                           int64_t n_vocab,
                                           int     rank,
                                           int32_t mask_token_id,
                                           int64_t k,
                                           float * __restrict__ part_val,
                                           int32_t * __restrict__ part_idx,
                                           float * scores) {
    extern __shared__ float w1s[];  // rank floats: the prev token's w1 row

    const int tid  = threadIdx.x;
    const int lane = tid & (DSPARK_WARP - 1);
    const int wid  = tid >> 5;

    const int64_t w1_off = (int64_t) prev[0] * (int64_t) rank;
    for (int r = tid; r < rank; r += DSPARK_BLK_THREADS) {
        w1s[r] = w1[w1_off + r];
    }
    __syncthreads();

    const int64_t base_off = k * n_vocab;

    float   bestv = -CUDART_INF_F;
    int32_t besti = INT32_MAX;

    const int64_t gw     = (int64_t) blockIdx.x * DSPARK_WARPS_PER_BLOCK + wid;
    const int64_t stride = (int64_t) gridDim.x * DSPARK_WARPS_PER_BLOCK;
    for (int64_t v = gw; v < n_vocab; v += stride) {
        if (v == mask_token_id) {
            if (lane == 0 && scores) {
                scores[v] = -CUDART_INF_F;
            }
            continue;
        }
        const float * w2row = w2 + v * (int64_t) rank;
        float         acc   = 0.0f;
        for (int r = lane; r < rank; r += DSPARK_WARP) {
            acc += w1s[r] * w2row[r];
        }
#pragma unroll
        for (int o = DSPARK_WARP / 2; o > 0; o >>= 1) {
            acc += __shfl_xor_sync(0xffffffffu, acc, o);
        }
        if (lane == 0) {
            const float logit = base[base_off + v] + acc;
            if (scores) {
                scores[v] = logit;
            }
            // strict >: rows are visited in increasing v, so the lowest index
            // wins ties (same as the host scalar/BLAS argmax).
            if (logit > bestv || (logit == bestv && v < besti)) {
                bestv = logit;
                besti = (int32_t) v;
            }
        }
    }

    __shared__ float   sv[DSPARK_WARPS_PER_BLOCK];
    __shared__ int32_t si[DSPARK_WARPS_PER_BLOCK];
    if (lane == 0) {
        sv[wid] = bestv;
        si[wid] = besti;
    }
    __syncthreads();

    if (tid == 0) {
        float   bv = sv[0];
        int32_t bi = si[0];
#pragma unroll
        for (int w = 1; w < DSPARK_WARPS_PER_BLOCK; ++w) {
            if (sv[w] > bv || (sv[w] == bv && si[w] < bi)) {
                bv = sv[w];
                bi = si[w];
            }
        }
        part_val[blockIdx.x] = bv;
        part_idx[blockIdx.x] = bi;
    }
}

__global__ void dspark_probability_partial(const float *   scores,
                                           const int32_t * out,
                                           int             k,
                                           int64_t         vocab,
                                           float *         partial) {
    __shared__ float sums[DSPARK_BLK_THREADS];
    const int        tid  = threadIdx.x;
    const float      best = scores[out[k]];
    float            sum  = 0.0f;
    for (int64_t v = (int64_t) blockIdx.x * blockDim.x + tid; v < vocab; v += (int64_t) gridDim.x * blockDim.x) {
        sum += expf(scores[v] - best);
    }
    sums[tid] = sum;
    __syncthreads();
    for (int step = blockDim.x / 2; step; step /= 2) {
        if (tid < step) {
            sums[tid] += sums[tid + step];
        }
        __syncthreads();
    }
    if (!tid) {
        partial[blockIdx.x] = sums[0];
    }
}

__global__ void dspark_probability_final(const float * partial, float * probabilities, int k) {
    __shared__ float sums[DSPARK_BLK_THREADS];
    const int        tid = threadIdx.x;
    float            sum = 0.0f;
    for (int i = tid; i < DSPARK_NBLOCKS; i += blockDim.x) {
        sum += partial[i];
    }
    sums[tid] = sum;
    __syncthreads();
    for (int step = blockDim.x / 2; step; step /= 2) {
        if (tid < step) {
            sums[tid] += sums[tid + step];
        }
        __syncthreads();
    }
    if (!tid) {
        probabilities[k] = 1.0f / sums[0];
    }
}

struct dspark_dcut_costs {
    float values[4];
};

__global__ void dspark_select_depth(const float * probabilities, int32_t * out, int n, dspark_dcut_costs costs) {
    float survival = 1.0f, expected = 1.0f, best_rate = -1.0f;
    int   retained = 1;
    for (int k = 0; k < n; ++k) {
        const float p = probabilities[k];
        if (!isfinite(p) || p < 0.0f || p > 1.000001f) {
            out[n] = -1;
            return;
        }
        survival *= fminf(p, 1.0f);
        expected += survival;
        const float rate = expected / costs.values[k];
        if (rate > best_rate) {
            best_rate = rate;
            retained  = k + 1;
        }
    }
    out[n] = retained;
}

// Reduce the per-block partials to the single global argmax, write it to
// out[k] and forward it into prev for the next position (the device-side
// chaining step). Single block.
__global__ void dspark_argmax_final(const float * __restrict__ part_val,
                                    const int32_t * __restrict__ part_idx,
                                    int     nparts,
                                    int64_t k,
                                    int32_t * __restrict__ out,
                                    int32_t * __restrict__ prev) {
    __shared__ float   sv[DSPARK_BLK_THREADS];
    __shared__ int32_t si[DSPARK_BLK_THREADS];

    const int tid = threadIdx.x;

    float   bv = -CUDART_INF_F;
    int32_t bi = 0;
    for (int i = tid; i < nparts; i += DSPARK_BLK_THREADS) {
        if (part_val[i] > bv || (part_val[i] == bv && part_idx[i] < bi)) {
            bv = part_val[i];
            bi = part_idx[i];
        }
    }
    sv[tid] = bv;
    si[tid] = bi;
    __syncthreads();

    for (int s = DSPARK_BLK_THREADS / 2; s > 0; s >>= 1) {
        if (tid < s) {
            if (sv[tid + s] > sv[tid] || (sv[tid + s] == sv[tid] && si[tid + s] < si[tid])) {
                sv[tid] = sv[tid + s];
                si[tid] = si[tid + s];
            }
        }
        __syncthreads();
    }

    if (tid == 0) {
        out[k]  = si[0];
        prev[0] = si[0];
    }
}

static bool dspark_markov_ensure_capacity(dspark_markov_cuda * ctx, int64_t rows) {
    if (rows <= ctx->base_cap_rows) {
        return true;
    }
    if (ctx->d_base) {
        cudaFree(ctx->d_base);
        ctx->d_base = nullptr;
    }
    if (ctx->h_base) {
        cudaFreeHost(ctx->h_base);
        ctx->h_base = nullptr;
    }
    if (ctx->d_out) {
        cudaFree(ctx->d_out);
        ctx->d_out = nullptr;
    }
    if (ctx->h_out) {
        cudaFreeHost(ctx->h_out);
        ctx->h_out = nullptr;
    }
    ctx->base_cap_rows = 0;

    DSPARK_CUDA_CHECK(cudaMalloc(&ctx->d_base, (size_t) rows * (size_t) ctx->n_vocab * sizeof(float)));
    DSPARK_CUDA_CHECK(
        cudaHostAlloc(&ctx->h_base, (size_t) rows * (size_t) ctx->n_vocab * sizeof(float), cudaHostAllocDefault));
    DSPARK_CUDA_CHECK(cudaMalloc(&ctx->d_out, (size_t) (rows + 1) * sizeof(int32_t)));
    DSPARK_CUDA_CHECK(cudaHostAlloc(&ctx->h_out, (size_t) (rows + 1) * sizeof(int32_t), cudaHostAllocDefault));
    ctx->base_cap_rows = rows;
    return true;
}

dspark_markov_cuda * dspark_markov_cuda_init(const float * w1,
                                             const float * w2,
                                             int64_t       n_vocab,
                                             int64_t       markov_rank,
                                             int32_t       mask_token_id) {
    if (w1 == nullptr || w2 == nullptr || n_vocab <= 1 || mask_token_id < 0 || mask_token_id >= n_vocab ||
        n_vocab > INT32_MAX || markov_rank <= 0 || markov_rank > INT32_MAX ||
        (uint64_t) n_vocab > SIZE_MAX / sizeof(float) / (uint64_t) markov_rank) {
        return nullptr;
    }

    int n_dev = 0;
    if (cudaGetDeviceCount(&n_dev) != cudaSuccess || n_dev <= 0) {
        fprintf(stderr, "dspark-markov: no CUDA device available\n");
        return nullptr;
    }

    dspark_markov_cuda * ctx = new dspark_markov_cuda();
    ctx->n_vocab             = n_vocab;
    ctx->rank                = markov_rank;
    ctx->mask_token_id       = mask_token_id;

    const size_t nbytes = (size_t) n_vocab * (size_t) markov_rank * sizeof(float);

    bool ok    = true;
    auto guard = [&](cudaError_t e, const char * what) {
        if (e != cudaSuccess) {
            fprintf(stderr, "dspark-markov init: %s: %s\n", what, cudaGetErrorString(e));
            ok = false;
        }
    };

    guard(cudaStreamCreate(&ctx->stream), "cudaStreamCreate");
    guard(cudaMalloc(&ctx->d_w1, nbytes), "cudaMalloc w1");
    guard(cudaMalloc(&ctx->d_w2, nbytes), "cudaMalloc w2");
    guard(cudaMalloc(&ctx->d_prev, sizeof(int32_t)), "cudaMalloc prev");
    guard(cudaMalloc(&ctx->d_part_val, DSPARK_NBLOCKS * sizeof(float)), "cudaMalloc part_val");
    guard(cudaMalloc(&ctx->d_part_idx, DSPARK_NBLOCKS * sizeof(int32_t)), "cudaMalloc part_idx");
    if (ok) {
        guard(cudaMemcpy(ctx->d_w1, w1, nbytes, cudaMemcpyHostToDevice), "H2D w1");
        guard(cudaMemcpy(ctx->d_w2, w2, nbytes, cudaMemcpyHostToDevice), "H2D w2");
    }

    if (!ok) {
        dspark_markov_cuda_free(ctx);
        return nullptr;
    }
    return ctx;
}

void dspark_markov_cuda_free(dspark_markov_cuda * ctx) {
    if (ctx == nullptr) {
        return;
    }
    if (ctx->d_w1) {
        cudaFree(ctx->d_w1);
    }
    if (ctx->d_w2) {
        cudaFree(ctx->d_w2);
    }
    if (ctx->d_prev) {
        cudaFree(ctx->d_prev);
    }
    if (ctx->d_scores) {
        cudaFree(ctx->d_scores);
    }
    if (ctx->d_prob) {
        cudaFree(ctx->d_prob);
    }
    if (ctx->d_part_val) {
        cudaFree(ctx->d_part_val);
    }
    if (ctx->d_part_idx) {
        cudaFree(ctx->d_part_idx);
    }
    if (ctx->d_base) {
        cudaFree(ctx->d_base);
    }
    if (ctx->d_out) {
        cudaFree(ctx->d_out);
    }
    if (ctx->h_base) {
        cudaFreeHost(ctx->h_base);
    }
    if (ctx->h_out) {
        cudaFreeHost(ctx->h_out);
    }
    if (ctx->stream) {
        cudaStreamDestroy(ctx->stream);
    }
    delete ctx;
}

static bool dspark_markov_cuda_run(dspark_markov_cuda * ctx,
                                   const float *        base_logits,
                                   int32_t              id_last,
                                   int32_t              n_use,
                                   int32_t *            out_ids,
                                   const float *        costs,
                                   int32_t *            retained,
                                   float *              probabilities) {
    if (ctx == nullptr || base_logits == nullptr || out_ids == nullptr || n_use <= 0 || id_last < 0 ||
        id_last >= ctx->n_vocab) {
        return false;
    }
    if (costs) {
        if (n_use > 4 || !retained) {
            return false;
        }
        for (int k = 0; k < n_use; ++k) {
            if (!std::isfinite(costs[k]) || costs[k] <= 0.0f) {
                return false;
            }
        }
        if (!ctx->d_scores) {
            DSPARK_CUDA_CHECK(cudaMalloc(&ctx->d_scores, (size_t) ctx->n_vocab * sizeof(float)));
        }
        if (!ctx->d_prob) {
            DSPARK_CUDA_CHECK(cudaMalloc(&ctx->d_prob, 4 * sizeof(float)));
        }
    }
    if (!dspark_markov_ensure_capacity(ctx, n_use)) {
        return false;
    }

    const int64_t V = ctx->n_vocab;
    const int     R = (int) ctx->rank;

    // Stage this round's base logits into pinned host memory, then one H2D.
    const size_t base_elems = (size_t) n_use * (size_t) V;
    memcpy(ctx->h_base, base_logits, base_elems * sizeof(float));
    DSPARK_CUDA_CHECK(
        cudaMemcpyAsync(ctx->d_base, ctx->h_base, base_elems * sizeof(float), cudaMemcpyHostToDevice, ctx->stream));

    // Seed the chained prev with the anchor token (prev for k == 0).
    DSPARK_CUDA_CHECK(cudaMemcpyAsync(ctx->d_prev, &id_last, sizeof(int32_t), cudaMemcpyHostToDevice, ctx->stream));

    const size_t shmem = (size_t) R * sizeof(float);
    for (int32_t k = 0; k < n_use; ++k) {
        dspark_gemv_argmax_partial<<<DSPARK_NBLOCKS, DSPARK_BLK_THREADS, shmem, ctx->stream>>>(
            ctx->d_w1, ctx->d_w2, ctx->d_base, ctx->d_prev, V, R, ctx->mask_token_id, (int64_t) k, ctx->d_part_val,
            ctx->d_part_idx, costs ? ctx->d_scores : nullptr);
        dspark_argmax_final<<<1, DSPARK_BLK_THREADS, 0, ctx->stream>>>(ctx->d_part_val, ctx->d_part_idx, DSPARK_NBLOCKS,
                                                                       (int64_t) k, ctx->d_out, ctx->d_prev);
        if (costs) {
            dspark_probability_partial<<<DSPARK_NBLOCKS, DSPARK_BLK_THREADS, 0, ctx->stream>>>(
                ctx->d_scores, ctx->d_out, k, V, ctx->d_part_val);
            dspark_probability_final<<<1, DSPARK_BLK_THREADS, 0, ctx->stream>>>(ctx->d_part_val, ctx->d_prob, k);
        }
    }
    if (costs) {
        dspark_dcut_costs table{};
        for (int k = 0; k < n_use; ++k) {
            table.values[k] = costs[k];
        }
        dspark_select_depth<<<1, 1, 0, ctx->stream>>>(ctx->d_prob, ctx->d_out, n_use, table);
        if (probabilities) {
            DSPARK_CUDA_CHECK(cudaMemcpyAsync(probabilities, ctx->d_prob, n_use * sizeof(float), cudaMemcpyDeviceToHost,
                                              ctx->stream));
        }
    }
    DSPARK_CUDA_CHECK(cudaGetLastError());

    DSPARK_CUDA_CHECK(cudaMemcpyAsync(ctx->h_out, ctx->d_out, (size_t) (n_use + (costs ? 1 : 0)) * sizeof(int32_t),
                                      cudaMemcpyDeviceToHost, ctx->stream));
    DSPARK_CUDA_CHECK(cudaStreamSynchronize(ctx->stream));

    memcpy(out_ids, ctx->h_out, (size_t) n_use * sizeof(int32_t));
    if (costs) {
        *retained = ctx->h_out[n_use];
        if (*retained < 1 || *retained > n_use) {
            return false;
        }
    }
    return true;
}

bool dspark_markov_cuda_resample(dspark_markov_cuda * ctx,
                                 const float *        base_logits,
                                 int32_t              id_last,
                                 int32_t              n_use,
                                 int32_t *            out_ids) {
    return dspark_markov_cuda_run(ctx, base_logits, id_last, n_use, out_ids, nullptr, nullptr, nullptr);
}

bool dspark_markov_cuda_dcut(dspark_markov_cuda * ctx,
                             const float *        base_logits,
                             int32_t              id_last,
                             int32_t              n_use,
                             int32_t *            out_ids,
                             const float *        costs,
                             int32_t *            retained,
                             float *              probabilities) {
    if (!costs) {
        return false;
    }
    return dspark_markov_cuda_run(ctx, base_logits, id_last, n_use, out_ids, costs, retained, probabilities);
}