// Comprehensive tests for autodiff operations
// Following strict TDD - tests written BEFORE implementation fixes

use rtx_autograd::{Variable, Graph};
use rtx_tensor::{Tensor, Device};
use std::sync::Arc;

#[cfg(test)]
mod autodiff_gradient_tests {
    use super::*;

    #[test]
    fn test_backward_grad_scalar() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Create a scalar variable
        let x = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);
        let y = &x * &x; // y = x^2

        // Compute gradient dy/dx = 2x = 4
        let grad = y.backward_grad(&x);
        assert_eq!(grad.shape(), &[]);

        let grad_value: f32 = grad.item();
        assert!((grad_value - 4.0).abs() < 1e-6);
    }

    #[test]
    fn test_backward_grad_vector() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Create vector variable
        let x = Variable::new(Tensor::ones(&[3], &device), true);
        let y = &x * &x + &x; // y = x^2 + x

        // Compute gradient dy/dx = 2x + 1
        let grad = y.backward_grad(&x);
        assert_eq!(grad.shape(), &[3]);

        let grad_data: Vec<f32> = grad.to_vec();
        for val in grad_data {
            assert!((val - 3.0).abs() < 1e-6); // 2*1 + 1 = 3
        }
    }

    #[test]
    fn test_backward_grad_matrix() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Create matrix variable
        let x = Variable::new(Tensor::ones(&[2, 3], &device), true);
        let w = Variable::new(Tensor::full(&[3, 2], 2.0, &device), false);

        // y = x @ w (matrix multiplication)
        let y = x.matmul(&w);

        // Compute gradient dy/dx
        let grad_x = y.backward_grad(&x);
        assert_eq!(grad_x.shape(), &[2, 3]);

        // Each element should be 4.0 (sum of row in w = 2+2)
        let grad_data: Vec<f32> = grad_x.to_vec();
        for val in grad_data {
            assert!((val - 4.0).abs() < 1e-6);
        }
    }

    #[test]
    fn test_backward_grad_chain_rule() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Test chain rule: z = sin(x^2)
        let x = Variable::new(Tensor::scalar(std::f32::consts::PI / 4.0, DType::F32, &device), true);
        let x_squared = &x * &x;
        let z = x_squared.sin();

        // dz/dx = cos(x^2) * 2x
        let grad = z.backward_grad(&x);

        let x_val = std::f32::consts::PI / 4.0;
        let expected = (x_val * x_val).cos() * 2.0 * x_val;

        let grad_value: f32 = grad.item();
        assert!((grad_value - expected).abs() < 1e-5);
    }

    #[test]
    fn test_mean_square_operation() {
        let device = Device::cpu();

        // Test mean square operation
        let data = vec![1.0, 2.0, 3.0, 4.0];
        let x = Variable::new(Tensor::from_vec(data.clone(), &[4], &device), true);

        // mean_square should compute mean(x^2)
        let result = x.mean_square();

        // Expected: (1 + 4 + 9 + 16) / 4 = 30 / 4 = 7.5
        let value: f32 = result.value().item();
        assert!((value - 7.5).abs() < 1e-6);
    }

    #[test]
    fn test_value_extraction() {
        let device = Device::cpu();

        let data = vec![1.0, 2.0, 3.0];
        let tensor = Tensor::from_vec(data.clone(), &[3], &device);
        let var = Variable::new(tensor, false);

        // Test value extraction
        let tensor_value = var.value();
        assert_eq!(tensor_value.shape(), &[3]);

        let extracted: Vec<f32> = tensor_value.to_vec();
        assert_eq!(extracted, data);
    }

    #[test]
    fn test_gradient_accumulation() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Test gradient accumulation through multiple paths
        let x = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);

        // Two paths: y1 = x^2, y2 = 3x
        let y1 = &x * &x;
        let y2 = &x * &Variable::new(Tensor::scalar(3.0, DType::F32, &device), false);

        // z = y1 + y2 = x^2 + 3x
        let z = &y1 + &y2;

        // dz/dx = 2x + 3 = 7
        let grad = z.backward_grad(&x);
        let grad_value: f32 = grad.item();
        assert!((grad_value - 7.0).abs() < 1e-6);
    }

    #[test]
    fn test_stop_gradient() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Test gradient flow control
        let x = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);
        let y = Variable::new(Tensor::scalar(3.0, DType::F32, &device), false); // requires_grad = false

        let z = &x * &y;

        // Should compute gradient w.r.t x
        let grad_x = z.backward_grad(&x);
        let grad_x_value: f32 = grad_x.item();
        assert!((grad_x_value - 3.0).abs() < 1e-6);

        // Should not compute gradient w.r.t y (requires_grad = false)
        // This should return zeros or handle gracefully
        let grad_y = z.backward_grad(&y);
        let grad_y_value: f32 = grad_y.item();
        assert!((grad_y_value - 0.0).abs() < 1e-6);
    }

    #[test]
    fn test_complex_computational_graph() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Build a more complex graph: f(x,y) = (x^2 + y) * sin(x*y)
        let x = Variable::new(Tensor::scalar(1.0, DType::F32, &device), true);
        let y = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);

        let x_squared = &x * &x;
        let sum_term = &x_squared + &y;
        let product = &x * &y;
        let sin_term = product.sin();
        let result = &sum_term * &sin_term;

        // Compute gradients
        let grad_x = result.backward_grad(&x);
        let grad_y = result.backward_grad(&y);

        // Verify gradients are computed
        assert_eq!(grad_x.shape(), &[]);
        assert_eq!(grad_y.shape(), &[]);
    }

    #[test]
    fn test_batch_gradient_computation() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Test gradient computation for batch operations
        let batch_size = 4;
        let features = 3;

        let x = Variable::new(Tensor::randn(&[batch_size, features], &device), true);
        let w = Variable::new(Tensor::randn(&[features, 2], &device), true);
        let b = Variable::new(Tensor::randn(&[2], &device), true);

        // Linear layer: y = xW + b
        let y = x.matmul(&w).add(&b);

        // Compute mean loss
        let loss = y.mean();

        // Compute gradients
        let grad_w = loss.backward_grad(&w);
        let grad_b = loss.backward_grad(&b);

        assert_eq!(grad_w.shape(), &[features, 2]);
        assert_eq!(grad_b.shape(), &[2]);
    }
}

