#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! GPU kernel Q-value parity tests.
//!
//! Validates that the CUDA experience collection kernel produces valid
//! outputs for all network variants:
//!   1. Standard dueling Q-network
//!   2. Distributional dueling Q-network (C51)
//!   3. NoisyNet dueling (non-deterministic -- verify noise injection works)
//!
//! These tests require a CUDA GPU.
//! Run with: `cargo test -p ml --test gpu_kernel_parity_test`

mod gpu_parity {
    use std::sync::Arc;

    use ml_core::device::MlDevice;
    use tracing::info;

    use ml::cuda_pipeline::gpu_experience_collector::{
        ExperienceCollectorConfig, GpuExperienceCollector,
    };
    use ml::cuda_pipeline::gpu_weights::extract_dueling_weights;
    use ml::dqn::dueling::{DuelingConfig, DuelingQNetwork};
    use ml::dqn::distributional_dueling::{
        DistributionalDuelingConfig, DistributionalDuelingQNetwork,
    };

    type CudaStream = cudarc::driver::CudaStream;
    type CudaSlice<T> = cudarc::driver::CudaSlice<T>;

    /// Returns (MlDevice, Arc<CudaStream>) if CUDA is available, None otherwise.
    fn try_cuda() -> Option<(MlDevice, Arc<CudaStream>)> {
        match MlDevice::cuda(0) {
            Ok(dev) => {
                let stream = dev.cuda_stream().expect("CUDA stream").clone();
                Some((dev, stream))
            }
            Err(_) => {
                tracing::warn!("CUDA not available, skipping GPU parity test");
                None
            }
        }
    }

    /// Download i32 actions from GPU to host.
    fn download_actions(stream: &Arc<CudaStream>, actions: &CudaSlice<i32>, n: usize) -> Vec<i32> {
        let mut host = vec![0_i32; n];
        stream.memcpy_dtoh(actions, &mut host).unwrap(); // test-only readback
        host
    }

    /// Download f32 rewards from GPU to host.
    fn download_f32(stream: &Arc<CudaStream>, buf: &CudaSlice<f32>, n: usize) -> Vec<f32> {
        let mut host = vec![0.0_f32; n];
        stream.memcpy_dtoh(buf, &mut host).unwrap(); // test-only readback
        host
    }

    fn dueling_config() -> DuelingConfig {
        DuelingConfig::new(54, 5, vec![256, 256], 128, 128)
    }

    fn dist_config() -> DistributionalDuelingConfig {
        DistributionalDuelingConfig::new(54, 5, 51, vec![256, 256], 128, 128)
    }

    /// Default network dims matching dueling_config/dist_config: shared=[256,256], value=128, adv=128
    const TEST_DIMS: (usize, usize, usize, usize) = (256, 256, 128, 128);
    /// Default kernel dims: state_dim=54 (tests use 51 market + 3 portfolio), market_dim=51, atoms_max=51
    const TEST_KERNEL_DIMS: (usize, usize, usize) = (54, 51, 51);

    /// Build synthetic market features on GPU (bf16, matching kernel signature).
    fn synthetic_market_data(
        total_bars: usize,
        stream: &Arc<CudaStream>,
    ) -> (CudaSlice<half::bf16>, CudaSlice<half::bf16>) {
        let market_len = total_bars * 51;
        let market_data: Vec<half::bf16> = (0..market_len)
            .map(|i| half::bf16::from_f32(((i as f32 * 0.7123 + 0.3).sin()) * 0.5))
            .collect();
        let mut market_buf = stream.alloc_zeros::<half::bf16>(market_len).unwrap();
        stream.memcpy_htod(&market_data, &mut market_buf).unwrap();

        let target_len = total_bars * 4;
        let mut target_data = vec![half::bf16::ZERO; target_len];
        for i in 0..total_bars {
            target_data[i * 4] = half::bf16::from_f32(100.0 + (i as f32 * 0.01));
            target_data[i * 4 + 1] = half::bf16::from_f32(100.5 + (i as f32 * 0.01));
            target_data[i * 4 + 2] = half::bf16::from_f32(99.5 + (i as f32 * 0.01));
            target_data[i * 4 + 3] = half::bf16::from_f32(1000.0);
        }
        let mut target_buf = stream.alloc_zeros::<half::bf16>(target_len).unwrap();
        stream.memcpy_htod(&target_data, &mut target_buf).unwrap();

        (market_buf, target_buf)
    }

    // -----------------------------------------------------------------------
    // Test 1: Standard dueling forward -- verify kernel produces valid outputs
    // -----------------------------------------------------------------------

