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

#[test]
fn test_onehot_encoder_creation() {
    let encoder = OneHotEncoder::new();
    assert!(!encoder.is_fitted());
    assert!(encoder.sparse());
    assert!(encoder.drop_first());

    let encoder = OneHotEncoder::with_params(false, false, None);
    assert!(!encoder.sparse());
    assert!(!encoder.drop_first());
}

#[test]
fn test_onehot_encoder_fit_simple() {
    let data = Tensor::from_slice(&[0.0, 1.0, 2.0, 1.0, 0.0], &[5, 1], &Device::cpu()).unwrap();
    let mut encoder = OneHotEncoder::new();

    let result = encoder.fit(&data);
    assert!(result.is_ok());
    assert!(encoder.is_fitted());
    assert_eq!(encoder.categories()[0].len(), 3); // [0, 1, 2]
}

#[test]
fn test_onehot_encoder_transform_simple() {
    let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
    let mut encoder = OneHotEncoder::with_params(false, false, None); // Dense, no drop

    encoder.fit(&data).unwrap();
    let encoded = encoder.transform(&data).unwrap();

    assert_eq!(encoded.shape(), &[3, 3]); // 3 samples, 3 categories
    let values: Vec<f32> = encoded.to_cpu().unwrap().iter().copied().collect();

    // Should be identity matrix
    let expected = vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0];
    for (actual, expected) in values.iter().zip(expected.iter()) {
        assert_abs_diff_eq!(actual, expected, epsilon = 1e-5);
    }
}

#[test]
fn test_onehot_encoder_transform_drop_first() {
    let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
    let mut encoder = OneHotEncoder::with_params(false, true, None); // Dense, drop first

    encoder.fit(&data).unwrap();
    let encoded = encoder.transform(&data).unwrap();

    assert_eq!(encoded.shape(), &[3, 2]); // 3 samples, 2 categories (dropped first)
    let values: Vec<f32> = encoded.to_cpu().unwrap().iter().copied().collect();

    // Should drop first column
    let expected = vec![0.0, 0.0, 1.0, 0.0, 0.0, 1.0];
    for (actual, expected) in values.iter().zip(expected.iter()) {
        assert_abs_diff_eq!(actual, expected, epsilon = 1e-5);
    }
}

#[test]
fn test_onehot_encoder_unknown_category() {
    let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
    let mut encoder = OneHotEncoder::new();

    encoder.fit(&data).unwrap();

    let unknown_data = Tensor::from_slice(&[3.0], &[1, 1], &Device::cpu()).unwrap(); // Category 3 not seen
    let result = encoder.transform(&unknown_data);

    // Should handle unknown categories based on handle_unknown parameter
    // Default might be 'error' or 'ignore'
    match result {
        Ok(_) => (), // If 'ignore' is default
        Err(e) => assert!(matches!(e, PreprocessingError::InvalidInput { .. })),
    }
}

#[test]
fn test_onehot_encoder_multi_feature() {
    let data =
        Tensor::from_slice(&[0.0, 10.0, 1.0, 20.0, 0.0, 10.0], &[3, 2], &Device::cpu()).unwrap();
    let mut encoder = OneHotEncoder::with_params(false, false, None);

    encoder.fit(&data).unwrap();
    let encoded = encoder.transform(&data).unwrap();

    // Feature 0: [0, 1] -> 2 categories
    // Feature 1: [10, 20] -> 2 categories
    // Total: 4 columns
    assert_eq!(encoded.shape(), &[3, 4]);
}

#[test]
fn test_onehot_encoder_inverse_transform() {
    let data = Tensor::from_slice(&[0.0, 1.0, 2.0, 1.0], &[4, 1], &Device::cpu()).unwrap();
    let mut encoder = OneHotEncoder::with_params(false, false, None);

    encoder.fit(&data).unwrap();
    let encoded = encoder.transform(&data).unwrap();
    let decoded = encoder.inverse_transform(&encoded).unwrap();

    let original_values: Vec<f32> = data.to_cpu().unwrap().iter().copied().collect();
    let decoded_values: Vec<f32> = decoded.to_cpu().unwrap().iter().copied().collect();

    for (orig, dec) in original_values.iter().zip(decoded_values.iter()) {
        assert_abs_diff_eq!(orig, dec, epsilon = 1e-5);
    }
}

#[test]
fn test_onehot_encoder_fit_transform() {
    let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
    let mut encoder = OneHotEncoder::new();

    let result1 = encoder.fit_transform(&data).unwrap();

    let mut encoder2 = OneHotEncoder::new();
    encoder2.fit(&data).unwrap();
    let result2 = encoder2.transform(&data).unwrap();

    assert_eq!(result1.shape(), result2.shape());
}

#[test]
fn test_onehot_encoder_reset() {
    let data = Tensor::from_slice(&[0.0, 1.0, 2.0], &[3, 1], &Device::cpu()).unwrap();
    let mut encoder = OneHotEncoder::new();

    encoder.fit(&data).unwrap();
    assert!(encoder.is_fitted());

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