/**
 * Training Guard + Q-Value Monitor CUDA Kernels
 *
 * Zero-synchronize safety checks for the DQN training hot path.
 * Kernels write results to mapped pinned memory (CU_MEMHOSTALLOC_DEVICEMAP);
 * __threadfence_system() ensures CPU visibility without memcpy_dtoh.
 * The Rust side uses double-buffered mapped buffers for pipelined readback.
 *
 * Kernels:
 *   training_guard_check       — NaN/Inf, loss-clip, gradient-collapse detection
 *   training_guard_accumulate  — epoch-boundary loss/grad running sums
 *   qvalue_stats_reduce        — batch-wide Q-value statistics (parallel reduction)
 *   qvalue_divergence_check    — single-sample Q-value divergence check
 */

/* ------------------------------------------------------------------ */
/*  Kernel 1: training_guard_check_and_accumulate (FUSED)             */
/*                                                                    */
/*  Launch: grid=(1,1,1), block=(1,1,1) — single-thread.             */
/*                                                                    */
/*  Combines guard check + accumulate into a single kernel launch.    */
/*  Reads GPU-resident loss and grad_norm scalars ONCE (not twice),   */
/*  performs safety checks, writes 7 floats to mapped pinned host     */
/*  buffer, AND accumulates to device-resident running sums.          */
/*                                                                    */
/*  Saves one kernel launch per training step (~2-4µs overhead each), */
/*  which compounds to 12-25ms/epoch at 6375 steps.                   */
/*                                                                    */
/*  Mapped output (7 floats, pinned host buffer):                     */
/*    [0] halt_nan         — 1.0 if loss or grad_norm is NaN/Inf      */
/*    [1] halt_loss_clip   — 1.0 if loss > clip_threshold             */
/*    [2] halt_grad_collapse — 1.0 if grad_norm < collapse_threshold  */
/*                             AND warmup == 0                        */
/*    [3] clipped_loss     — min(loss, clip_threshold)                */
/*    [4] raw_loss                                                     */
/*    [5] raw_grad_norm                                               */
/*    [6] reserved         — 0.0                                      */
/*                                                                    */
/*  acc_buf layout (3 floats, device memory):                         */
/*    [0] loss_sum                                                     */
/*    [1] grad_norm_sum                                               */
/*    [2] step_count (stored as float)                                */
/* ------------------------------------------------------------------ */
extern "C" __global__ void training_guard_check_and_accumulate(
    const float* __restrict__ loss_scalar,       /* GPU-resident scalar (1 float) */
    const float* __restrict__ grad_norm_scalar,  /* GPU-resident scalar (1 float) */
    float*       output,                         /* pinned host buffer (7 floats)  */
    float*       acc_buf,                        /* device accumulator (3 floats)  */
    float        clip_threshold,                 /* e.g. 1e6f                      */
    float        collapse_threshold,             /* e.g. 1e-7f                     */
    int          warmup                          /* 1 = in warmup, skip collapse   */
) {
    float loss     = loss_scalar[0];
    float grad_norm = grad_norm_scalar[0];

    /* ── Guard check ── */

    int loss_nan = isnan(loss) || isinf(loss);
    int grad_nan = isnan(grad_norm) || isinf(grad_norm);

    /* Use standard library functions for robust NaN/Inf detection */
    int halt_nan  = (loss_nan || grad_nan) ? 1 : 0;

    /* Loss clip */
    int halt_loss_clip = (!loss_nan && loss > clip_threshold) ? 1 : 0;

    /* Gradient collapse — only meaningful outside warmup */
    int halt_grad_collapse = 0;
    if (!warmup && !grad_nan) {
        halt_grad_collapse = (grad_norm < collapse_threshold) ? 1 : 0;
    }

    /* Clipped loss: NaN/Inf → passthrough raw (host will detect via halt_nan) */
    float clipped_loss = (!loss_nan)
                         ? fminf(loss, clip_threshold)
                         : loss;

    output[0] = (float)halt_nan;
    output[1] = (float)halt_loss_clip;
    output[2] = (float)halt_grad_collapse;
    output[3] = clipped_loss;
    output[4] = loss;
    output[5] = grad_norm;
    output[6] = 0.0f;  /* reserved */
    __threadfence_system(); /* ensure mapped writes visible to CPU */

    /* ── Accumulate ── */

    /* Skip NaN/Inf to keep running sums meaningful */
    if (!loss_nan) {
        acc_buf[0] += loss;
    }
    if (!grad_nan) {
        acc_buf[1] += grad_norm;
    }
    acc_buf[2] += 1.0f;
}

