#![allow(clippy::approx_constant)]
use approx::assert_abs_diff_eq;
use rtx_preprocessing::{PreprocessingError, TargetEncoder};
use rtx_tensor::{Device, Tensor};

#[test]
fn test_target_encoder_creation() {
    let encoder = TargetEncoder::new();
    assert!(!encoder.is_fitted());
    assert_abs_diff_eq!(encoder.smoothing(), 1.0, epsilon = 1e-5);

    let encoder = TargetEncoder::with_smoothing(5.0);
    assert_abs_diff_eq!(encoder.smoothing(), 5.0, epsilon = 1e-5);
}

#[test]
fn test_target_encoder_fit_simple() {
    let features = Tensor::from_slice(&[0.0, 1.0, 0.0, 1.0], &[4, 1], &Device::cpu()).unwrap();
    let targets = Tensor::from_slice(&[0.0, 1.0, 0.0, 1.0], &[4, 1], &Device::cpu()).unwrap();
    let mut encoder = TargetEncoder::new();

    let result = encoder.fit_with_target(&features, &targets);
    assert!(result.is_ok());
    assert!(encoder.is_fitted());
}

#[test]
fn test_target_encoder_transform_simple() {
    let features = Tensor::from_slice(&[0.0, 1.0, 0.0, 1.0], &[4, 1], &Device::cpu()).unwrap();
    let targets = Tensor::from_slice(&[0.0, 1.0, 0.0, 1.0], &[4, 1], &Device::cpu()).unwrap();
    let mut encoder = TargetEncoder::new();

    encoder.fit_with_target(&features, &targets).unwrap();
    let encoded = encoder.transform(&features).unwrap();

    let values = encoded
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();

    // Category 0: mean target = (0 + 0) / 2 = 0.0
    // Category 1: mean target = (1 + 1) / 2 = 1.0
    // With smoothing, values will be between global mean and category mean
    assert!(values[0] < values[1]); // Category 0 should have lower encoding than category 1
}

#[test]
fn test_target_encoder_smoothing_effect() {
    let features = Tensor::from_slice(&[0.0, 1.0], &[2, 1], &Device::cpu()).unwrap();
    let targets = Tensor::from_slice(&[0.0, 1.0], &[2, 1], &Device::cpu()).unwrap();

    // High smoothing (more regularization)
    let mut encoder_smooth = TargetEncoder::with_smoothing(100.0);
    encoder_smooth.fit_with_target(&features, &targets).unwrap();
    let encoded_smooth = encoder_smooth.transform(&features).unwrap();

    // Low smoothing (less regularization)
    let mut encoder_raw = TargetEncoder::with_smoothing(0.1);
    encoder_raw.fit_with_target(&features, &targets).unwrap();
    let encoded_raw = encoder_raw.transform(&features).unwrap();

    let smooth_values = encoded_smooth
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();
    let raw_values = encoded_raw
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();

    // High smoothing should be closer to global mean
    let diff_smooth = (smooth_values[1] - smooth_values[0]).abs();
    let diff_raw = (raw_values[1] - raw_values[0]).abs();
    assert!(diff_smooth < diff_raw);
}

#[test]
fn test_target_encoder_unknown_category() {
    let features = Tensor::from_slice(&[0.0, 1.0], &[2, 1], &Device::cpu()).unwrap();
    let targets = Tensor::from_slice(&[0.0, 1.0], &[2, 1], &Device::cpu()).unwrap();
    let mut encoder = TargetEncoder::new();

    encoder.fit_with_target(&features, &targets).unwrap();

    let unknown_features = Tensor::from_slice(&[2.0], &[1, 1], &Device::cpu()).unwrap();
    let encoded = encoder.transform(&unknown_features).unwrap();

    let values = encoded
        .to_cpu()
        .unwrap()
        .iter()
        .copied()
        .collect::<Vec<f32>>();
    // Unknown category should get global mean
    let global_mean = 0.5; // (0 + 1) / 2
    assert_abs_diff_eq!(values[0], global_mean, epsilon = 1e-1);
}

#[test]
fn test_target_encoder_multi_feature() {
    let features = Tensor::from_slice(&[0.0, 10.0, 1.0, 20.0], &[2, 2], &Device::cpu()).unwrap();
    let targets = Tensor::from_slice(&[0.0, 1.0], &[2, 1], &Device::cpu()).unwrap();
    let mut encoder = TargetEncoder::new();

    encoder.fit_with_target(&features, &targets).unwrap();
    let encoded = encoder.transform(&features).unwrap();

    assert_eq!(encoded.shape(), &[2, 2]);
}

#[test]
fn test_target_encoder_dimension_mismatch() {
    let features = Tensor::from_slice(&[0.0, 1.0], &[2, 1], &Device::cpu()).unwrap();
    let targets = Tensor::from_slice(&[0.0], &[1, 1], &Device::cpu()).unwrap(); // Wrong size
    let mut encoder = TargetEncoder::new();

    let result = encoder.fit_with_target(&features, &targets);
    assert!(result.is_err());
    assert!(matches!(
        result.unwrap_err(),
        PreprocessingError::DimensionMismatch { .. }
    ));
}

#[test]
fn test_target_encoder_reset() {
    let features = Tensor::from_slice(&[0.0, 1.0], &[2, 1], &Device::cpu()).unwrap();
    let targets = Tensor::from_slice(&[0.0, 1.0], &[2, 1], &Device::cpu()).unwrap();
    let mut encoder = TargetEncoder::new();

    encoder.fit_with_target(&features, &targets).unwrap();
    assert!(encoder.is_fitted());

    encoder.reset();
    assert!(!encoder.is_fitted());
}