/**
 * Fused epsilon-greedy action selection kernel with order routing + fill simulation.
 *
 * Given Q-values [batch_size, num_actions], selects actions entirely on GPU:
 * - With probability epsilon: random action (GPU LCG RNG)
 * - With probability 1-epsilon: argmax over Q-values
 * - Routes action through OrderRouter (spread/vol-dependent order type + urgency)
 * - Simulates fill (splitmix64 deterministic hash, order-type-dependent probability)
 * - Unfilled orders override to Flat (exposure=2)
 *
 * Reuses gpu_random(), route_order(), simulate_fill_check() from
 * common_device_functions.cuh (prepended via NVRTC).
 * Launch config: grid=(ceil(batch_size/256),1,1), block=(256,1,1).
 */

/**
 * Basic epsilon-greedy: outputs exposure index (0-4) only.
 * Used when fill simulation is not needed (e.g. inference, CPU routing fallback).
 */
extern "C" __global__ void epsilon_greedy_select(
    const float* __restrict__ q_values,  /* [batch_size, num_actions] */
    unsigned int* rng_states,            /* [batch_size] — persistent LCG state */
    unsigned int* actions_out,           /* [batch_size] — output action indices */
    const float epsilon,
    const int batch_size,
    const int num_actions,
    const float q_gap_threshold          /* min Q-gap for trade entry (0.0 = disabled) */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= batch_size) return;

    unsigned int rng = rng_states[idx];
    float r = gpu_random(&rng);

    unsigned int action;
    if (r < epsilon) {
        /* Random exploration */
        action = (unsigned int)(gpu_random(&rng) * (float)num_actions);
        if (action >= (unsigned int)num_actions) action = (unsigned int)(num_actions - 1);
    } else {
        /* Greedy: argmax over Q-values for this batch element */
        const float* q = q_values + idx * num_actions;
        float best_q = q[0];
        action = 0;
        for (int a = 1; a < num_actions; a++) {
            if (q[a] > best_q) {
                best_q = q[a];
                action = (unsigned int)a;
            }
        }
        /* Q-gap conviction filter: force flat when conviction is low */
        if (q_gap_threshold > 0.0f && num_actions > 2) {
            float q_flat = q[num_actions / 2]; /* center = Flat */
            if (best_q - q_flat < q_gap_threshold) {
                action = 2;
            }
        }
    }

    rng_states[idx] = rng;
    actions_out[idx] = action;
}

/**
 * Fused epsilon-greedy + order routing + fill simulation.
 *
 * Outputs post-fill action indices: if fill check fails, action is overridden
 * to Flat (exposure_idx=2). This eliminates the CPU readback + route_action()
 * + simulate_fill() roundtrip entirely.
 *
 * Parameters:
 *   step_offset — global training step offset for deterministic fill hash
 *   spread, median_spread — for order type selection
 *   volatility, median_vol — for urgency selection
 *   spread_bps — bid-ask spread in basis points (for fill cost)
 *   ioc_fill_prob, limit_fill_min, limit_fill_max — fill simulator config
 *   spread_cost_frac, spread_capture_frac — cost model config
 *
 * Outputs:
 *   actions_out — post-fill exposure indices (0-4), Flat=2 for unfilled
 *   fill_mask_out — 1 if filled, 0 if not (optional: NULL to skip)
 */
