// rl_confidence_gate.cu — override opening actions to Hold when the
// C51 distributional Q is insufficiently confident about the chosen
// action.
//
// For each batch where position==0 and the chosen action is an opening
// action (a0, a1, a5, a6):
//   1. Extract q_logits[b, a*, 0..Q_N_ATOMS]
//   2. Numerically-stable softmax → probs
//   3. μ = Σ probs[i] × atom_supports[i]
//   4. σ² = Σ probs[i] × (atom_supports[i] - μ)²
//   5. conf = clamp((μ - λ×√σ²) / σ_norm, 0, 1)
//   6. If conf < threshold → override actions[b] = 2 (Hold)
//
// One block per batch, single thread (Q_N_ATOMS=21 is small enough
// for sequential softmax + reduce). No shared memory needed.
//
// Per `feedback_no_atomicadd`: fired-count written by thread 0 via
// simple sequential count (b_size=1 in practice; multi-batch uses
// block-serial approach consistent with other diag counters).
// Per `feedback_cpu_is_read_only`: pure device-side.

#include <stdint.h>
#include <math.h>

#define Q_N_ATOMS                      21
#define N_ACTIONS                       11
#define ACTION_HOLD                     2
#define RL_CONF_GATE_THRESHOLD_INDEX    512
#define RL_CONF_GATE_LAMBDA_INDEX       513
#define RL_CONF_GATE_SIGMA_NORM_INDEX   514
#define RL_CONF_GATE_FIRED_COUNT_INDEX  515
#define RL_GATE_WARMUP_STEPS_INDEX     524
#define RL_STEP_COUNTER_ISV_INDEX      548
#define RL_HOLD_TARGET_FRAC_INDEX      575
#define RL_HOLD_FRAC_EMA_INDEX         576
#define RL_CONF_GATE_MAX_HOLD_FRAC_INDEX 584
// Phase 7b F5 (2026-06-05) Option A — interop with state-conditional
// action availability mask. When F5 master gate (slot 823) is engaged
// AND position is flat, F5 has DELIBERATELY forced the agent into an
// opening action by masking Hold (and other invalid actions) in
// pi_logits. Overriding that back to Hold here neutralizes F5's
// surfer→trend choice-set forcing and prevents F5-G1 from achieving
// `hold_frac_flat == 0`. The suppression is gated on F5-engaged AND
// flat — in-position safety overrides are preserved verbatim.
#define RL_F5_STATE_MASK_ENABLED_INDEX 823
#define THRESHOLD_MIN                  0.05f
#define THRESHOLD_MAX                  0.95f
#define HOLD_EMA_ALPHA                 0.1f
#define THRESHOLD_ADJUST_RATE          1.1f

