// rl_unit_state_update.cu — per-unit trade state machine with pyramid
// expansion.
//
// Runs AFTER each step's fill+extract sequence. Detects per-batch
// position transitions vs the prior step and updates per-unit
// buffers accordingly:
//
//   prev_pos == 0 AND curr_pos != 0  → OPEN: write unit slot 0
//   prev_pos != 0 AND curr_pos == 0  → CLOSE: deactivate all units
//   sign(prev) != sign(curr) AND both != 0 → REVERSE: close + open
//   sign(prev) == sign(curr) AND both != 0 → CONTINUE:
//       |curr| > |prev| → pyramid ADD: allocate next free slot
//       |curr| < |prev| → partial FLAT: deactivate oldest active slot
//       |curr| == |prev| → no structural change
//
// Per `pearl_first_observation_bootstrap`: prev_pos starts at sentinel
// 0; first non-zero curr_pos triggers OPEN via the standard branch.
//
// Per `feedback_no_atomicadd`: per-batch independent writes only.
// Per `feedback_cpu_is_read_only`: pure device-side update.
// Per `pearl_trade_level_vol_for_stop_distance`: trail initialized
// from `sqrt(realized_return_var_ema)` × ISV-driven k_init.

#include <stdint.h>

#define MAX_UNITS                           4
#define RL_TRAIL_K_INIT_INDEX               496
#define RL_MEAN_ABS_PNL_EMA_INDEX           423
#define RL_REWARD_SCALE_INDEX               406
#define RL_PYRAMID_ADD_COUNT_INDEX          507
#define RL_STEP_COUNTER_ISV_INDEX           548

// PosFlat layout (matches crates/ml-backtesting/src/lob/mod.rs::PosFlat):
//   offset 0: position_lots: i32
//   offset 4: vwap_entry:    f32
//   offset 8: realized_pnl:  f32
//   offset 12: peak_equity:  f32

extern "C" __global__ void rl_unit_state_update(
    const unsigned char* __restrict__ pos_state,    // [B * pos_bytes]
    float*               __restrict__ isv,
    int*                 __restrict__ prev_pos_lots,        // [B] tracking
    float*               __restrict__ unit_entry_price,     // [B * MAX_UNITS]
    int*                 __restrict__ unit_entry_step,      // [B * MAX_UNITS]
    int*                 __restrict__ unit_lots,            // [B * MAX_UNITS]
    float*               __restrict__ unit_initial_r,       // [B * MAX_UNITS]
    float*               __restrict__ unit_trail_distance,  // [B * MAX_UNITS]
    unsigned char*       __restrict__ unit_active,          // [B * MAX_UNITS]
    int*                 __restrict__ pyramid_units_count,  // [B]
    int                                b_size,
    int                                pos_bytes
) {
    const int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= b_size) return;

    const int current_step = (int)isv[RL_STEP_COUNTER_ISV_INDEX];

    const int curr_pos = *(const int*)(pos_state + b * pos_bytes + 0);
    const float vwap   = *(const float*)(pos_state + b * pos_bytes + 4);
    const int prev_pos = prev_pos_lots[b];

    const float k_init        = isv[RL_TRAIL_K_INIT_INDEX];
    const float mean_abs_pnl  = isv[RL_MEAN_ABS_PNL_EMA_INDEX];
    const float reward_scale  = isv[RL_REWARD_SCALE_INDEX];
    // mean_abs_pnl is in SCALED reward space. Convert to price space
    // by dividing out the reward_scale: raw_magnitude ≈ scaled / scale.
    // Trail distance must be in price units (compared against mid-entry).
    const float raw_magnitude = mean_abs_pnl / fmaxf(reward_scale, 1e-9f);
    const float bootstrap_d   = k_init * (raw_magnitude + 1e-8f);

    const int base = b * MAX_UNITS;

    const bool was_flat = (prev_pos == 0);
    const bool is_flat  = (curr_pos == 0);
    const bool reversed = (!was_flat && !is_flat && ((prev_pos > 0) != (curr_pos > 0)));

    int add_count = 0;

    if (is_flat && !was_flat) {
        // CLOSE: deactivate all units, zero lots.
        #pragma unroll
        for (int u = 0; u < MAX_UNITS; ++u) {
            unit_active[base + u] = 0;
            unit_lots[base + u]   = 0;
        }
        pyramid_units_count[b] = 0;
    } else if (!is_flat && (was_flat || reversed)) {
        // OPEN or REVERSE: clear all + populate slot 0.
        #pragma unroll
        for (int u = 0; u < MAX_UNITS; ++u) {
            unit_active[base + u] = 0;
            unit_lots[base + u]   = 0;
        }
        unit_entry_price[base + 0]    = vwap;
        unit_entry_step[base + 0]     = current_step;
        unit_lots[base + 0]           = curr_pos;
        unit_initial_r[base + 0]      = bootstrap_d;
        unit_trail_distance[base + 0] = bootstrap_d;
        unit_active[base + 0]         = 1;
        pyramid_units_count[b] = 1;
        add_count = 1;
    } else if (!is_flat && !was_flat && !reversed) {
        // CONTINUE: same direction.
        const int abs_curr = (curr_pos > 0) ? curr_pos : -curr_pos;
        const int abs_prev = (prev_pos > 0) ? prev_pos : -prev_pos;

        if (abs_curr > abs_prev) {
            // Pyramid ADD: lots grew. Allocate next free unit slot.
            int count = pyramid_units_count[b];
            if (count < MAX_UNITS) {
                const int slot = count;
                unit_entry_price[base + slot]    = vwap;
                unit_entry_step[base + slot]     = current_step;
                const int delta_lots = abs_curr - abs_prev;
                unit_lots[base + slot] = (curr_pos > 0) ? delta_lots : -delta_lots;
                unit_initial_r[base + slot]      = bootstrap_d;
                unit_trail_distance[base + slot] = bootstrap_d;
                unit_active[base + slot]         = 1;
                pyramid_units_count[b] = count + 1;
                add_count = 1;
            }
        } else if (abs_curr < abs_prev) {
            // Partial FLAT: lots shrank. Deactivate the oldest active
            // unit (lowest index among active slots).
            for (int u = 0; u < MAX_UNITS; ++u) {
                if (unit_active[base + u]) {
                    unit_active[base + u] = 0;
                    unit_lots[base + u]   = 0;
                    int count = pyramid_units_count[b];
                    if (count > 0) {
                        pyramid_units_count[b] = count - 1;
                    }
                    break;
                }
            }
        }
        // abs_curr == abs_prev: no structural change.
    }
    // (is_flat && was_flat): no-op, position remains flat.

    // Always advance prev_pos tracker.
    prev_pos_lots[b] = curr_pos;

    // Write add_count to ISV diag slot. b_size=1 in practice; for
    // multi-batch, only thread 0 writes (avoids atomicAdd).
    if (b == 0) {
        isv[RL_PYRAMID_ADD_COUNT_INDEX] = (float)add_count;
    }
}
