/**
 * IQL (Implicit Q-Learning) Value Network Training Kernel
 *
 * Kostrikov et al. 2021 -- trains V(s) with asymmetric expectile regression:
 *
 *   L_tau(u) = |tau - 1(u < 0)| * u^2,  where u = Q(s,a) - V(s)
 *
 * Architecture: 2-hidden-layer MLP with SiLU activation.
 *   state -> Linear(STATE_DIM, H) -> SiLU -> Linear(H, H) -> SiLU -> Linear(H, 1) -> V(s)
 *
 * Launch config:
 *   Forward + loss:  grid=(batch_size, 1, 1), block=(32, 1, 1)  -- one warp per sample
 *   Backward:        grid=(batch_size, 1, 1), block=(32, 1, 1)
 *   Adam:            grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1)
 *
 * Compile-time defines (injected via NVRTC dim_overrides):
 *   STATE_DIM          -- tensor-core-aligned state width (e.g. 48 or 56)
 *   VALUE_HIDDEN_DIM   -- hidden layer width (default 128)
 *   IQL_EXPECTILE_TAU  -- expectile parameter (default 0.7)
 *
 * Weight layout (flat f32 buffer, offsets computed on host):
 *   w1[H * STATE_DIM], b1[H], w2[H * H], b2[H], w3[1 * H], b3[1]
 *   TOTAL_PARAMS = H*STATE_DIM + H + H*H + H + H + 1
 */

#ifndef VALUE_HIDDEN_DIM
#define VALUE_HIDDEN_DIM 128
#endif

#ifndef IQL_EXPECTILE_TAU
#define IQL_EXPECTILE_TAU 0.7f
#endif

/* Parameter counts */
#define V_W1_SIZE  (VALUE_HIDDEN_DIM * STATE_DIM)
#define V_B1_SIZE  (VALUE_HIDDEN_DIM)
#define V_W2_SIZE  (VALUE_HIDDEN_DIM * VALUE_HIDDEN_DIM)
#define V_B2_SIZE  (VALUE_HIDDEN_DIM)
#define V_W3_SIZE  (VALUE_HIDDEN_DIM)       /* output layer: 1 × H */
#define V_B3_SIZE  (1)

#define V_TOTAL_PARAMS (V_W1_SIZE + V_B1_SIZE + V_W2_SIZE + V_B2_SIZE + V_W3_SIZE + V_B3_SIZE)

/* Offsets into the flat parameter buffer */
#define V_OFF_W1  0
#define V_OFF_B1  (V_OFF_W1 + V_W1_SIZE)
#define V_OFF_W2  (V_OFF_B1 + V_B1_SIZE)
#define V_OFF_B2  (V_OFF_W2 + V_W2_SIZE)
#define V_OFF_W3  (V_OFF_B2 + V_B2_SIZE)
#define V_OFF_B3  (V_OFF_W3 + V_W3_SIZE)

/* ------------------------------------------------------------------ */
/*  SiLU activation: x * sigmoid(x)                                    */
/* ------------------------------------------------------------------ */
__device__ __forceinline__ float silu(float x) {
    return x / (1.0f + expf(-x));
}

__device__ __forceinline__ float silu_grad(float x) {
    float s = 1.0f / (1.0f + expf(-x));
    return s * (1.0f + x * (1.0f - s));
}

/* ------------------------------------------------------------------ */
/*  Forward + Expectile Loss Kernel                                    */
/* ------------------------------------------------------------------ */
/**
 * Per-sample forward pass through V(s) MLP + expectile loss computation.
 *
 * One warp (32 threads) per sample. Warp-cooperative matvec for each layer.
 * Saves pre-activation values for backward pass.
 *
 * Inputs:
 *   states      [B, STATE_DIM]  -- input states
 *   q_values    [B]             -- target Q(s,a) from DQN
 *   params      [V_TOTAL_PARAMS] -- flat weight buffer
 *
 * Outputs:
 *   v_out       [B]             -- V(s) predictions
 *   loss_out    [B]             -- per-sample expectile loss
 *   total_loss  [1]             -- batch-mean loss (atomicAdd)
 *   save_pre1   [B, H]         -- pre-activation layer 1 (for backward)
 *   save_pre2   [B, H]         -- pre-activation layer 2 (for backward)
 *   save_h1     [B, H]         -- post-activation layer 1 (for backward)
 *   save_h2     [B, H]         -- post-activation layer 2 (for backward)
 */
