use rtx_rl::algorithms::{SAC, SACConfig};
use rtx_tensor::{DType, Device, Tensor};

#[tokio::test]
async fn test_sac_creation() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig {
        learning_rate: 3e-4,
        gamma: 0.99,
        tau: 0.005,
        alpha: 0.2,
        target_update_interval: 1,
        automatic_entropy_tuning: true,
        target_entropy: None,
        replay_buffer_size: 1_000_000,
        batch_size: 256,
    };

    let state_dim = 8;
    let action_dim = 3;
    let hidden_dim = 256;

    let sac = SAC::new(config, state_dim, action_dim, hidden_dim, device);
    assert_eq!(sac.state_dim(), state_dim);
    assert_eq!(sac.action_dim(), action_dim);
}

#[tokio::test]
#[ignore = "rtx-rl SAC actor tanh bounds incomplete"]
async fn test_sac_actor_forward() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig::default();

    let state_dim = 8;
    let action_dim = 3;
    let hidden_dim = 256;
    let batch_size = 32;

    let sac = SAC::new(config, state_dim, action_dim, hidden_dim, device.clone());

    let states = Tensor::randn(&[batch_size, state_dim], &device).unwrap();

    let (actions, log_probs) = sac
        .actor_forward(&states, false)
        .await
        .expect("Actor forward should work");

    assert_eq!(actions.shape(), &[batch_size, action_dim]);
    assert_eq!(log_probs.shape(), &[batch_size]);

    // Actions should be in [-1, 1] after tanh
    let action_data: Vec<f32> = actions.to_vec().expect("Should convert");
    for action in action_data {
        assert!(action >= -1.0 && action <= 1.0);
    }
}

#[tokio::test]
async fn test_sac_critic_forward() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig::default();
    let sac = SAC::new(config, 8, 3, 256, device.clone());

    let batch_size = 32;
    let states = Tensor::randn(&[batch_size, 8], &device).unwrap();
    let actions = Tensor::randn(&[batch_size, 3], &device).unwrap();

    let q_values = sac
        .critic_forward(&states, &actions, 0)
        .await
        .expect("Critic forward should work");

    assert_eq!(q_values.shape(), &[batch_size]);
}

#[tokio::test]
async fn test_sac_target_critic_forward() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig::default();
    let sac = SAC::new(config, 8, 3, 256, device.clone());

    let batch_size = 32;
    let states = Tensor::randn(&[batch_size, 8], &device).unwrap();
    let actions = Tensor::randn(&[batch_size, 3], &device).unwrap();

    let q_values = sac
        .target_critic_forward(&states, &actions, 0)
        .await
        .expect("Target critic forward should work");

    assert_eq!(q_values.shape(), &[batch_size]);
}

#[tokio::test]
#[ignore = "rtx-rl SAC loss shape mismatch"]
async fn test_sac_compute_actor_loss() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig::default();
    let sac = SAC::new(config, 8, 3, 256, device.clone());

    let batch_size = 64;
    let states = Tensor::randn(&[batch_size, 8], &device).unwrap();

    let actor_loss = sac
        .compute_actor_loss(&states)
        .await
        .expect("Should compute actor loss");

    assert_eq!(actor_loss.shape(), &[]);

    let loss_val: f32 = actor_loss.item().expect("Should get scalar");
    assert!(loss_val.is_finite());
}

#[tokio::test]
#[ignore = "rtx-tensor Bool dtype CPU copy not supported"]
async fn test_sac_compute_critic_loss() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig::default();
    let sac = SAC::new(config, 8, 3, 256, device.clone());

    let batch_size = 64;
    let states = Tensor::randn(&[batch_size, 8], &device).unwrap();
    let actions = Tensor::randn(&[batch_size, 3], &device).unwrap();
    let rewards = Tensor::randn(&[batch_size], &device).unwrap();
    let next_states = Tensor::randn(&[batch_size, 8], &device).unwrap();
    let dones = Tensor::zeros(&[batch_size], &device)
        .unwrap()
        .to_dtype(DType::Bool)
        .unwrap();

    let (critic1_loss, critic2_loss) = sac
        .compute_critic_loss(&states, &actions, &rewards, &next_states, &dones)
        .await
        .expect("Should compute critic losses");

    assert_eq!(critic1_loss.shape(), &[]);
    assert_eq!(critic2_loss.shape(), &[]);

    let loss1_val: f32 = critic1_loss.item().expect("Should get scalar");
    let loss2_val: f32 = critic2_loss.item().expect("Should get scalar");

    assert!(loss1_val >= 0.0); // MSE loss should be non-negative
    assert!(loss2_val >= 0.0);
}

