// Gather state vectors from pre-uploaded features + live portfolio state.
// Output: [n_windows, state_dim] tensor for model forward pass.
//
// State layout MATCHES the training experience_state_gather kernel:
//   [0 .. market_dim)                      : market features (42)
//   [market_dim .. market_dim+8)           : 8 portfolio features
//   [market_dim+8 .. market_dim+8+16)      : 16 multi-timeframe features
//   [market_dim+24 .. state_dim)           : zero-pad / OFI / alignment
//
// Portfolio state layout per window [8 floats]:
//   [0] value        - current portfolio value
//   [1] position     - current position size (-1.0 to +1.0)
//   [2] cash         - cash balance
//   [3] entry_price  - entry price of current position (0 if flat)
//   [4] max_equity   - peak equity for drawdown tracking
//   [5] step_pnl     - PnL this step (for reward)
//   [6] cum_return   - cumulative log return
//   [7] step_count   - number of completed steps
//
// 8 portfolio features (matches training experience_state_gather):
//   +0: position                       — raw position [-1, +1]
//   +1: unrealized_pnl / equity        — how is THIS trade doing?
//   +2: drawdown                       — (peak - equity) / peak [0, 1]
//   +3: hold_time / 100.0              — normalized holding duration (step_count as proxy)
//   +4: realized_pnl / equity          — cumulative return signal
//   +5: distance_to_floor              — (equity - floor) / equity [0, 1]
//   +6: trade_return                   — unrealized PnL / equity
//   +7: cash / equity                  — available capital ratio
//
// 16 multi-timeframe features (4 lookbacks × 4 features):
//   Lookback windows: 5, 15, 60, 240 bars
//   Per window: return, volatility, volume_trend, momentum