extern "C" __global__
void iql_forward_loss_kernel(
    const float* __restrict__ states,
    const float* __restrict__ q_values,
    const float* __restrict__ params,
    float* __restrict__ v_out,
    float* __restrict__ loss_out,
    float* __restrict__ total_loss,
    float* __restrict__ save_pre1,
    float* __restrict__ save_pre2,
    float* __restrict__ save_h1,
    float* __restrict__ save_h2,
    int batch_size
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;

    int lane = threadIdx.x;  /* 0..31 */

    const float* w1 = params + V_OFF_W1;
    const float* b1 = params + V_OFF_B1;
    const float* w2 = params + V_OFF_W2;
    const float* b2 = params + V_OFF_B2;
    const float* w3 = params + V_OFF_W3;
    const float* b3 = params + V_OFF_B3;

    const float* x = states + sample * STATE_DIM;

    /* ---- Layer 1: Linear(STATE_DIM -> H) + SiLU ---- */
    /* Each lane computes a subset of H outputs via warp reduction */
    float* pre1_ptr = save_pre1 + sample * VALUE_HIDDEN_DIM;
    float* h1_ptr   = save_h1   + sample * VALUE_HIDDEN_DIM;

    for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) {
        float acc = b1[j];
        for (int k = 0; k < STATE_DIM; k++) {
            acc += w1[j * STATE_DIM + k] * x[k];
        }
        pre1_ptr[j] = acc;
        h1_ptr[j]   = silu(acc);
    }
    __syncwarp();

    /* ---- Layer 2: Linear(H -> H) + SiLU ---- */
    float* pre2_ptr = save_pre2 + sample * VALUE_HIDDEN_DIM;
    float* h2_ptr   = save_h2   + sample * VALUE_HIDDEN_DIM;

    for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) {
        float acc = b2[j];
        for (int k = 0; k < VALUE_HIDDEN_DIM; k++) {
            acc += w2[j * VALUE_HIDDEN_DIM + k] * h1_ptr[k];
        }
        pre2_ptr[j] = acc;
        h2_ptr[j]   = silu(acc);
    }
    __syncwarp();

    /* ---- Output layer: Linear(H -> 1) ---- */
    /* Warp-reduce the dot product across lanes */
    float v_acc = 0.0f;
    for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) {
        v_acc += w3[k] * h2_ptr[k];
    }
    /* Warp reduce */
    for (int offset = 16; offset > 0; offset >>= 1) {
        v_acc += __shfl_down_sync(0xFFFFFFFF, v_acc, offset);
    }
    /* Lane 0 has the full sum */
    if (lane == 0) {
        float v_val = v_acc + b3[0];
        v_out[sample] = v_val;

        /* Expectile loss: L_tau(u) = |tau - 1(u<0)| * u^2 */
        float u = q_values[sample] - v_val;
        float weight = (u >= 0.0f) ? IQL_EXPECTILE_TAU : (1.0f - IQL_EXPECTILE_TAU);
        float sample_loss = weight * u * u;
        loss_out[sample] = sample_loss;

        /* Accumulate batch-mean loss */
        atomicAdd(total_loss, sample_loss / (float)batch_size);
    }
}

/* ------------------------------------------------------------------ */
/*  Backward Kernel                                                    */
/* ------------------------------------------------------------------ */
/**
 * Per-sample backward pass through V(s) MLP.
 *
 * Computes gradients of the expectile loss w.r.t. all parameters.
 * Uses atomicAdd to accumulate into the flat gradient buffer (safe for
 * concurrent samples since different samples write overlapping param indices).
 *
 * Gradient flow:
 *   dL/dV = -2 * weight * u   (where u = Q - V, weight = tau or 1-tau)
 *   dV/dw3[k] = h2[k]
 *   dV/db3    = 1
 *   ... backprop through SiLU and linear layers ...
 */
