//! Gradient budget and spectral norm validation tests.
//!
//! These tests validate the per-component gradient clipping infrastructure
//! and spectral normalization at the GPU kernel level, without requiring
//! training data. They construct a GpuDqnTrainer directly and exercise
//! the clipping/norm kernels with synthetic gradient data.

use std::sync::Arc;

use super::helpers::cuda_device;
use crate::cuda_pipeline::gpu_dqn_trainer::{GpuDqnTrainConfig, GpuDqnTrainer};
use crate::cuda_pipeline::gpu_weights::{
    DuelingWeightSet, BranchingWeightSet,
    DuelingWeightBacking, BranchingWeightBacking,
};

/// Small config for fast kernel-level tests (no training, just buffer ops).
fn test_config() -> GpuDqnTrainConfig {
    GpuDqnTrainConfig {
        state_dim: 16,
        shared_h1: 32,
        shared_h2: 32,
        value_h: 16,
        adv_h: 16,
        num_atoms: 11,
        v_min: -10.0,
        v_max: 10.0,
        branch_0_size: 4,
        branch_1_size: 3,
        branch_2_size: 3,
        batch_size: 8,
        max_grad_norm: 10.0,
        spectral_norm_sigma_max: 3.0,
        market_dim: 12,        // test state: 12 market + 4 portfolio = 16 state_dim
        bottleneck_dim: 0,     // disabled for spectral norm test (avoids w_s1 size mismatch)
        ..GpuDqnTrainConfig::default()
    }
}

/// Mirror of `compute_param_sizes`' `s1_input_dim` logic — when bottleneck
/// is off, w_s1's input dim is `STATE_DIM` (not `cfg.state_dim`, which is
/// vestigial in the bottleneck-off path). The test_config sets
/// `bottleneck_dim: 0` so this always returns STATE_DIM.
fn s1_input_dim_for_test(_cfg: &GpuDqnTrainConfig) -> usize {
    ml_core::state_layout::STATE_DIM
}

/// Allocate a DuelingWeightSet with random values (not zero — spectral norm
/// needs non-degenerate matrices for power iteration to converge).
/// Returns both the backing storage (keeps GPU memory alive) and the pointer view.
fn alloc_dueling(stream: &Arc<cudarc::driver::CudaStream>, cfg: &GpuDqnTrainConfig) -> (DuelingWeightBacking, DuelingWeightSet) {
    let na = cfg.num_atoms;
    let s1_in = s1_input_dim_for_test(cfg);
    let alloc = |n: usize| -> cudarc::driver::CudaSlice<f32> {
        // Fill with 0.1 (non-zero, non-degenerate for spectral norm)
        let bf16_data: Vec<f32> = vec![0.1; n];
        stream.clone_htod(&bf16_data).expect("alloc dueling weight")
    };
    // DuelingWeightBacking slices [0..12) map to flat-buffer indices via
    // `DUELING_FLAT_INDICES`. After the GRN refactor:
    //   slice[0] → flat[0] w_a_h_s1   [SH1, s1_in]
    //   slice[1] → flat[1] b_a_h_s1   [SH1]
    //   slice[2] → flat[2] w_b_h_s1   [2*SH1, SH1]  (Linear_b → GLU split)
    //   slice[3] → flat[3] b_b_h_s1   [2*SH1]
    //   slice[4..8) → flat[13..17) value head (post-GRN reshuffle)
    //   slice[8..12) → flat[17..21) branch 0 head
    // Sizes MUST match `compute_param_sizes` exactly or flatten_online_weights
    // rejects with "size mismatch" because flat-buf copies use the flat
    // sizes, not the slice sizes.
    let backing = DuelingWeightBacking {
        slices: [
            alloc(cfg.shared_h1 * s1_in), alloc(cfg.shared_h1),
            alloc((2 * cfg.shared_h1) * cfg.shared_h1), alloc(2 * cfg.shared_h1),
            alloc(cfg.value_h * cfg.shared_h2), alloc(cfg.value_h),
            alloc(na * cfg.value_h), alloc(na),
            // w_b0fc input dim grew to SH2+1 in SP14 (aux→Q wire); bias unchanged.
            alloc(cfg.adv_h * (cfg.shared_h2 + 1)), alloc(cfg.adv_h),
            alloc(cfg.branch_0_size * na * cfg.adv_h), alloc(cfg.branch_0_size * na),
        ],
    };
    let ws = backing.weight_set();
    (backing, ws)
}

