use tuphium::activations::Activation;
use tuphium::error::TuphiumError;
use tuphium::layers::{Layer, Module};

fn make_layer_relu() -> Layer {
    Layer::new(2, 3, Activation::ReLU).unwrap()
}

#[test]
fn forward_output_size_matches_out_features() {
    let mut layer = make_layer_relu();
    let output = layer.forward(&[1.0, 2.0]).unwrap();
    assert_eq!(output.len(), 3);
}

#[test]
fn forward_sigmoid_output_in_range() {
    let mut layer = Layer::new(4, 2, Activation::Sigmoid).unwrap();
    let output = layer.forward(&[0.5, -0.5, 1.0, -1.0]).unwrap();
    assert!(output.iter().all(|&v| v > 0.0 && v < 1.0));
}

#[test]
fn backward_grad_input_size_matches_in_features() {
    let mut layer = make_layer_relu();
    let output = layer.forward(&[1.0, 2.0]).unwrap();
    let grad_out = vec![1.0; output.len()];
    let grad_in = layer.backward(&grad_out).unwrap();
    assert_eq!(grad_in.len(), 2);
}

#[test]
fn params_and_grads_count_match() {
    let mut layer = Layer::new(3, 4, Activation::ReLU).unwrap();
    let p = layer.params().len();
    assert_eq!(p, 16);
    assert_eq!(layer.grads().len(), p);
}

#[test]
fn zero_grad_clears_accumulated_gradients() {
    let mut layer = Layer::new(2, 2, Activation::ReLU).unwrap();
    layer.forward(&[1.0, 1.0]).unwrap();
    layer.backward(&[1.0, 1.0]).unwrap();
    layer.zero_grad();
    assert!(layer.grads().iter().all(|&g| g == 0.0));
}

#[test]
fn backward_accumulates_gradients_across_calls() {
    let mut layer = Layer::new(2, 2, Activation::ReLU).unwrap();
    layer.weights.set(0, 0, 1.0).unwrap();
    layer.weights.set(1, 0, 1.0).unwrap();

    layer.forward(&[1.0, 0.0]).unwrap();
    layer.backward(&[1.0, 1.0]).unwrap();
    let grads_after_first = layer.grads();

    layer.forward(&[1.0, 0.0]).unwrap();
    layer.backward(&[1.0, 1.0]).unwrap();
    let grads_after_second = layer.grads();

    assert!(
        grads_after_first
            .iter()
            .zip(&grads_after_second)
            .any(|(a, b)| (b - a).abs() > 1e-6)
    );
}

#[test]
fn new_rejects_zero_dimensions() {
    let error = match Layer::new(0, 2, Activation::ReLU) {
        Ok(_) => panic!("expected Layer::new to fail"),
        Err(error) => error,
    };
    assert_eq!(
        error,
        TuphiumError::InvalidDimension {
            operation: "Layer::new",
            dimension: "input",
            value: 0,
            reason: "must be greater than zero",
        }
    );
}

#[test]
fn forward_returns_shape_mismatch() {
    let mut layer = Layer::new(2, 3, Activation::ReLU).unwrap();
    assert_eq!(
        layer.forward(&[1.0]).unwrap_err(),
        TuphiumError::ShapeMismatch {
            operation: "Layer::forward",
            expected: vec![2],
            actual: vec![1],
        }
    );
}

#[test]
fn backward_requires_forward_first() {
    let mut layer = Layer::new(2, 3, Activation::ReLU).unwrap();
    assert_eq!(
        layer.backward(&[1.0, 1.0, 1.0]).unwrap_err(),
        TuphiumError::InvalidState {
            operation: "Layer::backward",
            state: "forward must be called before backward",
        }
    );
}