extern "C" __global__ void epsilon_greedy_routed(
    const float* __restrict__ q_values,  /* [batch_size, num_actions] */
    unsigned int* rng_states,            /* [batch_size] */
    unsigned int* actions_out,           /* [batch_size] — post-fill exposure */
    int* fill_mask_out,                  /* [batch_size] — 1=filled, 0=not (nullable) */
    const float epsilon,
    const int batch_size,
    const int num_actions,
    const int step_offset,
    const float spread,
    const float median_spread,
    const float volatility,
    const float median_vol,
    const float spread_bps,
    const float ioc_fill_prob,
    const float limit_fill_min,
    const float limit_fill_max,
    const float spread_cost_frac,
    const float spread_capture_frac,
    const float q_gap_threshold
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= batch_size) return;

    unsigned int rng = rng_states[idx];
    float r = gpu_random(&rng);

    /* Step 1: Epsilon-greedy action selection (exposure index 0-4) */
    int exposure_idx;
    if (r < epsilon) {
        exposure_idx = (int)(gpu_random(&rng) * (float)num_actions);
        if (exposure_idx >= num_actions) exposure_idx = num_actions - 1;
    } else {
        const float* q = q_values + idx * num_actions;
        float best_q = q[0];
        exposure_idx = 0;
        for (int a = 1; a < num_actions; a++) {
            if (q[a] > best_q) {
                best_q = q[a];
                exposure_idx = a;
            }
        }
        /* Q-gap conviction filter */
        if (q_gap_threshold > 0.0f && num_actions > 2) {
            float q_flat = q[num_actions / 2];
            if (best_q - q_flat < q_gap_threshold) {
                exposure_idx = num_actions / 2;
            }
        }
    }

    /* Step 2: Order routing (spread/vol → order_type + urgency) */
    int order_type, urgency;
    route_order(spread, median_spread, volatility, median_vol,
                &order_type, &urgency);

    /* Step 3: Fill simulation (deterministic splitmix64 hash) */
    /* Normalized vol for fill probability calculation */
    float norm_vol = (median_vol > 0.0f) ? (volatility / median_vol) : 1.0f;
    float cost_adj;
    int filled = simulate_fill_check(
        order_type, urgency, norm_vol, spread_bps,
        step_offset + idx, exposure_idx,
        ioc_fill_prob, limit_fill_min, limit_fill_max,
        spread_cost_frac, spread_capture_frac,
        &cost_adj
    );

    /* Step 4: Override to Flat if not filled */
    int final_action = filled ? exposure_idx : (num_actions / 2);  /* center = Flat */

    rng_states[idx] = rng;
    actions_out[idx] = (unsigned int)final_action;
    if (fill_mask_out != NULL) {
        fill_mask_out[idx] = filled;
    }
}

/**
 * Branching DQN action selection: 3 independent epsilon-greedy heads.
 *
 * Each head independently selects via epsilon-greedy:
 *   - Exposure head: 5 actions (Short100/Short50/Flat/Long50/Long100)
 *   - Order head: 3 actions (LimitMaker/IoC/Market)
 *   - Urgency head: 3 actions (Patient/Normal/Aggressive)
 *
 * Composes factored action index: exposure * 9 + order * 3 + urgency (0-44).
 *
 * Launch: grid=(ceil(batch_size/256),1,1), block=(256,1,1).
 */