    #[test]
    fn test_gpu_dueling_forward_produces_valid_q_values() {
        let Some((_device, stream)) = try_cuda() else { return };

        let network = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();
        let target_network = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();

        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            network.store(),
            target_network.store(),
            None,
            100_000.0,
            0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        let total_bars = 2000;
        let (market_buf, target_buf) = synthetic_market_data(total_bars, &stream);

        let n_episodes = 4_usize;
        let timesteps = 50;
        let episode_starts: Vec<i32> = (0..n_episodes).map(|i| (i * 100) as i32).collect();

        let config = ExperienceCollectorConfig {
            n_episodes: n_episodes as i32,
            timesteps_per_episode: timesteps,
            total_bars: total_bars as i32,
            episode_length: timesteps,
            epsilon: 0.1,
            gamma: 0.99,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &config)
            .expect("GPU experience collection failed");

        let expected_total = n_episodes * timesteps as usize;

        // Download actions and rewards to host for validation
        let actions = download_actions(&stream, &batch.actions, expected_total);
        let rewards = download_f32(&stream, &batch.rewards, expected_total);

        assert_eq!(actions.len(), expected_total, "actions length mismatch");
        assert_eq!(rewards.len(), expected_total, "rewards length mismatch");

        for (i, &a) in actions.iter().enumerate() {
            assert!((0..5).contains(&a), "action[{i}] = {a} out of range [0, 4]");
        }

        for (i, &r) in rewards.iter().enumerate() {
            assert!(r.is_finite(), "reward[{i}] = {r} is not finite");
        }

        info!(expected_total, "Standard dueling kernel produced valid outputs");
    }

    // -----------------------------------------------------------------------
    // Test 2: Distributional dueling (C51) -- valid output test
    // -----------------------------------------------------------------------

    #[test]
    fn test_gpu_distributional_forward_produces_valid_q_values() {
        let Some((_device, stream)) = try_cuda() else { return };

        let network = DistributionalDuelingQNetwork::new(dist_config(), stream.clone()).unwrap();
        let target = DistributionalDuelingQNetwork::new(dist_config(), stream.clone()).unwrap();

        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            network.vars(),
            target.vars(),
            None,
            100_000.0,
            0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        let total_bars = 2000;
        let (market_buf, target_buf) = synthetic_market_data(total_bars, &stream);

        let n_episodes = 4_usize;
        let timesteps = 50;
        let episode_starts: Vec<i32> = (0..n_episodes).map(|i| (i * 100) as i32).collect();

        let config = ExperienceCollectorConfig {
            n_episodes: n_episodes as i32,
            timesteps_per_episode: timesteps,
            total_bars: total_bars as i32,
            episode_length: timesteps,
            epsilon: 0.1,
            gamma: 0.99,
            num_atoms: 51,
            v_min: -25.0,
            v_max: 25.0,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &config)
            .expect("GPU experience collection (C51) failed");

        let expected_total = n_episodes * timesteps as usize;
        let actions = download_actions(&stream, &batch.actions, expected_total);
        let rewards = download_f32(&stream, &batch.rewards, expected_total);

        assert_eq!(actions.len(), expected_total);

        for (i, &a) in actions.iter().enumerate() {
            assert!((0..5).contains(&a), "C51 action[{i}] = {a} out of range [0, 4]");
        }

        for (i, &r) in rewards.iter().enumerate() {
            assert!(r.is_finite(), "C51 reward[{i}] = {r} is not finite");
        }

        info!(expected_total, "C51 distributional kernel produced valid outputs");
    }

    // -----------------------------------------------------------------------
    // Test 3: Network vs GPU kernel action parity (standard dueling)
    // -----------------------------------------------------------------------

    #[test]
    fn test_candle_vs_kernel_q_value_parity() {
        let Some((_device, stream)) = try_cuda() else { return };

        let network = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();
        let target_network = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();

        // Network forward pass for a known state (host-side)
        let state_data: Vec<f32> = (0..54).map(|i| (i as f32 * 0.1).sin() * 0.5).collect();
        let q_values = network.forward_single(&state_data).unwrap();
        let num_actions = q_values.len();

        // Compute argmax and max Q on host (CPU)
        let (candle_argmax, candle_q_max) = q_values
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(idx, &val)| (idx, val))
            .unwrap();

        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            network.store(),
            target_network.store(),
            None,
            100_000.0,
            0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        // Market data with test state as every bar (bf16 to match kernel signature)
        let total_bars = 200;
        let market_len = total_bars * 51;
        let market_data: Vec<half::bf16> = (0..market_len)
            .map(|idx| {
                let bar = idx / 51;
                let f = idx % 51;
                half::bf16::from_f32(state_data[f.min(53)])
            })
            .collect();
        let mut market_buf = stream.alloc_zeros::<half::bf16>(market_len).unwrap();
        stream.memcpy_htod(&market_data, &mut market_buf).unwrap();