/* Legacy entry points — kept for backward compatibility, each delegates
   to the fused kernel's logic. New code should use the fused version. */

extern "C" __global__ void training_guard_check(
    const float* __restrict__ loss_scalar,
    const float* __restrict__ grad_norm_scalar,
    float*       output,
    float        clip_threshold,
    float        collapse_threshold,
    int          warmup
) {
    float loss     = loss_scalar[0];
    float grad_norm = grad_norm_scalar[0];
    int loss_nan = isnan(loss) || isinf(loss);
    int grad_nan = isnan(grad_norm) || isinf(grad_norm);
    int halt_nan  = (loss_nan || grad_nan) ? 1 : 0;
    int halt_loss_clip = (!loss_nan && loss > clip_threshold) ? 1 : 0;
    int halt_grad_collapse = 0;
    if (!warmup && !grad_nan) {
        halt_grad_collapse = (grad_norm < collapse_threshold) ? 1 : 0;
    }
    float clipped_loss = (!loss_nan) ? fminf(loss, clip_threshold) : loss;
    output[0] = (float)halt_nan;
    output[1] = (float)halt_loss_clip;
    output[2] = (float)halt_grad_collapse;
    output[3] = clipped_loss;
    output[4] = loss;
    output[5] = grad_norm;
    output[6] = 0.0f;
    __threadfence_system();
}

extern "C" __global__ void training_guard_accumulate(
    const float* __restrict__ loss_scalar,
    const float* __restrict__ grad_norm_scalar,
    float*       acc_buf
) {
    float loss      = loss_scalar[0];
    float grad_norm = grad_norm_scalar[0];
    if (!isnan(loss) && !isinf(loss)) {
        acc_buf[0] += loss;
    }
    if (!isnan(grad_norm) && !isinf(grad_norm)) {
        acc_buf[1] += grad_norm;
    }
    acc_buf[2] += 1.0f;
}

/* ------------------------------------------------------------------ */
/*  Kernel 3: qvalue_stats_reduce                                     */
/*                                                                    */
/*  Launch: grid=(1,1,1), block=(256,1,1).                            */
/*  Grid-stride loop handles batch_size > 256.                        */
/*                                                                    */
/*  Input: q_values[batch_size * num_actions] — row-major,            */
/*         row i holds Q-values for sample i across all actions.      */
/*                                                                    */
/*  For each sample finds max Q across actions, then reduces to:      */
/*    [0] q_min  — min of per-sample max-Q                           */
/*    [1] q_max  — max of per-sample max-Q                           */
/*    [2] q_mean — mean of per-sample max-Q                          */
/*    [3] mean_of_all — mean of ALL Q-values (collapse detection)    */
/* ------------------------------------------------------------------ */
extern "C" __global__ void qvalue_stats_reduce(
    const float* __restrict__ q_values,  /* [batch_size * num_actions]   */
    float*       output,                 /* pinned host buffer (4 floats) */
    int          batch_size,
    int          num_actions
) {
    __shared__ float s_min[256];
    __shared__ float s_max[256];
    __shared__ float s_sum[256];      /* sum of per-sample max-Q       */
    __shared__ float s_all_sum[256];  /* sum of ALL Q-values           */

    int tid = threadIdx.x;

    float local_min    =  1e30f;
    float local_max    = -1e30f;
    float local_sum    =  0.0f;
    float local_all    =  0.0f;

    /* Grid-stride loop: each thread processes multiple samples */
    for (int i = tid; i < batch_size; i += blockDim.x) {
        const float* row = q_values + i * num_actions;

        /* Find max Q across actions for this sample */
        float sample_max = -1e30f;
        for (int a = 0; a < num_actions; a++) {
            float qv = row[a];
            if (!isnan(qv) && !isinf(qv)) {
                if (qv > sample_max) sample_max = qv;
                local_all += qv;
            }
        }

        if (sample_max > -1e30f) {  /* at least one valid Q-value */
            if (sample_max < local_min) local_min = sample_max;
            if (sample_max > local_max) local_max = sample_max;
            local_sum += sample_max;
        }
    }

    s_min[tid]    = local_min;
    s_max[tid]    = local_max;
    s_sum[tid]    = local_sum;
    s_all_sum[tid] = local_all;
    __syncthreads();

    /* Parallel reduction phase 1: shared memory down to 32 threads */
    for (int stride = blockDim.x / 2; stride >= 32; stride >>= 1) {
        if (tid < stride) {
            s_min[tid]     = fminf(s_min[tid],     s_min[tid + stride]);
            s_max[tid]     = fmaxf(s_max[tid],     s_max[tid + stride]);
            s_sum[tid]    += s_sum[tid + stride];
            s_all_sum[tid] += s_all_sum[tid + stride];
        }
        __syncthreads();
    }

    /* Parallel reduction phase 2: warp-level reduction (no __syncthreads needed) */
    if (tid < 32) {
        float my_min     = s_min[tid];
        float my_max     = s_max[tid];
        float my_sum     = s_sum[tid];
        float my_all_sum = s_all_sum[tid];
        for (int mask = 16; mask >= 1; mask >>= 1) {
            my_min      = fminf(my_min, __shfl_xor_sync(0xFFFFFFFF, my_min, mask));
            my_max      = fmaxf(my_max, __shfl_xor_sync(0xFFFFFFFF, my_max, mask));
            my_sum     += __shfl_xor_sync(0xFFFFFFFF, my_sum, mask);
            my_all_sum += __shfl_xor_sync(0xFFFFFFFF, my_all_sum, mask);
        }
        if (tid == 0) {
            s_min[0]     = my_min;
            s_max[0]     = my_max;
            s_sum[0]     = my_sum;
            s_all_sum[0] = my_all_sum;
        }
    }
    __syncthreads();

    /* Thread 0 writes final results */
    if (tid == 0) {
        float n = (float)batch_size;
        float total_n = n * (float)num_actions;

        output[0] = (s_min[0] >  1e29f) ?  0.0f : s_min[0];   /* q_min        */
        output[1] = (s_max[0] < -1e29f) ?  0.0f : s_max[0];   /* q_max        */
        output[2] = (n > 0.0f)          ? s_sum[0] / n    : 0.0f; /* q_mean    */
        output[3] = (total_n > 0.0f)    ? s_all_sum[0] / total_n : 0.0f; /* mean_of_all */
        __threadfence_system(); /* ensure mapped writes visible to CPU */
    }
}

