/**
 * 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);
    float eps_bf = epsilon;

    unsigned int action;
    if (r < eps_bf) {
        /* Random exploration */
        float rand_val = gpu_random(&rng);
        action = (unsigned int)((float)rand_val * (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;
            }
        }
        /* Adaptive Q-gap conviction filter: threshold is a fraction of Q-range.
         * Auto-scales when v_min/v_max changes — conviction is relative, not absolute. */
        if (q_gap_threshold > 0.0f && num_actions > 2) {
            float q_flat = q[num_actions / 2]; /* center = Flat */
            /* Compute Q-range for adaptive scaling */
            float q_min_v = q[0], q_max_v = q[0];
            for (int a = 1; a < num_actions; a++) {
                if (q[a] > q_max_v) q_max_v = q[a];
                if (q[a] < q_min_v) q_min_v = q[a];
            }
            float q_range = fmaxf((q_max_v) - (q_min_v), 0.001f);
            float adaptive_gap = q_gap_threshold * q_range;
            if ((best_q) - (q_flat) < adaptive_gap) {
                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);
    float eps_bf = epsilon;

    /* Step 1: Epsilon-greedy action selection (exposure index 0-4) */
    int exposure_idx;
    if (r < eps_bf) {
        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;
            }
        }
        /* Adaptive Q-gap conviction filter: fraction of Q-range */
        if (q_gap_threshold > 0.0f && num_actions > 2) {
            float q_flat = q[num_actions / 2];
            float qmin = q[0], qmax = q[0];
            for (int a = 1; a < num_actions; a++) {
                if (q[a] > qmax) qmax = q[a];
                if (q[a] < qmin) qmin = q[a];
            }
            float qr = fmaxf((qmax) - (qmin), 0.001f);
            if ((best_q) - (q_flat) < q_gap_threshold * qr) {
                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, (float)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:
 *   - Direction head: 4 actions (Short/Hold/Long/Flat)
 *   - Order head: 3 actions (LimitMaker/IoC/Market)
 *   - Urgency head: 3 actions (Patient/Normal/Aggressive)
 *
 * Composes factored action index: direction * 9 + order * 3 + urgency (0-35).
 *
 * 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, 4] direction Q-values */
    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-35 */
    const float* __restrict__ per_sample_epsilon, /* [B] state-dependent epsilon */
    const int batch_size,
    const float q_gap_threshold,          /* min Q-gap for trade entry (0.0 = disabled) */
    const float* __restrict__ bonus_exposure, /* [4] UCB count bonus per direction 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];
    float eps = per_sample_epsilon[idx];

    /* Head 1: Direction (4 actions: Short/Hold/Long/Flat) -- Boltzmann softmax with UCB bonus.
     * Same pattern as order/urgency. Temperature = Q-gap → sharp when
     * Q-values differentiate, uniform when flat (no catastrophic flips). */
    int exposure;
    float r1 = gpu_random(&rng);
    if (r1 < eps) {
        exposure = (int)(gpu_random(&rng) * 4.0f);
        if (exposure >= 4) exposure = 3;
    } else {
        const float* qe = q_exposure + idx * 4;
        float q_max_e = qe[0] + ((bonus_exposure != NULL) ? bonus_exposure[0] : 0.0f);
        float q_min_e = q_max_e;
        float qe_adj[4];
        qe_adj[0] = q_max_e;
        for (int a = 1; a < 4; a++) {
            qe_adj[a] = qe[a] + ((bonus_exposure != NULL) ? bonus_exposure[a] : 0.0f);
            q_max_e = fmaxf(q_max_e, qe_adj[a]);
            q_min_e = fminf(q_min_e, qe_adj[a]);
        }
        float tau_exp = fmaxf(q_max_e - q_min_e, 0.01f);
        float sum_e = 0.0f;
        float exps_exp[4];
        for (int a = 0; a < 4; a++) {
            exps_exp[a] = expf((qe_adj[a] - q_max_e) / tau_exp);
            sum_e += exps_exp[a];
        }
        float re = gpu_random(&rng) * sum_e;
        float cum = 0.0f;
        exposure = 3;
        for (int a = 0; a < 4; a++) {
            cum += exps_exp[a];
            if (re < cum) { exposure = a; break; }
        }
    }

    /* Head 2: Order type (3 actions) -- Boltzmann softmax with UCB bonus */
    int order;
    float r2 = gpu_random(&rng);
    if (r2 < eps) {
        /* Random exploration */
        order = (int)(gpu_random(&rng) * 3.0f);
        if (order >= 3) order = 2;
    } else {
        /* Boltzmann softmax over order Q-values */
        const float* qo = q_order + idx * 3;
        float q_max_o = qo[0] + ((bonus_order != NULL) ? bonus_order[0] : 0.0f);
        float q_min_o = q_max_o;
        float qo_adj[3];
        qo_adj[0] = q_max_o;
        for (int a = 1; a < 3; a++) {
            qo_adj[a] = qo[a] + ((bonus_order != NULL) ? bonus_order[a] : 0.0f);
            q_max_o = fmaxf(q_max_o, qo_adj[a]);
            q_min_o = fminf(q_min_o, qo_adj[a]);
        }
        float tau_ord = fmaxf(q_max_o - q_min_o, 0.01f);
        float sum_e = 0.0f;
        float exps_ord[3];
        for (int a = 0; a < 3; a++) {
            exps_ord[a] = expf((qo_adj[a] - q_max_o) / tau_ord);
            sum_e += exps_ord[a];
        }
        float ro = gpu_random(&rng) * sum_e;
        float cum = 0.0f;
        order = 2;
        for (int a = 0; a < 3; a++) {
            cum += exps_ord[a];
            if (ro < cum) { order = a; break; }
        }
    }

    /* Head 3: Urgency (3 actions) -- Boltzmann softmax with UCB bonus */
    int urgency;
    float r3 = gpu_random(&rng);
    if (r3 < eps) {
        /* Random exploration */
        urgency = (int)(gpu_random(&rng) * 3.0f);
        if (urgency >= 3) urgency = 2;
    } else {
        /* Boltzmann softmax over urgency Q-values */
        const float* qu = q_urgency + idx * 3;
        float q_max_u = qu[0] + ((bonus_urgency != NULL) ? bonus_urgency[0] : 0.0f);
        float q_min_u = q_max_u;
        float qu_adj[3];
        qu_adj[0] = q_max_u;
        for (int a = 1; a < 3; a++) {
            qu_adj[a] = qu[a] + ((bonus_urgency != NULL) ? bonus_urgency[a] : 0.0f);
            q_max_u = fmaxf(q_max_u, qu_adj[a]);
            q_min_u = fminf(q_min_u, qu_adj[a]);
        }
        float tau_urg = fmaxf(q_max_u - q_min_u, 0.01f);
        float sum_e = 0.0f;
        float exps_urg[3];
        for (int a = 0; a < 3; a++) {
            exps_urg[a] = expf((qu_adj[a] - q_max_u) / tau_urg);
            sum_e += exps_urg[a];
        }
        float ru = gpu_random(&rng) * sum_e;
        float cum = 0.0f;
        urgency = 2;
        for (int a = 0; a < 3; a++) {
            cum += exps_urg[a];
            if (ru < cum) { urgency = a; break; }
        }
    }

    /* 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
 *
 * 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;
}

/**
 * Fill f32 array with a scalar value.
 * Used for GPU-side epsilon fill — eliminates per-step HtoD memcpy.
 * Launch: grid=(ceil(n/256),1,1), block=(256,1,1).
 */
extern "C" __global__ void fill_f32(
    float* __restrict__ out,
    float value,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = value;
}