extern "C" __global__ void branching_action_select(
    const float* __restrict__ q_exposure,  /* [batch_size, 5] */
    const float* __restrict__ q_order,     /* [batch_size, 3] */
    const float* __restrict__ q_urgency,   /* [batch_size, 3] */
    unsigned int* rng_states,             /* [batch_size] */
    unsigned int* actions_out,            /* [batch_size] — factored index 0-44 */
    const float epsilon,
    const int batch_size,
    const float q_gap_threshold,          /* min Q-gap for trade entry (0.0 = disabled) */
    const float* __restrict__ bonus_exposure, /* [5] UCB count bonus per exposure action (NULL = disabled) */
    const float* __restrict__ bonus_order,    /* [3] UCB count bonus per order action (NULL = disabled) */
    const float* __restrict__ bonus_urgency   /* [3] UCB count bonus per urgency action (NULL = disabled) */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= batch_size) return;

    unsigned int rng = rng_states[idx];

    /* Head 1: Exposure (5 actions) — with UCB bonus + Q-gap conviction filter */
    int exposure;
    float r1 = gpu_random(&rng);
    if (r1 < epsilon) {
        exposure = (int)(gpu_random(&rng) * 5.0f);
        if (exposure >= 5) exposure = 4;
    } else {
        const float* qe = q_exposure + idx * 5;
        float best = qe[0] + ((bonus_exposure != NULL) ? bonus_exposure[0] : 0.0f);
        exposure = 0;
        for (int a = 1; a < 5; a++) {
            float q_plus_bonus = qe[a] + ((bonus_exposure != NULL) ? bonus_exposure[a] : 0.0f);
            if (q_plus_bonus > best) { best = q_plus_bonus; exposure = a; }
        }
        /* Q-gap filter: force flat when conviction is low.
         * Branching 5-action: S100=0, S50=1, Flat=2, L50=3, L100=4.
         * Compare best Q (without bonus) against Q(Flat) — conviction is
         * about Q-value quality, not exploration bonus. */
        if (q_gap_threshold > 0.0f) {
            float q_flat = qe[2]; /* index 2 = Flat (center of 5-action branching) */
            float q_best_raw = qe[exposure]; /* raw Q without bonus */
            if (q_best_raw - q_flat < q_gap_threshold) {
                exposure = 2;
            }
        }
    }

    /* Head 2: Order type (3 actions) — with UCB bonus */
    int order;
    float r2 = gpu_random(&rng);
    if (r2 < epsilon) {
        order = (int)(gpu_random(&rng) * 3.0f);
        if (order >= 3) order = 2;
    } else {
        const float* qo = q_order + idx * 3;
        float best = qo[0] + ((bonus_order != NULL) ? bonus_order[0] : 0.0f);
        order = 0;
        for (int a = 1; a < 3; a++) {
            float q_plus_bonus = qo[a] + ((bonus_order != NULL) ? bonus_order[a] : 0.0f);
            if (q_plus_bonus > best) { best = q_plus_bonus; order = a; }
        }
    }

    /* Head 3: Urgency (3 actions) — with UCB bonus */
    int urgency;
    float r3 = gpu_random(&rng);
    if (r3 < epsilon) {
        urgency = (int)(gpu_random(&rng) * 3.0f);
        if (urgency >= 3) urgency = 2;
    } else {
        const float* qu = q_urgency + idx * 3;
        float best = qu[0] + ((bonus_urgency != NULL) ? bonus_urgency[0] : 0.0f);
        urgency = 0;
        for (int a = 1; a < 3; a++) {
            float q_plus_bonus = qu[a] + ((bonus_urgency != NULL) ? bonus_urgency[a] : 0.0f);
            if (q_plus_bonus > best) { best = q_plus_bonus; urgency = a; }
        }
    }

    /* Compose factored action: exposure * 9 + order * 3 + urgency */
    unsigned int factored = (unsigned int)(exposure * 9 + order * 3 + urgency);

    rng_states[idx] = rng;
    actions_out[idx] = factored;
}

/**
 * Batch convert exposure indices (0-4) to factored action indices (0-44).
 *
 * Applies OrderRouter::route() logic on GPU via route_order() device function:
 *   factored = exposure * 9 + order_type * 3 + urgency
 *
 * Same arithmetic as the fused epsilon_greedy_routed kernel but as a standalone
 * kernel for post-hoc routing of epsilon_greedy_select output.
 *
 * Launch: grid=(ceil(N/256),1,1), block=(256,1,1).
 */
extern "C" __global__ void batch_route_exposure_to_factored(
    const unsigned int* __restrict__ exposure_actions,  /* [N] 0-4 */
    unsigned int* __restrict__ factored_actions,         /* [N] 0-44 output */
    float spread,
    float median_spread,
    float volatility,
    float median_volatility,
    int N
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= N) return;

    /* Route order: determine order_type and urgency from market conditions */
    int order_type, urgency;
    route_order(spread, median_spread, volatility, median_volatility,
                &order_type, &urgency);

    unsigned int exp = exposure_actions[i];
    factored_actions[i] = exp * 9u + (unsigned int)order_type * 3u + (unsigned int)urgency;
}
