// Simple standalone test for LogCoshLoss
use std::process::exit;

// Mock types for standalone testing
#[derive(Debug, Clone)]
struct Device;

impl Device {
    fn cpu() -> Self { Device }
}

#[derive(Debug, Clone)]
struct Tensor {
    data: Vec<f32>,
    shape: Vec<usize>,
    device: Device,
}

impl Tensor {
    fn from_data(data: Vec<f32>, shape: Vec<usize>, _device: &Device) -> Result<Self, String> {
        if data.len() != shape.iter().product::<usize>() {
            return Err("Shape mismatch".to_string());
        }
        Ok(Tensor { data, shape, device: Device::cpu() })
    }
    
    fn shape(&self) -> &[usize] { &self.shape }
    fn device(&self) -> &Device { &self.device }
    fn to_cpu(&self) -> Result<Vec<f32>, String> { Ok(self.data.clone()) }
    fn numel(&self) -> usize { self.data.len() }
    
    fn sum(&self, _dim: Option<usize>) -> Result<Self, String> {
        let sum: f32 = self.data.iter().sum();
        Ok(Tensor::from_data(vec![sum], vec![], &Device::cpu())?)
    }
    
    fn div_scalar(&self, scalar: f32) -> Result<Self, String> {
        let result_data: Vec<f32> = self.data.iter().map(|x| x / scalar).collect();
        Ok(Tensor::from_data(result_data, self.shape.clone(), &self.device)?)
    }
}

impl std::ops::Sub for &Tensor {
    type Output = Tensor;
    
    fn sub(self, other: &Tensor) -> Tensor {
        assert_eq!(self.data.len(), other.data.len());
        let result_data: Vec<f32> = self.data.iter().zip(other.data.iter())
            .map(|(a, b)| a - b)
            .collect();
        Tensor::from_data(result_data, self.shape.clone(), &self.device).unwrap()
    }
}

type Result<T> = std::result::Result<T, String>;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Reduction {
    None,
    Mean,
    Sum,
}

impl Default for Reduction {
    fn default() -> Self { Self::Mean }
}

trait Loss {
    fn forward(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor>;
    fn reduction(&self) -> Reduction;
    fn supports_backprop(&self) -> bool { true }
    fn name(&self) -> &'static str;
}

// LogCoshLoss implementation
#[derive(Debug, Clone)]
struct LogCoshLoss {
    reduction: Reduction,
}

impl LogCoshLoss {
    fn new() -> Self {
        Self {
            reduction: Reduction::Mean,
        }
    }
    
    fn with_reduction(mut self, reduction: Reduction) -> Self {
        self.reduction = reduction;
        self
    }
    
    fn stable_logcosh(&self, x: f32) -> f32 {
        let abs_x = x.abs();
        
        if abs_x < 12.0 {
            x.cosh().ln()
        } else {
            abs_x - 2.0f32.ln()
        }
    }
    
    fn compute_logcosh_elementwise(&self, errors: &Tensor) -> Result<Tensor> {
        let error_data = errors.to_cpu()?;
        
        let mut result_data = Vec::new();
        for &error in &error_data {
            let logcosh_value = self.stable_logcosh(error);
            result_data.push(logcosh_value);
        }
        
        Tensor::from_data(result_data, errors.shape().to_vec(), errors.device())
    }
    
    fn apply_reduction(&self, losses: &Tensor) -> Result<Tensor> {
        match self.reduction {
            Reduction::None => Ok(losses.clone()),
            Reduction::Mean => {
                let sum_result = losses.sum(None)?;
                let numel = losses.numel() as f32;
                sum_result.div_scalar(numel)
            },
            Reduction::Sum => losses.sum(None),
        }
    }
}

impl Loss for LogCoshLoss {
    fn forward(&self, predictions: &Tensor, targets: &Tensor) -> Result<Tensor> {
        // Validate input shapes
        if predictions.shape() != targets.shape() {
            return Err(format!(
                "Shape mismatch: predictions {:?} vs targets {:?}",
                predictions.shape(),
                targets.shape()
            ));
        }

        // Compute element-wise errors
        let errors = predictions - targets;
        
        // Compute LogCosh loss element-wise
        let logcosh_values = self.compute_logcosh_elementwise(&errors)?;
        
        // Apply reduction
        self.apply_reduction(&logcosh_values)
    }

    fn reduction(&self) -> Reduction {
        self.reduction
    }

    fn supports_backprop(&self) -> bool {
        true
    }

    fn name(&self) -> &'static str {
        "LogCoshLoss"
    }
}