extern "C" __global__
void iql_backward_kernel(
    const float* __restrict__ states,
    const float* __restrict__ q_values,
    const float* __restrict__ v_out,
    const float* __restrict__ params,
    const float* __restrict__ save_pre1,
    const float* __restrict__ save_pre2,
    const float* __restrict__ save_h1,
    const float* __restrict__ save_h2,
    float* __restrict__ grads,
    int batch_size
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;

    int lane = threadIdx.x;

    const float* w1 = params + V_OFF_W1;
    const float* w2 = params + V_OFF_W2;
    const float* w3 = params + V_OFF_W3;

    float* gw1 = grads + V_OFF_W1;
    float* gb1 = grads + V_OFF_B1;
    float* gw2 = grads + V_OFF_W2;
    float* gb2 = grads + V_OFF_B2;
    float* gw3 = grads + V_OFF_W3;
    float* gb3 = grads + V_OFF_B3;

    const float* x    = states   + sample * STATE_DIM;
    const float* h1   = save_h1  + sample * VALUE_HIDDEN_DIM;
    const float* h2   = save_h2  + sample * VALUE_HIDDEN_DIM;
    const float* pre1 = save_pre1 + sample * VALUE_HIDDEN_DIM;
    const float* pre2 = save_pre2 + sample * VALUE_HIDDEN_DIM;

    float v_val = v_out[sample];
    float q_val = q_values[sample];

    /* dL/dV = -2 * weight * (Q - V) / batch_size */
    float u = q_val - v_val;
    float weight = (u >= 0.0f) ? IQL_EXPECTILE_TAU : (1.0f - IQL_EXPECTILE_TAU);
    float dldv = -2.0f * weight * u / (float)batch_size;

    /* ---- Output layer gradient: dL/dw3, dL/db3 ---- */
    /* dV/dw3[k] = h2[k],  dV/db3 = 1 */
    for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) {
        atomicAdd(&gw3[k], dldv * h2[k]);
    }
    if (lane == 0) {
        atomicAdd(&gb3[0], dldv);
    }

    /* ---- Backprop through layer 2 ---- */
    /* dL/dh2[k] = dldv * w3[k] */
    /* dL/dpre2[k] = dL/dh2[k] * silu'(pre2[k]) */
    for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) {
        float dh2_j = dldv * w3[j];
        float dpre2_j = dh2_j * silu_grad(pre2[j]);

        /* dL/db2[j] = dpre2_j */
        atomicAdd(&gb2[j], dpre2_j);

        /* dL/dw2[j, k] = dpre2_j * h1[k] */
        for (int k = 0; k < VALUE_HIDDEN_DIM; k++) {
            atomicAdd(&gw2[j * VALUE_HIDDEN_DIM + k], dpre2_j * h1[k]);
        }
    }
    __syncwarp();

    /* ---- Backprop through layer 1 ---- */
    /* dL/dh1[k] = sum_j(dL/dpre2[j] * w2[j, k]) */
    for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) {
        float dh1_k = 0.0f;
        for (int j = 0; j < VALUE_HIDDEN_DIM; j++) {
            float dh2_j = dldv * w3[j];
            float dpre2_j = dh2_j * silu_grad(pre2[j]);
            dh1_k += dpre2_j * w2[j * VALUE_HIDDEN_DIM + k];
        }
        float dpre1_k = dh1_k * silu_grad(pre1[k]);

        /* dL/db1[k] = dpre1_k */
        atomicAdd(&gb1[k], dpre1_k);

        /* dL/dw1[k, d] = dpre1_k * x[d] */
        for (int d = 0; d < STATE_DIM; d++) {
            atomicAdd(&gw1[k * STATE_DIM + d], dpre1_k * x[d]);
        }
    }
}

/* ------------------------------------------------------------------ */
/*  Adam Update Kernel                                                 */
/* ------------------------------------------------------------------ */
/**
 * AdamW optimizer step over the flat parameter buffer.
 *
 * grid = ceil(V_TOTAL_PARAMS / 256), block = 256.
 * Each thread updates one parameter.
 *
 * Supports gradient clipping via max_grad_norm (applied before Adam step
 * using the pre-computed grad_norm scalar).
 */
