/**
 * GPU-resident curiosity forward model training kernel.
 *
 * Trains the curiosity forward model (2-layer MLP) entirely on GPU with zero
 * CPU involvement. Experience data (states, actions, next_states) is already
 * on GPU from the experience collector. Weights are updated in-place via Adam.
 *
 * Architecture: [CUR_INPUT=MARKET_DIM+3] -> [CUR_HIDDEN=128] LeakyReLU(0.01) -> [CUR_OUTPUT=MARKET_DIM]
 * Total params: 128*(MARKET_DIM+3) + 128 + MARKET_DIM*128 + MARKET_DIM
 *
 * Requires common_device_functions.cuh to be prepended for CUR_INPUT/CUR_HIDDEN/CUR_OUTPUT.
 */

/* Total trainable parameters */
#define CUR_TOTAL_PARAMS (CUR_HIDDEN * CUR_INPUT + CUR_HIDDEN + CUR_OUTPUT * CUR_HIDDEN + CUR_OUTPUT)

/* ------------------------------------------------------------------ */
/*  Kernel 1: Build shifted next_states buffer                        */
/* ------------------------------------------------------------------ */

/**
 * Copy states[i + state_dim] -> next_states[i] for i in [0, N * state_dim).
 * Shifts the states buffer by one timestep to produce next_states.
 * Episode boundary noise (at most n_episodes samples out of total) is
 * negligible for a tiny auxiliary model.
 */
extern "C" __global__ void curiosity_shift_states(
    const float* __restrict__ states,
    float* __restrict__ next_states,
    int N,
    int state_dim
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    int total = N * state_dim;
    if (i < total) {
        next_states[i] = states[i + state_dim];
    }
}

/* ------------------------------------------------------------------ */
/*  Kernel 2: Zero gradient buffers                                   */
/* ------------------------------------------------------------------ */

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

/* ------------------------------------------------------------------ */
/*  Kernel 3: Forward + backward pass, accumulate gradients           */
/* ------------------------------------------------------------------ */

/* Warp-level sum reduction via butterfly shuffle (all lanes get result) */
__device__ __forceinline__ float warp_sum_cur(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    return val;
}

/**
 * One thread per sample. Computes forward pass, MSE loss, and backpropagates
 * gradients via atomicAdd with warp-level pre-reduction (32x fewer atomics).
 * Threads within a warp process different samples but accumulate gradients
 * to the same weight indices, so warp reduction before atomicAdd is valid.
 */
