//! Pre-compile all ml-alpha CUDA kernels into arch-specific cubins.
//!
//! Per `feedback_no_nvrtc.md`: no runtime kernel compilation.
//! Per `pearl_build_rs_rerun_if_env_changed.md`: every `std::env::var`
//! is paired with `cargo:rerun-if-env-changed`.

use std::path::{Path, PathBuf};
use std::process::Command;

const KERNELS: &[&str] = &[
    "mamba2_alpha_kernel",          // Mamba2 SSM scan kernel (used by PerceptionTrainer's encoder prefix)
    "snap_feature_assemble",
    "cfc_step",
    "multi_horizon_heads",
    "projection",
    "bce_loss_multi_horizon",        // Kendall σ-weighted multi-horizon BCE (axis A)
    "adamw_step",
    "grad_norm",
    "horizon_lambda",                // ISV-driven per-horizon gradient scaler (EMA + lambda)
    "layer_norm",                    // Phase 1: trunk pre-CfC normalisation
    "variable_selection",            // Phase 2D: TFT-style per-feature gating
    "attention_pool",                // Phase 3: single-Q learned content summary at CfC k=0
    "reduce_axis0",                  // Phase B: cross-batch param-grad reducer
    "output_smoothness",             // CRT.train: per-horizon adjacent-position prob-jitter penalty
    "smoothness_lambda_controller",  // CRT.train: ISV-driven λ controller anchored on h30 jitter
    "gpu_log_ring",                  // GPU diagnostic log ring — tick kernel + log_record helper
    "bucket_transition_kernels",     // Per-horizon CfC Phase 1→2 transition: tau_sort, bucket_assign, bucket_iqr, channels_in_bucket, heads_compact, zero_off_bucket (ALPHA fix 2026-05-21)
    "cfc_step_per_branch",           // Per-horizon CfC Phase 2: fused per-(batch, branch) fwd + bwd over [25,25,25,25,28] buckets
    "heads_block_diagonal_fwd",      // Per-horizon CfC Phase 2: heads w_skip projection with compact ragged storage (640→128 floats)
    "aux_trunk",                     // SDD-3 Layer B3: smaller single-bucket CfC trunk (AUX_HIDDEN=64) for outcome-supervision (D-labels)
    "aux_heads",                     // SDD-3 Layer B4: per-direction linear regression heads on AuxTrunk output (long + short, N_AUX_HORIZONS each)
    "aux_loss",                      // SDD-3 Layer B4: Huber loss + grad for aux trade-outcome regression targets (NaN-masked)
    "aux_vec_add",                   // SDD-3 Layer B5: element-wise dst += src for aux→encoder gradient accumulation (lifted stop-grad)
    "dqn_distributional_q",          // RL Phase C: C51 distributional Q-head fwd + Bellman TD bwd for integrated RL trainer
    "rl_gamma_controller",           // RL Phase C: ISV controller emitting γ to ISV[RL_GAMMA_INDEX=400]
    "rl_target_tau_controller",      // RL Phase C: ISV controller emitting τ to ISV[RL_TARGET_TAU_INDEX=401]
    "ppo_clipped_surrogate",         // RL Phase D: PPO clipped-surrogate + entropy bonus + value MSE fwd/bwd
    "rl_ppo_clip_controller",        // RL Phase D: ISV controller emitting ε to ISV[RL_PPO_CLIP_INDEX=402]
    "rl_entropy_coef_controller",    // RL Phase D: ISV controller emitting entropy bonus weight to ISV[RL_ENTROPY_COEF_INDEX=403]
    "v_head_fwd_bwd",                // RL Phase E.2: scalar V(s) head fwd + MSE bwd (linear layer; per-batch scratch + reduce_axis0)
    "grad_h_accumulate",             // RL Phase E.2: element-wise grad_h_encoder += λ × grad_h_head accumulator (one head at a time, serialised by stream)
    "bellman_target_projection",     // RL Phase E.2-DEFER: C51 categorical projection of Bellman target Z(s_{t+1}, a*) onto the discrete support, reads γ from ISV[400]; replaces host-side build_synthetic_bellman_target stand-in
    "rl_lr_controller",              // RL Phase E.2-DEFER: per-head learning-rate ISV emitter — bootstraps ISV[412..417] with 1e-3 (BCE/Q/π/V/aux); replaces hardcoded PHASE_E2_DEFAULT_LR
    "rl_rollout_steps_controller",   // RL Phase E.3b: rollout-buffer-length ISV emitter — emits ISV[RL_N_ROLLOUT_STEPS_INDEX=404] from var(advantage)/|mean A| EMA; bootstraps 2048
    "rl_per_alpha_controller",       // RL Phase E.3b: PER priority-exponent ISV emitter — emits ISV[RL_PER_ALPHA_INDEX=405] from TD-error kurtosis EMA; bootstraps 0.6
    "rl_reward_scale_controller",    // RL Phase R1 (rebuild): reward-standardisation scale ISV emitter — emits ISV[RL_REWARD_SCALE_INDEX=406] from mean |realized_pnl_usd| EMA; bootstraps 1.0
    "ema_update_on_done",            // RL Phase R3: generic done-gated EMA producer (slot-parameterised) for closed-trade-magnitude EMAs (mean_abs_pnl, q_divergence, td_kurtosis)
    "ema_update_per_step",           // RL Phase R3: generic per-step EMA producer (slot-parameterised) for continuous EMAs (kl_pi, entropy_observed, advantage_var_ratio, trade_duration)
    "compute_advantage_return",      // RL Phase R3: element-wise A_t = r + γ(1-done)·V(s_{t+1}) − V(s_t), R_t = r + γ(1-done)·V(s_{t+1}); reads γ from ISV[400]
    "rl_action_kernel",              // RL Phase R4: Thompson sampler over C51 atoms; one block per batch, N_ACTIONS threads; per-batch xorshift32 PRNG state; replaces host Thompson loop per feedback_cpu_is_read_only
    "argmax_expected_q",             // RL Phase R4: argmax over expected Q per action; Bellman-target argmax (Double-DQN); pairs with rl_action_kernel per pearl_thompson_for_distributional_action_selection
    "log_pi_at_action",              // RL Phase R4: per-batch log π(action_b) via log-softmax + lookup; PPO importance-ratio path
    "dqn_target_soft_update",        // RL Phase R5: element-wise target[i] = (1-τ)·target + τ·current, reads τ from ISV[401]; closes defect #4 (no target-net soft update in flawed branch)
    "extract_realized_pnl_delta",    // RL Phase R6: GPU-pure reward + done extraction from device Pos array; replaces host read_pos loop per feedback_cpu_is_read_only
    "apply_reward_scale",            // RL Phase R6: element-wise rewards *= ISV[RL_REWARD_SCALE_INDEX=406]; closes the F.3b host roundtrip
    "actions_to_market_targets",     // RL Phase R6: 9-action grid → LobSim market_targets[B*2] on device; replaces host submit_market loop per feedback_cpu_is_read_only
    "abs_copy",                      // RL Phase R7a: element-wise dst[b] = fabsf(src[b]); feeds |reward| into ema_update_on_done for the MEAN_ABS_PNL_EMA slot
    "rl_var_over_abs_mean_streaming",// EMA-streaming var/|mean| (folds across STEPS, fixes b_size=1 → ISV[421] feeding rl_rollout_steps
    "rl_kurtosis_streaming",         // EMA-streaming kurtosis M4/M2² (folds across STEPS, fixes b_size=1) → ISV[422] feeding rl_per_alpha
    "rl_kl_approx_b",                // Schulman-style KL = mean(log π_old − log π_new) → kl_pi_ema (ISV[419]) feeding rl_ppo_clip
    "rl_l2_diff_norm",               // ‖W_online − W_target‖₂ → q_divergence_ema (ISV[418]) feeding rl_target_tau
    "rl_step_counter_update",        // per-batch trade-duration counter + done-gated emit → mean_trade_duration_ema (ISV[417]) feeding rl_gamma
    "rl_l2_norm",                    // ‖x‖₂ single-buffer reduction → q/pi/v grad-norm EMAs (ISV[424..427]) feeding rl_lr_controller
    // (entropy_observed_ema, ISV[420], feeds rl_entropy_coef directly via ema_update_per_step's internal mean reduce on entropy_d — no separate kernel needed.)
    "rl_ppo_ratio_clamp_controller", // RL R9: PPO ratio clamp ceiling at ISV[440], anchored on ε at ISV[402] — bounds catastrophic unclipped-branch surrogate
    "ppo_log_ratio_abs_max_b",       // RL R9 diag: per-batch max|log π_new − log π_old| → ISV[441]; surfaces ratio-clamp activity in diag JSONL
    "rl_streaming_clamp_init",       // RL R9: device-side seeder for streaming-kernel output clamp ceilings (ISV[447], ISV[448]) — no HtoD
    "rl_isv_write",                  // RL R9: generic single-slot ISV seeder for tunable design constants — no HtoD
    "rl_q_pi_agree_b",               // RL R9 audit: per-batch fraction (argmax Q == argmax π) → ISV[407] EMA; wires previously-dead slot
    "rl_pi_action_kernel",           // audit Option B: π drives action selection via multinomial sampling from softmax(pi_logits); Q becomes pure critic
    "rl_reward_clamp_controller",    // audit 2026-05-24: adaptive [-LOSS, +WIN] clamp from positive-tail EMA; replaces static [-3, +1] that crushed winning-trade signal in rmgm5
    "rl_atom_support_update",        // audit 2026-05-24 followup: refreshes atom_supports_d from ISV V_MIN/V_MAX so C51 atom span adapts with reward clamp (Q learning was capped at V_MAX=1.0)
    "rl_kl_reference_grad",
    "rl_q_pi_distill_grad",          // audit 2026-05-24 vj5f6 followup: KL(softmax(E_Q/τ) || π_new) gradient ADDED to pi_grad_logits — couples Q's improved C51 calibration to action selection (was decoupled per Option B)
    "action_entropy_per_step",       // POST-gate action entropy EMA for SAC α/τ co-tuning
    "rl_drawdown_stop",              // per-step drawdown penalty + hard stop-loss
    "rl_q_distill_lambda_controller",// audit 2026-05-24 rljzl followup: adaptive λ_distill via Schulman bounded step on KL_EMA vs target
    "rl_unit_state_update",          // SP20 P1+P5 audit fix: per-unit trade state machine — detects open/close/reverse transitions, sets up unit slot 0 entry+trail
    "rl_trail_mutate",               // SP20 P1+P5 audit fix (a7/a8 dead): TrailTighten/TrailLoosen mutate unit_trail_distance bounded MIN/MAX, symmetric reciprocal adjust
    "rl_trail_stop_check",           // SP20 P1+P5 audit fix: per-unit trail breach check; OVERRIDE action to FlatFromLong/Short on breach (close routes through existing flat plumbing)
    "rl_frd_fwd",                    // SP20 P3: Forward-Return-Distribution head fwd — 2-layer MLP [HIDDEN_DIM → FRD_HIDDEN_DIM → FRD_N_HORIZONS × FRD_N_ATOMS]; ReLU hidden cached for bwd; softmax + CE happen in bwd
    "rl_frd_softmax_ce_grad",        // SP20 P3 F.3a: per-(batch, horizon) softmax + CE loss + dL/dlogits; 1 block per (b, h), 21 threads; label = -1 sentinel masks the row
    "rl_frd_layer2_bwd",             // SP20 P3 F.3b: FRD head layer-2 backward — dW2 (per-batch scratch), db2 (per-batch scratch), dhidden (per-batch overwrite); 1 block per batch, 64 threads
    "rl_frd_layer1_bwd",             // SP20 P3 F.3c: FRD head layer-1 backward — dW1 (per-batch scratch), db1 (per-batch scratch), dh_t (per-batch overwrite); applies ReLU mask via cached post-ReLU hidden; 1 block per batch, 128 threads
    "rl_position_heat_check",        // SP20 P6: position heat cap — force-flat when |position_lots| exceeds ISV-driven max; last defense before actions_to_market_targets
    "rl_confidence_gate",            // P8: override opening actions to Hold when C51 distributional Q is insufficiently confident (LCB < threshold)
    "rl_frd_gate",                   // P9: override opening actions to Hold when FRD head predicts insufficient favorable probability mass
    "rl_recent_outcome_update",      // P10: per-batch signed outcome EMA for anti-martingale sizing
    "rl_trade_context_update",       // P1: per-batch trade-arc features (time_in_trade, unrealized_R, pos_mag, entry_dist)
    "rl_multires_features_update",   // P0: per-batch multi-resolution streaming features (3 horizons × 4 features)
    "rl_encoder_context_broadcast",  // P2: broadcast per-batch context (16 dims) into encoder input [B,K,56] at cols 40-55
    "rl_gate_threshold_controller",  // adaptive gate thresholds from trade frequency (dones EMA)
    "rl_reward_shaping",             // surfer-philosophy: entry cost + short-hold penalty + hold bonus
    "rl_min_hold_check",             // hard minimum hold time — override close actions to Hold before N steps
    "rl_asymmetric_trail_decay",     // auto-tighten losers, auto-widen winners — structural P&L asymmetry
    "rl_session_risk_check",         // session-level loss limit circuit breaker
    "rl_iqn_forward",                // IQN distributional Q-head: quantile embedding + action-value projection; complementary to C51
    "rl_iqn_loss",                   // IQN quantile Huber loss: ρ_τ(δ) = |τ - 1(δ<0)| × Huber(δ, κ=1.0); forward + backward
    "rl_iqn_backward",               // IQN backward through forward pass: grad_output → grad_w_out/b_out/w_embed/b_embed per-batch scratch
    "rl_ensemble_action_value",      // C51+IQN ensemble: E_ensemble = α×E_C51 + (1-α)×E_IQN; α from ISV[544]
    "rl_noisy_linear_forward",       // NoisyNet: factored noisy linear forward — y = (mu_w + sigma_w ⊙ eps_w) × x + (mu_b + sigma_b ⊙ eps_b); state-dependent exploration for C51/IQN final projection
    "rl_noisy_linear_backward",      // NoisyNet: factored noisy linear backward — grad_mu_w/sigma_w/mu_b/sigma_b per-batch scratch for reduce_axis0
    "rl_sample_tau",                 // CUDA graph prereq: device-side xorshift32 tau ~ U(0,1) for IQN; replaces host ChaCha8 + mapped-pinned upload
    "rl_sample_noise",               // CUDA graph prereq: device-side factored noise f(rand) for NoisyLinear; replaces host ChaCha8 + mapped-pinned upload
    "rl_write_u64",                  // CUDA graph prereq: single-thread u64 scalar write for device-resident ts_ns (graph-captured kernels read via pointer)
    "rl_fused_reward_pipeline",      // P3: 7→1 fused per-batch reward extraction + shaping + EMA + outcome update
    "rl_per_push_ring",              // GPU PER: coalesced n-step ring write + flush decision (Grid=B, Block=128)
    "rl_per_push_flush",             // GPU PER: coordinated coalesced replay write with prefix-sum slot allocation (Grid=B, Block=128)
    "rl_per_sample",                 // GPU PER: stratified proportional sampling via sum-tree walk + gather
    "rl_per_update_priority",        // GPU PER: write |TD|^α to tree leaves + block-wide max reduction
    "rl_per_tree_rebuild",           // GPU PER: bottom-up parallel sum-tree rebuild (no atomics)
    "rl_hindsight_track",            // HER Phase 1: per-step mid-price ring + peak tracking for backward hindsight
    "rl_hindsight_inject",           // HER Phase 2: backward inject — synthetic replay push on done if peak >> actual
    "rl_hindsight_forward",          // HER Phase 3: forward continuation — evaluates closed trades after lookahead
    "rl_increment_step",             // device-resident step counter bump (ISV[548] += 1.0); graph-safe prereq — removes scalar current_step from all downstream kernel args
    "rl_fused_controllers",          // fused kernel: 10 RL ISV controllers in one launch (gamma, tau, ppo_clip, entropy_coef, rollout_steps, per_alpha, reward_scale, ppo_ratio_clamp, gate_threshold, q_distill_lambda) — saves 9 kernel launch overheads (~40-80μs/step)
    "rl_popart_normalize",           // PopArt: Welford-EMA reward normalization (replaces apply_reward_scale)
    "rl_popart_v_correct",           // PopArt: V-head output correction after stats shift
    "rl_spectral_norm",              // Spectral norm: power iteration σ_max estimate + W rescale
    "rl_spectral_decouple",          // Spectral decoupling: L2 penalty on logit magnitudes
    "rl_q_bias_correction",          // Q-bias: EMA of (Q_pred - actual_return) → Bellman correction
    "rl_per_branch_lr",              // Per-branch LR: adaptive per-head learning rate scaling
    "rl_outcome_fwd",                // Outcome aux: linear forward h_t → 3-class logits
    "rl_outcome_ce",                 // Outcome aux: softmax CE loss + gradient (masked by sentinel -1)
    "rl_outcome_label",              // Outcome aux: assign labels from reward/done (Profit/Timeout/Loss)
    "rl_outcome_bwd",                // Outcome aux: backward through linear layer → grad_W/b/h_t
    "rl_outcome_fused",              // Outcome aux: fused fwd + CE + bwd — eliminates 2 global round-trips (logits, grad_logits kept in smem)
    "rl_curriculum_weights",         // E8: per-segment difficulty-weighted softmax from Sharpe → PER weights
    "rl_adversarial_boost",          // Adversarial: boost PER priority for negative-reward transitions
    "rl_outcome_bwd",                // Outcome aux: single linear layer backward — dW (per-batch), db (per-batch), dh_t (per-batch); 1 block per batch, 128 threads
    "snapshot_aos_to_soa",           // AoS→SoA scatter: one thread per snapshot reads contiguous Mbp10RawInput, writes into 10 SoA device buffers; replaces host nested loops + 10 DtoD copies
    "gpu_sample_and_gather",         // GPU-resident batch sampler: random file+anchor sampling + AoS→SoA gather from pre-uploaded dataset; eliminates ALL per-step CPU data loading
];

