/**
 * DQN utility kernels — gradient norm reduction and AdamW update.
 *
 * These two kernels form the optimizer step in the DQN training pipeline.
 * They operate on the flat parameter / gradient buffers produced by the
 * cuBLAS backward pass and are compiled together with
 * common_device_functions.cuh (which also provides the f32_to_bf16_kernel
 * and bf16_to_f32_kernel converters needed by the same NVRTC module).
 *
 * No #include of common_device_functions.cuh is needed here — the build
 * function in gpu_dqn_trainer.rs prepends it via string concatenation
 * before passing the combined source to NVRTC.
 *
 * Extracted from dqn_training_kernel.cu (lines 1207-1293).
 */

/* ══════════════════════════════════════════════════════════════════════
 * GRADIENT NORM KERNEL (Phase 3a)
 *
 * Computes gradient L2 norm (sum of squares) via warp + block reduction
 * and atomicAdd into a single output float. Must be followed by
 * dqn_adam_update_kernel which reads the completed norm.
 *
 * Launch config: grid=(ceil(TOTAL_PARAMS/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_grad_norm_kernel(
    const float* __restrict__ grads,      /* [TOTAL_PARAMS] accumulated gradients */
    float* __restrict__ out_grad_norm,    /* [1] output: sum of squares */
    int total_params
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    float g = (idx < total_params) ? grads[idx] : 0.0f;
    float g2 = g * g;

    /* Warp-level reduction via shuffle (no shared memory) */
    for (int offset = 16; offset > 0; offset >>= 1)
        g2 += __shfl_xor_sync(0xFFFFFFFF, g2, offset);

    /* Block-level cross-warp reduction via shared memory.
     * Pad warp_sums to 2 elements per warp (stride=2) to avoid bank conflicts:
     * warp 0 → bank 0, warp 1 → bank 2, ... warp 7 → bank 14.
     * Without padding, sequential warp writes to warp_sums[0..7] map to
     * banks 0..7 which is conflict-free for writes, but the __syncthreads()
     * + read-back pattern benefits from padding to avoid false sharing
     * between adjacent cache lines on H100's 128-byte L1 lines. */
    __shared__ float warp_sums[16];  /* 8 warps × 2 stride (padded) */
    int warp_id = threadIdx.x / 32;
    int warp_lane = threadIdx.x % 32;
    if (warp_lane == 0) warp_sums[warp_id * 2] = g2;
    __syncthreads();

    /* First warp reduces across warps using shuffles (pure register path) */
    if (warp_id == 0) {
        float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0)
            atomicAdd(out_grad_norm, val);
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * ADAM UPDATE KERNEL (Phase 3b)
 *
 * Reads the COMPLETED gradient L2 norm from dqn_grad_norm_kernel,
 * applies gradient clipping, then AdamW parameter update.
 * Trivially parallel: 1 thread per parameter element.
 *
 * Launch config: grid=(ceil(TOTAL_PARAMS/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_adam_update_kernel(
    float* __restrict__ params,           /* [TOTAL_PARAMS] flattened parameters */
    const float* __restrict__ grads,      /* [TOTAL_PARAMS] accumulated gradients */
    float* __restrict__ m,                /* [TOTAL_PARAMS] first moment (Adam) */
    float* __restrict__ v,                /* [TOTAL_PARAMS] second moment (Adam) */
    const float* __restrict__ grad_norm_sq, /* [1] completed sum of squares */
    float lr,
    float beta1,
    float beta2,
    float epsilon,
    float weight_decay,
    float max_grad_norm,
    const int* __restrict__ t_ptr,        /* Adam step counter on device (for CUDA Graph) */
    int total_params
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_params) return;

    int t = *t_ptr;
    float g = grads[idx];

    /* Clip using the COMPLETED norm (no race — separate kernel launch) */
    float norm = sqrtf(*grad_norm_sq + 1e-12f);
    float clip_scale = (norm > max_grad_norm) ? (max_grad_norm / norm) : 1.0f;
    float clipped_g = g * clip_scale;

    /* Adam update */
    float beta1_t = 1.0f - powf(beta1, (float)t);
    float beta2_t = 1.0f - powf(beta2, (float)t);

    float m_i = beta1 * m[idx] + (1.0f - beta1) * clipped_g;
    float v_i = beta2 * v[idx] + (1.0f - beta2) * clipped_g * clipped_g;
    m[idx] = m_i;
    v[idx] = v_i;

    float m_hat = m_i / beta1_t;
    float v_hat = v_i / beta2_t;

    /* AdamW weight decay (decoupled) */
    params[idx] -= lr * (m_hat / (sqrtf(v_hat) + epsilon) + weight_decay * params[idx]);
}

/* ══════════════════════════════════════════════════════════════════════
 * SAXPY KERNEL
 *
 * y[i] += alpha * x[i]   for i = 0..n-1
 *
 * Used by IQN trunk gradient to apply auxiliary SGD correction to the
 * shared trunk weights. Runs outside the CUDA Graph.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_saxpy_kernel(
    float* __restrict__ y,
    const float* __restrict__ x,
    float alpha,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) y[i] += alpha * x[i];
}

/* ══════════════════════════════════════════════════════════════════════
 * ZERO KERNEL
 *
 * buf[i] = 0.0f   for i = 0..n-1
 *
 * Zeroes a section of the gradient buffer before IQN trunk backward
 * accumulation (backward_fc_layer uses beta=1.0 for grad accumulation).
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_zero_kernel(
    float* __restrict__ buf,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) buf[i] = 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * REGIME-ADAPTIVE PER SCALING KERNEL
 *
 * Scales td_errors by regime similarity before PER priority update.
 * Samples from similar market regimes (ADX/CUSUM) get higher priority.
 *
 * Uses STATE_DIM from common_device_functions.cuh for state stride.
 *
 * Launch config: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_regime_scale_kernel(
    float* __restrict__ td_errors,          /* [B] scaled in-place */
    const float* __restrict__ states,       /* [B, STATE_DIM] */
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    /* Read target regime from first sample (GPU-only, no CPU readback).
     * ADX at feature index 40, CUSUM at index 41. */
    float target_adx   = states[0 * STATE_DIM + 40];
    float target_cusum = states[0 * STATE_DIM + 41];

    float sample_adx   = states[i * STATE_DIM + 40];
    float sample_cusum = states[i * STATE_DIM + 41];

    float adx_diff   = sample_adx   - target_adx;
    float cusum_diff = sample_cusum - target_cusum;
    float dist_sq = adx_diff * adx_diff + cusum_diff * cusum_diff;

    float temperature = 0.3f;
    float sim = expf(-dist_sq / (temperature * temperature + 1e-8f));
    /* Clamp to [0.5, 2.0]: never fully suppress, max 2x upweight */
    float scale = 0.5f + 1.5f * sim;

    td_errors[i] *= scale;
}

