//! Tests for PPO Hidden State Management
//!
//! Validates LSTM hidden state initialization, propagation, and reset logic.

use candle_core::{DType, Device, Tensor};
use ml::ppo::hidden_state_manager::HiddenStateManager;

#[test]
fn test_hidden_state_initialization() -> Result<(), Box<dyn std::error::Error>> {
    // Test that hidden states are initialized to zeros
    let device = Device::cuda_if_available(0)?;
    let num_layers = 2;
    let batch_size = 4;
    let hidden_dim = 128;

    let manager = HiddenStateManager::new(num_layers, batch_size, hidden_dim, &device)?;

    let (policy_h, policy_c) = manager.get_policy_state();
    let (value_h, value_c) = manager.get_value_state();

    // Verify shapes
    assert_eq!(policy_h.shape().dims(), &[num_layers, batch_size, hidden_dim]);
    assert_eq!(policy_c.shape().dims(), &[num_layers, batch_size, hidden_dim]);
    assert_eq!(value_h.shape().dims(), &[num_layers, batch_size, hidden_dim]);
    assert_eq!(value_c.shape().dims(), &[num_layers, batch_size, hidden_dim]);

    // Verify all zeros
    let policy_h_sum = policy_h.sum_all()?.to_scalar::<f32>()?;
    let policy_c_sum = policy_c.sum_all()?.to_scalar::<f32>()?;
    let value_h_sum = value_h.sum_all()?.to_scalar::<f32>()?;
    let value_c_sum = value_c.sum_all()?.to_scalar::<f32>()?;

    assert_eq!(policy_h_sum, 0.0, "Policy hidden state should be zeros");
    assert_eq!(policy_c_sum, 0.0, "Policy cell state should be zeros");
    assert_eq!(value_h_sum, 0.0, "Value hidden state should be zeros");
    assert_eq!(value_c_sum, 0.0, "Value cell state should be zeros");

    Ok(())
}

#[test]
fn test_hidden_state_propagation() -> Result<(), Box<dyn std::error::Error>> {
    // Test that hidden states carry across timesteps
    let device = Device::cuda_if_available(0)?;
    let num_layers = 2;
    let batch_size = 4;
    let hidden_dim = 128;

    let mut manager = HiddenStateManager::new(num_layers, batch_size, hidden_dim, &device)?;

    // Create new states with non-zero values
    let new_h = Tensor::ones(&[num_layers, batch_size, hidden_dim], DType::F32, &device)?;
    let new_c = Tensor::ones(&[num_layers, batch_size, hidden_dim], DType::F32, &device)?.affine(2.0, 0.0)?;

    // Update policy states
    manager.update_policy_state(new_h.clone(), new_c.clone())?;

    // Verify states were updated
    let (policy_h, policy_c) = manager.get_policy_state();
    let policy_h_sum = policy_h.sum_all()?.to_scalar::<f32>()?;
    let policy_c_sum = policy_c.sum_all()?.to_scalar::<f32>()?;

    let expected_h_sum = (num_layers * batch_size * hidden_dim) as f32;
    let expected_c_sum = expected_h_sum * 2.0;

    assert!(
        (policy_h_sum - expected_h_sum).abs() < 0.01,
        "Policy hidden state should be updated. Expected {}, got {}",
        expected_h_sum,
        policy_h_sum
    );
    assert!(
        (policy_c_sum - expected_c_sum).abs() < 0.01,
        "Policy cell state should be updated. Expected {}, got {}",
        expected_c_sum,
        policy_c_sum
    );

    // Update value states
    manager.update_value_state(new_h, new_c)?;

    // Verify value states were updated
    let (value_h, value_c) = manager.get_value_state();
    let value_h_sum = value_h.sum_all()?.to_scalar::<f32>()?;
    let value_c_sum = value_c.sum_all()?.to_scalar::<f32>()?;

    assert!(
        (value_h_sum - expected_h_sum).abs() < 0.01,
        "Value hidden state should be updated. Expected {}, got {}",
        expected_h_sum,
        value_h_sum
    );
    assert!(
        (value_c_sum - expected_c_sum).abs() < 0.01,
        "Value cell state should be updated. Expected {}, got {}",
        expected_c_sum,
        value_c_sum
    );

    Ok(())
}

