rustyml 0.14.0

//! Integration tests for the optimizer objects: SGD, Adam, AdamW, RMSprop, and AdaGrad.
//!
//! Coverage:
//! - Constructor validation: learning_rate or epsilon at or below 0, NaN, or Inf. Adam and
//!   RMSprop betas outside [0, 1), with boundary values 0.0 (valid) and 1.0 (invalid) checked
//!   explicitly.
//! - End-to-end convergence: each optimizer drives a single Dense layer's MSE loss strictly
//!   down over a fixed number of epochs on a fixed, seeded regression problem.
//! - Multi-layer convergence: a 2-Dense-layer net with Adam verifies per-layer state buffers
//!   are allocated correctly (both layers updated, loss falls).
//!
//! Gradient correctness lives in `gradient_check.rs`.

use approx::assert_abs_diff_eq;
use ndarray::{Array, Array2, ArrayD};
use rustyml::error::Error;
use rustyml::neural_network::Tensor;
use rustyml::neural_network::layers::activation::linear::Linear;
use rustyml::neural_network::layers::dense::Dense;
use rustyml::neural_network::layers::layer_weight::LayerWeight;
use rustyml::neural_network::layers::regularization::normalization::batch_normalization::BatchNormalization;
use rustyml::neural_network::losses::mean_squared_error::MeanSquaredError;
use rustyml::neural_network::optimizers::AdaGrad;
use rustyml::neural_network::optimizers::Adam;
use rustyml::neural_network::optimizers::AdamW;
use rustyml::neural_network::optimizers::RMSprop;
use rustyml::neural_network::optimizers::SGD;
use rustyml::neural_network::sequential::Sequential;
use rustyml::neural_network::traits::{Layer, Optimizer};

// Helper: simple regression problem

/// Fixed, deterministic (x, y) pair for a tiny 1-input -> 1-output regression
///
/// Target y = 2*x over 4 samples. With identity weights (w=1, b=0), the initial MSE is 1.875
fn regression_data() -> (Tensor, Tensor) {
    let x = Array::from_shape_vec((4, 1), vec![0.5_f32, 1.0, 1.5, 2.0])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![1.0_f32, 2.0, 3.0, 4.0])
        .unwrap()
        .into_dyn();
    (x, y)
}

/// Dense(1->1, Linear) layer with weight=1, bias=0, acting as a passthrough
///
/// Gives a known starting loss of 1.875 on the regression_data() problem
fn identity_dense() -> Dense {
    let mut layer = Dense::new(1, 1, Linear::new()).unwrap();
    let w = Array::from_shape_vec((1, 1), vec![1.0_f32]).unwrap();
    let b = Array::from_shape_vec((1, 1), vec![0.0_f32]).unwrap();
    layer.set_weights(w, b).unwrap();
    layer
}

/// MSE loss for `model.predict(x)` against `y`: mean((pred - y)^2)
///
/// Matches MeanSquaredError::compute_loss so the comparison is like-for-like
fn eval_mse(model: &Sequential, x: &Tensor, y: &Tensor) -> f32 {
    let pred = model.predict(x).unwrap();
    let diff = &pred - y;
    let sq: Tensor = diff.mapv(|v| v * v);
    sq.sum() / sq.len() as f32
}

// SGD: constructor validation