extern "C" __global__ void curiosity_forward_backward(
    const float* __restrict__ states,       /* [N, state_dim]            */
    const int*   __restrict__ actions,       /* [N]                       */
    const float* __restrict__ next_states,   /* [N, state_dim]            */
    const float* __restrict__ w1,            /* [CUR_HIDDEN, CUR_INPUT]   */
    const float* __restrict__ b1,            /* [CUR_HIDDEN]              */
    const float* __restrict__ w2,            /* [CUR_OUTPUT, CUR_HIDDEN]  */
    const float* __restrict__ b2,            /* [CUR_OUTPUT]              */
    float* __restrict__ grad_w1,             /* [CUR_HIDDEN, CUR_INPUT]   */
    float* __restrict__ grad_b1,             /* [CUR_HIDDEN]              */
    float* __restrict__ grad_w2,             /* [CUR_OUTPUT, CUR_HIDDEN]  */
    float* __restrict__ grad_b2,             /* [CUR_OUTPUT]              */
    int N,
    int state_dim
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= N) return;

    const float* state = states + tid * state_dim;
    int action_idx = actions[tid];
    const float* next_state = next_states + tid * state_dim;

    /* ---- Forward pass ---- */

    /* Build input: first MARKET_DIM state features + 3-class action one-hot */
    float input[CUR_INPUT];
    for (int i = 0; i < MARKET_DIM; i++) input[i] = state[i];
    input[MARKET_DIM + 0] = 0.0f;
    input[MARKET_DIM + 1] = 0.0f;
    input[MARKET_DIM + 2] = 0.0f;

    /* Action to category one-hot — decode exposure from factored action.
     * Factored action = exposure_idx * (b1*b2) + order_idx * b2 + urgency_idx.
     * With b0=9, b1=3, b2=3: exposure_idx = action / 9.
     * Categories: Short (idx 0-3), Flat (idx 4), Long (idx 5-8). */
    int exposure_idx = action_idx / (DQN_ORDER_ACTIONS * DQN_URGENCY_ACTIONS);
    int category;
    if (exposure_idx < DQN_NUM_ACTIONS / 2)       category = 0;  /* Short */
    else if (exposure_idx == DQN_NUM_ACTIONS / 2)  category = 1;  /* Flat */
    else                                           category = 2;  /* Long */
    input[MARKET_DIM + category] = 1.0f;

    /* Layer 1: pre_act = w1 * input + b1, hidden = LeakyReLU(pre_act, alpha=0.01) */
    float pre_act[CUR_HIDDEN];
    float hidden[CUR_HIDDEN];
    for (int h = 0; h < CUR_HIDDEN; h++) {
        float sum = b1[h];
        for (int i = 0; i < CUR_INPUT; i++) {
            sum += w1[h * CUR_INPUT + i] * input[i];
        }
        pre_act[h] = sum;
        hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum;
    }

    /* Layer 2: pred = w2 * hidden + b2 (no activation) */
    float pred[CUR_OUTPUT];
    for (int o = 0; o < CUR_OUTPUT; o++) {
        float sum = b2[o];
        for (int h = 0; h < CUR_HIDDEN; h++) {
            sum += w2[o * CUR_HIDDEN + h] * hidden[h];
        }
        pred[o] = sum;
    }

    /* ---- Loss: MSE(pred, next_state[:MARKET_DIM]) ---- */
    /* d_loss/d_pred[o] = 2 * (pred[o] - next_state[o]) / CUR_OUTPUT */
    float d_pred[CUR_OUTPUT];
    float inv_out = 2.0f / (float)CUR_OUTPUT;
    for (int o = 0; o < CUR_OUTPUT; o++) {
        d_pred[o] = (pred[o] - next_state[o]) * inv_out;
    }

    /* ---- Backward: Layer 2 (warp-reduced atomics) ---- */
    float d_hidden[CUR_HIDDEN];
    for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = 0.0f;

    for (int o = 0; o < CUR_OUTPUT; o++) {
        {
            float val = warp_sum_cur(d_pred[o]);
            if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b2[o], val);
        }
        for (int h = 0; h < CUR_HIDDEN; h++) {
            {
                float val = warp_sum_cur(d_pred[o] * hidden[h]);
                if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w2[o * CUR_HIDDEN + h], val);
            }
            d_hidden[h] += w2[o * CUR_HIDDEN + h] * d_pred[o];
        }
    }

    /* ---- Backward: LeakyReLU ---- */
    for (int h = 0; h < CUR_HIDDEN; h++) {
        if (pre_act[h] <= 0.0f) d_hidden[h] *= 0.01f;
    }

    /* ---- Backward: Layer 1 (warp-reduced atomics) ---- */
    for (int h = 0; h < CUR_HIDDEN; h++) {
        {
            float val = warp_sum_cur(d_hidden[h]);
            if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b1[h], val);
        }
        for (int i = 0; i < CUR_INPUT; i++) {
            {
                float val = warp_sum_cur(d_hidden[h] * input[i]);
                if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w1[h * CUR_INPUT + i], val);
            }
        }
    }
}

/* ------------------------------------------------------------------ */
/*  Kernel 4: Adam optimizer step (single param group, legacy)        */
/* ------------------------------------------------------------------ */

/**
 * One thread per parameter. Applies Adam update with bias correction.
 * Gradients are divided by batch_size to compute mean gradient.
 */
extern "C" __global__ void curiosity_adam_step(
    float* __restrict__ params,       /* weights/biases to update      */
    const float* __restrict__ grads,  /* accumulated gradients         */
    float* __restrict__ m,            /* first moment                  */
    float* __restrict__ v,            /* second moment                 */
    int num_params,
    int batch_size,                   /* divide grads by this for mean */
    float lr,
    float beta1,
    float beta2,
    float eps,
    int step                          /* 1-based step counter          */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= num_params) return;

    float g = grads[i] / (float)batch_size;
    m[i] = beta1 * m[i] + (1.0f - beta1) * g;
    v[i] = beta2 * v[i] + (1.0f - beta2) * g * g;
    float m_hat = m[i] / (1.0f - powf(beta1, (float)step));
    float v_hat = v[i] / (1.0f - powf(beta2, (float)step));
    params[i] -= lr * m_hat / (sqrtf(v_hat) + eps);
}