#[test]
fn test_hidden_state_reset_on_done() -> Result<(), Box<dyn std::error::Error>> {
    // Test that states reset when episodes end
    let device = Device::cuda_if_available(0)?;
    let num_layers = 2;
    let batch_size = 4;
    let hidden_dim = 128;

    let mut manager = HiddenStateManager::new(num_layers, batch_size, hidden_dim, &device)?;

    // Set states to non-zero
    let ones = Tensor::ones(&[num_layers, batch_size, hidden_dim], DType::F32, &device)?;
    manager.update_policy_state(ones.clone(), ones.clone())?;
    manager.update_value_state(ones.clone(), ones.clone())?;

    // Create done mask: environments 0 and 2 are done
    let done_mask = Tensor::new(&[1u8, 0u8, 1u8, 0u8], &device)?;

    // Reset states for done environments
    manager.reset_on_done(&done_mask)?;

    // Verify policy states
    let (policy_h, policy_c) = manager.get_policy_state();
    let policy_h_data = policy_h.flatten_all()?.to_vec1::<f32>()?;
    let policy_c_data = policy_c.flatten_all()?.to_vec1::<f32>()?;

    // Check that environments 0 and 2 are reset (all zeros)
    // and environments 1 and 3 retain their values (all ones)
    for layer in 0..num_layers {
        for env in 0..batch_size {
            for dim in 0..hidden_dim {
                let idx = layer * batch_size * hidden_dim + env * hidden_dim + dim;
                let expected = if env == 0 || env == 2 { 0.0 } else { 1.0 };
                assert!(
                    (policy_h_data[idx] - expected).abs() < 0.01,
                    "Policy hidden state at [{}, {}, {}] should be {}. Got {}",
                    layer,
                    env,
                    dim,
                    expected,
                    policy_h_data[idx]
                );
                assert!(
                    (policy_c_data[idx] - expected).abs() < 0.01,
                    "Policy cell state at [{}, {}, {}] should be {}. Got {}",
                    layer,
                    env,
                    dim,
                    expected,
                    policy_c_data[idx]
                );
            }
        }
    }

    // Verify value states similarly
    let (value_h, value_c) = manager.get_value_state();
    let value_h_data = value_h.flatten_all()?.to_vec1::<f32>()?;
    let value_c_data = value_c.flatten_all()?.to_vec1::<f32>()?;

    for layer in 0..num_layers {
        for env in 0..batch_size {
            for dim in 0..hidden_dim {
                let idx = layer * batch_size * hidden_dim + env * hidden_dim + dim;
                let expected = if env == 0 || env == 2 { 0.0 } else { 1.0 };
                assert!(
                    (value_h_data[idx] - expected).abs() < 0.01,
                    "Value hidden state at [{}, {}, {}] should be {}. Got {}",
                    layer,
                    env,
                    dim,
                    expected,
                    value_h_data[idx]
                );
                assert!(
                    (value_c_data[idx] - expected).abs() < 0.01,
                    "Value cell state at [{}, {}, {}] should be {}. Got {}",
                    layer,
                    env,
                    dim,
                    expected,
                    value_c_data[idx]
                );
            }
        }
    }

    Ok(())
}

#[test]
fn test_hidden_state_batching() -> Result<(), Box<dyn std::error::Error>> {
    // Test that manager handles multiple parallel environments correctly
    let device = Device::cuda_if_available(0)?;
    let num_layers = 2;
    let batch_size = 8; // 8 parallel environments
    let hidden_dim = 64;

    let mut manager = HiddenStateManager::new(num_layers, batch_size, hidden_dim, &device)?;

    // Create different values for different environments
    let mut h_data = vec![0.0f32; num_layers * batch_size * hidden_dim];
    let mut c_data = vec![0.0f32; num_layers * batch_size * hidden_dim];

    for layer in 0..num_layers {
        for env in 0..batch_size {
            for dim in 0..hidden_dim {
                let idx = layer * batch_size * hidden_dim + env * hidden_dim + dim;
                h_data[idx] = (env + 1) as f32; // Environment 0 = 1.0, env 1 = 2.0, etc.
                c_data[idx] = (env + 1) as f32 * 10.0; // Environment 0 = 10.0, env 1 = 20.0, etc.
            }
        }
    }

    let new_h = Tensor::from_vec(h_data.clone(), &[num_layers, batch_size, hidden_dim], &device)?;
    let new_c = Tensor::from_vec(c_data.clone(), &[num_layers, batch_size, hidden_dim], &device)?;

    // Update states
    manager.update_policy_state(new_h.clone(), new_c.clone())?;

    // Verify each environment has correct values
    let (policy_h, policy_c) = manager.get_policy_state();
    let retrieved_h = policy_h.flatten_all()?.to_vec1::<f32>()?;
    let retrieved_c = policy_c.flatten_all()?.to_vec1::<f32>()?;

    for layer in 0..num_layers {
        for env in 0..batch_size {
            for dim in 0..hidden_dim {
                let idx = layer * batch_size * hidden_dim + env * hidden_dim + dim;
                let expected_h = (env + 1) as f32;
                let expected_c = (env + 1) as f32 * 10.0;

                assert!(
                    (retrieved_h[idx] - expected_h).abs() < 0.01,
                    "Hidden state for env {} should be {}. Got {}",
                    env,
                    expected_h,
                    retrieved_h[idx]
                );
                assert!(
                    (retrieved_c[idx] - expected_c).abs() < 0.01,
                    "Cell state for env {} should be {}. Got {}",
                    env,
                    expected_c,
                    retrieved_c[idx]
                );
            }
        }
    }

    Ok(())
}