fn alloc_branching(stream: &Arc<cudarc::driver::CudaStream>, cfg: &GpuDqnTrainConfig) -> (BranchingWeightBacking, BranchingWeightSet) {
    let na = cfg.num_atoms;
    let alloc = |n: usize| -> cudarc::driver::CudaSlice<f32> {
        let bf16_data: Vec<f32> = vec![0.1; n];
        stream.clone_htod(&bf16_data).expect("alloc branching weight")
    };
    // Branch 1 (magnitude) is direction-conditioned: input dim = SH2 + branch_0_size.
    // Branches 2-3 (order, urgency) are OFI-conditioned: input dim = SH2 + 3
    // (3 OFI features per branch — see concat_ofi_features kernel).
    let b1_input = cfg.shared_h2 + cfg.branch_0_size;
    let b2_input = cfg.shared_h2 + 3;
    let b3_input = cfg.shared_h2 + 3;
    let backing = BranchingWeightBacking {
        slices: [
            alloc(cfg.adv_h * b1_input), alloc(cfg.adv_h),
            alloc(cfg.branch_1_size * na * cfg.adv_h), alloc(cfg.branch_1_size * na),
            alloc(cfg.adv_h * b2_input), alloc(cfg.adv_h),
            alloc(cfg.branch_2_size * na * cfg.adv_h), alloc(cfg.branch_2_size * na),
            alloc(cfg.adv_h * b3_input), alloc(cfg.adv_h),
            alloc(cfg.branch_3_size * na * cfg.adv_h), alloc(cfg.branch_3_size * na),
        ],
    };
    let ws = backing.weight_set();
    (backing, ws)
}

/// Spectral norm on all 10 weight matrices must not panic and must
/// produce finite weights after normalization.
#[test]
fn test_spectral_norm_all_heads_no_panic() -> anyhow::Result<()> {
    let dev = cuda_device();
    let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
    let cfg = test_config();
    let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?;
    let (d_backing, dueling) = alloc_dueling(&stream, &cfg);
    let (b_backing, branching) = alloc_branching(&stream, &cfg);

    // Run spectral norm — should not panic with non-zero weights.
    // Spectral norm now writes directly to params_bf16; unflatten syncs
    // the individual weight slices back so we can read them.
    trainer.apply_spectral_norm(&dueling, &branching)?;
    trainer.unflatten_online_weights(&dueling, &branching)?;

    // Verify a few weights are still finite after normalization (f32 -> f32 readback)
    let mut w_s1_bf16 = vec![0.0_f32; cfg.shared_h1 * s1_input_dim_for_test(&cfg)];
    stream.memcpy_dtoh(&d_backing.slices[0], &mut w_s1_bf16)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let w_s1_host: Vec<f32> = w_s1_bf16.to_vec();
    assert!(
        w_s1_host.iter().all(|v| v.is_finite()),
        "W_s1 has non-finite values after spectral norm"
    );

    let mut w_v1_bf16 = vec![0.0_f32; cfg.value_h * cfg.shared_h2];
    stream.memcpy_dtoh(&d_backing.slices[4], &mut w_v1_bf16)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let w_v1_host: Vec<f32> = w_v1_bf16.to_vec();
    assert!(
        w_v1_host.iter().all(|v| v.is_finite()),
        "W_v1 has non-finite values after spectral norm"
    );

    let mut w_bo1_bf16 = vec![0.0_f32; cfg.adv_h * (cfg.shared_h2 + cfg.branch_0_size)];
    stream.memcpy_dtoh(&b_backing.slices[0], &mut w_bo1_bf16)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let w_bo1_host: Vec<f32> = w_bo1_bf16.to_vec();
    assert!(
        w_bo1_host.iter().all(|v| v.is_finite()),
        "W_bo1 has non-finite values after spectral norm"
    );

    Ok(())
}