// Cache bust v31 — five new reduce / derive kernels populate the input
// EMAs for the previously-frozen controllers (entropy_coef,
// rollout_steps, per_alpha, ppo_clip, target_tau, gamma). Each kernel
// is a single-block reduction (tree-reduce, grid-stride, or per-batch
// state update). The trainer launches each one immediately after its
// source signal is populated:
//   * `rl_var_over_abs_mean_b`   after compute_advantage_return
//   * `rl_kurtosis_b`            after dqn_distributional_q_bwd
//   * `rl_kl_approx_b`           after PPO surrogate forward
//   * `rl_l2_diff_norm`          after target-net soft update
//   * `rl_step_counter_update`   after extract_realized_pnl_delta
// The scalar output is then consumed by ema_update_per_step (continuous
// EMAs) or ema_update_on_done (done-gated EMAs) at the right ISV slot.
// `entropy_observed_ema` reuses ema_update_per_step's built-in
// per-batch mean reduce on entropy_d directly.

fn main() {
    println!("cargo:rerun-if-changed=build.rs");
    // Track shared headers so .cuh / .h edits trigger rebuilds of every
    // .cu that #includes them. Without these, an edit to a helper header
    // leaves a stale cubin.
    println!("cargo:rerun-if-changed=cuda/gpu_log_ids.h");
    println!("cargo:rerun-if-changed=cuda/gpu_log_helpers.cuh");

    println!("cargo:rerun-if-env-changed=CARGO_FEATURE_CUDA");
    if std::env::var("CARGO_FEATURE_CUDA").is_err() {
        eprintln!("  ml-alpha: cuda feature disabled, skipping kernel build");
        return;
    }

    println!("cargo:rerun-if-env-changed=CUDA_HOME");
    println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP");

    let nvcc = match find_nvcc() {
        Some(p) => p,
        None => {
            eprintln!("  ml-alpha: nvcc not found, skipping kernel build (set CUDA_HOME or install CUDA toolkit)");
            return;
        }
    };

    let out = PathBuf::from(std::env::var("OUT_DIR").expect("OUT_DIR not set by cargo"));

    // Detect GPU arch: CUDA_COMPUTE_CAP env > nvidia-smi query > default 86
    let arch = detect_arch(&nvcc);
    eprintln!("  ml-alpha: compiling kernels for sm_{arch}");

    for k in KERNELS {
        let src = PathBuf::from(format!("cuda/{k}.cu"));
        if !src.exists() {
            eprintln!("  ml-alpha: skipping {k} — source not yet present");
            continue;
        }
        println!("cargo:rerun-if-changed={}", src.display());
        let cubin = out.join(format!("{k}.cubin"));
        compile(&nvcc, &src, &cubin, &arch);
    }
}

