//! Reward v8 unit tests — verify mathematical properties of each component.
//! CPU-side tests that validate the CUDA kernel logic produces correct results.

#[cfg(test)]
mod tests {
    use std::f64::consts::PI;

    fn asymmetric_soft_clamp(x: f32) -> f32 {
        if x >= 0.0 { x.min(10.0) } else { -10.0 * (1.0 - (x / 10.0).exp()) }
    }

    #[test]
    fn test_asymmetric_soft_clamp_properties() {
        assert!((asymmetric_soft_clamp(0.0)).abs() < 1e-6, "f(0) = 0");
        assert!((asymmetric_soft_clamp(5.0) - 5.0).abs() < 1e-6, "f(5) = 5 (linear)");
        assert!((asymmetric_soft_clamp(15.0) - 10.0).abs() < 1e-6, "f(15) = 10 (capped)");
        assert!((asymmetric_soft_clamp(-5.0) - (-3.935)).abs() < 0.01, "f(-5) ≈ -3.93");
        assert!((asymmetric_soft_clamp(-15.0) - (-7.769)).abs() < 0.01, "f(-15) ≈ -7.77");
        assert!(asymmetric_soft_clamp(3.0) > asymmetric_soft_clamp(2.0), "monotonic positive");
        assert!(asymmetric_soft_clamp(-2.0) > asymmetric_soft_clamp(-3.0), "monotonic negative");
        assert!(asymmetric_soft_clamp(-5.0).abs() < asymmetric_soft_clamp(5.0).abs(), "asymmetric");
    }

    #[test]
    fn test_kelly_prior_neutral() {
        let prior_wins = 2.0_f64;
        let prior_losses = 2.0;
        let prior_sum_wins = 0.01;
        let prior_sum_losses = 0.01;
        let eff_total = prior_wins + prior_losses;
        let win_rate = prior_wins / eff_total;
        let avg_win = prior_sum_wins / prior_wins;
        let avg_loss = prior_sum_losses / prior_losses;
        let payoff = avg_win / avg_loss;
        let kelly_f = (payoff * win_rate - (1.0 - win_rate)) / payoff;
        assert!(kelly_f.abs() < 0.01, "Kelly with prior only should be ~0, got {kelly_f}");
    }

    #[test]
    fn test_cosine_epsilon_schedule() {
        let eps_start = 0.3_f64;
        let eps_end = 0.02;
        let total = 20.0;
        let e0 = eps_end + 0.5 * (eps_start - eps_end) * (1.0 + (PI * 0.0 / total).cos());
        assert!((e0 - eps_start).abs() < 0.001, "Epoch 0 should be eps_start, got {e0}");
        let ef = eps_end + 0.5 * (eps_start - eps_end) * (1.0 + PI.cos());
        assert!((ef - eps_end).abs() < 0.001, "Final epoch should be eps_end, got {ef}");
        let em = eps_end + 0.5 * (eps_start - eps_end) * (1.0 + (PI * 0.5).cos());
        assert!(em > eps_end && em < eps_start, "Mid epsilon {em} between bounds");
    }

    #[test]
    fn test_td_lambda_degenerates() {
        let rewards = [1.0_f64, 2.0, 3.0];
        let q_next = [10.0, 20.0, 30.0];
        let gamma = 0.95;
        let g1 = rewards[0] + gamma * q_next[0];
        assert!((g1 - 10.5).abs() < 0.01, "1-step return should be 10.5, got {g1}");
        let g_mc = rewards[0] + gamma * rewards[1] + gamma * gamma * rewards[2]
                 + gamma * gamma * gamma * q_next[2];
        let expected = 1.0 + 0.95 * 2.0 + 0.9025 * 3.0 + 0.857375 * 30.0;
        assert!((g_mc - expected).abs() < 0.01, "MC return should be {expected}, got {g_mc}");
    }