impl Default for LogCoshLoss {
    fn default() -> Self {
        Self::new()
    }
}

fn main() {
    // Test LogCoshLoss implementation
    println!("Testing LogCoshLoss implementation...");
    
    let device = Device::cpu();
    let loss = LogCoshLoss::new();
    
    // Test 1: Zero error
    println!("\nTest 1: Zero error");
    let predictions = Tensor::from_data(vec![1.0, 2.0, 3.0], vec![3], &device).unwrap();
    let targets = predictions.clone();
    
    match loss.forward(&predictions, &targets) {
        Ok(result) => {
            println!("Loss value: {}", result.data[0]);
            if result.data[0].abs() < 1e-6 {
                println!("✓ Zero error test passed");
            } else {
                println!("✗ Zero error test failed: expected ~0.0, got {}", result.data[0]);
            }
        }
        Err(e) => {
            println!("✗ Zero error test failed: {}", e);
            exit(1);
        }
    }
    
    // Test 2: Small error (should be approximately quadratic)
    println!("\nTest 2: Small error");
    let predictions = Tensor::from_data(vec![0.0], vec![1], &device).unwrap();
    let targets = Tensor::from_data(vec![0.1], vec![1], &device).unwrap();
    
    match loss.forward(&predictions, &targets) {
        Ok(result) => {
            let expected_quad = 0.5 * 0.1 * 0.1; // x²/2
            let actual_logcosh = 0.1f32.cosh().ln();
            println!("Loss value: {}", result.data[0]);
            println!("Expected quadratic approximation: {}", expected_quad);
            println!("Actual log(cosh(0.1)): {}", actual_logcosh);
            
            if (result.data[0] - actual_logcosh).abs() < 1e-6 {
                println!("✓ Small error test passed");
            } else {
                println!("✗ Small error test failed");
            }
        }
        Err(e) => {
            println!("✗ Small error test failed: {}", e);
            exit(1);
        }
    }
    
    // Test 3: Large error (should be approximately linear)
    println!("\nTest 3: Large error");
    let predictions = Tensor::from_data(vec![0.0], vec![1], &device).unwrap();
    let targets = Tensor::from_data(vec![10.0], vec![1], &device).unwrap();
    
    match loss.forward(&predictions, &targets) {
        Ok(result) => {
            let expected_linear = 10.0 - 2.0f32.ln(); // |x| - log(2)
            println!("Loss value: {}", result.data[0]);
            println!("Expected linear approximation: {}", expected_linear);
            
            if (result.data[0] - expected_linear).abs() < 1e-2 {
                println!("✓ Large error test passed");
            } else {
                println!("✗ Large error test failed");
            }
        }
        Err(e) => {
            println!("✗ Large error test failed: {}", e);
            exit(1);
        }
    }
    
    // Test 4: Reduction modes
    println!("\nTest 4: Reduction modes");
    let predictions = Tensor::from_data(vec![0.0, 1.0, 2.0, 3.0], vec![4], &device).unwrap();
    let targets = Tensor::from_data(vec![0.5, 0.5, 3.0, 2.0], vec![4], &device).unwrap();
    
    let loss_mean = LogCoshLoss::new().with_reduction(Reduction::Mean);
    let loss_sum = LogCoshLoss::new().with_reduction(Reduction::Sum);
    let loss_none = LogCoshLoss::new().with_reduction(Reduction::None);
    
    let mean_result = loss_mean.forward(&predictions, &targets).unwrap();
    let sum_result = loss_sum.forward(&predictions, &targets).unwrap();
    let none_result = loss_none.forward(&predictions, &targets).unwrap();
    
    println!("Mean result: {}", mean_result.data[0]);
    println!("Sum result: {}", sum_result.data[0]);
    println!("None result: {:?}", none_result.data);
    
    // Check that sum = mean * num_elements
    let expected_sum = mean_result.data[0] * 4.0;
    if (sum_result.data[0] - expected_sum).abs() < 1e-6 {
        println!("✓ Reduction modes test passed");
    } else {
        println!("✗ Reduction modes test failed");
    }
    
    // Test 5: Utility functions
    println!("\nTest 5: Utility functions");
    let loss = LogCoshLoss::new();
    
    assert!(loss.is_in_quadratic_region(0.5));
    assert!(!loss.is_in_quadratic_region(2.0));
    assert!(loss.is_in_linear_region(5.0));
    assert!(!loss.is_in_linear_region(0.5));
    
    println!("✓ Utility functions test passed");
    
    println!("\n🎉 All tests passed! LogCoshLoss implementation is working correctly.");
}