/**
 * MSE loss kernel on expected Q-values for DQN warmup.
 *
 * Computes Mean Squared Error between online expected Q-value and Bellman
 * target expected Q-value. Used during warmup before switching to C51
 * cross-entropy, because MSE provides 10-100x stronger gradient signal
 * when per-bar returns are ~0.001 and C51 distributions are nearly identical.
 *
 * Same input signature as c51_loss_batched so the two kernels are
 * interchangeable in the CUDA Graph capture.
 *
 * Grid  = (batch_size, 1, 1)
 * Block = (BLOCK_THREADS=256, 1, 1)
 *
 * Requires #define: NUM_ATOMS, V_MIN, V_MAX, BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE
 */

#ifndef NUM_ATOMS
#define NUM_ATOMS 51
#endif
#ifndef BRANCH_0_SIZE
#define BRANCH_0_SIZE 9
#endif
#ifndef BRANCH_1_SIZE
#define BRANCH_1_SIZE 3
#endif
#ifndef BRANCH_2_SIZE
#define BRANCH_2_SIZE 3
#endif
#ifndef V_MIN
#define V_MIN (-25.0f)
#endif
#ifndef V_MAX
#define V_MAX (25.0f)
#endif
#define DELTA_Z ((V_MAX - V_MIN) / (float)(NUM_ATOMS - 1))

#define BRANCH_0_ATOMS (BRANCH_0_SIZE * NUM_ATOMS)
#define BRANCH_1_ATOMS (BRANCH_1_SIZE * NUM_ATOMS)
#define BRANCH_2_ATOMS (BRANCH_2_SIZE * NUM_ATOMS)
#define MAX_BRANCH_SIZE 9
#define NUM_BRANCHES 3
#define NUM_WARPS 8
#define BLOCK_THREADS 256

/* ── Shared memory layout (same as c51_loss_kernel.cu) ─────────────── */
#define SHMEM_SUPPORT_OFF      0
#define SHMEM_VAL_OFF          (SHMEM_SUPPORT_OFF   + NUM_ATOMS)
#define SHMEM_ADV_OFF          (SHMEM_VAL_OFF        + NUM_ATOMS)
#define SHMEM_LP_OFF           (SHMEM_ADV_OFF        + MAX_BRANCH_SIZE * NUM_ATOMS)
#define SHMEM_PROJ_OFF         (SHMEM_LP_OFF         + NUM_ATOMS)
#define SHMEM_CURRENT_LP_OFF   (SHMEM_PROJ_OFF       + NUM_ATOMS)
#define SHMEM_REDUCE_OFF       (SHMEM_CURRENT_LP_OFF + NUM_ATOMS)
#define SHMEM_TOTAL_FLOATS     (SHMEM_REDUCE_OFF     + NUM_WARPS)

/* ── Device helpers (same as c51_loss_kernel.cu) ───────────────────── */

__device__ __forceinline__ float block_reduce_max(
    float val, float* __restrict__ shmem_reduce, int tid
) {
    unsigned int mask = 0xFFFFFFFF;
    int lane = tid & 31;
    int warp = tid >> 5;
    for (int offset = 16; offset > 0; offset >>= 1)
        val = fmaxf(val, __shfl_xor_sync(mask, val, offset));
    if (lane == 0) shmem_reduce[warp] = val;
    __syncthreads();
    float result = -1e30f;
    if (tid < NUM_WARPS) result = shmem_reduce[tid];
    if (tid < 32) {
        for (int offset = NUM_WARPS >> 1; offset > 0; offset >>= 1)
            result = fmaxf(result, __shfl_xor_sync(mask, result, offset));
        if (tid == 0) shmem_reduce[0] = result;
    }
    __syncthreads();
    return shmem_reduce[0];
}