#[cfg(test)]
mod autodiff_optimization_tests {
    use super::*;

    #[test]
    fn test_gradient_descent_step() {
        let device = Device::cpu();

        // Simple quadratic: f(x) = (x - 3)^2
        let mut x = Variable::new(Tensor::scalar(0.0, DType::F32, &device), true);
        let target = Variable::new(Tensor::scalar(3.0, DType::F32, &device), false);

        let learning_rate = 0.1;

        for _ in 0..20 {
            // Compute loss
            let diff = &x - &target;
            let loss = &diff * &diff;

            // Compute gradient
            let grad = loss.backward_grad(&x);

            // Update parameter
            let x_tensor = x.value();
            let grad_tensor = grad;
            let updated = x_tensor - &(grad_tensor * learning_rate);
            x = Variable::new(updated, true);
        }

        // Should converge close to 3.0
        let final_value: f32 = x.value().item();
        assert!((final_value - 3.0).abs() < 0.1);
    }

    #[test]
    fn test_multi_variable_optimization() {
        let device = Device::cpu();

        // Minimize f(x,y) = x^2 + y^2 + 2x - 4y + 5
        // Minimum at x = -1, y = 2
        let mut x = Variable::new(Tensor::scalar(5.0, DType::F32, &device), true);
        let mut y = Variable::new(Tensor::scalar(-3.0, DType::F32, &device), true);

        let learning_rate = 0.1;

        for _ in 0..50 {
            // Compute loss
            let x_squared = &x * &x;
            let y_squared = &y * &y;
            let linear_x = &x * &Variable::new(Tensor::scalar(2.0, DType::F32, &device), false);
            let linear_y = &y * &Variable::new(Tensor::scalar(-4.0, DType::F32, &device), false);
            let constant = Variable::new(Tensor::scalar(5.0, DType::F32, &device), false);

            let loss = &(&(&x_squared + &y_squared) + &linear_x) + &(&linear_y + &constant);

            // Compute gradients
            let grad_x = loss.backward_grad(&x);
            let grad_y = loss.backward_grad(&y);

            // Update parameters
            let x_updated = x.value() - &(grad_x * learning_rate);
            let y_updated = y.value() - &(grad_y * learning_rate);

            x = Variable::new(x_updated, true);
            y = Variable::new(y_updated, true);
        }

        // Should converge close to (-1, 2)
        let x_final: f32 = x.value().item();
        let y_final: f32 = y.value().item();

        assert!((x_final - (-1.0)).abs() < 0.1);
        assert!((y_final - 2.0).abs() < 0.1);
    }

