//! Minimal test to understand Var vs Tensor gradient flow with scatter_add

use candle_core::{Device, DType, Tensor, Var};
use ml::MLError;

#[test]
fn test_scatter_add_tensor_vs_var() -> Result<(), MLError> {
    let device = Device::cuda_if_available(0)?;

    println!("\n=== Testing Tensor::zeros + scatter_add ===");
    {
        // Input that should have gradients
        let input = Tensor::ones((2, 3), DType::F32, &device)?;

        // Base for scatter (using Tensor::zeros)
        let base = Tensor::zeros((2, 3), DType::F32, &device)?;
        let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;

        // Scatter
        let result = base.scatter_add(&indices, &input, 1)?;

        // Compute loss and backward
        let loss = result.sum_all()?;
        let grads = loss.backward()?;

        let has_grads = grads.get(&input).is_some();
        println!("Tensor::zeros -> has gradients: {}", has_grads);
    }

    println!("\n=== Testing Var::zeros + scatter_add ===");
    {
        // Input that should have gradients
        let input = Tensor::ones((2, 3), DType::F32, &device)?;

        // Base for scatter (using Var::zeros)
        let base_var = Var::zeros((2, 3), DType::F32, &device)?;
        let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;

        // Scatter - need to convert Var to Tensor for scatter_add
        let base_tensor = base_var.as_tensor();
        let result = base_tensor.scatter_add(&indices, &input, 1)?;

        // Compute loss and backward
        let loss = result.sum_all()?;
        let grads = loss.backward()?;

        let has_grads = grads.get(&input).is_some();
        println!("Var::zeros -> has gradients: {}", has_grads);
    }

    println!("\n=== Testing Var::from_tensor (input) + scatter_add ===");
    {
        // Input wrapped in Var
        let input_tensor = Tensor::ones((2, 3), DType::F32, &device)?;
        let input_var = Var::from_tensor(&input_tensor)?;

        // Base for scatter (regular Tensor)
        let base = Tensor::zeros((2, 3), DType::F32, &device)?;
        let indices = Tensor::new(&[[0i64, 1i64, 2i64], [0i64, 1i64, 2i64]], &device)?;

        // Scatter - use Var as Tensor
        let result = base.scatter_add(&indices, &input_var.as_tensor(), 1)?;

        // Compute loss and backward
        let loss = result.sum_all()?;
        let grads = loss.backward()?;

        let has_grads_tensor = grads.get(&input_tensor).is_some();
        let has_grads_var = grads.get(input_var.as_tensor()).is_some();
        println!("Var(input) -> has gradients on tensor: {}", has_grads_tensor);
        println!("Var(input) -> has gradients on var: {}", has_grads_var);
    }

    Ok(())
}