extern "C" __global__
void iql_adam_kernel(
    float* __restrict__ params,
    float* __restrict__ grads,
    float* __restrict__ m,          /* first moment */
    float* __restrict__ v,          /* second moment */
    const float* __restrict__ grad_norm_buf,
    float lr,
    float beta1,
    float beta2,
    float eps,
    float weight_decay,
    float max_grad_norm,
    int t,                          /* Adam step (1-indexed) */
    int total_params
)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total_params) return;

    /* Gradient clipping */
    float gn = grad_norm_buf[0];
    float clip_scale = (gn > max_grad_norm && gn > 0.0f) ? (max_grad_norm / gn) : 1.0f;
    float g = grads[i] * clip_scale;

    /* AdamW: decoupled weight decay */
    float p = params[i];
    p -= lr * weight_decay * p;

    /* Adam moment updates */
    float m_i = beta1 * m[i] + (1.0f - beta1) * g;
    float v_i = beta2 * v[i] + (1.0f - beta2) * g * g;

    /* Bias correction */
    float m_hat = m_i / (1.0f - powf(beta1, (float)t));
    float v_hat = v_i / (1.0f - powf(beta2, (float)t));

    /* Parameter update */
    p -= lr * m_hat / (sqrtf(v_hat) + eps);

    params[i] = p;
    m[i] = m_i;
    v[i] = v_i;
    grads[i] = 0.0f;  /* zero gradient for next step */
}

/* ------------------------------------------------------------------ */
/*  Gradient Norm Kernel                                               */
/* ------------------------------------------------------------------ */
/**
 * Compute L2 norm of the gradient buffer via parallel reduction.
 *
 * grid = ceil(V_TOTAL_PARAMS / 256), block = 256.
 * Each block reduces 256 elements, atomicAdds the partial sum into grad_norm_buf[0].
 * Caller must zero grad_norm_buf[0] before launch.
 */
extern "C" __global__
void iql_grad_norm_kernel(
    const float* __restrict__ grads,
    float* __restrict__ grad_norm_buf,
    int total_params
)
{
    __shared__ float sdata[256];

    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    /* Load and square */
    float val = (i < total_params) ? grads[i] : 0.0f;
    sdata[tid] = val * val;
    __syncthreads();

    /* Tree reduction within block */
    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) {
            sdata[tid] += sdata[tid + s];
        }
        __syncthreads();
    }

    /* Block 0 writes partial sum */
    if (tid == 0) {
        atomicAdd(grad_norm_buf, sqrtf(sdata[0]));
    }
}

/* ------------------------------------------------------------------ */
/*  Forward-Only Kernel (for inference / advantage computation)        */
/* ------------------------------------------------------------------ */
/**
 * V(s) forward pass only -- no loss, no backward.
 *
 * Used for computing advantages A(s,a) = Q(s,a) - V(s) at inference time.
 * Same architecture as forward_loss but skips loss and activation saves.
 */
extern "C" __global__
void iql_forward_kernel(
    const float* __restrict__ states,
    const float* __restrict__ params,
    float* __restrict__ v_out,
    int batch_size
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;

    int lane = threadIdx.x;

    const float* w1 = params + V_OFF_W1;
    const float* b1 = params + V_OFF_B1;
    const float* w2 = params + V_OFF_W2;
    const float* b2 = params + V_OFF_B2;
    const float* w3 = params + V_OFF_W3;
    const float* b3 = params + V_OFF_B3;

    const float* x = states + sample * STATE_DIM;

    /* Shared memory for intermediate activations (reused per layer) */
    __shared__ float sh1[VALUE_HIDDEN_DIM];
    __shared__ float sh2[VALUE_HIDDEN_DIM];

    /* Layer 1: SiLU */
    for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) {
        float acc = b1[j];
        for (int k = 0; k < STATE_DIM; k++) {
            acc += w1[j * STATE_DIM + k] * x[k];
        }
        sh1[j] = silu(acc);
    }
    __syncwarp();

    /* Layer 2: SiLU */
    for (int j = lane; j < VALUE_HIDDEN_DIM; j += 32) {
        float acc = b2[j];
        for (int k = 0; k < VALUE_HIDDEN_DIM; k++) {
            acc += w2[j * VALUE_HIDDEN_DIM + k] * sh1[k];
        }
        sh2[j] = silu(acc);
    }
    __syncwarp();

    /* Output: dot product + bias */
    float v_acc = 0.0f;
    for (int k = lane; k < VALUE_HIDDEN_DIM; k += 32) {
        v_acc += w3[k] * sh2[k];
    }
    for (int offset = 16; offset > 0; offset >>= 1) {
        v_acc += __shfl_down_sync(0xFFFFFFFF, v_acc, offset);
    }
    if (lane == 0) {
        v_out[sample] = v_acc + b3[0];
    }
}