extern "C" __global__ void rl_confidence_gate(
    int*                 __restrict__ actions,        // [B] IN/OUT
    const float*         __restrict__ q_logits,       // [B × N_ACTIONS × Q_N_ATOMS]
    const float*         __restrict__ atom_supports,  // [Q_N_ATOMS]
    const unsigned char* __restrict__ pos_state,      // [B × pos_bytes]
    float*               __restrict__ isv,
    int                                b_size,
    int                                pos_bytes
) {
    const int b = blockIdx.x;
    if (b >= b_size) return;

    const int current_step = (int)isv[RL_STEP_COUNTER_ISV_INDEX];
    const int warmup = (int)isv[RL_GATE_WARMUP_STEPS_INDEX];
    if (current_step < warmup) return;

    // Exploration slots: the first `min_explore` batch elements are
    // NEVER gated. This guarantees a minimum fraction of trading
    // actions for Q to learn from, breaking the self-reinforcing
    // Hold trap where gate→Hold→Q learns Hold→conf=0→gate.
    const float max_hold = isv[RL_CONF_GATE_MAX_HOLD_FRAC_INDEX];
    const int min_explore = (int)((float)b_size * (1.0f - max_hold));
    if (b < min_explore) return;

    const int action = actions[b];
    if (action == ACTION_HOLD) return;

    // Gate is opening-only: a non-flat position has already cleared
    // the confidence check, so overriding it to Hold would suppress
    // legitimate adds/exits. pos_state[0..4] = position_lots:i32.
    const int position_lots =
        *reinterpret_cast<const int*>(pos_state + b * pos_bytes);
    if (position_lots != 0) return;

    const float threshold  = isv[RL_CONF_GATE_THRESHOLD_INDEX];
    const float lambda     = isv[RL_CONF_GATE_LAMBDA_INDEX];
    const float sigma_norm = isv[RL_CONF_GATE_SIGMA_NORM_INDEX];

    // q_logits layout: [B, N_ACTIONS, Q_N_ATOMS]
    const float* logits = q_logits + b * N_ACTIONS * Q_N_ATOMS + action * Q_N_ATOMS;

    // Numerically-stable softmax.
    float max_l = logits[0];
    for (int i = 1; i < Q_N_ATOMS; ++i) {
        if (logits[i] > max_l) max_l = logits[i];
    }
    float sum_exp = 0.0f;
    float probs[Q_N_ATOMS];
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        probs[i] = expf(logits[i] - max_l);
        sum_exp += probs[i];
    }
    const float inv_sum = 1.0f / sum_exp;
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        probs[i] *= inv_sum;
    }

    // μ and σ² from the categorical distribution.
    float mu = 0.0f;
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        mu += probs[i] * atom_supports[i];
    }
    float var = 0.0f;
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        const float diff = atom_supports[i] - mu;
        var += probs[i] * diff * diff;
    }
    const float sigma = sqrtf(var + 1e-8f);

    // Lower confidence bound, shifted relative to atom span so conf=0
    // means "worst possible" and conf=1 means "certain at V_MAX".
    // Without the V_MIN shift, any distribution with σ > μ gives conf=0
    // (permanent block for near-uniform Q at training start).
    const float v_min = atom_supports[0];
    const float v_max = atom_supports[Q_N_ATOMS - 1];
    const float span = v_max - v_min + 1e-8f;
    const float lcb = (mu - v_min) - lambda * sigma;
    const float conf = fmaxf(0.0f, fminf(lcb / span, 1.0f));

    // Phase 7b F5 (2026-06-05) Option A — suppress Hold override when
    // F5 is engaged and we are flat. F5 has masked Hold (and other
    // invalid actions) in pi_logits, so the sampled action is a
    // deliberate F5-allowed opening. The early-return at line above
    // (`position_lots != 0 → return`) already restricts us to flat
    // here, so checking only the F5 master gate is sufficient.
    const bool f5_engaged = (isv[RL_F5_STATE_MASK_ENABLED_INDEX] > 0.5f);
    if (conf < threshold && !f5_engaged) {
        actions[b] = ACTION_HOLD;
        isv[RL_CONF_GATE_FIRED_COUNT_INDEX] += 1.0f;
    }

    // ── Adaptive threshold controller (block 0 only). ───────────────
    // Count Hold actions in the batch (post-gate), compute hold_frac,
    // EMA it, and adjust threshold to maintain the target Hold fraction.
    // Runs once per step via the single-block launch.
    if (b == 0) {
        // Count post-gate Hold actions across the batch.
        int hold_count = 0;
        for (int i = 0; i < b_size; i++) {
            if (actions[i] == ACTION_HOLD) hold_count++;
        }
        const float hold_frac = (float)hold_count / (float)b_size;

        // EMA of Hold fraction.
        const float prev_ema = isv[RL_HOLD_FRAC_EMA_INDEX];
        const float new_ema = (prev_ema == 0.0f)
            ? hold_frac
            : (1.0f - HOLD_EMA_ALPHA) * prev_ema + HOLD_EMA_ALPHA * hold_frac;
        isv[RL_HOLD_FRAC_EMA_INDEX] = new_ema;

        // Schulman-style bounded adjustment: if Hold fraction is below
        // target, raise threshold (gate more aggressively). If above
        // target, lower threshold (allow more trading).
        // Symmetric rates — the old 100× asymmetry caused the threshold
        // to ratchet up and never recover.
        const float hold_target = isv[RL_HOLD_TARGET_FRAC_INDEX];
        float new_threshold = threshold;
        if (new_ema < hold_target * 0.9f) {
            new_threshold = threshold * THRESHOLD_ADJUST_RATE;
        } else if (new_ema > hold_target * 1.1f) {
            new_threshold = threshold / THRESHOLD_ADJUST_RATE;
        }
        new_threshold = fmaxf(THRESHOLD_MIN, fminf(new_threshold, THRESHOLD_MAX));
        isv[RL_CONF_GATE_THRESHOLD_INDEX] = new_threshold;
    }
}