extern "C" __global__ void gather_states(
    const float* __restrict__ features,      // [n_windows * max_len * feat_dim]
    const float* __restrict__ portfolio,     // [n_windows * 8]
    float*       states_out,                 // [n_windows * state_dim]
    int          n_windows,
    int          max_len,
    int          feat_dim,
    int          state_dim,
    int          current_step,
    float        initial_capital,
    float        spread_cost,
    int          ofi_dim                     // 0 = no OFI, 8 = standard OFI features
) {
    // Shared memory tile for coalesced portfolio reads — 8 floats per thread
    __shared__ float shmem_portfolio[256 * 8];

    int w = blockIdx.x * blockDim.x + threadIdx.x;
    int local_tid = threadIdx.x;

    // Cooperative tile load: all threads load their portfolio slice in one
    // coalesced warp transaction, then read from shared memory below.
    int ps = w * 8;
    if (w < n_windows) {
        for (int i = 0; i < 8; i++) {
            shmem_portfolio[local_tid * 8 + i] = __ldg(&portfolio[ps + i]);
        }
    }
    __syncthreads();

    if (w >= n_windows) return;

    int feat_base = (w * max_len + current_step) * feat_dim;
    int out_base  = w * state_dim;

    // Read portfolio state from shared memory tile
    float value      = shmem_portfolio[local_tid * 8 + 0];
    float position   = shmem_portfolio[local_tid * 8 + 1];
    float cash       = shmem_portfolio[local_tid * 8 + 2];
    float entry_price = shmem_portfolio[local_tid * 8 + 3];
    float max_equity = shmem_portfolio[local_tid * 8 + 4];
    float step_count = shmem_portfolio[local_tid * 8 + 7];

    // Determine market_dim: if OFI is present, market features are feat_dim - ofi_dim
    int market_dim = (ofi_dim > 0) ? (feat_dim - ofi_dim) : feat_dim;

    // ── 1. Market features [0 .. market_dim) ── //
    int i = 0;
    for (; i + 3 < market_dim; i += 4) {
        states_out[out_base + i]     = __ldg(&features[feat_base + i]);
        states_out[out_base + i + 1] = __ldg(&features[feat_base + i + 1]);
        states_out[out_base + i + 2] = __ldg(&features[feat_base + i + 2]);
        states_out[out_base + i + 3] = __ldg(&features[feat_base + i + 3]);
    }
    for (; i < market_dim; i++) {
        states_out[out_base + i] = __ldg(&features[feat_base + i]);
    }

    // ── 2. Portfolio features [market_dim .. market_dim+8) ── //
    // Match training experience_state_gather exactly.
    float equity = (value > 1.0f) ? value : 1.0f;

    // Drawdown from peak
    float drawdown = (max_equity > 1.0f)
                   ? (max_equity - value) / max_equity
                   : 0.0f;
    if (drawdown < 0.0f) drawdown = 0.0f;

    // Unrealized PnL: position * (current_close - entry_price)
    // Use feature[0] as close price proxy (same as training kernel)
    float close_now = __ldg(&features[feat_base]);
    float unrealized_pnl = (entry_price > 0.0f && position != 0.0f)
                         ? position * (close_now - entry_price)
                         : 0.0f;

    // Realized PnL: derive from equity state
    // In backtest, realized PnL = value - initial_capital (approximate)
    float realized_pnl = value - initial_capital;

    // Trade return (unrealized / equity)
    float trade_return = unrealized_pnl / equity;

    // Capital floor distance (75% of peak)
    float floor = max_equity * 0.75f;
    float floor_dist = (equity > floor && equity > 1.0f)
                     ? (equity - floor) / equity
                     : 0.0f;

    int pf_base = market_dim;
    if (pf_base + 7 < state_dim) {
        states_out[out_base + pf_base + 0] = position;                    // raw position
        states_out[out_base + pf_base + 1] = unrealized_pnl / equity;     // trade P&L signal
        states_out[out_base + pf_base + 2] = drawdown;                    // risk: how deep are we?
        states_out[out_base + pf_base + 3] = step_count / 100.0f;         // how long in trade?
        states_out[out_base + pf_base + 4] = realized_pnl / equity;       // session P&L
        states_out[out_base + pf_base + 5] = floor_dist;                  // distance to game over
        states_out[out_base + pf_base + 6] = trade_return;                // this trade's return
        states_out[out_base + pf_base + 7] = cash / equity;               // available capital
    }

    // ── 3. Multi-timeframe features [market_dim+8 .. market_dim+8+16) ── //
    // 4 lookback windows × 4 features = 16 GPU-native features.
    // Matches training experience_state_gather exactly.
    const int lookbacks[4] = {5, 15, 60, 240};
    int mtf_base = market_dim + 8;

    // Global bar index for this window at current step
    int window_bar_base = w * max_len;

    for (int lb = 0; lb < 4; lb++) {
        int N_lb = lookbacks[lb];
        int past_step = current_step - N_lb;
        int slot = mtf_base + lb * 4;

        if (past_step >= 0 && slot + 3 < state_dim) {
            int now_feat_off  = (window_bar_base + current_step) * feat_dim;
            int past_feat_off = (window_bar_base + past_step) * feat_dim;
            float close_cur  = __ldg(&features[now_feat_off]);
            float close_past = __ldg(&features[past_feat_off]);

            // Return over N bars
            float ret = (close_past > 0.0f) ? (close_cur - close_past) / close_past : 0.0f;
            float scaled_ret = ret * 100.0f;
            states_out[out_base + slot + 0] = fmaxf(-10.0f, fminf(10.0f, scaled_ret));

            // Volatility: scan high/low over window (from close changes)
            float max_val = close_cur;
            float min_val = close_cur;
            float vol_sum = 0.0f;
            int vol_count = 0;
            for (int j = past_step; j <= current_step; j++) {
                int j_off = (window_bar_base + j) * feat_dim;
                float v = __ldg(&features[j_off]);
                if (v > max_val) max_val = v;
                if (v < min_val) min_val = v;
                // Volume proxy: feature index 4
                if (feat_dim > 4) {
                    vol_sum += __ldg(&features[j_off + 4]);
                    vol_count++;
                }
            }
            float range = (close_cur > 0.0f) ? (max_val - min_val) / close_cur : 0.0f;
            float scaled_range = range * 100.0f;
            states_out[out_base + slot + 1] = fmaxf(0.0f, fminf(10.0f, scaled_range));

            // Volume trend: current vs average
            float avg_vol = (vol_count > 0) ? vol_sum / (float)vol_count : 1.0f;
            float cur_vol = (feat_dim > 4) ? __ldg(&features[now_feat_off + 4]) : 1.0f;
            float vol_ratio = (avg_vol > 0.0f) ? cur_vol / avg_vol : 1.0f;
            states_out[out_base + slot + 2] = fmaxf(0.0f, fminf(5.0f, vol_ratio));

            // Momentum: position within range [0=bottom, 1=top]
            float range_size = max_val - min_val;
            states_out[out_base + slot + 3] = (range_size > 0.0f)
                          ? (close_cur - min_val) / range_size
                          : 0.5f;
        } else {
            // Not enough history — zero pad
            if (slot + 3 < state_dim) {
                states_out[out_base + slot + 0] = 0.0f;
                states_out[out_base + slot + 1] = 0.0f;
                states_out[out_base + slot + 2] = 0.0f;
                states_out[out_base + slot + 3] = 0.0f;
            }
        }
    }

    // ── 4. OFI features (if present) ── //
    if (ofi_dim > 0) {
        int ofi_out_base = market_dim + 8 + 16;
        // OFI features are stored after market features in the feature vector
        for (i = 0; i + 3 < ofi_dim; i += 4) {
            states_out[out_base + ofi_out_base + i]     = __ldg(&features[feat_base + market_dim + i]);
            states_out[out_base + ofi_out_base + i + 1] = __ldg(&features[feat_base + market_dim + i + 1]);
            states_out[out_base + ofi_out_base + i + 2] = __ldg(&features[feat_base + market_dim + i + 2]);
            states_out[out_base + ofi_out_base + i + 3] = __ldg(&features[feat_base + market_dim + i + 3]);
        }
        for (; i < ofi_dim; i++) {
            states_out[out_base + ofi_out_base + i] = __ldg(&features[feat_base + market_dim + i]);
        }
    }

    // ── 5. Zero-pad remaining [filled .. state_dim) ── //
    int filled = market_dim + 8 + 16 + ofi_dim;
    for (i = filled; i + 3 < state_dim; i += 4) {
        states_out[out_base + i]     = 0.0f;
        states_out[out_base + i + 1] = 0.0f;
        states_out[out_base + i + 2] = 0.0f;
        states_out[out_base + i + 3] = 0.0f;
    }
    for (; i < state_dim; i++) {
        states_out[out_base + i] = 0.0f;
    }

    // Suppress unused variable warning
    (void)spread_cost;
}