/* ══════════════════════════════════════════════════════════════════════
 * SHRINK-AND-PERTURB KERNEL (GPU-native noise generation)
 *
 * params[i] = alpha * params[i] + (1-alpha) * sigma * noise_i
 *
 * Generates Gaussian noise via Box-Muller transform using per-element
 * LCG PRNG seeded from the element index + step counter. No CPU noise
 * generation or HtoD upload needed.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_shrink_perturb_kernel(
    float* __restrict__ params,
    float alpha,
    float sigma,
    int n,
    unsigned int seed            /* changes per call to produce different noise */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;

    /* Per-element LCG PRNG: hash(seed, index) → uniform random */
    unsigned int state = seed ^ (unsigned int)(i * 2654435761u);
    state = state * 1664525u + 1013904223u;
    float u1 = (float)(state >> 8) / 16777216.0f;  /* (0, 1) */
    state = state * 1664525u + 1013904223u;
    float u2 = (float)(state >> 8) / 16777216.0f;

    /* Box-Muller: uniform → Gaussian N(0, sigma) */
    u1 = fmaxf(u1, 1e-7f);  /* prevent log(0) */
    float noise = sigma * sqrtf(-2.0f * logf(u1)) * cosf(6.283185307f * u2);

    /* Shrink-and-Perturb: blend old weights with noise */
    params[i] = alpha * params[i] + (1.0f - alpha) * noise;
}

