// Per-window metrics reduction kernel — native F32.
#include "state_layout.cuh"
// One block per window. Threads cooperate to reduce step_returns.
//
// Output per window [14 f32 values written to metrics_out]:
//   [0] sharpe_ratio (annualized)
//   [1] total_pnl (multiplicative compounding)
//   [2] max_drawdown (exact sequential scan)
//   [3] sortino_ratio
//   [4] win_rate
//   [5] total_trades
//   [6] var_95
//   [7] cvar_95
//   [8] calmar_ratio
//   [9] omega_ratio
//  [10] buy_count
//  [11] sell_count
//  [12] hold_count
//  [13] unique_action_count

extern "C" __global__ void compute_backtest_metrics(
    const float* __restrict__ step_returns,   // [n_windows * max_len]
    const float* __restrict__ portfolio_state, // [n_windows * 8]
    const int*   __restrict__ window_lens,    // [n_windows]
    const int*   __restrict__ actions_history, // [n_windows * max_len]
    float* metrics_out,                       // [n_windows * 14]
    int          n_windows,
    int          max_len,
    float        annualization_factor,        // f32 scalar
    int          num_actions,
    int          order_actions,
    int          urgency_actions
) {
    int w = blockIdx.x;
    if (w >= n_windows) return;

    int wlen = window_lens[w];
    int tid = threadIdx.x;
    int stride = blockDim.x;
    int base = w * max_len;

    // Shared memory layout (6 reduction arrays + sort scratch):
    extern __shared__ float shmem[];
    float* s_sum     = shmem;
    float* s_sq_sum  = shmem + stride;
    float* s_down_sq = shmem + 2*stride;
    float* s_buys    = shmem + 3*stride;
    float* s_sells   = shmem + 4*stride;
    float* s_holds   = shmem + 5*stride;
    float* s_sorted  = shmem + 6*stride;     // [4096] for sort + boundary data

    // Per-thread local accumulators — native F32
    float local_sum = 0.0f;
    float local_sq  = 0.0f;
    float local_down = 0.0f;
    float local_cum = 1.0f;
    int local_buys = 0, local_sells = 0, local_holds = 0;
    int local_action_mask = 0;

    // Boundary-aware trade tracking.
    // Sentinel = -2 (was -1). After the trade-key change to signed_dir
    // {-1=Short, 0=Hold/Flat, +1=Long}, -1 is a legitimate direction, so
    // "no data in this chunk" must use a value outside the signed_dir range.
    int bnd_first_action = -2;
    int bnd_last_action = -2;
    float bnd_prefix_return = 0.0f;
    float bnd_suffix_return = 0.0f;
    int bnd_complete_trades = 0;
    int bnd_complete_wins = 0;
    int bnd_num_changes = 0;
    int bnd_cur_action = -1;
    float bnd_cur_return = 0.0f;

    // Consecutive chunks for correct drawdown tracking
    int chunk_size = (wlen + stride - 1) / stride;
    int chunk_start = tid * chunk_size;
    int chunk_end = chunk_start + chunk_size;
    if (chunk_end > wlen) chunk_end = wlen;

    /* Trade boundary key: DIRECTION only, not (direction × magnitude).
     *
     * Previously the boundary key was `exp_idx = act / (order×urgency)` which
     * decodes to (dir × mag). That counted Long-Half → Long-Full as a NEW
     * trade even though the position sign didn't change. Each "new" trade
     * absorbed the bar-of-change tx_cost as its first step_return, biasing
     * win-rate downward asymmetrically vs. training's win-rate (which counts
     * trades only on position-sign cycles, see experience_kernels.cu:1592).
     *
     * 4-branch action encoding: action = dir*27 + mag*9 + ord*3 + urg.
     * Direction values: Short(0), Hold(1), Long(2), Flat(3).
     * dir_idx = action / (mag×order×urgency) = action / 27 = exp_idx / 3.
     *
     * For trade-boundary purposes we collapse Hold(1) and Flat(3) — both
     * mean "no exposure" — so a Long → Hold → Long is also one trade cycle.
     * This matches training's "position sign change" semantics. */
    for (int i = chunk_start; i < chunk_end; i++) {
        float r = step_returns[base + i];
        local_sum += r;
        local_sq  += r * r;
        if (r < 0.0f) local_down += r * r;

        local_cum *= (1.0f + r);

        int act = actions_history[base + i];
        /* Skip unwritten slots (act < 0): after `done_flags[w]=1` (capital
         * floor breach), backtest_env_step early-returns without writing
         * actions_history for the remaining slots in [done_step, max_len).
         * `reset_evaluation_state` initialises those to -1 sentinel. Without
         * this guard, C integer division (`-1 / 9 = 0`) decoded unwritten
         * slots as Short Quarter (action 0), inflating sell_count and
         * collapsing active_frac/dir_entropy to a measurement artifact
         * rather than real model behaviour. step_returns at these slots are
         * still zero (memset_zeros'd, never written), so summing them with
         * `r=0` is a no-op and remains correct. */
        if (act < 0) continue;
        int exp_idx = act / (order_actions * urgency_actions);
        int raw_dir = exp_idx / 3;  /* 0=Short, 1=Hold, 2=Long, 3=Flat */
        /* Collapse Hold(1) and Flat(3) → no-exposure (-1).
         * Keep Short(0) and Long(2) as distinct sign states. */
        int signed_dir = (raw_dir == DIR_SHORT) ? -1 :
                         (raw_dir == DIR_LONG)  ?  1 :
                                                   0;  /* Hold/Flat = no exposure */

        /* Action-distribution counters: now reflect ACTUAL directions, not
         * legacy 9-action exposure-band thresholds. Counts: long, short,
         * flat-or-hold (treated as "neutral"). */
        if (signed_dir > 0)      local_buys++;
        else if (signed_dir < 0) local_sells++;
        else                     local_holds++;

        if (exp_idx >= 0 && exp_idx < num_actions)
            local_action_mask |= (1 << exp_idx);

        if (i == chunk_start) {
            bnd_first_action = signed_dir;
            bnd_cur_action = signed_dir;
        }

        if (signed_dir != bnd_cur_action) {
            bnd_num_changes++;
            if (bnd_num_changes == 1) {
                bnd_prefix_return = bnd_cur_return;
            } else {
                bnd_complete_trades++;
                if (bnd_cur_return > 0.0f) bnd_complete_wins++;
            }
            bnd_cur_return = 0.0f;
            bnd_cur_action = signed_dir;
        }

        bnd_cur_return += r;
        bnd_last_action = signed_dir;
    }

    bnd_suffix_return = bnd_cur_return;
    if (bnd_num_changes == 0) {
        bnd_prefix_return = bnd_cur_return;
    }

    // Store to shared memory for parallel reduction
    s_sum[tid] = local_sum;
    s_sq_sum[tid] = local_sq;
    s_down_sq[tid] = local_down;
    s_buys[tid] = (float)local_buys;
    s_sells[tid] = (float)local_sells;
    s_holds[tid] = (float)local_holds;
    __syncthreads();

    // Block-level parallel reduction — native f32 += f32
    for (int s = stride / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_sum[tid] += s_sum[tid + s];
            s_sq_sum[tid] += s_sq_sum[tid + s];
            s_down_sq[tid] += s_down_sq[tid + s];
            s_buys[tid] += s_buys[tid + s];
            s_sells[tid] += s_sells[tid + s];
            s_holds[tid] += s_holds[tid + s];
        }
        __syncthreads();
    }

    // Equity product reduction (multiplicative)
    s_sorted[tid] = local_cum;
    __syncthreads();
    for (int s = stride / 2; s > 0; s >>= 1) {
        if (tid < s) {
            s_sorted[tid] *= s_sorted[tid + s];
        }
        __syncthreads();
    }

    // Export boundary trade data to shared memory.
    {
        float* s_bnd = &s_sorted[stride];
        s_bnd[0 * stride + tid] = (float)bnd_first_action;
        s_bnd[1 * stride + tid] = (float)bnd_last_action;
        s_bnd[2 * stride + tid] = bnd_prefix_return;
        s_bnd[3 * stride + tid] = bnd_suffix_return;
        s_bnd[4 * stride + tid] = (float)bnd_complete_trades;
        s_bnd[5 * stride + tid] = (float)bnd_complete_wins;
        s_bnd[6 * stride + tid] = (float)bnd_num_changes;
    }
    __syncthreads();

    // Thread 0: final metrics from fully reduced values
    if (tid == 0) {
        float n = (float)wlen;
        float mean = s_sum[0] / n;
        float var = s_sq_sum[0] / n - mean * mean;
        float std_val = sqrtf(fmaxf(var, 1e-10f));
        float down_std = sqrtf(fmaxf(s_down_sq[0] / n, 1e-10f));

        // Exact max drawdown via sequential scan
        float exact_equity = 1.0f;
        float exact_peak = 1.0f;
        float exact_max_dd = 0.0f;
        for (int i = 0; i < wlen; i++) {
            float r = step_returns[base + i];
            exact_equity *= (1.0f + r);
            exact_peak = fmaxf(exact_peak, exact_equity);
            float dd = (exact_peak - exact_equity) / fmaxf(exact_peak, 1e-8f);
            exact_max_dd = fmaxf(exact_max_dd, dd);
        }

        // Boundary stitching: exact trade count + win rate
        float* s_bnd = &s_sorted[stride];
        int total_trades = 0;
        int total_wins = 0;

        for (int t = 0; t < stride; t++) {
            total_trades += (int)s_bnd[4 * stride + t];
            total_wins   += (int)s_bnd[5 * stride + t];
        }

        /* Boundary stitching across CUDA blocks.
         * fa/la are now signed_dir values {-1, 0, +1}; sentinel "no data"
         * is -2. Replaced `< 0` checks with `< -1` (only -2 means empty). */
        float open_return = 0.0f;
        int open_action = -2;

        for (int t = 0; t < stride; t++) {
            int fa = (int)s_bnd[0 * stride + t];
            int la = (int)s_bnd[1 * stride + t];
            float pr = s_bnd[2 * stride + t];
            float sr = s_bnd[3 * stride + t];
            int nc = (int)s_bnd[6 * stride + t];

            if (fa < -1) continue;  /* -2 = no data in this block */

            if (open_action < -1) {
                open_action = fa;
                open_return = 0.0f;
            } else if (fa != open_action) {
                total_trades++;
                if (open_return > 0.0f) total_wins++;
                open_action = fa;
                open_return = 0.0f;
            }

            if (nc == 0) {
                open_return += pr;
            } else {
                open_return += pr;
                total_trades++;
                if (open_return > 0.0f) total_wins++;
                open_action = la;
                open_return = sr;
            }
        }

        if (open_action >= -1) {
            total_trades++;
            if (open_return > 0.0f) total_wins++;
        }

        int out_base = w * 14;
        metrics_out[out_base + 0] = (mean / std_val) * annualization_factor;
        metrics_out[out_base + 1] = s_sorted[0] - 1.0f;
        metrics_out[out_base + 2] = exact_max_dd;
        metrics_out[out_base + 3] = (mean / down_std) * annualization_factor;

        float total_trades_f = (float)total_trades;
        metrics_out[out_base + 4] = (total_trades > 0)
            ? (float)total_wins / total_trades_f : 0.0f;
        metrics_out[out_base + 5] = total_trades_f;

        metrics_out[out_base + 10] = s_buys[0];
        metrics_out[out_base + 11] = s_sells[0];
        metrics_out[out_base + 12] = s_holds[0];
    }

    // Unique action count via bitmask OR-reduction
    {
        int* s_action_mask = (int*)s_sorted;
        s_action_mask[tid] = local_action_mask;
        __syncthreads();
        for (int s = stride / 2; s > 0; s >>= 1) {
            if (tid < s) s_action_mask[tid] |= s_action_mask[tid + s];
            __syncthreads();
        }
        if (tid == 0) {
            int out_base = w * 14;
            metrics_out[out_base + 13] = (float)__popc(s_action_mask[0]);
        }
    }

    // VaR, CVaR, Calmar, Omega via bitonic sort
    int sort_len = wlen < 4096 ? wlen : 4096;
    int sample_step = (wlen + sort_len - 1) / sort_len;
    for (int i = tid; i < sort_len; i += stride) {
        int src = i * sample_step;
        s_sorted[i] = (src < wlen) ? step_returns[base + src] : 0.0f;
    }

    int padded_len = 1;
    while (padded_len < sort_len) padded_len <<= 1;

    float f32_inf = 1e30f;
    for (int i = sort_len + tid; i < padded_len; i += stride) {
        s_sorted[i] = f32_inf;
    }
    __syncthreads();

    // Bitonic sort — ascending, native f32 comparisons
    for (int k = 2; k <= padded_len; k <<= 1) {
        for (int j = k >> 1; j > 0; j >>= 1) {
            for (int i = tid; i < padded_len; i += stride) {
                int ixj = i ^ j;
                if (ixj > i) {
                    bool ascending = ((i & k) == 0);
                    if ((ascending  && s_sorted[i] > s_sorted[ixj]) ||
                        (!ascending && s_sorted[i] < s_sorted[ixj])) {
                        float tmp = s_sorted[i];
                        s_sorted[i] = s_sorted[ixj];
                        s_sorted[ixj] = tmp;
                    }
                }
            }
            __syncthreads();
        }
    }

    // Thread 0: extended metrics from sorted array
    if (tid == 0) {
        int out_base = w * 14;

        int var_idx = (int)(0.05f * (float)sort_len);
        if (var_idx < 0) var_idx = 0;
        if (var_idx >= sort_len) var_idx = sort_len - 1;
        float var_95 = s_sorted[var_idx];

        int cvar_count = var_idx + 1;
        float cvar_sum = 0.0f;
        for (int i = 0; i < cvar_count; i++) {
            cvar_sum += s_sorted[i];
        }
        float cvar_95 = cvar_sum / (float)cvar_count;

        float daily_mean = s_sum[0] / (float)wlen;
        float max_dd = metrics_out[out_base + 2];
        float calmar = (max_dd > 0.001f)
            ? (daily_mean * annualization_factor * annualization_factor) / max_dd
            : 0.0f;
        calmar = fmaxf(-100.0f, fminf(100.0f, calmar));

        float gain_sum = 0.0f, loss_sum = 0.0f;
        for (int i = 0; i < sort_len; i++) {
            if (s_sorted[i] > 0.0f) gain_sum += s_sorted[i];
            else                     loss_sum -= s_sorted[i];
        }
        float omega = (loss_sum > 0.001f)
            ? gain_sum / loss_sum : 0.0f;
        omega = fmaxf(0.0f, fminf(100.0f, omega));

        metrics_out[out_base + 6] = var_95;
        metrics_out[out_base + 7] = cvar_95;
        metrics_out[out_base + 8] = calmar;
        metrics_out[out_base + 9] = omega;
    }
}
