use tuphium::error::TuphiumError;
use tuphium::losses::{Loss, BinaryCrossEntropy, MeanSquaredError, CrossEntropy};

#[test]
fn mse_forward_and_backward_work() {
    let loss = MeanSquaredError;
    assert_eq!(loss.forward(&[1.0, 0.0], &[1.0, 0.0]), Ok(0.0));
    assert!((loss.forward(&[0.0], &[2.0]).unwrap() - 4.0).abs() < 1e-5);
    let g = loss.backward(&[0.0], &[1.0]).unwrap();
    assert!((g[0] - 2.0).abs() < 1e-5);
}

#[test]
fn bce_forward_and_backward_work() {
    let loss = BinaryCrossEntropy;
    let v = loss.forward(&[1.0, 0.0], &[0.9999, 0.0001]).unwrap();
    assert!(v < 0.01);
    let g = loss.backward(&[1.0], &[0.3]).unwrap();
    assert!(g[0] < 0.0);
}

#[test]
fn ce_forward_and_backward_work() {
    let loss = CrossEntropy;
    let v = loss
        .forward(&[1.0, 0.0, 0.0], &[0.9999, 0.00005, 0.00005])
        .unwrap();
    assert!(v < 0.01);
    let g = loss.backward(&[0.0, 1.0], &[0.3, 0.7]).unwrap();
    assert_eq!(g[0], 0.0);
}

#[test]
fn loss_works_as_trait_object() {
    let losses: Vec<Box<dyn Loss>> = vec![Box::new(MeanSquaredError), Box::new(BinaryCrossEntropy)];
    for loss in &losses {
        assert!(loss.forward(&[1.0], &[0.8]).unwrap() >= 0.0);
        assert_eq!(loss.backward(&[1.0], &[0.8]).unwrap().len(), 1);
    }
}

#[test]
fn forward_returns_shape_mismatch() {
    let loss = MeanSquaredError;
    assert_eq!(
        loss.forward(&[1.0, 2.0], &[1.0]).unwrap_err(),
        TuphiumError::ShapeMismatch {
            operation: "MeanSquaredError::forward",
            expected: vec![2],
            actual: vec![1],
        }
    );
}

#[test]
fn forward_returns_empty_data() {
    let loss = BinaryCrossEntropy;
    assert_eq!(
        loss.forward(&[], &[]).unwrap_err(),
        TuphiumError::EmptyData {
            operation: "BinaryCrossEntropy::forward",
            target: "loss inputs".to_string(),
        }
    );
}

#[test]
fn backward_returns_non_finite_value() {
    let loss = CrossEntropy;
    assert!(matches!(
        loss.backward(&[1.0], &[f32::NAN]).unwrap_err(),
        TuphiumError::NonFiniteValue {
            operation: "CrossEntropy::backward",
            ..
        }
    ));
}
