/* state_kl_moment_match — per-dim Gaussian moment-match KL between
 * train-state and val-state batches. Summed over a configurable
 * dimension range (typically the OFI block — bits that should be
 * similar across train/val for the same instrument).
 *
 * Plan 3 Task 7 C.3.
 *
 * Formula per dim d:
 *   μ_t, σ²_t = mean+var of train_state[*, d]
 *   μ_v, σ²_v = mean+var of val_state[*, d]
 *   kl_d      = 0.5 × (log(σ²_v/σ²_t) + (σ²_t + (μ_t − μ_v)²) / σ²_v − 1)
 * Total KL = sum over monitored dims.
 *
 * Adaptive amp formula (no tuned constants):
 *   prev_kl_ema = isv[kl_ema_idx]                         (trailing baseline)
 *   sharpe      = clamp(isv[sharpe_ema_idx], -2, 2)
 *   alpha       = alpha_base × (1 + 0.5 × |sharpe|)       (Task 3 convention)
 *   new_kl_ema  = (1 − alpha) × prev_kl_ema + alpha × current_kl
 *
 *   ratio       = new_kl_ema / max(1e-6, prev_kl_ema)
 *   amp_target  = 1 + clamp(ratio − 1, 0, 1)              (∈ [1, 2])
 *   prev_amp    = isv[amp_idx]
 *   new_amp     = (1 − alpha) × prev_amp + alpha × amp_target
 *
 * Consumer sites (B.1 opp_cost, B.2 bonus) multiply the bounded
 * `kl_amp` ∈ [1, 2] into the reward terms — amplifying Flat-trap
 * escape pressure when train/val distributions diverge.
 *
 * Single-block kernel, one thread per monitored dim. Each thread
 * reduces ALL N elements for its dim sequentially. Cold-path —
 * runs ONCE per validation epoch (acceptable cost). Uses smem
 * for the per-dim KL block reduction. No atomicAdd, no DtoH.
 *
 * Kernel-internal trailing-EMA-of-self pattern per
 * feedback_isv_for_adaptive_bounds.md — ratio is computed against
 * the previous EMA value before it is overwritten, so no extra
 * threshold ISV slot is needed. */

#include <cuda_runtime.h>

extern "C" __global__ void state_kl_moment_match(
    const float* __restrict__ train_state,    /* [n_train, d_total] (state_dim_padded stride) */
    const float* __restrict__ val_state,      /* [n_val,   d_total] (state_dim_padded stride) */
    int          n_train,
    int          n_val,
    int          d_total,
    int          monitored_dims_begin,
    int          monitored_dims_end,
    float*       __restrict__ isv_out,        /* ISV bus [ISV_TOTAL_DIM] */
    int          kl_ema_idx,                  /* = STATE_KL_TRAIN_VAL_EMA_INDEX */
    int          amp_idx,                     /* = STATE_KL_AMPLIFICATION_INDEX */
    int          sharpe_ema_idx,              /* = SHARPE_EMA_INDEX = 22 */
    float        alpha_base                   /* base EMA rate (e.g. 0.05) */
) {
    if (blockIdx.x != 0) return;
    const int d_monitored = monitored_dims_end - monitored_dims_begin;
    if (d_monitored <= 0) return;
    if (threadIdx.x >= d_monitored) return;
    const int d = monitored_dims_begin + threadIdx.x;

    /* Per-dim moments — train */
    float sum_t = 0.0f, sumsq_t = 0.0f;
    for (int i = 0; i < n_train; ++i) {
        const float x = train_state[(long long)i * d_total + d];
        sum_t   += x;
        sumsq_t += x * x;
    }
    const float n_t_f  = fmaxf(1.0f, (float)n_train);
    const float mean_t = sum_t / n_t_f;
    const float var_t  = fmaxf(1e-6f, sumsq_t / n_t_f - mean_t * mean_t);

    /* Per-dim moments — val */
    float sum_v = 0.0f, sumsq_v = 0.0f;
    for (int i = 0; i < n_val; ++i) {
        const float x = val_state[(long long)i * d_total + d];
        sum_v   += x;
        sumsq_v += x * x;
    }
    const float n_v_f  = fmaxf(1.0f, (float)n_val);
    const float mean_v = sum_v / n_v_f;
    const float var_v  = fmaxf(1e-6f, sumsq_v / n_v_f - mean_v * mean_v);

    /* Per-dim KL (numerically clamped to ≥ 0) */
    const float diff = mean_t - mean_v;
    const float kl_d = 0.5f * (logf(var_v / var_t) + (var_t + diff * diff) / var_v - 1.0f);

    extern __shared__ float smem[];
    smem[threadIdx.x] = fmaxf(0.0f, kl_d);
    __syncthreads();

    if (threadIdx.x == 0) {
        float total_kl = 0.0f;
        for (int t = 0; t < d_monitored; ++t) total_kl += smem[t];

        const float prev_kl  = isv_out[kl_ema_idx];
        const float sharpe   = fmaxf(-2.0f, fminf(2.0f, isv_out[sharpe_ema_idx]));
        const float alpha    = alpha_base * (1.0f + 0.5f * fabsf(sharpe));
        const float new_kl   = (1.0f - alpha) * prev_kl + alpha * total_kl;

        /* Adaptive amp: ratio of new EMA to prev EMA (trailing-self pattern).
         * When matched, ratio≈1 → target=1 → amp drifts toward 1.
         * When KL doubles, ratio≈2 → target=2 → amp drifts toward 2.
         * Bounded [1, 2] by construction. */
        const float ratio      = new_kl / fmaxf(1e-6f, prev_kl);
        const float amp_target = 1.0f + fminf(1.0f, fmaxf(0.0f, ratio - 1.0f));
        const float prev_amp   = isv_out[amp_idx];
        const float new_amp    = (1.0f - alpha) * prev_amp + alpha * amp_target;

        /* Commit AFTER computing ratio so the trailing-self pattern works. */
        isv_out[kl_ema_idx] = new_kl;
        isv_out[amp_idx]    = new_amp;
    }
}
