rustytorch/crates/training/rtx-automeasure/tests/integration_test.rs

use rtx_automeasure::{AutoMLAgent, AutoMLConfig, AutoMLResult, OptimizationObjective, TaskType};
use rtx_tensor::{Device, Tensor};

#[tokio::test]
async fn test_automl_agent_creation() -> AutoMLResult<()> {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(60)
        .with_objective(OptimizationObjective::Accuracy)
        .with_cv_folds(3);

    let agent = AutoMLAgent::new(config)?;
    assert!(!agent.get_id().is_empty());

    Ok(())
}

#[tokio::test]
async fn test_automl_fit_basic() -> AutoMLResult<()> {
    let device = Device::cpu();
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(10)
        .with_cv_folds(3);

    let mut agent = AutoMLAgent::new(config)?;

    let x_train = Tensor::randn(&[100, 8], &device)?;
    let y_train = Tensor::zeros(&[100], &device)?;

    let pipeline = agent.fit(&x_train, &y_train).await?;

    assert!(!pipeline.get_models().is_empty());
    assert!(!pipeline.get_best_model().is_empty());
    assert!(pipeline.get_validation_score() >= 0.0);

    Ok(())
}

#[tokio::test]
async fn test_automl_predict() -> AutoMLResult<()> {
    let device = Device::cpu();
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(10);

    let mut agent = AutoMLAgent::new(config)?;

    let x_train = Tensor::randn(&[80, 6], &device)?;
    let y_train = Tensor::zeros(&[80], &device)?;

    let pipeline = agent.fit(&x_train, &y_train).await?;

    let x_test = Tensor::randn(&[20, 6], &device)?;
    let predictions = agent.predict(&pipeline, &x_test).await?;

    assert_eq!(predictions.shape()[0], 20);

    Ok(())
}

#[tokio::test]
async fn test_automl_regression() -> AutoMLResult<()> {
    let device = Device::cpu();
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Regression)
        .with_time_budget(10)
        .with_objective(OptimizationObjective::MSE);

    let mut agent = AutoMLAgent::new(config)?;

    let x_train = Tensor::randn(&[60, 5], &device)?;
    let y_train = Tensor::randn(&[60], &device)?;

    let pipeline = agent.fit(&x_train, &y_train).await?;

    assert!(!pipeline.get_best_model().is_empty());

    Ok(())
}

#[tokio::test]
async fn test_automl_leaderboard() -> AutoMLResult<()> {
    let device = Device::cpu();
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(10);

    let mut agent = AutoMLAgent::new(config)?;

    let x_train = Tensor::randn(&[100, 10], &device)?;
    let y_train = Tensor::zeros(&[100], &device)?;

    let _pipeline = agent.fit(&x_train, &y_train).await?;

    let leaderboard = agent.get_leaderboard();
    assert!(!leaderboard.is_empty());

    // Leaderboard should be sorted by score
    for i in 1..leaderboard.len() {
        assert!(leaderboard[i - 1].score >= leaderboard[i].score);
    }

    Ok(())
}

#[tokio::test]
async fn test_automl_feature_importance() -> AutoMLResult<()> {
    let device = Device::cpu();
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(10);

    let mut agent = AutoMLAgent::new(config)?;

    let x_train = Tensor::randn(&[100, 8], &device)?;
    let y_train = Tensor::zeros(&[100], &device)?;

    let pipeline = agent.fit(&x_train, &y_train).await?;

    let importance = agent.get_feature_importance(&pipeline)?;
    assert!(!importance.is_empty());

    Ok(())
}

#[tokio::test]
async fn test_automl_progress_tracking() -> AutoMLResult<()> {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(60);

    let agent = AutoMLAgent::new(config)?;

    let progress = agent.get_progress();
    assert_eq!(progress.elapsed_time_seconds, 0.0);
    assert_eq!(progress.completion_percentage, 0.0);

    Ok(())
}

#[tokio::test]
async fn test_automl_config_validation() {
    // Test invalid time budget
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(0);

    let result = AutoMLAgent::new(config);
    assert!(result.is_err());

    // Test invalid cv_folds
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_cv_folds(0);

    let result = AutoMLAgent::new(config);
    assert!(result.is_err());
}

#[tokio::test]
async fn test_automl_pipeline_serialization() -> AutoMLResult<()> {
    let device = Device::cpu();
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(10);

    let mut agent = AutoMLAgent::new(config)?;

    let x_train = Tensor::randn(&[60, 4], &device)?;
    let y_train = Tensor::zeros(&[60], &device)?;

    let pipeline = agent.fit(&x_train, &y_train).await?;

    // Test serialization
    let json = pipeline.to_json()?;
    assert!(!json.is_empty());

    // Test deserialization
    let restored = rtx_automeasure::AutoMLPipeline::from_json(&json)?;
    assert_eq!(restored.get_best_model(), pipeline.get_best_model());

    Ok(())
}

#[tokio::test]
async fn test_automl_with_different_objectives() -> AutoMLResult<()> {
    let device = Device::cpu();

    let objectives = vec![
        OptimizationObjective::Accuracy,
        OptimizationObjective::F1Score,
        OptimizationObjective::Precision,
        OptimizationObjective::Recall,
    ];

    for objective in objectives {
        let config = AutoMLConfig::new()
            .with_task_type(TaskType::Classification)
            .with_time_budget(10)
            .with_objective(objective);

        let mut agent = AutoMLAgent::new(config)?;

        let x_train = Tensor::randn(&[50, 5], &device)?;
        let y_train = Tensor::zeros(&[50], &device)?;

        let pipeline = agent.fit(&x_train, &y_train).await?;
        assert!(!pipeline.get_best_model().is_empty());
    }

    Ok(())
}

#[tokio::test]
async fn test_automl_data_validation() {
    let device = Device::cpu();
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Classification)
        .with_time_budget(10);

    let mut agent = AutoMLAgent::new(config).unwrap();

    // Test mismatched shapes
    let x_train = Tensor::randn(&[100, 8], &device).unwrap();
    let y_train = Tensor::zeros(&[50], &device).unwrap(); // Wrong size

    let result = agent.fit(&x_train, &y_train).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_config_builder_pattern() {
    let config = AutoMLConfig::new()
        .with_task_type(TaskType::Regression)
        .with_time_budget(120)
        .with_memory_budget(8 * 1024 * 1024 * 1024)
        .with_cv_folds(5)
        .with_objective(OptimizationObjective::RMSE);

    assert_eq!(config.task_type, TaskType::Regression);
    assert_eq!(config.time_budget_seconds, 120);
    assert_eq!(config.cv_folds, 5);
}