//! TDD tests for Variable gradient computation
//! Following strict TDD approach - implementing gradient functionality

use rtx_science::types::Variable;
use rtx_tensor::{Device, Tensor};

#[test]
fn test_backward_grad_with_inputs() {
    let device = Device::cpu();
    let data = Tensor::ones(&[2, 3], &device).unwrap();
    let variable = Variable::new(data, true);

    let inputs = Tensor::zeros(&[2, 2], &device).unwrap();

    // Should compute gradient with respect to inputs at index 0
    let grad = variable
        .backward_grad_with_inputs(&inputs, Some(0))
        .unwrap();
    assert_eq!(grad.value().shape().dims(), &[2, 3]);
}

#[test]
fn test_backward_grad_returns_variable() {
    let device = Device::cpu();
    let data = Tensor::ones(&[2, 3], &device).unwrap();
    let variable = Variable::new(data, true);

    let inputs = Tensor::zeros(&[2, 2], &device).unwrap();

    // Should return a Variable, not ()
    let grad = variable
        .backward_grad_with_inputs(&inputs, Some(0))
        .unwrap();
    assert!(grad.requires_grad() == false); // Gradients don't require gradients
}

#[test]
fn test_simple_backward() {
    let device = Device::cpu();
    let data = Tensor::ones(&[1], &device).unwrap();
    let variable = Variable::new(data, true);

    // Simple backward should succeed
    variable.backward().unwrap();
}