/* ------------------------------------------------------------------ */
/*  Kernel 4b: Fused Adam step — all 4 param groups in one launch     */
/* ------------------------------------------------------------------ */

/**
 * Processes w1, b1, w2, b2 in a single kernel launch. Each thread
 * determines which parameter group it belongs to via offset comparison.
 * Eliminates 3 kernel launch overheads (~150 µs/step on H100).
 */
extern "C" __global__ void curiosity_adam_step_fused(
    float* __restrict__ p0, const float* __restrict__ g0, float* __restrict__ m0, float* __restrict__ v0, int n0,
    float* __restrict__ p1, const float* __restrict__ g1, float* __restrict__ m1, float* __restrict__ v1, int n1,
    float* __restrict__ p2, const float* __restrict__ g2, float* __restrict__ m2, float* __restrict__ v2, int n2,
    float* __restrict__ p3, const float* __restrict__ g3, float* __restrict__ m3, float* __restrict__ v3, int n3,
    int batch_size,
    float lr, float beta1, float beta2, float eps, int step
) {
    int total = n0 + n1 + n2 + n3;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total) return;

    /* Determine which param group and local offset */
    float* p; const float* g; float* mm; float* vv;
    int local_i;
    if (idx < n0) {
        p = p0; g = g0; mm = m0; vv = v0; local_i = idx;
    } else if (idx < n0 + n1) {
        p = p1; g = g1; mm = m1; vv = v1; local_i = idx - n0;
    } else if (idx < n0 + n1 + n2) {
        p = p2; g = g2; mm = m2; vv = v2; local_i = idx - n0 - n1;
    } else {
        p = p3; g = g3; mm = m3; vv = v3; local_i = idx - n0 - n1 - n2;
    }

    float grad = g[local_i] / (float)batch_size;
    mm[local_i] = beta1 * mm[local_i] + (1.0f - beta1) * grad;
    vv[local_i] = beta2 * vv[local_i] + (1.0f - beta2) * grad * grad;
    float m_hat = mm[local_i] / (1.0f - powf(beta1, (float)step));
    float v_hat = vv[local_i] / (1.0f - powf(beta2, (float)step));
    p[local_i] -= lr * m_hat / (sqrtf(v_hat) + eps);
}

/* ------------------------------------------------------------------ */
/*  Kernel 5: Fully fused zero + forward/backward + Adam              */
/* ------------------------------------------------------------------ */

/**
 * Fuses gradient zeroing, forward+backward pass, and Adam optimizer
 * update into a single kernel launch. Uses an atomic block-arrival
 * counter for grid-wide synchronization between the fwd/bwd phase
 * (sample-parallel) and the Adam phase (parameter-parallel).
 *
 * Phase 1 (all threads): Zero gradient buffers via grid-stride loop,
 *   then __syncthreads() within each block.
 * Phase 2 (all threads): Forward + backward pass (one sample per thread),
 *   accumulate gradients via warp-reduced atomicAdd.
 * Phase 3 (last-arriving block only): After all blocks complete phase 2,
 *   the last block to arrive (detected via atomic counter) applies
 *   Adam update across all parameters in a grid-stride loop.
 *
 * Saves 4 memset dispatches + 1 kernel launch = 5 fewer GPU dispatches
 * per training step (~30-50 µs on H100 at high training frequency).
 */