__device__ __forceinline__ float block_reduce_sum(
    float val, float* __restrict__ shmem_reduce, int tid
) {
    unsigned int mask = 0xFFFFFFFF;
    int lane = tid & 31;
    int warp = tid >> 5;
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(mask, val, offset);
    if (lane == 0) shmem_reduce[warp] = val;
    __syncthreads();
    float result = 0.0f;
    if (tid < NUM_WARPS) result = shmem_reduce[tid];
    if (tid < 32) {
        for (int offset = NUM_WARPS >> 1; offset > 0; offset >>= 1)
            result += __shfl_xor_sync(mask, result, offset);
        if (tid == 0) shmem_reduce[0] = result;
    }
    __syncthreads();
    return shmem_reduce[0];
}

/**
 * Block-cooperative log_softmax, then expected Q from log-probs.
 * Also writes softmax probs to save_probs if save_probs != NULL.
 */
__device__ float block_softmax_expected_q(
    float* __restrict__ shmem_logits,      /* [NUM_ATOMS] in shmem (modified in-place) */
    const float* __restrict__ shmem_support, /* [NUM_ATOMS] z_j */
    float* __restrict__ shmem_reduce,      /* [NUM_WARPS] scratch */
    float* __restrict__ save_probs_out,    /* [NUM_ATOMS] global mem or NULL */
    int tid
) {
    /* Pass 1: block-max */
    float local_max = -1e30f;
    for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
        local_max = fmaxf(local_max, shmem_logits[i]);
    float block_max = block_reduce_max(local_max, shmem_reduce, tid);

    /* Pass 2: sum(exp(x - max)) */
    float local_sum = 0.0f;
    for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
        local_sum += expf(shmem_logits[i] - block_max);
    float block_sum = block_reduce_sum(local_sum, shmem_reduce, tid);
    float log_sum_exp = logf(block_sum + 1e-10f);

    /* Pass 3: softmax probs + expected Q */
    /* Write probs to shmem_logits in-place */
    for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS) {
        float lp = shmem_logits[i] - block_max - log_sum_exp;
        float prob = expf(lp);
        shmem_logits[i] = prob;  /* overwrite logits with probs */
        if (save_probs_out) save_probs_out[i] = prob;
    }
    __syncthreads();

    /* E[Q] = sum_j(prob_j * z_j) */
    float local_eq = 0.0f;
    for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
        local_eq += shmem_logits[j] * shmem_support[j];
    return block_reduce_sum(local_eq, shmem_reduce, tid);
}