#[tokio::test]
#[ignore = "rtx-rl SAC temperature implementation incomplete"]
async fn test_sac_compute_temperature_loss() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig {
        automatic_entropy_tuning: true,
        target_entropy: Some(-3.0),
        ..SACConfig::default()
    };
    let sac = SAC::new(config, 8, 3, 256, device.clone());

    let batch_size = 64;
    let log_probs = Tensor::randn(&[batch_size], &device)
        .unwrap()
        .scalar_mul(-2.0)
        .unwrap(); // Negative log probs

    let temperature_loss = sac
        .compute_temperature_loss(&log_probs)
        .await
        .expect("Should compute temperature loss");

    assert_eq!(temperature_loss.shape(), &[]);

    let loss_val: f32 = temperature_loss.item().expect("Should get scalar");
    assert!(loss_val.is_finite());
}

#[tokio::test]
#[ignore = "rtx-tensor Bool dtype CPU copy not supported"]
async fn test_sac_soft_update() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig {
        tau: 0.1, // Larger tau for more noticeable updates
        ..SACConfig::default()
    };
    let mut sac = SAC::new(config, 4, 2, 64, device.clone());

    // Get initial target parameters
    let initial_target_params = sac
        .get_target_critic_params()
        .await
        .expect("Should get initial params");

    // Modify main critic parameters by training on some dummy data
    let batch_size = 32;
    let states = Tensor::randn(&[batch_size, 4], &device).unwrap();
    let actions = Tensor::randn(&[batch_size, 2], &device).unwrap();
    let rewards = Tensor::ones(&[batch_size], &device).unwrap();
    let next_states = Tensor::randn(&[batch_size, 4], &device).unwrap();
    let dones = Tensor::zeros(&[batch_size], &device)
        .unwrap()
        .to_dtype(DType::Bool)
        .unwrap();

    // Perform one update step
    let _metrics = sac
        .update(&states, &actions, &rewards, &next_states, &dones)
        .await
        .expect("Update should succeed");

    // Verify target networks were soft-updated
    let updated_target_params = sac
        .get_target_critic_params()
        .await
        .expect("Should get updated params");

    // Target parameters should have changed but not completely
    assert_ne!(initial_target_params.len(), 0);
    assert_eq!(initial_target_params.len(), updated_target_params.len());
}

#[tokio::test]
#[ignore = "rtx-rl SAC deterministic action incomplete"]
async fn test_sac_deterministic_action() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig::default();
    let sac = SAC::new(config, 4, 2, 64, device.clone());

    let state = Tensor::randn(&[1, 4], &device).unwrap();

    let (action1, _) = sac
        .actor_forward(&state, true)
        .await
        .expect("Should get deterministic action");
    let (action2, _) = sac
        .actor_forward(&state, true)
        .await
        .expect("Should get deterministic action");

    // Deterministic actions should be identical for same state
    let action1_data: Vec<f32> = action1.to_vec().expect("Should convert");
    let action2_data: Vec<f32> = action2.to_vec().expect("Should convert");

    for (a1, a2) in action1_data.iter().zip(action2_data.iter()) {
        assert!((a1 - a2).abs() < 1e-6);
    }
}

#[tokio::test]
async fn test_sac_stochastic_action() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig::default();
    let sac = SAC::new(config, 4, 2, 64, device.clone());

    let state = Tensor::randn(&[1, 4], &device).unwrap();

    let (action1, _) = sac
        .actor_forward(&state, false)
        .await
        .expect("Should get stochastic action");
    let (action2, _) = sac
        .actor_forward(&state, false)
        .await
        .expect("Should get stochastic action");

    // Stochastic actions should be different for same state
    let action1_data: Vec<f32> = action1.to_vec().expect("Should convert");
    let action2_data: Vec<f32> = action2.to_vec().expect("Should convert");

    let mut different = false;
    for (a1, a2) in action1_data.iter().zip(action2_data.iter()) {
        if (a1 - a2).abs() > 1e-3 {
            different = true;
            break;
        }
    }
    assert!(different, "Stochastic actions should be different");
}

#[tokio::test]
#[ignore = "rtx-tensor Bool dtype CPU copy not supported"]
async fn test_sac_full_update_cycle() {
    let device = Device::cuda(0).unwrap_or_else(|_| Device::cpu());
    let config = SACConfig {
        batch_size: 32,
        target_update_interval: 1,
        ..SACConfig::default()
    };
    let mut sac = SAC::new(config, 4, 2, 64, device.clone());

    let batch_size = 32;
    let states = Tensor::randn(&[batch_size, 4], &device).unwrap();
    let actions = Tensor::randn(&[batch_size, 2], &device).unwrap();
    let rewards = Tensor::randn(&[batch_size], &device).unwrap();
    let next_states = Tensor::randn(&[batch_size, 4], &device).unwrap();
    let dones = Tensor::zeros(&[batch_size], &device)
        .unwrap()
        .to_dtype(DType::Bool)
        .unwrap();

    let metrics = sac
        .update(&states, &actions, &rewards, &next_states, &dones)
        .await
        .expect("Full update should succeed");

    assert!(metrics.actor_loss.is_finite());
    assert!(metrics.critic1_loss.is_finite());
    assert!(metrics.critic2_loss.is_finite());
    assert!(metrics.temperature_loss.is_finite());
    assert!(metrics.alpha > 0.0);
    assert!(metrics.target_entropy.is_finite());
}