use tuphium::error::TuphiumError;
use tuphium::optimizers::{Optimizer, SGD};

#[test]
fn sgd_update() {
    let mut opt = SGD { lr: 0.1 };
    let mut params = vec![1.0, 2.0, 3.0];
    let grads = vec![0.5, -1.0, 2.0];
    opt.update(&mut params, &grads).unwrap();
    assert!((params[0] - 0.95).abs() == 0.0);
    assert!((params[1] - 2.10).abs() == 0.0);
    assert!((params[2] - 2.80).abs() == 0.0);
}

#[test]
fn sgd_zero_lr() {
    let mut opt = SGD { lr: 0.0 };
    let mut params = vec![1.0, 2.0];
    let grads = vec![10.0, -20.0];
    opt.update(&mut params, &grads).unwrap();
    assert!((params[0] - 1.0).abs() == 0.0);
    assert!((params[1] - 2.0).abs() == 0.0);
}

#[test]
fn sgd_neg_lr() {
    let mut opt = SGD { lr: -0.1 };
    let mut params = vec![1.0];
    let grads = vec![5.0];
    opt.update(&mut params, &grads).unwrap();
    assert!((params[0] - 1.5).abs() == 0.0);
}

#[test]
fn sgd_empty() {
    let mut opt = SGD { lr: 1.0 };
    let mut params: Vec<f32> = vec![];
    let grads: Vec<f32> = vec![];
    opt.update(&mut params, &grads).unwrap();
    assert!(params.is_empty());
}

#[test]
fn update_returns_length_mismatch() {
    let mut opt = SGD { lr: 0.1 };
    let mut params = vec![1.0, 2.0];
    let grads = vec![1.0];
    let error = opt.update(&mut params, &grads).unwrap_err();
    assert_eq!(
        error,
        TuphiumError::LengthMismatch {
            operation: "SGD::update",
            expected: 2,
            actual: 1,
        }
    );
}

#[test]
fn update_returns_non_finite_value() {
    let mut opt = SGD { lr: f32::INFINITY };
    let mut params = vec![1.0];
    let grads = vec![1.0];
    let error = opt.update(&mut params, &grads).unwrap_err();
    assert_eq!(
        error,
        TuphiumError::NonFiniteValue {
            operation: "SGD::update",
            value: f32::INFINITY,
        }
    );
}