    #[test]
    fn test_micro_reward_adaptive_scale() {
        let base = 0.001_f64;
        let vol_calm = 0.0025_f64;
        let adaptive_calm = base / (vol_calm / 0.005_f64).sqrt().max(0.5);
        assert!(adaptive_calm > base, "Calm scale ({adaptive_calm}) > base ({base})");
        let vol_volatile = 0.02_f64;
        let adaptive_volatile = base / (vol_volatile / 0.005_f64).sqrt().max(0.5);
        assert!(adaptive_volatile < base, "Volatile scale ({adaptive_volatile}) < base ({base})");
    }

    #[test]
    fn test_exposure_aux_gradient_unique() {
        let b0 = 3;  // 4-branch: direction branch now 3 (Short/Flat/Long)
        let target = 1; // Flat
        let prob = 1.0 / b0 as f64;
        let grads: Vec<f64> = (0..b0).map(|a| prob - if a == target { 1.0 } else { 0.0 }).collect();
        assert!(grads[target] < 0.0, "Target gets negative gradient");
        for a in 0..b0 {
            if a != target { assert!(grads[a] > 0.0, "Non-target {a} gets positive gradient"); }
        }
        assert!((grads[target] - grads[0]).abs() > 0.1, "Target differs from non-target");
    }

    #[test]
    fn test_hindsight_relabel_improves_reward() {
        let entry = 100.0_f64;
        let futures = [101.0, 103.0, 102.0];
        let position = 1.0;
        let actual_pnl = position * (futures[0] - entry);
        let best_pnl = futures.iter().map(|&p| position * (p - entry)).fold(0.0_f64, f64::max);
        assert!((best_pnl - 3.0).abs() < 0.01, "Best PnL should be 3.0");
        assert!(best_pnl >= actual_pnl, "Hindsight >= actual");
    }

    #[test]
    fn test_curriculum_ordering() {
        let trending = 0.02_f64;
        let choppy = 0.001;
        let score_trending = 1.0 / (trending + 0.001);
        let score_choppy = 1.0 / (choppy + 0.001);
        assert!(score_trending < score_choppy, "Trending (easy) < choppy (hard)");
    }

    #[test]
    fn test_popart_normalization() {
        let rewards = vec![1.0_f64, 2.0, 3.0, 4.0, 5.0];
        let mean = rewards.iter().sum::<f64>() / rewards.len() as f64;
        let var = rewards.iter().map(|r| (r - mean).powi(2)).sum::<f64>() / rewards.len() as f64;
        let std = var.sqrt();
        let normalized: Vec<f64> = rewards.iter().map(|r| (r - mean) / std).collect();
        let norm_mean = normalized.iter().sum::<f64>() / normalized.len() as f64;
        let norm_var = normalized.iter().map(|r| (r - norm_mean).powi(2)).sum::<f64>() / normalized.len() as f64;
        assert!(norm_mean.abs() < 1e-10, "Mean should be 0, got {norm_mean}");
        assert!((norm_var - 1.0).abs() < 0.01, "Variance should be 1, got {norm_var}");
    }

    // ── 4-Branch Architecture Tests ─────────────────────────────────────

    #[test]
    fn test_4branch_action_encoding() {
        // All 81 actions encode/decode correctly
        for dir in 0..3_usize {
            for mag in 0..3_usize {
                for order in 0..3_usize {
                    for urgency in 0..3_usize {
                        let action = dir * 27 + mag * 9 + order * 3 + urgency;
                        assert_eq!(action / 27, dir);
                        assert_eq!((action % 27) / 9, mag);
                        assert_eq!((action % 9) / 3, order);
                        assert_eq!(action % 3, urgency);
                        assert!(action < 81);
                    }
                }
            }
        }
    }

