// crates/ml/src/cuda_pipeline/dd_state_kernel.cu
//
// SP15 Phase 1.3 — per-step drawdown state tracking.
// SP15 Phase 1.3.b-followup (2026-05-07) — per-env redesign (Path A → Path B).
//
// READ-ONLY on the position-state buffer's equity fields. Reads existing
// PS_PEAK_EQUITY (slot 7) and PS_PREV_EQUITY (slot 9) — both maintained
// by `experience_env_step` (`experience_kernels.cu:3473-3475`, which
// writes ps[PS_PEAK_EQUITY]=peak_equity and
// ps[PS_PREV_EQUITY]=new_portfolio_value at the end of each env step).
// This kernel does NOT recompute equity and does NOT write back to
// PS_PEAK_EQUITY / PS_PREV_EQUITY — the env-step kernel is the canonical
// equity writer.
//
// Phase 1.3.b-followup migration: the original Phase 1.3.b atomic fix
// (Path A) chose **env 0 as the canonical DD observable** and wrote
// directly to 6 ISV scalar slots [401..407). That choice was wrong for
// production where each env has its own portfolio + DD trajectory: when
// env-3 was in 30% DD and env-0 was at ATH, the per-(i,t) reward shaping
// looked up env-0's `dd_pct=0` and skipped recovery shaping for env-3's
// transitions.
//
// Path B (this kernel): per-env grid `[n_envs, 1, 1]` × `[1, 1, 1]`. Each
// thread computes its own env's DD trajectory and writes 6 scalars to
// the per-env tile buffer `dd_state_per_env[env_id * 6 + k]` for
// `k ∈ {DD_CURRENT, DD_MAX, DD_RECOVERY_BARS, DD_PERSISTENCE, CALMAR,
// DD_PCT}`. Each thread's per-bar walk is sequential (the running max,
// recovery counter, and persistence counter are stateful per-env), but
// envs are independent so the grid parallelises across them.
//
// The 6 ISV scalar slots [401..407) are now populated by a separate
// reduction kernel (`dd_state_reduce_kernel`) that mean-aggregates
// across envs for HEALTH_DIAG diagnostic reporting + max-aggregates
// DD_PERSISTENCE into the new `DD_PERSISTENCE_MAX_INDEX = 443` slot for
// the plasticity-injection trigger (one set of advantage weights →
// global firing semantics → max-aggregate is the only correct rule
// when ANY env exceeding the threshold should arm the gate).
//
// Per-env tile layout (offsets within each env's [6] sub-array):
//   tile[env*6 + 0] = DD_CURRENT          (current_dd = max(0, (peak-equity)/peak))
//   tile[env*6 + 1] = DD_MAX              (running max of DD_CURRENT in this fold)
//   tile[env*6 + 2] = DD_RECOVERY_BARS    (bars since last new HWM)
//   tile[env*6 + 3] = DD_PERSISTENCE      (twin counter to recovery)
//   tile[env*6 + 4] = CALMAR              (max(dd_max, 1e-4) — calmar denominator floor)
//   tile[env*6 + 5] = DD_PCT              (clip(current_dd / max(dd_budget, 1e-4), 0, 1))
//
// Layout fingerprint contract: PS_PEAK_EQUITY and PS_PREV_EQUITY are
// `#define`d in state_layout.cuh — use the macros so the kernel tracks
// any future portfolio-state reshuffle automatically.
//
// Allowed-write rule: the kernel writes ONLY to the per-env tile (which
// has a FoldReset arm `sp15_dd_state_per_env`) and never touches the
// position state buffer or the ISV scalar slots. The ISV scalar slots
// are written exclusively by `dd_state_reduce_kernel`. No
// pos_state ownership conflict; no double-writer for the ISV slots.

#include "state_layout.cuh"

// PS_STRIDE is defined in state_layout.cuh (== 43 floats per env). Env i
// offset is therefore `i * PS_STRIDE`. The Rust constant
// `PORTFOLIO_STRIDE` mirrors this same value in
// `gpu_experience_collector.rs` (cross-checked at constructor init);
// the layout-fingerprint check guards drift.

extern "C" __global__ void dd_state_kernel(
    float                    dd_budget,
    int                      n_envs,
    float* __restrict__      dd_state_per_env,  /* [n_envs * 6] tile output */
    const float* __restrict__ pos_state         /* [n_envs * PS_STRIDE]    */
) {
    /* Per-env grid: one thread per env, one block per env (mirrors the
     * single-thread / single-block shape of the original Path A kernel
     * but lifted into a per-env loop via blockIdx.x). The state-machine
     * inside the thread is sequential per env — running max, recovery
     * counter, persistence counter — so we cannot parallelise across
     * the [n_envs * 6] output dimension. Envs ARE independent of each
     * other, so the gridDim.x = n_envs parallelism is straightforward. */
    if (threadIdx.x != 0) return;
    const int env_id = blockIdx.x;
    if (env_id >= n_envs) return;

    /* Read prev_equity (PS_PREV_EQUITY = slot 9) and peak (PS_PEAK_EQUITY
     * = slot 7) from this env's row of the portfolio state buffer. Both
     * are written by experience_env_step earlier this step — by launching
     * this kernel immediately after env_step on the same stream we get
     * the live values for every env. */
    const float prev_equity = pos_state[env_id * PS_STRIDE + PS_PREV_EQUITY];
    const float peak        = pos_state[env_id * PS_STRIDE + PS_PEAK_EQUITY];

    /* current_dd = max(0, (peak − equity) / peak), guarded against peak ≤ 0. */
    const float current_dd = (peak > 0.0f)
        ? fmaxf(0.0f, (peak - prev_equity) / peak)
        : 0.0f;

    const int base = env_id * 6;

    /* Recovery + persistence counters: zero on a fresh new-high-water-mark
     * step, increment otherwise. We detect a new HWM by checking whether
     * current_dd is exactly zero (peak == prev_equity in the no-DD case).
     * env_step already updated peak to max(prev_peak, prev_equity) before
     * we read; if equity ≥ peak then current_dd == 0 ↔ new HWM (or flat
     * pre-trade equity), which is the canonical "no drawdown right now"
     * condition the recovery/persistence counters must zero on. */
    float recovery_bars;
    float persistence_bars;
    if (current_dd <= 0.0f) {
        recovery_bars   = 0.0f;
        persistence_bars = 0.0f;
    } else {
        recovery_bars   = dd_state_per_env[base + 2] + 1.0f;
        persistence_bars = dd_state_per_env[base + 3] + 1.0f;
    }

    /* Running max of current_dd. */
    const float prior_max = dd_state_per_env[base + 1];
    const float new_max   = (current_dd > prior_max) ? current_dd : prior_max;

    /* dd_pct = clip(current_dd / max(dd_budget, 1e-4), 0, 1). */
    const float dd_pct = fminf(1.0f, current_dd / fmaxf(dd_budget, 1e-4f));

    /* CALMAR slot stores the floored max_dd denominator; host (post-
     * reduction) composes calmar = mean_pnl / value-here. Floor
     * eliminates the saturation-at-100 artifact that the legacy host-
     * side composer exhibited when max_dd ≈ 0. */
    const float calmar_denom = fmaxf(new_max, 1e-4f);

    dd_state_per_env[base + 0] = current_dd;
    dd_state_per_env[base + 1] = new_max;
    dd_state_per_env[base + 2] = recovery_bars;
    dd_state_per_env[base + 3] = persistence_bars;
    dd_state_per_env[base + 4] = calmar_denom;
    dd_state_per_env[base + 5] = dd_pct;

    __threadfence_system();
}
