// rl_iqn_loss.cu — IQN quantile Huber loss (forward + backward).
//
// Computes the quantile regression loss between online and target IQN
// outputs for the action taken:
//
//   δ_{ij} = target_q[b, j, a*] - online_q[b, i, a*]
//   ρ_{τ_i}(δ_{ij}) = |τ_i - 1(δ_{ij} < 0)| × Huber(δ_{ij}, κ=1.0)
//   L_b = (1 / N_TAU_ONLINE) × (1 / N_TAU_TARGET) × Σ_i Σ_j ρ_{τ_i}(δ_{ij})
//
// where:
//   Huber(δ, κ) = 0.5 × δ² / κ       if |δ| ≤ κ
//               = |δ| - 0.5 × κ       otherwise
//
// Backward: ∂L/∂online_q[b, i, a*] is computed and written to grad_out.
//
// Block layout:
//   grid = (B, 1, 1); block = (N_TAU_ONLINE_MAX, 1, 1).
//   One block per batch entry. Each thread handles one online quantile
//   index i, reducing across all target quantile indices j.
//   Cross-batch loss aggregation uses block tree-reduce into shared
//   memory (no atomicAdd per `feedback_no_atomicadd.md`), with block 0
//   writing the final scalar loss.
//
// κ (Huber threshold) = 1.0 is standard for IQN (Dabney et al. 2018).

#define HIDDEN_DIM       128
#define N_ACTIONS        11
#define N_TAU_MAX        64
#define KAPPA            1.0f

// ─────────────────────────────────────────────────────────────────────
// rl_iqn_loss_fwd:
//   Compute per-batch quantile Huber loss and per-online-quantile gradients.
//
// Inputs:
//   online_q        [B, N_TAU_ONLINE, N_ACTIONS] — online IQN Q values
//   target_q        [B, N_TAU_TARGET, N_ACTIONS] — target IQN Q values
//   tau_online      [B, N_TAU_ONLINE]            — online quantile fractions
//   actions_taken   [B]                          — action index per batch
//   B                                           — batch size
//   N_TAU_ONLINE                                — number of online quantiles
//   N_TAU_TARGET                                — number of target quantiles
// Outputs:
//   loss_per_batch  [B]                         — per-batch scalar loss
//   grad_online_q   [B, N_TAU_ONLINE, N_ACTIONS] — gradient w.r.t. online Q
//                   (only the action-taken slice is non-zero)
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_iqn_loss_fwd(
    const float* __restrict__ online_q,       // [B, N_TAU_ONLINE, N_ACTIONS]
    const float* __restrict__ target_q,       // [B, N_TAU_TARGET, N_ACTIONS]
    const float* __restrict__ tau_online,     // [B, N_TAU_ONLINE]
    const int*   __restrict__ actions_taken,  // [B]
    int          B,
    int          N_TAU_ONLINE,
    int          N_TAU_TARGET,
    float*       __restrict__ loss_per_batch, // [B]
    float*       __restrict__ grad_online_q   // [B, N_TAU_ONLINE, N_ACTIONS]
) {
    const int batch = blockIdx.x;
    const int i = threadIdx.x;  // online quantile index
    if (batch >= B) return;
    if (i >= N_TAU_ONLINE) return;

    const int act = actions_taken[batch];

    // Online Q value for this quantile and action
    const float q_online_i = online_q[batch * N_TAU_ONLINE * N_ACTIONS + i * N_ACTIONS + act];
    const float tau_i = tau_online[batch * N_TAU_ONLINE + i];

    float loss_sum = 0.0f;
    float grad_sum = 0.0f;

    // Sum over all target quantiles j
    for (int j = 0; j < N_TAU_TARGET; ++j) {
        const float q_target_j = target_q[batch * N_TAU_TARGET * N_ACTIONS + j * N_ACTIONS + act];
        const float delta = q_target_j - q_online_i;
        const float abs_delta = fabsf(delta);

        // Huber loss
        float huber;
        float huber_grad;  // dHuber/d(online_q) = -dHuber/d(delta)
        if (abs_delta <= KAPPA) {
            huber = 0.5f * delta * delta / KAPPA;
            huber_grad = -delta / KAPPA;
        } else {
            huber = abs_delta - 0.5f * KAPPA;
            huber_grad = (delta < 0.0f) ? 1.0f : -1.0f;
        }

        // Quantile weight: |tau_i - 1(delta < 0)|
        float indicator = (delta < 0.0f) ? 1.0f : 0.0f;
        float weight = fabsf(tau_i - indicator);

        loss_sum += weight * huber;
        // Gradient: d/d(online_q) of weight * huber
        // weight doesn't depend on online_q (it depends on sign of delta
        // which is a subgradient boundary). Standard IQN gradient:
        grad_sum += weight * huber_grad;
    }

    // Normalize by N_TAU_TARGET (the N_TAU_ONLINE normalization is applied
    // when writing to loss_per_batch after the shared-mem reduce below)
    float loss_i = loss_sum / (float)N_TAU_TARGET;
    float grad_i = grad_sum / ((float)N_TAU_ONLINE * (float)N_TAU_TARGET);

    // Write gradient for this online quantile's action slot
    // Zero-init the non-taken-action gradient slots
    const int grad_base = batch * N_TAU_ONLINE * N_ACTIONS + i * N_ACTIONS;
    for (int a = 0; a < N_ACTIONS; ++a) {
        grad_online_q[grad_base + a] = (a == act) ? grad_i : 0.0f;
    }

    // Block tree-reduce to get per-batch loss (sum over online quantiles,
    // then divide by N_TAU_ONLINE)
    __shared__ float s_loss[N_TAU_MAX];
    s_loss[i] = loss_i;
    __syncthreads();

    // Tree reduce within the block
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (i < stride && (i + stride) < N_TAU_ONLINE) {
            s_loss[i] += s_loss[i + stride];
        }
        __syncthreads();
    }

    if (i == 0) {
        loss_per_batch[batch] = s_loss[0] / (float)N_TAU_ONLINE;
    }
}