#[test]
fn sgd_rejects_invalid_learning_rate() {
    for lr in [0.0_f32, -0.1, f32::INFINITY, f32::NAN] {
        assert!(
            matches!(
                SGD::new(lr, 0.0, false, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for learning_rate={lr:?}"
        );
    }
}

#[test]
fn sgd_accepts_valid_learning_rate() {
    assert!(SGD::new(0.01, 0.0, false, 0.0).is_ok());
    assert!(SGD::new(1.0, 0.0, false, 0.0).is_ok());
    // smallest positive finite f32
    assert!(SGD::new(f32::MIN_POSITIVE, 0.0, false, 0.0).is_ok());
}

// Adam: constructor validation

#[test]
fn adam_rejects_invalid_learning_rate() {
    for lr in [0.0_f32, -1e-3, f32::INFINITY, f32::NAN] {
        assert!(
            matches!(
                Adam::new(lr, 0.9, 0.999, 1e-8, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for learning_rate={lr:?}"
        );
    }
}

/// beta1 must lie in [0, 1): 0.0 (inclusive lower) is accepted. 1.0 (exclusive upper),
/// out-of-range values, and NaN are rejected
#[test]
fn adam_beta1_bounds() {
    assert!(
        Adam::new(0.001, 0.0, 0.999, 1e-8, 0.0).is_ok(),
        "beta1=0.0 (inclusive lower bound) is accepted"
    );
    for beta1 in [1.0_f32, 1.1, -0.1, f32::NAN] {
        assert!(
            matches!(
                Adam::new(0.001, beta1, 0.999, 1e-8, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for beta1={beta1:?}"
        );
    }
}

/// beta2 must lie in [0, 1): 0.0 (inclusive lower) is accepted. 1.0 (exclusive upper)
/// and NaN are rejected
#[test]
fn adam_beta2_bounds() {
    assert!(
        Adam::new(0.001, 0.9, 0.0, 1e-8, 0.0).is_ok(),
        "beta2=0.0 (inclusive lower bound) is accepted"
    );
    for beta2 in [1.0_f32, f32::NAN] {
        assert!(
            matches!(
                Adam::new(0.001, 0.9, beta2, 1e-8, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for beta2={beta2:?}"
        );
    }
}

#[test]
fn adam_rejects_invalid_epsilon() {
    for eps in [0.0_f32, -1e-8, f32::NAN, f32::INFINITY] {
        assert!(
            matches!(
                Adam::new(0.001, 0.9, 0.999, eps, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for epsilon={eps:?}"
        );
    }
}

#[test]
fn adam_accepts_valid_hyperparameters() {
    assert!(Adam::new(0.001, 0.9, 0.999, 1e-8, 0.0).is_ok());
    // typical alternative: small beta1
    assert!(Adam::new(0.01, 0.5, 0.9, 1e-6, 0.0).is_ok());
}

// RMSprop: constructor validation

#[test]
fn rmsprop_rejects_invalid_learning_rate() {
    for lr in [0.0_f32, -0.01, f32::INFINITY, f32::NAN] {
        assert!(
            matches!(
                RMSprop::new(lr, 0.9, 1e-8, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for learning_rate={lr:?}"
        );
    }
}

/// rho must lie in [0, 1): 0.0 (inclusive lower) is accepted. 1.0 (exclusive upper),
/// out-of-range values, and NaN are rejected
#[test]
fn rmsprop_rho_bounds() {
    assert!(
        RMSprop::new(0.01, 0.0, 1e-8, 0.0).is_ok(),
        "rho=0.0 (inclusive lower bound) is accepted"
    );
    for rho in [1.0_f32, 1.5, -0.5, f32::NAN] {
        assert!(
            matches!(
                RMSprop::new(0.01, rho, 1e-8, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for rho={rho:?}"
        );
    }
}

#[test]
fn rmsprop_rejects_invalid_epsilon() {
    for eps in [0.0_f32, f32::NAN, f32::INFINITY] {
        assert!(
            matches!(
                RMSprop::new(0.01, 0.9, eps, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for epsilon={eps:?}"
        );
    }
}

#[test]
fn rmsprop_accepts_valid_hyperparameters() {
    assert!(RMSprop::new(0.001, 0.9, 1e-8, 0.0).is_ok());
    assert!(RMSprop::new(0.01, 0.95, 1e-5, 0.0).is_ok());
}

// AdaGrad: constructor validation

#[test]
fn adagrad_rejects_invalid_learning_rate() {
    for lr in [0.0_f32, -0.01, f32::INFINITY, f32::NAN] {
        assert!(
            matches!(
                AdaGrad::new(lr, 1e-8, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for learning_rate={lr:?}"
        );
    }
}

#[test]
fn adagrad_rejects_invalid_epsilon() {
    for eps in [0.0_f32, -1e-8, f32::NAN, f32::INFINITY] {
        assert!(
            matches!(
                AdaGrad::new(0.01, eps, 0.0),
                Err(Error::InvalidParameter { .. })
            ),
            "expected InvalidParameter for epsilon={eps:?}"
        );
    }
}

#[test]
fn adagrad_accepts_valid_hyperparameters() {
    assert!(AdaGrad::new(0.01, 1e-8, 0.0).is_ok());
    assert!(AdaGrad::new(0.001, 1e-5, 0.0).is_ok());
}

// Known initial-loss sanity check

/// Untrained Dense(w=1, b=0) predicts y_hat = x, giving MSE = 1.875 on the fixture
#[test]
fn identity_dense_initial_mse_is_1_875() {
    let (x, y) = regression_data();
    let mut model = Sequential::new();
    model.add(identity_dense()).compile(
        SGD::new(0.01, 0.0, false, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    let mse = eval_mse(&model, &x, &y);
    assert_abs_diff_eq!(mse, 1.875_f32, epsilon = 1e-5);
}

// End-to-end convergence: each optimizer drives loss down over 20 epochs

/// SGD: loss after 20 epochs is strictly below the initial loss (1.875)
#[test]
fn sgd_single_layer_loss_decreases_over_20_epochs() {
    let (x, y) = regression_data();
    let initial_mse = 1.875_f32;

    let mut model = Sequential::new();
    // Plain SGD needs lr < 2 / lambda_max(Hessian) to converge here, about 2 / 5.5, or about
    // 0.36. A rate of 0.5 overshoots and diverges. This test uses 0.1, which reduces the loss
    // steadily.
    model.add(identity_dense()).compile(
        SGD::new(0.1, 0.0, false, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    let mse_before = eval_mse(&model, &x, &y);
    assert_abs_diff_eq!(mse_before, initial_mse, epsilon = 1e-5);

    model.fit(&x, &y, 20).unwrap();

    let mse_after = eval_mse(&model, &x, &y);
    assert!(
        mse_after < mse_before,
        "SGD: loss should decrease; before={mse_before}, after={mse_after}"
    );
}

/// Adam: loss after 20 epochs is strictly below the initial loss (1.875)
#[test]
fn adam_single_layer_loss_decreases_over_20_epochs() {
    let (x, y) = regression_data();
    let initial_mse = 1.875_f32;

    let mut model = Sequential::new();
    model.add(identity_dense()).compile(
        Adam::new(0.1, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    let mse_before = eval_mse(&model, &x, &y);
    assert_abs_diff_eq!(mse_before, initial_mse, epsilon = 1e-5);

    model.fit(&x, &y, 20).unwrap();

    let mse_after = eval_mse(&model, &x, &y);
    assert!(
        mse_after < mse_before,
        "Adam: loss should decrease; before={mse_before}, after={mse_after}"
    );
}

/// RMSprop: loss after 20 epochs is strictly below the initial loss (1.875)
#[test]
fn rmsprop_single_layer_loss_decreases_over_20_epochs() {
    let (x, y) = regression_data();
    let initial_mse = 1.875_f32;

    let mut model = Sequential::new();
    model.add(identity_dense()).compile(
        RMSprop::new(0.1, 0.9, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    let mse_before = eval_mse(&model, &x, &y);
    assert_abs_diff_eq!(mse_before, initial_mse, epsilon = 1e-5);

    model.fit(&x, &y, 20).unwrap();

    let mse_after = eval_mse(&model, &x, &y);
    assert!(
        mse_after < mse_before,
        "RMSprop: loss should decrease; before={mse_before}, after={mse_after}"
    );
}

/// AdaGrad: loss after 20 epochs is strictly below the initial loss (1.875)
#[test]
fn adagrad_single_layer_loss_decreases_over_20_epochs() {
    let (x, y) = regression_data();
    let initial_mse = 1.875_f32;

    let mut model = Sequential::new();
    model.add(identity_dense()).compile(
        AdaGrad::new(0.5, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    let mse_before = eval_mse(&model, &x, &y);
    assert_abs_diff_eq!(mse_before, initial_mse, epsilon = 1e-5);

    model.fit(&x, &y, 20).unwrap();

    let mse_after = eval_mse(&model, &x, &y);
    assert!(
        mse_after < mse_before,
        "AdaGrad: loss should decrease; before={mse_before}, after={mse_after}"
    );
}

// Multi-layer convergence with Adam (verifies per-layer state-buffer allocation)

/// Adam on Dense(1->4) -> Dense(4->1) allocates moment buffers for both layers. Loss falls:
/// below the initial value after 20 epochs, and not above the 20-epoch value after 40
#[test]
fn adam_two_layer_loss_decreases_and_buffers_allocated_correctly() {
    let (x, y) = regression_data();

    // Seed the dense layers and the fit-time shuffle so the test is deterministic
    // and never flakes on a pathological Xavier init
    const SEED: u64 = 0;
    let build_model = || -> Sequential {
        let layer1 = Dense::new(1, 4, Linear::new())
            .unwrap()
            .with_random_state(SEED);
        let layer2 = Dense::new(4, 1, Linear::new())
            .unwrap()
            .with_random_state(SEED);

        let mut model = Sequential::new_with_seed(SEED);
        model.add(layer1).add(layer2).compile(
            Adam::new(0.05, 0.9, 0.999, 1e-8, 0.0).unwrap(),
            MeanSquaredError::new(),
        );
        model
    };

    let mut model = build_model();
    let mse_initial = eval_mse(&model, &x, &y);

    model.fit(&x, &y, 20).unwrap();
    let mse_after_20 = eval_mse(&model, &x, &y);

    // 20 more epochs (total 40)
    model.fit(&x, &y, 20).unwrap();
    let mse_after_40 = eval_mse(&model, &x, &y);

    assert!(
        mse_after_20 < mse_initial,
        "Adam 2-layer: loss after 20 epochs ({mse_after_20}) should be < initial ({mse_initial})"
    );

    // Small tolerance rather than strict inequality, to absorb occasional plateaus
    assert!(
        mse_after_40 <= mse_after_20 + 1e-4,
        "Adam 2-layer: loss after 40 epochs ({mse_after_40}) should not be greater than after 20 ({mse_after_20})"
    );
}

// Multi-layer convergence: one test per remaining optimizer

/// SGD on a 2-layer net (1->4->1): loss falls and converges below 0.1 over 50 epochs
#[test]
fn sgd_two_layer_loss_decreases() {
    const SEED: u64 = 0;
    let (x, y) = regression_data();

    let layer1 = Dense::new(1, 4, Linear::new())
        .unwrap()
        .with_random_state(SEED);
    let layer2 = Dense::new(4, 1, Linear::new())
        .unwrap()
        .with_random_state(SEED);

    let mut model = Sequential::new();
    model.add(layer1).add(layer2).compile(
        SGD::new(0.05, 0.0, false, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    let mse_before = eval_mse(&model, &x, &y);
    model.fit(&x, &y, 50).unwrap();
    let mse_after = eval_mse(&model, &x, &y);

    assert!(
        mse_after < mse_before,
        "SGD 2-layer: loss should decrease; before={mse_before}, after={mse_after}"
    );
    assert!(
        mse_after < 0.1,
        "SGD 2-layer: loss should converge near 0; after={mse_after}"
    );
}

/// RMSprop on a 2-layer net (1->4->1): loss falls and converges below 0.1 over 150 epochs
#[test]
fn rmsprop_two_layer_loss_decreases() {
    const SEED: u64 = 0;
    let (x, y) = regression_data();

    let layer1 = Dense::new(1, 4, Linear::new())
        .unwrap()
        .with_random_state(SEED);
    let layer2 = Dense::new(4, 1, Linear::new())
        .unwrap()
        .with_random_state(SEED);

    let mut model = Sequential::new();
    model.add(layer1).add(layer2).compile(
        RMSprop::new(0.01, 0.9, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    let mse_before = eval_mse(&model, &x, &y);
    model.fit(&x, &y, 150).unwrap();
    let mse_after = eval_mse(&model, &x, &y);

    assert!(
        mse_after < mse_before,
        "RMSprop 2-layer: loss should decrease; before={mse_before}, after={mse_after}"
    );
    assert!(
        mse_after < 0.1,
        "RMSprop 2-layer: loss should converge near 0; after={mse_after}"
    );
}

/// AdaGrad on a 2-layer net (1->4->1): loss falls over 30 epochs
#[test]
fn adagrad_two_layer_loss_decreases() {
    const SEED: u64 = 0;
    let (x, y) = regression_data();

    let layer1 = Dense::new(1, 4, Linear::new())
        .unwrap()
        .with_random_state(SEED);
    let layer2 = Dense::new(4, 1, Linear::new())
        .unwrap()
        .with_random_state(SEED);

    let mut model = Sequential::new();
    model.add(layer1).add(layer2).compile(
        AdaGrad::new(0.5, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    let mse_before = eval_mse(&model, &x, &y);
    model.fit(&x, &y, 30).unwrap();
    let mse_after = eval_mse(&model, &x, &y);

    assert!(
        mse_after < mse_before,
        "AdaGrad 2-layer: loss should decrease; before={mse_before}, after={mse_after}"
    );
}

// Numerical value: SGD 1-step weight update on known weights

/// 1 SGD step (lr=0.01) on w=1, b=0 with x=2, y=6 yields y_hat=2.40 after refit
#[test]
fn sgd_one_step_weight_update_matches_hand_calculation() {
    let x = Array::from_shape_vec((1, 1), vec![2.0_f32])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((1, 1), vec![6.0_f32])
        .unwrap()
        .into_dyn();

    // w=1, b=0  ->  y_hat=2,  loss=(2-6)^2/1 = 16
    let w = Array::from_shape_vec((1, 1), vec![1.0_f32]).unwrap();
    let b = Array::from_shape_vec((1, 1), vec![0.0_f32]).unwrap();
    let mut layer = Dense::new(1, 1, Linear::new()).unwrap();
    layer.set_weights(w, b).unwrap();

    let mut model = Sequential::new();
    model.add(layer).compile(
        SGD::new(0.01, 0.0, false, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    // initial prediction: y_hat = 2.0
    let pred_before = model.predict(&x).unwrap();
    let val_before = *pred_before.iter().next().unwrap();
    assert_abs_diff_eq!(val_before, 2.0_f32, epsilon = 1e-6);

    model.fit(&x, &y, 1).unwrap();

    // after 1 SGD step: w_new=1.16, b_new=0.08, so y_hat = 2.0*1.16 + 0.08 = 2.40
    let pred_after = model.predict(&x).unwrap();
    let val_after = *pred_after.iter().next().unwrap();
    assert_abs_diff_eq!(val_after, 2.40_f32, epsilon = 1e-4);
}

// Clip-by-global-norm (opt-in via the `with_global_clipnorm` builder)

/// Reads the (weight, bias) scalars of a model whose first layer is a 1x1 Dense
fn dense_wb(model: &Sequential) -> (f32, f32) {
    match &model.get_weights()[0] {
        LayerWeight::Dense(d) => (d.weight[[0, 0]], d.bias[[0, 0]]),
        _ => panic!("expected Dense layer weights"),
    }
}

/// 1 clipped SGD step on the same w=1, b=0, x=2, y=6 problem as the hand-calc test above.
/// Clipping scales every gradient by max_norm / global_norm, preserving direction
#[test]
fn clip_by_global_norm_scales_sgd_step() {
    let x = Array::from_shape_vec((1, 1), vec![2.0_f32])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((1, 1), vec![6.0_f32])
        .unwrap()
        .into_dyn();

    let w = Array::from_shape_vec((1, 1), vec![1.0_f32]).unwrap();
    let b = Array::from_shape_vec((1, 1), vec![0.0_f32]).unwrap();
    let mut layer = Dense::new(1, 1, Linear::new()).unwrap();
    layer.set_weights(w, b).unwrap();

    let max_norm = 8.0_f32;
    let mut model = Sequential::new();
    model.add(layer).compile(
        SGD::new(0.01, 0.0, false, 0.0)
            .unwrap()
            .with_global_clipnorm(max_norm)
            .unwrap(),
        MeanSquaredError::new(),
    );

    model.fit(&x, &y, 1).unwrap();

    // grad_w=-16, grad_b=-8, so global_norm = sqrt(16^2 + 8^2), which is sqrt(320).
    // scale = max_norm / global_norm = 8 / sqrt(320).
    // Unclipped deltas were +0.16 (w) and +0.08 (b).
    let scale = max_norm / 320.0_f32.sqrt();
    let (w_new, b_new) = dense_wb(&model);
    assert_abs_diff_eq!(w_new, 1.0 + 0.16 * scale, epsilon = 1e-5);
    assert_abs_diff_eq!(b_new, 0.0 + 0.08 * scale, epsilon = 1e-5);
    // Clipping scales both parameters by the same factor, so direction is preserved
    assert_abs_diff_eq!((w_new - 1.0) / (b_new), 0.16 / 0.08, epsilon = 1e-4);
}

/// A clip threshold above the global gradient norm (sqrt(320) ~= 17.9) leaves the update identical
/// to plain SGD: w_new=1.16, b_new=0.08
#[test]
fn clip_by_global_norm_above_norm_is_noop() {
    let x = Array::from_shape_vec((1, 1), vec![2.0_f32])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((1, 1), vec![6.0_f32])
        .unwrap()
        .into_dyn();

    let w = Array::from_shape_vec((1, 1), vec![1.0_f32]).unwrap();
    let b = Array::from_shape_vec((1, 1), vec![0.0_f32]).unwrap();
    let mut layer = Dense::new(1, 1, Linear::new()).unwrap();
    layer.set_weights(w, b).unwrap();

    let mut model = Sequential::new();
    model.add(layer).compile(
        SGD::new(0.01, 0.0, false, 0.0)
            .unwrap()
            .with_global_clipnorm(100.0)
            .unwrap(),
        MeanSquaredError::new(),
    );

    model.fit(&x, &y, 1).unwrap();

    let (w_new, b_new) = dense_wb(&model);
    assert_abs_diff_eq!(w_new, 1.16_f32, epsilon = 1e-5);
    assert_abs_diff_eq!(b_new, 0.08_f32, epsilon = 1e-5);
}

/// `with_global_clipnorm` rejects non-positive or non-finite thresholds, and accepts a valid
/// positive one. Constructing without it leaves clipping disabled
#[test]
fn new_rejects_invalid_global_clipnorm() {
    for bad in [0.0_f32, -1.0, f32::NAN, f32::INFINITY] {
        assert!(
            matches!(
                SGD::new(0.01, 0.0, false, 0.0)
                    .unwrap()
                    .with_global_clipnorm(bad),
                Err(Error::InvalidParameter { .. })
            ),
            "global_clipnorm {bad} should be rejected"
        );
    }
    assert!(
        SGD::new(0.01, 0.0, false, 0.0)
            .unwrap()
            .with_global_clipnorm(5.0)
            .is_ok()
    );
    assert!(SGD::new(0.01, 0.0, false, 0.0).is_ok());
}

// SGD: momentum, weight decay, and LR scheduling (integration)

/// `set_learning_rate` retunes the step: doubling lr before 1 SGD step doubles the weight delta.
/// Reuses the w=1, b=0, x=2, y=6 problem (grad_w=-16, grad_b=-8)
#[test]
fn set_learning_rate_scales_the_step() {
    let x = Array::from_shape_vec((1, 1), vec![2.0_f32])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((1, 1), vec![6.0_f32])
        .unwrap()
        .into_dyn();
    let mut layer = Dense::new(1, 1, Linear::new()).unwrap();
    layer
        .set_weights(
            Array::from_shape_vec((1, 1), vec![1.0_f32]).unwrap(),
            Array::from_shape_vec((1, 1), vec![0.0_f32]).unwrap(),
        )
        .unwrap();

    let mut model = Sequential::new();
    model.add(layer).compile(
        SGD::new(0.01, 0.0, false, 0.0).unwrap(),
        MeanSquaredError::new(),
    );
    model.set_learning_rate(0.02); // double the configured 0.01
    model.fit(&x, &y, 1).unwrap();

    let (w_new, b_new) = dense_wb(&model);
    assert_abs_diff_eq!(w_new, 1.0 + 0.02 * 16.0, epsilon = 1e-5); // 1.32
    assert_abs_diff_eq!(b_new, 0.02 * 8.0, epsilon = 1e-5); // 0.16
}

/// Decoupled weight decay shrinks the parameter by (1 - lr*wd) before the gradient step.
/// With w=1, wd=0.5, lr=0.01: w := 1*(1 - 0.005) - 0.01*(-16) = 0.995 + 0.16 = 1.155
#[test]
fn sgd_decoupled_weight_decay_shrinks_param() {
    let x = Array::from_shape_vec((1, 1), vec![2.0_f32])
        .unwrap()
        .into_dyn();
    let y = Array::from_shape_vec((1, 1), vec![6.0_f32])
        .unwrap()
        .into_dyn();
    let mut layer = Dense::new(1, 1, Linear::new()).unwrap();
    layer
        .set_weights(
            Array::from_shape_vec((1, 1), vec![1.0_f32]).unwrap(),
            Array::from_shape_vec((1, 1), vec![0.0_f32]).unwrap(),
        )
        .unwrap();

    let mut model = Sequential::new();
    model.add(layer).compile(
        SGD::new(0.01, 0.0, false, 0.5).unwrap(),
        MeanSquaredError::new(),
    );
    model.fit(&x, &y, 1).unwrap();

    let (w_new, b_new) = dense_wb(&model);
    assert_abs_diff_eq!(w_new, 0.995 + 0.16, epsilon = 1e-5); // 1.155
    assert_abs_diff_eq!(b_new, 0.08, epsilon = 1e-5); // b=0, decay no-op
}

// Weight decay applies to weights only, not to biases or to normalization gamma and beta.
// These run 2 layers through an identical forward and backward pass and differ only in
// `weight_decay`, so decay explains any divergence. The gradient values cancel out of the
// comparison: weight_decay shrinks `value` by `(1 - lr*wd)` before the same gradient step.

/// Runs a Dense(2->2, Linear) with fixed weights and bias through 1 forward and backward pass
/// and 1 SGD step at the given `weight_decay`. The upstream gradient is fixed and nonzero.
/// Returns the resulting (weights, bias). Both decay settings see identical gradients.
fn dense_after_one_sgd_step(
    w0: &Array2<f32>,
    b0: &Array2<f32>,
    lr: f32,
    weight_decay: f32,
) -> (Array2<f32>, Array2<f32>) {
    let mut layer = Dense::new(2, 2, Linear::new()).unwrap();
    layer.set_weights(w0.clone(), b0.clone()).unwrap();
    let x = Array::from_shape_vec((1, 2), vec![1.0_f32, 2.0])
        .unwrap()
        .into_dyn();
    let _ = layer.forward(&x).unwrap();
    let grad_out = Array::from_shape_vec((1, 2), vec![0.7_f32, -1.3])
        .unwrap()
        .into_dyn();
    layer.backward(&grad_out).unwrap();

    let mut opt = SGD::new(lr, 0.0, false, weight_decay).unwrap();
    opt.step();
    opt.update(&mut layer, 1.0);
    match layer.get_weights() {
        LayerWeight::Dense(d) => ((*d.weight).clone(), (*d.bias).clone()),
        _ => panic!("expected Dense weights"),
    }
}

/// Decoupled weight decay shrinks Dense weights by exactly `lr*wd*w0`, but leaves the bias
/// byte-identical with and without decay
#[test]
fn weight_decay_decays_dense_weights_but_skips_bias() {
    let w0 = Array::from_shape_vec((2, 2), vec![1.0_f32, -2.0, 3.0, -4.0]).unwrap();
    let b0 = Array::from_shape_vec((1, 2), vec![0.5_f32, -1.5]).unwrap();
    let (lr, wd) = (0.1_f32, 0.5_f32);

    let (w_plain, b_plain) = dense_after_one_sgd_step(&w0, &b0, lr, 0.0);
    let (w_decay, b_decay) = dense_after_one_sgd_step(&w0, &b0, lr, wd);

    // Bias is excluded from weight decay -> identical with and without it
    for i in 0..2 {
        assert_abs_diff_eq!(b_decay[[0, i]], b_plain[[0, i]], epsilon = 1e-6);
    }
    // Weights are decayed by exactly lr*wd*w0 relative to the no-decay step
    for i in 0..2 {
        for j in 0..2 {
            assert_abs_diff_eq!(
                w_decay[[i, j]],
                w_plain[[i, j]] - lr * wd * w0[[i, j]],
                epsilon = 1e-6
            );
        }
    }
    // Guard against a vacuous pass: decay must have actually moved the weights
    assert!(
        (w_decay[[0, 0]] - w_plain[[0, 0]]).abs() > 1e-4,
        "weight decay should change the weights"
    );
}

/// Runs a BatchNormalization layer (gamma=1, beta=0) through 1 training forward and backward
/// pass and 1 SGD step at the given `weight_decay`. The upstream gradient is fixed and nonzero.
/// Returns the resulting (gamma, beta).
fn batchnorm_gamma_beta_after_one_sgd_step(weight_decay: f32) -> (ArrayD<f32>, ArrayD<f32>) {
    let mut bn = BatchNormalization::new(vec![2, 3], 0.9, 1e-5).unwrap();
    bn.set_training_if_mode_dependent(true);
    let x = Array::from_shape_vec((2, 3), vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0])
        .unwrap()
        .into_dyn();
    let _ = bn.forward(&x).unwrap();
    let grad_out = Array::from_shape_vec((2, 3), vec![0.5_f32, -0.5, 1.0, -1.0, 0.25, -0.25])
        .unwrap()
        .into_dyn();
    bn.backward(&grad_out).unwrap();

    let mut opt = SGD::new(0.1, 0.0, false, weight_decay).unwrap();
    opt.step();
    opt.update(&mut bn, 1.0);
    match bn.get_weights() {
        LayerWeight::BatchNormalization(w) => ((*w.gamma).clone(), (*w.beta).clone()),
        _ => panic!("expected BatchNormalization weights"),
    }
}

/// Normalization scale and shift (gamma and beta) are excluded from weight decay: a non-zero
/// `weight_decay` produces a byte-identical update to no decay at all
#[test]
fn weight_decay_skips_batchnorm_gamma_and_beta() {
    let (g_plain, b_plain) = batchnorm_gamma_beta_after_one_sgd_step(0.0);
    let (g_decay, b_decay) = batchnorm_gamma_beta_after_one_sgd_step(0.5);

    assert_eq!(g_plain.shape(), g_decay.shape());
    for (p, d) in g_plain.iter().zip(g_decay.iter()) {
        assert_abs_diff_eq!(*p, *d, epsilon = 1e-6);
    }
    for (p, d) in b_plain.iter().zip(b_decay.iter()) {
        assert_abs_diff_eq!(*p, *d, epsilon = 1e-6);
    }
    // Guard against a vacuous pass. The gradient step must actually update gamma.
    // The "identical" check above then compares moved values, not 2 untouched 1.0 arrays.
    assert!(
        g_plain.iter().any(|&v| (v - 1.0).abs() > 1e-5),
        "gamma should have a non-trivial gradient update"
    );
}

// Adam (classic coupled L2 weight decay) vs AdamW (decoupled weight decay)

/// Runs a Dense(2->2, Linear) with fixed weights and bias through 1 forward and backward pass
/// and 1 step of the given optimizer. The upstream gradient is fixed and nonzero. Returns the
/// resulting (weights, bias). Generic over the optimizer, so Adam and AdamW share the harness.
fn dense_weights_after_one_step<O: Optimizer>(
    mut opt: O,
    w0: &Array2<f32>,
    b0: &Array2<f32>,
) -> (Array2<f32>, Array2<f32>) {
    let mut layer = Dense::new(2, 2, Linear::new()).unwrap();
    layer.set_weights(w0.clone(), b0.clone()).unwrap();
    let x = Array::from_shape_vec((1, 2), vec![1.0_f32, 2.0])
        .unwrap()
        .into_dyn();
    let _ = layer.forward(&x).unwrap();
    let grad_out = Array::from_shape_vec((1, 2), vec![0.7_f32, -1.3])
        .unwrap()
        .into_dyn();
    layer.backward(&grad_out).unwrap();

    opt.step();
    opt.update(&mut layer, 1.0);
    match layer.get_weights() {
        LayerWeight::Dense(d) => ((*d.weight).clone(), (*d.bias).clone()),
        _ => panic!("expected Dense weights"),
    }
}

/// With `weight_decay == 0.0`, Adam and AdamW are the same algorithm: identical weights and bias
#[test]
fn adam_equals_adamw_without_weight_decay() {
    let w0 = Array::from_shape_vec((2, 2), vec![1.0_f32, -2.0, 3.0, -4.0]).unwrap();
    let b0 = Array::from_shape_vec((1, 2), vec![0.5_f32, -1.5]).unwrap();

    let (w_adam, b_adam) =
        dense_weights_after_one_step(Adam::new(0.1, 0.9, 0.999, 1e-8, 0.0).unwrap(), &w0, &b0);
    let (w_adamw, b_adamw) =
        dense_weights_after_one_step(AdamW::new(0.1, 0.9, 0.999, 1e-8, 0.0).unwrap(), &w0, &b0);

    for (a, w) in w_adam.iter().zip(w_adamw.iter()) {
        assert_abs_diff_eq!(*a, *w, epsilon = 1e-7);
    }
    for (a, w) in b_adam.iter().zip(b_adamw.iter()) {
        assert_abs_diff_eq!(*a, *w, epsilon = 1e-7);
    }
}

/// Coupled L2 (Adam) flows through the moments and the adaptive denominator. Decoupled decay
/// (AdamW) does not, so the 2 schemes diverge on the weights when weight_decay is non-zero
#[test]
fn adam_l2_and_adamw_decoupled_differ_with_weight_decay() {
    let w0 = Array::from_shape_vec((2, 2), vec![1.0_f32, -2.0, 3.0, -4.0]).unwrap();
    let b0 = Array::from_shape_vec((1, 2), vec![0.5_f32, -1.5]).unwrap();
    let wd = 0.5_f32;

    let (w_adam, b_adam) =
        dense_weights_after_one_step(Adam::new(0.1, 0.9, 0.999, 1e-8, wd).unwrap(), &w0, &b0);
    let (w_adamw, b_adamw) =
        dense_weights_after_one_step(AdamW::new(0.1, 0.9, 0.999, 1e-8, wd).unwrap(), &w0, &b0);

    // Weights differ between the 2 decay schemes
    assert!(
        w_adam
            .iter()
            .zip(w_adamw.iter())
            .any(|(a, w)| (a - w).abs() > 1e-5),
        "coupled (Adam) and decoupled (AdamW) weight decay should produce different weights"
    );
    // Bias is excluded from weight decay in both, so it updates identically
    for (a, w) in b_adam.iter().zip(b_adamw.iter()) {
        assert_abs_diff_eq!(*a, *w, epsilon = 1e-7);
    }
}

/// AdamW drives MSE strictly down on the seeded regression problem (with a non-zero decoupled
/// weight_decay active)
#[test]
fn adamw_single_layer_loss_decreases_over_20_epochs() {
    let (x, y) = regression_data();

    let mut model = Sequential::new();
    model.add(identity_dense()).compile(
        AdamW::new(0.1, 0.9, 0.999, 1e-8, 0.01).unwrap(),
        MeanSquaredError::new(),
    );

    let mse_before = eval_mse(&model, &x, &y);
    model.fit(&x, &y, 20).unwrap();
    let mse_after = eval_mse(&model, &x, &y);
    assert!(
        mse_after < mse_before,
        "AdamW: loss should decrease; before={mse_before}, after={mse_after}"
    );
}

/// AdamW routes through the same validators as Adam: rejects out-of-range betas or negative
/// weight_decay, accepts valid hyperparameters
#[test]
fn adamw_validates_hyperparameters() {
    assert!(matches!(
        AdamW::new(0.001, 1.0, 0.999, 1e-8, 0.0),
        Err(Error::InvalidParameter { .. })
    ));
    assert!(matches!(
        AdamW::new(0.001, 0.9, 0.999, 1e-8, -0.1),
        Err(Error::InvalidParameter { .. })
    ));
    assert!(AdamW::new(0.001, 0.9, 0.999, 1e-8, 0.01).is_ok());
    assert!(
        AdamW::new(0.001, 0.9, 0.999, 1e-8, 0.0)
            .unwrap()
            .with_global_clipnorm(1.0)
            .is_ok()
    );
}

/// SGD with momentum still drives MSE strictly down on the seeded regression problem
#[test]
fn sgd_momentum_loss_decreases() {
    let (x, y) = regression_data();
    let mut model = Sequential::new();
    model.add(identity_dense()).compile(
        SGD::new(0.05, 0.9, true, 0.0).unwrap(),
        MeanSquaredError::new(),
    );
    let before = eval_mse(&model, &x, &y);
    model.fit(&x, &y, 5).unwrap();
    let after = eval_mse(&model, &x, &y);
    assert!(
        after < before,
        "SGD+momentum should reduce loss; before={before}, after={after}"
    );
}

/// Negative momentum or weight_decay values are rejected
#[test]
fn new_rejects_negative_momentum_and_weight_decay() {
    assert!(matches!(
        SGD::new(0.01, -0.1, false, 0.0),
        Err(Error::InvalidParameter { .. })
    ));
    assert!(matches!(
        SGD::new(0.01, 0.0, false, -0.1),
        Err(Error::InvalidParameter { .. })
    ));
    assert!(matches!(
        Adam::new(0.001, 0.9, 0.999, 1e-8, -0.1),
        Err(Error::InvalidParameter { .. })
    ));
}

// learning rate: readable as well as writable

/// Every optimizer round-trips learning_rate(): it reports the value it was built with, and the
/// value it was last set to. This lets a schedule read the current rate instead of a stale copy
#[test]
fn every_optimizer_reports_its_current_learning_rate() {
    fn round_trip(mut opt: impl Optimizer, configured: f32) {
        assert_abs_diff_eq!(opt.learning_rate(), configured, epsilon = 0.0_f32);

        // The decay a scheduler writes: derived from the current value, not from a copy
        let decayed = opt.learning_rate() * 0.1;
        opt.set_learning_rate(decayed);
        assert_abs_diff_eq!(opt.learning_rate(), configured * 0.1, epsilon = 1e-9);
    }

    round_trip(SGD::new(0.05, 0.9, false, 0.0).unwrap(), 0.05);
    round_trip(Adam::new(0.003, 0.9, 0.999, 1e-8, 0.0).unwrap(), 0.003);
    round_trip(AdamW::new(0.003, 0.9, 0.999, 1e-8, 0.01).unwrap(), 0.003);
    round_trip(RMSprop::new(0.002, 0.9, 1e-8, 0.0).unwrap(), 0.002);
    round_trip(AdaGrad::new(0.02, 1e-8, 0.0).unwrap(), 0.02);
}