/* ══════════════════════════════════════════════════════════════════════
 * RELU MASK KERNEL (for IQN trunk gradient — separate from backward module)
 *
 * dx[i] *= (activation[i] > 0.0f)
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void dqn_relu_mask_kernel(
    float* __restrict__ dx,
    const float* __restrict__ activation,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    if (activation[i] <= 0.0f) dx[i] = 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * SPECTRAL NORM POWER ITERATION KERNEL
 *
 * One step of power iteration for spectral normalization:
 *   v_new = W^T u / ||W^T u||
 *   u_new = W v_new / ||W v_new||
 *   sigma = u_new^T (W v_new)
 *
 * Then scales the weight matrix: W[i] /= max(1.0, sigma / sigma_max)
 *
 * For small matrices (256x256), this runs efficiently with a single block.
 * Uses shared memory for the matmul + reduction.
 *
 * Launch: grid=(1,1,1), block=(256,1,1)
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void spectral_norm_kernel(
    float* __restrict__ W,        /* [out_dim, in_dim] weight matrix — scaled in-place */
    float* __restrict__ u,        /* [out_dim] left singular vector (persistent) */
    float* __restrict__ v,        /* [in_dim] right singular vector (persistent) */
    int out_dim,
    int in_dim,
    float sigma_max               /* clip sigma above this (typically 1.0) */
) {
    __shared__ float shmem[256];  /* scratch for reductions */
    int tid = threadIdx.x;
    int bd = blockDim.x;          /* typically 256 */
    int n_total = out_dim * in_dim;

    /* ── Step 1: v_new = W^T u ── (strided for dims > blockDim) */
    /* Each thread handles multiple v elements via stride loop. */
    for (int col = tid; col < in_dim; col += bd) {
        float val = 0.0f;
        for (int row = 0; row < out_dim; row++)
            val += W[row * in_dim + col] * u[row];
        v[col] = val;  /* unnormalized — normalize below */
    }
    __syncthreads();

    /* ── Normalize v_new: ||v||₂ via parallel reduction ── */
    float local_v2 = 0.0f;
    for (int col = tid; col < in_dim; col += bd)
        local_v2 += v[col] * v[col];
    shmem[tid] = local_v2;
    __syncthreads();
    for (int s = bd / 2; s > 0; s >>= 1) {
        if (tid < s) shmem[tid] += shmem[tid + s];
        __syncthreads();
    }
    float v_norm = sqrtf(shmem[0] + 1e-12f);
    for (int col = tid; col < in_dim; col += bd)
        v[col] /= v_norm;
    __syncthreads();

    /* ── Step 2: u_new = W v_new ── (strided for dims > blockDim) */
    for (int row = tid; row < out_dim; row += bd) {
        float val = 0.0f;
        for (int col = 0; col < in_dim; col++)
            val += W[row * in_dim + col] * v[col];
        u[row] = val;  /* unnormalized — sigma = ||u_new|| */
    }
    __syncthreads();

    /* ── Sigma = ||W v_new|| = ||u_new|| (before normalizing) ── */
    float local_u2 = 0.0f;
    for (int row = tid; row < out_dim; row += bd)
        local_u2 += u[row] * u[row];
    shmem[tid] = local_u2;
    __syncthreads();
    for (int s = bd / 2; s > 0; s >>= 1) {
        if (tid < s) shmem[tid] += shmem[tid + s];
        __syncthreads();
    }
    float sigma = sqrtf(shmem[0] + 1e-12f);

    /* Normalize u_new */
    for (int row = tid; row < out_dim; row += bd)
        u[row] /= (sigma + 1e-12f);
    __syncthreads();

    /* ── Step 3: Scale W if sigma > sigma_max ── */
    if (sigma < 1e-6f) return;
    float scale = (sigma > sigma_max) ? (sigma_max / sigma) : 1.0f;
    if (scale < 1.0f) {
        for (int i = tid; i < n_total; i += bd)
            W[i] *= scale;
    }
}