/* ══════════════════════════════════════════════════════════════════════
 * MAIN KERNEL: mse_loss_batched
 *
 * For each sample b, each branch d:
 *   1. Dueling on online logits → softmax → E[Q_online]
 *   2. Double DQN: online-next argmax → target dueling → softmax → E[Q_target]
 *   3. Bellman target: target_Q = reward + gamma * (1-done) * E[Q_target]
 *   4. TD error: td = E[Q_online] - target_Q
 *   5. MSE loss: 0.5 * td^2 (averaged over branches, IS-weighted)
 *
 * Saves: softmax probs → save_current_lp
 *        td_error + E[Q] per branch → save_projected (for grad kernel)
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void mse_loss_batched(
    /* ── Online network outputs on current STATES ─────────────────── */
    const float* __restrict__ on_value_logits,    /* [B, NUM_ATOMS] */
    const float* __restrict__ on_adv_logits_b0,   /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
    const float* __restrict__ on_adv_logits_b1,   /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
    const float* __restrict__ on_adv_logits_b2,   /* [B, BRANCH_2_SIZE * NUM_ATOMS] */

    /* ── Target network outputs on NEXT_STATES ────────────────────── */
    const float* __restrict__ tg_value_logits,    /* [B, NUM_ATOMS] */
    const float* __restrict__ tg_adv_logits_b0,   /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
    const float* __restrict__ tg_adv_logits_b1,   /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
    const float* __restrict__ tg_adv_logits_b2,   /* [B, BRANCH_2_SIZE * NUM_ATOMS] */

    /* ── Online network outputs on NEXT_STATES (Double DQN selector) */
    const float* __restrict__ on_next_value_logits,   /* [B, NUM_ATOMS] */
    const float* __restrict__ on_next_adv_logits_b0,  /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
    const float* __restrict__ on_next_adv_logits_b1,  /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
    const float* __restrict__ on_next_adv_logits_b2,  /* [B, BRANCH_2_SIZE * NUM_ATOMS] */

    /* ── Batch data ───────────────────────────────────────────────── */
    const int*   __restrict__ actions,            /* [B] factored action indices 0-44 */
    const float* __restrict__ rewards,            /* [B] */
    const float* __restrict__ dones,              /* [B] */
    const float* __restrict__ is_weights,         /* [B] PER importance-sampling weights */

    /* ── Outputs ──────────────────────────────────────────────────── */
    float* __restrict__ per_sample_loss,          /* [B] IS-weighted loss per sample */
    float* __restrict__ td_errors,               /* [B] unweighted, for PER priority update */
    float* __restrict__ total_loss,              /* [1] batch mean loss (atomicAdd) */

    /* ── Saved tensors for backward pass ─────────────────────────── */
    float* __restrict__ save_current_lp,          /* [B, NUM_BRANCHES, NUM_ATOMS] online probs */
    float* __restrict__ save_projected,           /* [B, NUM_BRANCHES, NUM_ATOMS] [td, E_Q, 0..] */

    /* ── Config ───────────────────────────────────────────────────── */
    float gamma,
    int batch_size
) {
    extern __shared__ float shmem[];

    float* shmem_support    = shmem + SHMEM_SUPPORT_OFF;
    float* shmem_val        = shmem + SHMEM_VAL_OFF;
    float* shmem_adv        = shmem + SHMEM_ADV_OFF;
    float* shmem_lp         = shmem + SHMEM_LP_OFF;
    float* shmem_proj       = shmem + SHMEM_PROJ_OFF;
    float* shmem_current_lp = shmem + SHMEM_CURRENT_LP_OFF;
    float* shmem_reduce     = shmem + SHMEM_REDUCE_OFF;

    int sample_id = blockIdx.x;
    int tid = threadIdx.x;

    if (sample_id >= batch_size) return;

    /* ── Precompute C51 support atoms once per block ─────────────── */
    for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
        shmem_support[j] = V_MIN + (float)j * DELTA_Z;
    __syncthreads();

    /* ── Decode factored action into per-branch indices ──────────── */
    int factored_action = actions[sample_id];
    int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE;
    if (factored_action < 0 || factored_action >= max_action)
        factored_action = 0;
    int a0 = factored_action / (BRANCH_1_SIZE * BRANCH_2_SIZE);
    int a1 = (factored_action / BRANCH_2_SIZE) % BRANCH_1_SIZE;
    int a2 = factored_action % BRANCH_2_SIZE;

    const int branch_action[NUM_BRANCHES] = { a0, a1, a2 };
    const int branch_sizes[NUM_BRANCHES]  = { BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE };

    const float* on_adv_row[NUM_BRANCHES] = {
        on_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
        on_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
        on_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
    };
    const float* tg_adv_row[NUM_BRANCHES] = {
        tg_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
        tg_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
        tg_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
    };
    const float* on_next_adv_row[NUM_BRANCHES] = {
        on_next_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
        on_next_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
        on_next_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
    };

    const float* on_val_row      = on_value_logits      + (long long)sample_id * NUM_ATOMS;
    const float* tg_val_row      = tg_value_logits      + (long long)sample_id * NUM_ATOMS;
    const float* on_next_val_row = on_next_value_logits + (long long)sample_id * NUM_ATOMS;

    float reward    = rewards[sample_id];
    float done      = dones[sample_id];
    float is_weight = is_weights[sample_id];

    float total_mse = 0.0f;
    float total_abs_td = 0.0f;

    for (int d = 0; d < NUM_BRANCHES; d++) {
        int n_d    = branch_sizes[d];
        int a_d    = branch_action[d];
        int n_atoms = n_d * NUM_ATOMS;

        long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * NUM_ATOMS;

        /* ═══════════════════════════════════════════════════════════
         * STEP a: Online E[Q] for the taken action a_d
         * ═══════════════════════════════════════════════════════════ */

        /* Load online value logits */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_val[j] = on_val_row[j];
        __syncthreads();

        /* Load advantage logits */
        for (int i = tid; i < n_atoms; i += BLOCK_THREADS)
            shmem_adv[i] = on_adv_row[d][i];
        __syncthreads();

        /* Dueling for action a_d: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j]) */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * NUM_ATOMS + j];
            a_mean /= (float)n_d;
            shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * NUM_ATOMS + j] - a_mean;
        }
        __syncthreads();

        /* softmax + E[Q], saving probs to save_current_lp */
        float online_eq = block_softmax_expected_q(
            shmem_lp, shmem_support, shmem_reduce,
            save_current_lp + save_off, tid
        );

        /* ═══════════════════════════════════════════════════════════
         * STEP b: Double DQN — argmax action on online_next network
         * ═══════════════════════════════════════════════════════════ */

        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_val[j] = on_next_val_row[j];
        __syncthreads();

        for (int i = tid; i < n_atoms; i += BLOCK_THREADS)
            shmem_adv[i] = on_next_adv_row[d][i];
        __syncthreads();

        /* Per-atom mean advantage for dueling */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * NUM_ATOMS + j];
            a_mean /= (float)n_d;
            shmem_proj[j] = a_mean;
        }
        __syncthreads();

        __shared__ float eq_per_action[MAX_BRANCH_SIZE];
        __shared__ int best_next_a_shared;

        for (int a = 0; a < n_d; a++) {
            for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
                shmem_lp[j] = shmem_val[j] + shmem_adv[a * NUM_ATOMS + j] - shmem_proj[j];
            __syncthreads();

            float eq = block_softmax_expected_q(
                shmem_lp, shmem_support, shmem_reduce,
                (float*)0, tid  /* no save needed for action selection */
            );
            if (tid == 0) eq_per_action[a] = eq;
            __syncthreads();
        }

        if (tid == 0) {
            int best = 0;
            for (int a = 1; a < n_d; a++)
                if (eq_per_action[a] > eq_per_action[best]) best = a;
            best_next_a_shared = best;
        }
        __syncthreads();
        int best_next_a = best_next_a_shared;

        /* ═══════════════════════════════════════════════════════════
         * STEP c: Target E[Q] for best_next action
         * ═══════════════════════════════════════════════════════════ */

        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_val[j] = tg_val_row[j];
        __syncthreads();

        for (int i = tid; i < n_atoms; i += BLOCK_THREADS)
            shmem_adv[i] = tg_adv_row[d][i];
        __syncthreads();

        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * NUM_ATOMS + j];
            a_mean /= (float)n_d;
            shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * NUM_ATOMS + j] - a_mean;
        }
        __syncthreads();

        float target_eq = block_softmax_expected_q(
            shmem_lp, shmem_support, shmem_reduce,
            (float*)0, tid
        );

        /* ═══════════════════════════════════════════════════════════
         * STEP d: Bellman target + MSE
         * ═══════════════════════════════════════════════════════════ */
        float target_q = reward + gamma * (1.0f - done) * target_eq;
        float td = online_eq - target_q;
        float mse = 0.5f * td * td;

        total_mse += mse;
        total_abs_td += fabsf(td);

        /* Save td_error and E[Q] for gradient kernel.
         * Layout: save_projected[save_off + 0] = td_error
         *         save_projected[save_off + 1] = E[Q_online]
         *         save_projected[save_off + 2..] = 0 (padding) */
        if (tid == 0) {
            save_projected[save_off + 0] = td;
            save_projected[save_off + 1] = online_eq;
        }
        /* Zero padding */
        for (int j = tid + 2; j < NUM_ATOMS; j += BLOCK_THREADS)
            save_projected[save_off + j] = 0.0f;
        __syncthreads();

    }  /* end branch loop */

    /* ── Aggregate: average over branches, apply IS weight ───────── */
    float avg_mse = total_mse / (float)NUM_BRANCHES;
    float avg_td  = total_abs_td / (float)NUM_BRANCHES;

    if (tid == 0) {
        float weighted_loss = avg_mse * is_weight;
        per_sample_loss[sample_id] = weighted_loss;
        td_errors[sample_id]       = avg_td;
        atomicAdd(total_loss, weighted_loss / (float)batch_size);
    }
}