/// Spectral norm must actually constrain the spectral norm of weight matrices.
/// After applying spectral norm with σ_max=3.0, the L2 operator norm
/// of each weight matrix should be ≤ σ_max.
#[test]
fn test_spectral_norm_constrains_operator_norm() -> anyhow::Result<()> {
    let dev = cuda_device();
    let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
    let cfg = test_config();
    let mut trainer = GpuDqnTrainer::new(stream.clone(), cfg.clone())?;

    // Create weights with large values — spectral norm should constrain them
    let alloc_large = |n: usize| -> cudarc::driver::CudaSlice<f32> {
        let bf16_data: Vec<f32> = vec![5.0; n];
        stream.clone_htod(&bf16_data).expect("alloc large weight")
    };
    let na = cfg.num_atoms;
    let s1_in = s1_input_dim_for_test(&cfg);
    // See alloc_dueling above for the DUELING_FLAT_INDICES slot → flat-
    // buffer layout mapping. slices[2]/[3] are w_b_h_s1/b_b_h_s1 [2*SH1,
    // SH1] / [2*SH1] in the GRN-expanded layout, not w_s2/b_s2.
    let d_backing = DuelingWeightBacking {
        slices: [
            alloc_large(cfg.shared_h1 * s1_in), alloc_large(cfg.shared_h1),
            alloc_large((2 * cfg.shared_h1) * cfg.shared_h1), alloc_large(2 * cfg.shared_h1),
            alloc_large(cfg.value_h * cfg.shared_h2), alloc_large(cfg.value_h),
            alloc_large(na * cfg.value_h), alloc_large(na),
            // w_b0fc input dim grew to SH2+1 in SP14 (aux→Q wire); bias unchanged.
            alloc_large(cfg.adv_h * (cfg.shared_h2 + 1)), alloc_large(cfg.adv_h),
            alloc_large(cfg.branch_0_size * na * cfg.adv_h), alloc_large(cfg.branch_0_size * na),
        ],
    };
    let dueling = d_backing.weight_set();
    // Branch 1 (magnitude) is direction-conditioned: input dim = SH2 + branch_0_size.
    // Branches 2-3 (order, urgency) are OFI-conditioned: input dim = SH2 + 3.
    let b1_input = cfg.shared_h2 + cfg.branch_0_size;
    let b2_input = cfg.shared_h2 + 3;
    let b3_input = cfg.shared_h2 + 3;
    let b_backing = BranchingWeightBacking {
        slices: [
            alloc_large(cfg.adv_h * b1_input), alloc_large(cfg.adv_h),
            alloc_large(cfg.branch_1_size * na * cfg.adv_h), alloc_large(cfg.branch_1_size * na),
            alloc_large(cfg.adv_h * b2_input), alloc_large(cfg.adv_h),
            alloc_large(cfg.branch_2_size * na * cfg.adv_h), alloc_large(cfg.branch_2_size * na),
            alloc_large(cfg.adv_h * b3_input), alloc_large(cfg.adv_h),
            alloc_large(cfg.branch_3_size * na * cfg.adv_h), alloc_large(cfg.branch_3_size * na),
        ],
    };
    let branching = b_backing.weight_set();

    // Run spectral norm multiple times (power iteration converges over iterations).
    // Spectral norm writes directly to params_bf16; unflatten syncs back to
    // individual slices for readback. Only needed once after the final iteration.
    for _ in 0..5 {
        trainer.apply_spectral_norm(&dueling, &branching)?;
    }
    trainer.unflatten_online_weights(&dueling, &branching)?;

    // Check W_s1 [shared_h1, s1_input_dim]: Frobenius norm of a rank-1
    // matrix with σ_max=3.0 would be 3.0. A uniform matrix of 5.0 has
    // σ = 5*sqrt(rows*cols), after spectral norm it should have σ ≤ 3.0.
    // Proxy check: Frobenius norm should be significantly reduced from initial.
    //
    // s1_input_dim is STATE_DIM (128) when bottleneck is off — NOT
    // cfg.state_dim (which is vestigial in the bottleneck-off path).
    // See `s1_input_dim_for_test` for the canonical formula.
    let rows = cfg.shared_h1;
    let cols = s1_input_dim_for_test(&cfg);
    let mut w_bf16 = vec![0.0_f32; rows * cols];
    stream.memcpy_dtoh(&d_backing.slices[0], &mut w_bf16)
        .map_err(|e| anyhow::anyhow!("{e}"))?;
    let w_host: Vec<f32> = w_bf16.to_vec();

    let frobenius: f32 = w_host.iter().map(|x| x * x).sum::<f32>().sqrt();
    let initial_frobenius = 5.0 * (rows as f32 * cols as f32).sqrt();

    assert!(
        frobenius < initial_frobenius * 0.5,
        "Frobenius norm should be significantly reduced: {frobenius:.2} vs initial {initial_frobenius:.2}"
    );

    // The spectral norm (largest singular value) should be ≤ σ_max.
    // For a uniform matrix after spectral norm, the spectral norm equals
    // σ_max if the original σ > σ_max. We can estimate σ ≈ Frobenius / sqrt(min(m,n)).
    let approx_sigma = frobenius / (rows.min(cols) as f32).sqrt();
    assert!(
        approx_sigma < cfg.spectral_norm_sigma_max * 1.5, // tolerance for approximation
        "Approximate spectral norm {approx_sigma:.2} should be near σ_max={:.1}",
        cfg.spectral_norm_sigma_max
    );

    Ok(())
}

/// Dynamic budget: auxiliary fractions must leave room for C51.
/// C51 gets 1.0 - sum(active auxiliary budgets). When all are active: C51=70%.
#[test]
fn test_budget_fractions_leave_room_for_c51() {
    // Auxiliary budgets from fused_training.rs
    let cql: f32 = 0.15;
    let iqn: f32 = 0.10;
    let ens: f32 = 0.05;
    let aux_total = cql + iqn + ens;

    // C51 gets the remainder
    let c51 = 1.0 - aux_total;
    assert!(c51 > 0.5, "C51 should get at least 50%, got {c51}");
    assert!(
        (c51 + aux_total - 1.0).abs() < 1e-6,
        "Total budget must equal 1.0, got {}", c51 + aux_total
    );

    // When no auxiliaries: C51 gets 100%
    let c51_solo: f32 = 1.0 - 0.0 - 0.0 - 0.0;
    assert!((c51_solo - 1.0).abs() < 1e-6, "Solo C51 should be 1.0");
}
