use ml::dqn::agent::TradingState;
use ml::dqn::reward::{RewardConfig, RewardFunction};
use ml::dqn::action_space::{FactoredAction, ExposureLevel, OrderType, Urgency};

#[test]
fn test_pnl_reward_is_normalized_not_raw_dollars() {
    // Create reward function with default config
    let mut reward_fn = RewardFunction::new(RewardConfig::default());

    // Scenario: $100,000 portfolio gains $500 (0.5% return)
    // Portfolio value representation: 100000 → 100500 (raw dollars)
    let current_state = TradingState {
        price_features: vec![100.0; 4],  // 4 price features
        technical_indicators: vec![0.0; 121],  // 121 technical indicators
        market_features: vec![100.0; 0],  // No market features (0 length)
        portfolio_features: vec![100_000.0, 0.0, 0.01],  // [value=$100k, position=0, spread=0.01]
    };

    let next_state = TradingState {
        price_features: vec![100.0; 4],
        technical_indicators: vec![0.0; 121],
        market_features: vec![100.0; 0],
        portfolio_features: vec![100_500.0, 0.0, 0.01],  // [value=$100.5k, position=0, spread=0.01]
    };

    // Execute BUY action (triggers P&L calculation)
    let action = FactoredAction {
        exposure: ExposureLevel::Long50,
        order: OrderType::Market,  // Fixed: 'order' not 'order_type'
        urgency: Urgency::Normal,
    };

    let reward = reward_fn.calculate_reward(
        action,
        &current_state,
        &next_state,
        &vec![],  // No recent actions (diversity penalty = 0)
    ).unwrap();

    // Extract P&L component (reward = pnl_weight * pnl_reward - risk - cost - diversity)
    // With default config: pnl_weight=1.0, risk_weight=0.1, cost_weight=0.05
    // For flat position (position=0): risk_penalty=0, cost_penalty=0
    // So reward ≈ pnl_reward (before diversity adjustment)

    // CRITICAL TEST: If bug exists, pnl_reward = $500 (raw dollars)
    // If fixed, pnl_reward = 0.005 (normalized percentage)

    // The reward includes transaction costs and other penalties, but the P&L component
    // should be the dominant term. For a 0.5% gain ($500), normalized reward should be ~0.005
    // If the bug exists, we'd see a reward around $500 (massive value)

    let reward_f64 = reward.to_string().parse::<f64>().unwrap();

    // CRITICAL ASSERTION: Reward must be in normalized range (-1.0 to 1.0), NOT raw dollars
    // Before fix: reward would be ~10 (raw $500 / some scaling factor)
    // After fix: reward should be in [-1, 1] range (normalized percentage)
    assert!(
        reward_f64.abs() < 1.0,
        "P&L reward must be normalized! Got {}, expected small normalized value. \
         Value > 1.0 indicates raw dollar amount (Wave 16Q bug still exists)!",
        reward_f64
    );

    // Note: Final reward includes transaction costs (-0.05 weight * spread penalty)
    // and diversity penalty (-0.1 for empty action history), so it will be negative
    // even with positive P&L. The key test is that it's normalized (< 1.0), not in
    // raw dollar amounts (> 1.0).
    println!("✅ Test passed: Reward {} is in normalized range (< 1.0), confirming fix", reward_f64);
}

#[test]
fn test_pnl_reward_handles_loss_correctly() {
    let mut reward_fn = RewardFunction::new(RewardConfig::default());

    // Scenario: $100,000 portfolio loses $2,000 (2% loss)
    let current_state = TradingState {
        price_features: vec![100.0; 4],
        technical_indicators: vec![0.0; 121],
        market_features: vec![100.0; 0],
        portfolio_features: vec![100_000.0, 0.0, 0.01],  // [value=$100k, position=0, spread=0.01]
    };

    let next_state = TradingState {
        price_features: vec![100.0; 4],
        technical_indicators: vec![0.0; 121],
        market_features: vec![100.0; 0],
        portfolio_features: vec![98_000.0, 0.0, 0.01],  // [value=$98k, position=0, spread=0.01]
    };

    let action = FactoredAction {
        exposure: ExposureLevel::Long50,
        order: OrderType::Market,
        urgency: Urgency::Normal,
    };

    let reward = reward_fn.calculate_reward(
        action,
        &current_state,
        &next_state,
        &vec![],
    ).unwrap();

    let reward_f64 = reward.to_string().parse::<f64>().unwrap();

    // Critical test: Reward must be normalized (< 1.0), not raw dollars (would be -2000 or -10 with bug)
    assert!(
        reward_f64.abs() < 1.0,
        "Loss reward must be normalized! Got {}, expected normalized value < 1.0",
        reward_f64
    );

    // Note: Reward will be more negative than -2% due to transaction costs and diversity penalty
    // The key is that it's in normalized range, not raw dollar amounts
    println!("✅ Test passed: Loss reward {} is in normalized range (< 1.0), confirming fix", reward_f64);
}

#[test]
fn test_large_portfolio_value_is_normalized() {
    let mut reward_fn = RewardFunction::new(RewardConfig::default());

    // Scenario: $10 MILLION portfolio gains $50,000 (0.5% return)
    // This tests that even large dollar amounts get normalized
    let current_state = TradingState {
        price_features: vec![100.0; 4],
        technical_indicators: vec![0.0; 121],
        market_features: vec![100.0; 0],
        portfolio_features: vec![10_000_000.0, 0.0, 0.01],  // $10M
    };

    let next_state = TradingState {
        price_features: vec![100.0; 4],
        technical_indicators: vec![0.0; 121],
        market_features: vec![100.0; 0],
        portfolio_features: vec![10_050_000.0, 0.0, 0.01],  // $10.05M (+$50k)
    };

    let action = FactoredAction {
        exposure: ExposureLevel::Long50,
        order: OrderType::Market,
        urgency: Urgency::Normal,
    };

    let reward = reward_fn.calculate_reward(
        action,
        &current_state,
        &next_state,
        &vec![],
    ).unwrap();

    let reward_f64 = reward.to_string().parse::<f64>().unwrap();

    // CRITICAL: Even with $50k raw gain, normalized reward must be in [-1, 1] range
    // If bug exists, we'd see reward around 50,000 or 1,000 (massive values)
    // After fix: reward should be in normalized range (< 1.0)
    assert!(
        reward_f64.abs() < 1.0,
        "Large portfolio P&L must be normalized! Got {}, which indicates \
         raw dollar bug! Expected normalized value < 1.0",
        reward_f64
    );

    println!("✅ Test passed: Large portfolio reward {} is in normalized range (< 1.0), confirming fix works for large values", reward_f64);
}