fn detect_arch(_nvcc: &Path) -> String {
    // 1. Explicit env override
    if let Ok(cap) = std::env::var("CUDA_COMPUTE_CAP") {
        return cap;
    }
    // 2. Query the GPU on this machine
    if let Ok(output) = Command::new("nvidia-smi")
        .args(["--query-gpu=compute_cap", "--format=csv,noheader"])
        .output()
    {
        if output.status.success() {
            let s = String::from_utf8_lossy(&output.stdout);
            let cap = s.trim().replace('.', "");
            if !cap.is_empty() {
                return cap;
            }
        }
    }
    // 3. Default to sm_86 (RTX 3050 Ti local dev)
    "86".to_string()
}

fn compile(nvcc: &Path, src: &Path, cubin: &Path, arch: &str) {
    let status = Command::new(nvcc)
        .args([
            "-cubin",
            &format!("-arch=sm_{arch}"),
            "-O3",
            "--use_fast_math",
            "--ftz=true",
            "--fmad=true",
            "-o",
            cubin.to_str().unwrap(),
            src.to_str().unwrap(),
        ])
        .status()
        .unwrap_or_else(|e| panic!("nvcc spawn failed for {}: {e}", src.display()));
    if !status.success() {
        panic!(
            "nvcc failed for {} (exit {})",
            src.display(),
            status.code().unwrap_or(-1)
        );
    }
    eprintln!(
        "  ml-alpha: compiled {} -> {} (sm_{arch})",
        src.display(),
        cubin.display()
    );
}

fn find_nvcc() -> Option<PathBuf> {
    if let Ok(home) = std::env::var("CUDA_HOME") {
        let p = PathBuf::from(home).join("bin/nvcc");
        if p.exists() {
            return Some(p);
        }
    }
    for cand in ["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] {
        let p = PathBuf::from(cand);
        if p.exists() {
            return Some(p);
        }
    }
    Command::new("nvcc")
        .arg("--version")
        .output()
        .ok()
        .filter(|o| o.status.success())
        .map(|_| PathBuf::from("nvcc"))
}