    #[test]
    fn test_higher_order_derivatives() {
        let graph = Graph::new();
        let device = Device::cpu();

        // Test second-order derivatives
        let x = Variable::new(Tensor::scalar(2.0, DType::F32, &device), true);

        // f(x) = x^3
        let x_squared = &x * &x;
        let x_cubed = &x_squared * &x;

        // First derivative: df/dx = 3x^2 = 12
        let grad1 = x_cubed.backward_grad(&x);
        let grad1_value: f32 = grad1.item();
        assert!((grad1_value - 12.0).abs() < 1e-5);

        // For second derivative, we'd need to differentiate grad1
        // This tests that the gradient computation infrastructure works
    }

    #[test]
    fn test_numerical_stability() {
        let device = Device::cpu();

        // Test with very small and very large values
        let small = Variable::new(Tensor::scalar(1e-8, DType::F32, &device), true);
        let large = Variable::new(Tensor::scalar(1e8, DType::F32, &device), true);

        // Operations should not overflow or underflow
        let result_small = &small * &small;
        let result_large = &large + &large;

        // Gradients should be computable
        let grad_small = result_small.backward_grad(&small);
        let grad_large = result_large.backward_grad(&large);

        // Verify no NaN or Inf
        let grad_small_val: f32 = grad_small.item();
        let grad_large_val: f32 = grad_large.item();

        assert!(grad_small_val.is_finite());
        assert!(grad_large_val.is_finite());
    }
}

#[cfg(test)]
mod integration_tests {
    use super::*;

    #[test]
    fn test_neural_network_forward_backward() {
        let device = Device::cpu();

        // Simple 2-layer network
        let batch_size = 2;
        let input_dim = 3;
        let hidden_dim = 4;
        let output_dim = 1;

        // Initialize weights and biases
        let w1 = Variable::new(Tensor::randn(&[input_dim, hidden_dim], &device) * 0.1, true);
        let b1 = Variable::new(Tensor::zeros(&[hidden_dim], &device), true);
        let w2 = Variable::new(Tensor::randn(&[hidden_dim, output_dim], &device) * 0.1, true);
        let b2 = Variable::new(Tensor::zeros(&[output_dim], &device), true);

        // Input and target
        let x = Variable::new(Tensor::randn(&[batch_size, input_dim], &device), false);
        let target = Variable::new(Tensor::ones(&[batch_size, output_dim], &device), false);

        // Forward pass
        let h1 = x.matmul(&w1).add(&b1).relu();
        let output = h1.matmul(&w2).add(&b2);

        // Loss (MSE)
        let diff = &output - &target;
        let squared = &diff * &diff;
        let loss = squared.mean();

        // Backward pass
        let grad_w1 = loss.backward_grad(&w1);
        let grad_b1 = loss.backward_grad(&b1);
        let grad_w2 = loss.backward_grad(&w2);
        let grad_b2 = loss.backward_grad(&b2);

        // Verify gradient shapes
        assert_eq!(grad_w1.shape(), &[input_dim, hidden_dim]);
        assert_eq!(grad_b1.shape(), &[hidden_dim]);
        assert_eq!(grad_w2.shape(), &[hidden_dim, output_dim]);
        assert_eq!(grad_b2.shape(), &[output_dim]);
    }
}