        let target_len = total_bars * 4;
        let mut target_data = vec![half::bf16::ZERO; target_len];
        for i in 0..total_bars {
            target_data[i * 4] = half::bf16::from_f32(100.0);
            target_data[i * 4 + 1] = half::bf16::from_f32(100.5);
            target_data[i * 4 + 2] = half::bf16::from_f32(99.5);
            target_data[i * 4 + 3] = half::bf16::from_f32(1000.0);
        }
        let mut target_buf = stream.alloc_zeros::<half::bf16>(target_len).unwrap();
        stream.memcpy_htod(&target_data, &mut target_buf).unwrap();

        let episode_starts = vec![0_i32];
        let config = ExperienceCollectorConfig {
            n_episodes: 1,
            timesteps_per_episode: 1,
            total_bars: total_bars as i32,
            episode_length: 1,
            epsilon: 0.0,
            gamma: 0.99,
            count_bonus_coefficient: 0.0,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &config)
            .unwrap();

        let actions = download_actions(&stream, &batch.actions, 1);
        let rewards = download_f32(&stream, &batch.rewards, 1);
        let kernel_action = actions[0] as usize;

        info!(candle_argmax, candle_q_max, kernel_action, "Network vs kernel argmax comparison");
        info!(reward = rewards[0], "Kernel reward");

        // The kernel assembles state as [51 market features, 3 portfolio features]
        // while the network test uses the raw 54-dim test state. Portfolio at t=0 = [100000, 0, 0]
        // differs from the test state's last 3 dims, so argmax may differ for close
        // Q-values. Verify the kernel at least produces valid actions.
        assert!(
            (0..5).contains(&actions[0]),
            "Kernel action {} out of valid range [0, 4]",
            actions[0]
        );

        // Compute Q-gap on host: max Q minus second-best Q
        let mut sorted_q = q_values.clone();
        sorted_q.sort_by(|a, b| b.partial_cmp(a).unwrap());
        let second_best = sorted_q[1.min(sorted_q.len() - 1)];
        let q_gap = candle_q_max - second_best;