/* ------------------------------------------------------------------ */
/*  Kernel 4: qvalue_divergence_check                                 */
/*                                                                    */
/*  Launch: grid=(1,1,1), block=(1,1,1) — single-thread.             */
/*                                                                    */
/*  Single-sample Q-value divergence check.                           */
/*  Computes min/max/mean/variance of Q-values for one sample,        */
/*  then checks whether |min| or |max| exceeds divergence_threshold.  */
/*                                                                    */
/*  Output (5 floats, pinned host buffer):                            */
/*    [0] q_min                                                        */
/*    [1] q_max                                                        */
/*    [2] q_mean                                                       */
/*    [3] q_variance                                                   */
/*    [4] divergence_detected — 1.0 if |q_min| or |q_max| > threshold */
/* ------------------------------------------------------------------ */
extern "C" __global__ void qvalue_divergence_check(
    const float* __restrict__ q_values,       /* [num_actions] for one sample  */
    float*       output,                      /* pinned host buffer (5 floats) */
    int          num_actions,
    float        divergence_threshold         /* e.g. 1e4f                     */
) {
    float q_min  =  1e30f;
    float q_max  = -1e30f;
    float q_sum  =  0.0f;
    float q_sq   =  0.0f;
    int   valid  =  0;

    for (int a = 0; a < num_actions; a++) {
        float qv = q_values[a];
        if (isnan(qv) || isinf(qv)) continue;

        if (qv < q_min) q_min = qv;
        if (qv > q_max) q_max = qv;
        q_sum += qv;
        q_sq  += qv * qv;
        valid++;
    }

    float q_mean     = 0.0f;
    float q_variance = 0.0f;
    if (valid > 0) {
        q_mean = q_sum / (float)valid;
        float mean_sq = q_sq / (float)valid;
        q_variance = mean_sq - q_mean * q_mean;
        if (q_variance < 0.0f) q_variance = 0.0f;  /* numerical floor */
    }

    /* Sentinel: if no valid values, zero out extremes */
    if (valid == 0) {
        q_min = 0.0f;
        q_max = 0.0f;
    }

    int diverged = (fabsf(q_min) > divergence_threshold ||
                    fabsf(q_max) > divergence_threshold) ? 1 : 0;

    output[0] = q_min;
    output[1] = q_max;
    output[2] = q_mean;
    output[3] = q_variance;
    output[4] = (float)diverged;
    __threadfence_system(); /* ensure mapped writes visible to CPU */
}
