#!/usr/bin/env python3
"""
Comprehensive test of the edaflow.ml subpackage functionality.
"""

print("🔬 Testing ML subpackage comprehensive functionality...")

# Import required libraries
from edaflow import ml
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.datasets import make_classification
import pandas as pd
import numpy as np

# 1. Data preparation
print("📊 1. Preparing test data...")
X, y = make_classification(
    n_samples=300, 
    n_features=10, 
    n_classes=2, 
    random_state=42
)
df = pd.DataFrame(X, columns=[f'feature_{i}' for i in range(10)])
df['target'] = y
print("✅ 1. Data preparation complete")

# 2. ML experiment setup
print("🔧 2. Setting up ML experiment...")
experiment_data = ml.setup_ml_experiment(df, 'target', verbose=False)
print("✅ 2. ML experiment setup complete")

# 3. Model training
print("🤖 3. Training models...")
models = {
    'RandomForest': RandomForestClassifier(random_state=42, n_estimators=10),
    'LogisticRegression': LogisticRegression(random_state=42, max_iter=200),
    'GradientBoosting': GradientBoostingClassifier(random_state=42, n_estimators=10)
}

# Train models
trained_models = {}
for name, model in models.items():
    trained_models[name] = model.fit(experiment_data['X_train'], experiment_data['y_train'])
print("✅ 3. Models trained successfully")

# 4. Model comparison
print("⚖️ 4. Comparing models...")
try:
    comparison = ml.compare_models(
        trained_models,
        experiment_data['X_train'],
        experiment_data['X_val'], 
        experiment_data['y_train'],
        experiment_data['y_val'],
        verbose=False
    )
    print("✅ 4. Model comparison complete")
    print(f"📈 Comparison results shape: {comparison.shape}")
except Exception as e:
    print(f"❌ Model comparison failed: {e}")

# 5. Hyperparameter optimization
print("🎯 5. Testing hyperparameter optimization...")
try:
    param_distributions = {
        'n_estimators': [10, 20],
        'max_depth': [3, 5]
    }
    
    results = ml.optimize_hyperparameters(
        model=RandomForestClassifier(random_state=42),
        param_distributions=param_distributions,
        X_train=experiment_data['X_train'],
        y_train=experiment_data['y_train'],
        method='grid',
        cv=3,
        verbose=False
    )
    print("✅ 5. Hyperparameter optimization complete")
    print(f"🏆 Best score: {results['best_score']:.4f}")
except Exception as e:
    print(f"❌ Hyperparameter optimization failed: {e}")

# 6. Learning curves
print("📈 6. Testing learning curves...")
try:
    fig = ml.plot_learning_curves(
        RandomForestClassifier(random_state=42, n_estimators=10),
        experiment_data['X_train'],
        experiment_data['y_train'],
        cv=3
    )
    print("✅ 6. Learning curves generated")
except Exception as e:
    print(f"❌ Learning curves failed: {e}")

# 7. Artifacts (basic test)
print("💾 7. Testing artifacts functionality...")
try:
    # Save model artifacts
    test_model = RandomForestClassifier(random_state=42, n_estimators=10)
    test_model.fit(experiment_data['X_train'], experiment_data['y_train'])
    
    # Create basic experiment config and metrics
    config = {
        'target_column': 'target',
        'problem_type': 'classification',
        'train_size': len(experiment_data['X_train'])
    }
    
    metrics = {
        'accuracy': 0.95,
        'f1_score': 0.94
    }
    
    ml.save_model_artifacts(
        model=test_model,
        model_name='test_model',
        experiment_config=config,
        performance_metrics=metrics
    )
    print("✅ 7. Artifacts functionality working")
except Exception as e:
    print(f"❌ Artifacts failed: {e}")

print("\n🎉 ML subpackage comprehensive test complete!")
print("🚀 All major functionality is working properly!")