    #[test]
    fn test_4branch_position_mapping() {
        let max_pos = 2.0_f64;
        // direction: 0=Short(-1), 1=Flat(0), 2=Long(+1)
        // magnitude: 0=Small(0.25), 1=Medium(0.5), 2=Full(1.0)
        // Short × Full = -2.0
        assert!(((- 1.0_f64) * 1.0 * max_pos - (-2.0)).abs() < 1e-10);
        // Flat × anything = 0
        assert!((0.0_f64 * 0.25 * max_pos).abs() < 1e-10);
        assert!((0.0_f64 * 1.0 * max_pos).abs() < 1e-10);
        // Long × Quarter = +0.5
        assert!((1.0_f64 * 0.25 * max_pos - 0.5).abs() < 1e-10);
        // Long × Half = +1.0
        assert!((1.0_f64 * 0.5 * max_pos - 1.0).abs() < 1e-10);
        // Long × Full = +2.0
        assert!((1.0_f64 * 1.0 * max_pos - 2.0).abs() < 1e-10);
        // Short × Small = -0.5
        assert!(((- 1.0_f64) * 0.25 * max_pos - (-0.5)).abs() < 1e-10);
    }

    #[test]
    fn test_flat_shortcut() {
        // Flat direction (idx=1) × any magnitude = 0 position
        for mag in 0..3_usize {
            let dir_fraction = 0.0_f64; // Flat
            let mag_scale = match mag { 0 => 0.25, 1 => 0.5, _ => 1.0 };
            assert!((dir_fraction * mag_scale).abs() < 1e-10,
                "Flat × mag={mag} must be zero");
        }
    }

    #[test]
    fn test_magnitude_epsilon_equals_direction() {
        let eps_start = 0.20_f64;
        let eps_end = 0.02;
        let eps_old_mag = eps_end + 0.1 * (eps_start - eps_end);
        let eps_new_mag = eps_start;
        assert!(eps_new_mag > eps_old_mag * 4.0,
            "New magnitude epsilon ({eps_new_mag}) should be >4× old ({eps_old_mag})");
        assert!((eps_new_mag - eps_start).abs() < 1e-10,
            "Magnitude epsilon should equal direction epsilon");
    }

    #[test]
    fn test_counterfactual_weight() {
        let cf_weight = 0.3_f64;
        let flat_fraction = 0.58;
        let effective = cf_weight * flat_fraction + 1.0 * (1.0 - flat_fraction);
        assert!(effective > 0.5, "Effective gradient should be >50%: {effective}");
        assert!(effective < 1.0, "Counterfactual must reduce magnitude gradient: {effective}");
        let amp = 1.0 / effective;
        assert!((amp - 1.68).abs() < 0.1, "Expected ~1.68× amplification, got {amp}");
    }

    #[test]
    fn test_advantage_standardization() {
        // Raw advantages: Small=10.0, Medium=1.0, Full=0.5 (collapse scenario)
        let advantages = [10.0_f64, 1.0, 0.5];
        let mean = advantages.iter().sum::<f64>() / 3.0;
        let centered: Vec<f64> = advantages.iter().map(|a| a - mean).collect();
        let variance = centered.iter().map(|c| c * c).sum::<f64>() / 3.0;
        let std = variance.sqrt();
        let standardized: Vec<f64> = centered.iter().map(|c| c / (std + 1e-6)).collect();
        // After standardization, values are O(1)
        assert!(standardized[0] < 2.0, "Small should be <2.0 after std, got {}", standardized[0]);
        assert!(standardized[2] > -2.0, "Full should be >-2.0 after std, got {}", standardized[2]);
        // Relative ordering preserved
        assert!(standardized[0] > standardized[1], "Small > Medium ordering preserved");
        assert!(standardized[1] > standardized[2], "Medium > Full ordering preserved");
        // Gap is bounded
        let gap = standardized[0] - standardized[2];
        assert!(gap < 3.0, "Max gap should be <3.0, got {gap}");
    }

    #[test]
    fn test_magnitude_micro_reward_scaling() {
        let dir_reward = 0.01_f64;
        let adaptive_scale = 1.0;
        let micro_vol = 0.005;
        let small_reward = adaptive_scale * dir_reward * 0.25 / micro_vol;
        let full_reward = adaptive_scale * dir_reward * 1.00 / micro_vol;
        assert!((full_reward / small_reward - 4.0).abs() < 1e-10,
            "Full/Small reward ratio should be 4.0, got {}", full_reward / small_reward);
    }
}