        if q_gap > 0.5 {
            assert_eq!(
                kernel_action, candle_argmax,
                "Kernel action {kernel_action} != network argmax {candle_argmax} \
                 despite Q gap {q_gap:.3} -- forward pass likely wrong"
            );
        } else {
            info!(
                q_gap, kernel_action, candle_argmax,
                "Q-values too close for strict parity -- portfolio feature perturbation can flip argmax"
            );
        }
    }

    // -----------------------------------------------------------------------
    // Test 4: NoisyNet exploration produces different actions than clean
    // -----------------------------------------------------------------------

    #[test]
    fn test_gpu_noisy_net_exploration_differs_from_clean() {
        let Some((_device, stream)) = try_cuda() else { return };

        let network = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();
        let target = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();

        let total_bars = 2000;
        let (market_buf, target_buf) = synthetic_market_data(total_bars, &stream);

        let n_episodes = 32_usize;
        let timesteps = 100;
        let episode_starts: Vec<i32> = (0..n_episodes).map(|i| (i * 50) as i32).collect();

        // Run WITHOUT noise
        let mut collector_clean = GpuExperienceCollector::new(
            stream.clone(), network.store(), target.store(), None,
            100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        let config_clean = ExperienceCollectorConfig {
            n_episodes: n_episodes as i32,
            timesteps_per_episode: timesteps,
            total_bars: total_bars as i32,
            episode_length: timesteps,
            epsilon: 0.0,
            gamma: 0.99,
            count_bonus_coefficient: 0.0,
            ..Default::default()
        };

        let batch_clean = collector_clean
            .collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &config_clean)
            .unwrap();

        // Run WITH noise
        let mut collector_noisy = GpuExperienceCollector::new(
            stream.clone(), network.store(), target.store(), None,
            100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        let config_noisy = ExperienceCollectorConfig {
            n_episodes: n_episodes as i32,
            timesteps_per_episode: timesteps,
            total_bars: total_bars as i32,
            episode_length: timesteps,
            epsilon: 0.0,
            gamma: 0.99,
            noisy_sigma_init: 0.5,
            count_bonus_coefficient: 0.0,
            ..Default::default()
        };

        let batch_noisy = collector_noisy
            .collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &config_noisy)
            .unwrap();

        let total = n_episodes * timesteps as usize;
        let actions_clean = download_actions(&stream, &batch_clean.actions, total);
        let actions_noisy = download_actions(&stream, &batch_noisy.actions, total);

        let diff_count = (0..total)
            .filter(|&i| actions_clean[i] != actions_noisy[i])
            .count();

        let diff_pct = (diff_count as f64 / total as f64) * 100.0;
        info!(diff_count, total, diff_pct, "NoisyNet action divergence from clean");

        assert!(
            diff_count > 0,
            "NoisyNet produced ZERO different actions vs clean -- noise injection is broken"
        );
        assert!(
            diff_pct > 2.0,
            "NoisyNet only changed {diff_pct:.1}% of actions -- noise may be too weak"
        );
    }

    // -----------------------------------------------------------------------
    // Test 5: Weight sync roundtrip
    // -----------------------------------------------------------------------

    #[test]
    fn test_gpu_weight_extraction_and_sync() {
        let Some((_device, stream)) = try_cuda() else { return };

        let network = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();
        let target = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();

        let weights = extract_dueling_weights(network.store(), &stream)
            .expect("Weight extraction failed");

        // Download and verify w_s1 is non-zero and finite (bf16 weights)
        let mut shared_0_w = vec![half::bf16::ZERO; 256 * 54];
        stream.memcpy_dtoh(&weights.w_s1, &mut shared_0_w).unwrap(); // test-only readback
        assert!(!shared_0_w.iter().all(|&x| x == half::bf16::ZERO), "w_s1 is all zeros");
        assert!(shared_0_w.iter().all(|x| x.to_f32().is_finite()), "w_s1 has NaN/Inf");

        // Verify sync works
        let mut collector = GpuExperienceCollector::new(
            stream.clone(), network.store(), target.store(), None,
            100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        collector.sync_online_weights(network.store()).expect("Online sync failed");
        collector.sync_target_weights(target.store()).expect("Target sync failed");

        info!("Weight extraction and sync roundtrip complete");
    }

    // -----------------------------------------------------------------------
    // Test 6: C51 distributional weight shapes and RMSNorm extraction
    // -----------------------------------------------------------------------

    #[test]
    fn test_gpu_distributional_weight_shapes() {
        let Some((_device, stream)) = try_cuda() else { return };

        let network = DistributionalDuelingQNetwork::new(dist_config(), stream.clone()).unwrap();

        let weights = extract_dueling_weights(network.vars(), &stream)
            .expect("Distributional weight extraction failed");

        // value_out: [51, 128] = 6528 elements (bf16 weights)
        let mut value_out_w = vec![half::bf16::ZERO; 51 * 128];
        stream.memcpy_dtoh(&weights.w_v2, &mut value_out_w).unwrap(); // test-only readback
        assert!(value_out_w.iter().any(|&x| x != half::bf16::ZERO), "w_v2 all zeros");

        // advantage_out: [255, 128] = 32640 elements (bf16 weights)
        let mut adv_out_w = vec![half::bf16::ZERO; 255 * 128];
        stream.memcpy_dtoh(&weights.w_a2, &mut adv_out_w).unwrap(); // test-only readback
        assert!(adv_out_w.iter().any(|&x| x != half::bf16::ZERO), "w_a2 all zeros");

        // RMSNorm gamma should exist and be ~1.0
        let rmsnorm = ml::cuda_pipeline::gpu_weights::extract_rmsnorm_weights(
            network.vars(), &stream,
        )
        .expect("RMSNorm extraction failed");

        assert!(rmsnorm.is_some(), "Distributional network should have RMSNorm weights");

        let rmsnorm = rmsnorm.unwrap();
        let mut gamma_s0 = vec![half::bf16::ZERO; 256];
        stream.memcpy_dtoh(&rmsnorm.gamma_s0, &mut gamma_s0).unwrap(); // test-only readback

        let mean_gamma: f32 = gamma_s0.iter().map(|x| x.to_f32()).sum::<f32>() / gamma_s0.len() as f32;
        assert!(
            (mean_gamma - 1.0).abs() < 0.01,
            "RMSNorm gamma_s0 mean = {mean_gamma} (expected ~1.0)"
        );

        info!(mean_gamma, "Distributional weights verified: value_out [51,128], adv_out [255,128]");
    }

    // -----------------------------------------------------------------------
    // Test 7: Multiple kernel launches produce finite results
    // -----------------------------------------------------------------------

    #[test]
    fn test_gpu_kernel_repeated_launches_stable() {
        let Some((_device, stream)) = try_cuda() else { return };

        let network = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();
        let target = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();

        let total_bars = 2000;
        let (market_buf, target_buf) = synthetic_market_data(total_bars, &stream);

        let n_episodes = 4_usize;
        let timesteps = 20;
        let episode_starts: Vec<i32> = (0..n_episodes).map(|i| (i * 100) as i32).collect();

        let cfg = ExperienceCollectorConfig {
            n_episodes: n_episodes as i32,
            timesteps_per_episode: timesteps,
            total_bars: total_bars as i32,
            episode_length: timesteps,
            epsilon: 0.0,
            gamma: 0.99,
            count_bonus_coefficient: 0.0,
            ..Default::default()
        };

        let mut collector = GpuExperienceCollector::new(
            stream.clone(), network.store(), target.store(), None,
            100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        let expected_total = n_episodes * timesteps as usize;

        for run in 0..5 {
            let batch = collector
                .collect_experiences_gpu(&market_buf, &target_buf, &episode_starts, &cfg)
                .unwrap();

            let rewards = download_f32(&stream, &batch.rewards, expected_total);

            for (i, &r) in rewards.iter().enumerate() {
                assert!(r.is_finite(), "Run {run} reward[{i}] = {r} is not finite");
            }
        }

        info!("5 consecutive kernel launches: all produce finite rewards");
    }

    // -----------------------------------------------------------------------
    // Test 8: Standard vs C51 both produce valid output
    // -----------------------------------------------------------------------

    #[test]
    fn test_gpu_distributional_vs_standard_both_valid() {
        let Some((_device, stream)) = try_cuda() else { return };

        let std_net = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();
        let std_target = DuelingQNetwork::new(dueling_config(), stream.clone()).unwrap();
        let dist_net = DistributionalDuelingQNetwork::new(dist_config(), stream.clone()).unwrap();
        let dist_target = DistributionalDuelingQNetwork::new(dist_config(), stream.clone()).unwrap();

        let total_bars = 2000;
        let (market_buf, target_buf) = synthetic_market_data(total_bars, &stream);
        let n_episodes = 4_usize;
        let timesteps = 50;
        let episode_starts: Vec<i32> = (0..n_episodes).map(|i| (i * 100) as i32).collect();
        let expected_total = n_episodes * timesteps as usize;

        let mut std_collector = GpuExperienceCollector::new(
            stream.clone(), std_net.store(), std_target.store(), None,
            100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        let std_batch = std_collector
            .collect_experiences_gpu(
                &market_buf, &target_buf, &episode_starts,
                &ExperienceCollectorConfig {
                    n_episodes: n_episodes as i32,
                    timesteps_per_episode: timesteps,
                    total_bars: total_bars as i32,
                    episode_length: timesteps,
                    ..Default::default()
                },
            )
            .unwrap();

        let mut dist_collector = GpuExperienceCollector::new(
            stream.clone(), dist_net.vars(), dist_target.vars(), None,
            100_000.0, 0.01, 0.05, TEST_DIMS, TEST_KERNEL_DIMS, 8, 50,
        ).unwrap();

        let dist_batch = dist_collector
            .collect_experiences_gpu(
                &market_buf, &target_buf, &episode_starts,
                &ExperienceCollectorConfig {
                    n_episodes: n_episodes as i32,
                    timesteps_per_episode: timesteps,
                    total_bars: total_bars as i32,
                    episode_length: timesteps,
                    num_atoms: 51,
                    v_min: -25.0,
                    v_max: 25.0,
                    ..Default::default()
                },
            )
            .unwrap();

        let std_actions = download_actions(&stream, &std_batch.actions, expected_total);
        let dist_actions = download_actions(&stream, &dist_batch.actions, expected_total);
        let std_rewards = download_f32(&stream, &std_batch.rewards, expected_total);
        let dist_rewards = download_f32(&stream, &dist_batch.rewards, expected_total);

        assert_eq!(std_actions.len(), dist_actions.len());

        // Both should produce finite rewards
        assert!(std_rewards.iter().all(|r| r.is_finite()), "Standard rewards have NaN/Inf");
        assert!(dist_rewards.iter().all(|r| r.is_finite()), "C51 rewards have NaN/Inf");

        // Both should produce valid actions
        assert!(std_actions.iter().all(|&a| (0..5).contains(&a)), "Standard actions out of range");
        assert!(dist_actions.iter().all(|&a| (0..5).contains(&a)), "C51 actions out of range");

        info!("Standard and C51 both produced valid outputs");
    }
}