extern "C" __global__ void curiosity_fused_zero_fwd_bwd_adam(
    const float* __restrict__ states,       /* [N, state_dim]            */
    const int*   __restrict__ actions,       /* [N]                       */
    const float* __restrict__ next_states,   /* [N, state_dim]            */
    /* Weights (updated in-place by Adam in phase 3) */
    float* __restrict__ w1,                  /* [CUR_HIDDEN, CUR_INPUT]   */
    float* __restrict__ b1,                  /* [CUR_HIDDEN]              */
    float* __restrict__ w2,                  /* [CUR_OUTPUT, CUR_HIDDEN]  */
    float* __restrict__ b2,                  /* [CUR_OUTPUT]              */
    /* Gradient accumulators (zeroed in phase 1, accumulated in phase 2) */
    float* __restrict__ grad_w1,             /* [CUR_HIDDEN, CUR_INPUT]   */
    float* __restrict__ grad_b1,             /* [CUR_HIDDEN]              */
    float* __restrict__ grad_w2,             /* [CUR_OUTPUT, CUR_HIDDEN]  */
    float* __restrict__ grad_b2,             /* [CUR_OUTPUT]              */
    /* Adam first moment */
    float* __restrict__ adam_m_w1,           /* [CUR_HIDDEN * CUR_INPUT]  */
    float* __restrict__ adam_m_b1,           /* [CUR_HIDDEN]              */
    float* __restrict__ adam_m_w2,           /* [CUR_OUTPUT * CUR_HIDDEN] */
    float* __restrict__ adam_m_b2,           /* [CUR_OUTPUT]              */
    /* Adam second moment */
    float* __restrict__ adam_v_w1,           /* [CUR_HIDDEN * CUR_INPUT]  */
    float* __restrict__ adam_v_b1,           /* [CUR_HIDDEN]              */
    float* __restrict__ adam_v_w2,           /* [CUR_OUTPUT * CUR_HIDDEN] */
    float* __restrict__ adam_v_b2,           /* [CUR_OUTPUT]              */
    /* Atomic block-arrival counter (must be zeroed before launch) */
    int* __restrict__ block_counter,
    /* Scalar parameters */
    int N,                                   /* number of training samples */
    int state_dim,
    int batch_size,                          /* == N, for gradient averaging */
    float lr, float beta1, float beta2, float eps,
    int adam_step                             /* 1-based step counter */
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;

    /* Gradient buffers are pre-zeroed by host-side memset_zeros on the same
     * CUDA stream before this kernel launch. Stream ordering guarantees the
     * memsets complete before this kernel starts. No in-kernel zeroing needed.
     *
     * The previous in-kernel Phase 1 zeroing had an inter-block race:
     * __syncthreads() only syncs within a block, so fast blocks could start
     * atomicAdd-ing gradients while slow blocks were still zeroing. */

    int w1_len = CUR_HIDDEN * CUR_INPUT;
    int b1_len = CUR_HIDDEN;
    int w2_len = CUR_OUTPUT * CUR_HIDDEN;
    int total_grad_elems = CUR_TOTAL_PARAMS;

    /* ================================================================ */
    /*  PHASE 1: Forward + backward pass (one sample per thread)        */
    /* ================================================================ */
    if (tid < N) {
        const float* state = states + tid * state_dim;
        int action_idx = actions[tid];
        const float* next_state = next_states + tid * state_dim;

        /* ---- Forward pass ---- */

        /* Build input: first MARKET_DIM state features + 3-class action one-hot */
        float input[CUR_INPUT];
        for (int i = 0; i < MARKET_DIM; i++) input[i] = state[i];
        input[MARKET_DIM + 0] = 0.0f;
        input[MARKET_DIM + 1] = 0.0f;
        input[MARKET_DIM + 2] = 0.0f;

        /* Action to category one-hot */
        int category;
        if (action_idx <= 1)      category = 0;  /* Short100/Short50 */
        else if (action_idx == 2) category = 1;  /* Flat */
        else                      category = 2;  /* Long50/Long100 */
        input[MARKET_DIM + category] = 1.0f;

        /* Layer 1: pre_act = w1 * input + b1, hidden = LeakyReLU(pre_act, 0.01) */
        float pre_act[CUR_HIDDEN];
        float hidden[CUR_HIDDEN];
        for (int h = 0; h < CUR_HIDDEN; h++) {
            float sum = b1[h];
            for (int i = 0; i < CUR_INPUT; i++) {
                sum += w1[h * CUR_INPUT + i] * input[i];
            }
            pre_act[h] = sum;
            hidden[h] = (sum > 0.0f) ? sum : 0.01f * sum;
        }

        /* Layer 2: pred = w2 * hidden + b2 (no activation) */
        float pred[CUR_OUTPUT];
        for (int o = 0; o < CUR_OUTPUT; o++) {
            float sum = b2[o];
            for (int h = 0; h < CUR_HIDDEN; h++) {
                sum += w2[o * CUR_HIDDEN + h] * hidden[h];
            }
            pred[o] = sum;
        }

        /* ---- Loss: MSE(pred, next_state[:MARKET_DIM]) ---- */
        float d_pred[CUR_OUTPUT];
        float inv_out = 2.0f / (float)CUR_OUTPUT;
        for (int o = 0; o < CUR_OUTPUT; o++) {
            d_pred[o] = (pred[o] - next_state[o]) * inv_out;
        }

        /* ---- Backward: Layer 2 (warp-reduced atomics) ---- */
        float d_hidden[CUR_HIDDEN];
        for (int h = 0; h < CUR_HIDDEN; h++) d_hidden[h] = 0.0f;

        for (int o = 0; o < CUR_OUTPUT; o++) {
            {
                float val = warp_sum_cur(d_pred[o]);
                if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b2[o], val);
            }
            for (int h = 0; h < CUR_HIDDEN; h++) {
                {
                    float val = warp_sum_cur(d_pred[o] * hidden[h]);
                    if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w2[o * CUR_HIDDEN + h], val);
                }
                d_hidden[h] += w2[o * CUR_HIDDEN + h] * d_pred[o];
            }
        }

        /* ---- Backward: LeakyReLU ---- */
        for (int h = 0; h < CUR_HIDDEN; h++) {
            if (pre_act[h] <= 0.0f) d_hidden[h] *= 0.01f;
        }

        /* ---- Backward: Layer 1 (warp-reduced atomics) ---- */
        for (int h = 0; h < CUR_HIDDEN; h++) {
            {
                float val = warp_sum_cur(d_hidden[h]);
                if ((threadIdx.x & 31) == 0) atomicAdd(&grad_b1[h], val);
            }
            for (int i = 0; i < CUR_INPUT; i++) {
                {
                    float val = warp_sum_cur(d_hidden[h] * input[i]);
                    if ((threadIdx.x & 31) == 0) atomicAdd(&grad_w1[h * CUR_INPUT + i], val);
                }
            }
        }
    }

    /* ================================================================ */
    /*  Grid-wide barrier via atomic block-arrival counter              */
    /* ================================================================ */
    /* Ensure all global memory writes (gradient atomicAdds) from this
     * block are visible to all other blocks before signalling arrival. */
    __threadfence();
    __syncthreads();

    /* Thread 0 of each block increments the arrival counter.
     * The last block to arrive (counter == gridDim.x - 1) proceeds
     * to phase 3. All other blocks exit. */
    __shared__ int is_last_block;
    if (threadIdx.x == 0) {
        int arrived = atomicAdd(block_counter, 1);
        is_last_block = (arrived == (int)gridDim.x - 1) ? 1 : 0;
    }
    __syncthreads();
    if (!is_last_block) return;

    /* ================================================================ */
    /*  PHASE 3: Adam optimizer update (last block, grid-stride loop)   */
    /* ================================================================ */
    /* The last-arriving block processes all CUR_TOTAL_PARAMS parameters
     * using blockDim.x threads in a stride loop. This avoids a second
     * kernel launch entirely. */
    int block_tid = threadIdx.x;
    int block_size = blockDim.x;

    for (int i = block_tid; i < total_grad_elems; i += block_size) {
        /* Determine which param group and local offset */
        float* p; float* g; float* mm; float* vv;
        int local_i;
        if (i < w1_len) {
            p = w1; g = grad_w1; mm = adam_m_w1; vv = adam_v_w1; local_i = i;
        } else if (i < w1_len + b1_len) {
            p = b1; g = grad_b1; mm = adam_m_b1; vv = adam_v_b1; local_i = i - w1_len;
        } else if (i < w1_len + b1_len + w2_len) {
            p = w2; g = grad_w2; mm = adam_m_w2; vv = adam_v_w2; local_i = i - w1_len - b1_len;
        } else {
            p = b2; g = grad_b2; mm = adam_m_b2; vv = adam_v_b2; local_i = i - w1_len - b1_len - w2_len;
        }

        float grad = g[local_i] / (float)batch_size;
        mm[local_i] = beta1 * mm[local_i] + (1.0f - beta1) * grad;
        vv[local_i] = beta2 * vv[local_i] + (1.0f - beta2) * grad * grad;
        float m_hat = mm[local_i] / (1.0f - powf(beta1, (float)adam_step));
        float v_hat = vv[local_i] / (1.0f - powf(beta2, (float)adam_step));
        p[local_i] -= lr * m_hat / (sqrtf(v_hat) + eps);
    }
}
