#!/usr/bin/env rust-script
//! sklearn Classifier Migration Example - Transpiled Output
//!
//! This shows the expected Rust/Aprender code after Batuta transpilation.
//!
//! Demonstrates:
//! - sklearn → Aprender ML algorithm mapping
//! - Training, evaluation, and inference pipeline
//! - Automatic backend selection for compute-intensive operations
//!
//! Run with: cargo run --example sklearn_classifier_output

use std::collections::HashMap;

// Aprender imports (would be used in real implementation)
// use aprender::{LogisticRegression, StandardScaler, TrainTestSplit};
// use aprender::metrics::{accuracy_score, confusion_matrix, classification_report};

/// Iris dataset (embedded for demonstration)
struct IrisDataset {
    data: Vec<Vec<f64>>,
    target: Vec<usize>,
    feature_names: Vec<String>,
    target_names: Vec<String>,
}

/// Load the iris dataset
fn load_data() -> (Vec<Vec<f64>>, Vec<usize>, Vec<String>, Vec<String>) {
    // Simplified dataset (first 30 samples for demo)
    let dataset = IrisDataset {
        data: vec![
            vec![5.1, 3.5, 1.4, 0.2], // setosa
            vec![4.9, 3.0, 1.4, 0.2],
            vec![4.7, 3.2, 1.3, 0.2],
            vec![4.6, 3.1, 1.5, 0.2],
            vec![5.0, 3.6, 1.4, 0.2],
            vec![7.0, 3.2, 4.7, 1.4], // versicolor
            vec![6.4, 3.2, 4.5, 1.5],
            vec![6.9, 3.1, 4.9, 1.5],
            vec![5.5, 2.3, 4.0, 1.3],
            vec![6.5, 2.8, 4.6, 1.5],
            vec![6.3, 3.3, 6.0, 2.5], // virginica
            vec![5.8, 2.7, 5.1, 1.9],
            vec![7.1, 3.0, 5.9, 2.1],
            vec![6.3, 2.9, 5.6, 1.8],
            vec![6.5, 3.0, 5.8, 2.2],
            // ... (full dataset would have 150 samples)
        ],
        target: vec![0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2],
        feature_names: vec![
            "sepal length (cm)".to_string(),
            "sepal width (cm)".to_string(),
            "petal length (cm)".to_string(),
            "petal width (cm)".to_string(),
        ],
        target_names: vec![
            "setosa".to_string(),
            "versicolor".to_string(),
            "virginica".to_string(),
        ],
    };

    (
        dataset.data,
        dataset.target,
        dataset.feature_names,
        dataset.target_names,
    )
}

/// Split data into training and test sets
fn split_data(
    X: &[Vec<f64>],
    y: &[usize],
    test_size: f64,
) -> (Vec<Vec<f64>>, Vec<Vec<f64>>, Vec<usize>, Vec<usize>) {
    let n = X.len();
    let test_count = (n as f64 * test_size).round() as usize;
    let train_count = n - test_count;

    // Simple split (in practice, use stratified sampling)
    let X_train = X[..train_count].to_vec();
    let X_test = X[train_count..].to_vec();
    let y_train = y[..train_count].to_vec();
    let y_test = y[train_count..].to_vec();

    (X_train, X_test, y_train, y_test)
}

/// StandardScaler implementation
struct StandardScaler {
    mean: Vec<f64>,
    std: Vec<f64>,
}

impl StandardScaler {
    fn new() -> Self {
        Self {
            mean: Vec::new(),
            std: Vec::new(),
        }
    }

    fn fit(&mut self, X: &[Vec<f64>]) {
        let n_samples = X.len();
        let n_features = X[0].len();

        // Compute mean for each feature
        self.mean = (0..n_features)
            .map(|j| {
                let sum: f64 = X.iter().map(|row| row[j]).sum();
                sum / n_samples as f64
            })
            .collect();

        // Compute std for each feature
        self.std = (0..n_features)
            .map(|j| {
                let variance: f64 = X
                    .iter()
                    .map(|row| (row[j] - self.mean[j]).powi(2))
                    .sum::<f64>()
                    / n_samples as f64;
                variance.sqrt()
            })
            .collect();
    }

    fn transform(&self, X: &[Vec<f64>]) -> Vec<Vec<f64>> {
        X.iter()
            .map(|row| {
                row.iter()
                    .enumerate()
                    .map(|(j, &val)| (val - self.mean[j]) / self.std[j])
                    .collect()
            })
            .collect()
    }
}

/// Simple logistic regression (placeholder - real impl would be in Aprender)
struct LogisticRegression {
    coefficients: Vec<Vec<f64>>,
    intercept: Vec<f64>,
    classes: Vec<usize>,
}

impl LogisticRegression {
    fn new() -> Self {
        Self {
            coefficients: Vec::new(),
            intercept: Vec::new(),
            classes: Vec::new(),
        }
    }

    fn fit(&mut self, X: &[Vec<f64>], y: &[usize]) {
        // Simplified training (real implementation uses L-BFGS)
        let n_features = X[0].len();
        let n_classes = *y.iter().max().unwrap() + 1;

        self.classes = (0..n_classes).collect();
        self.coefficients = vec![vec![0.1; n_features]; n_classes];
        self.intercept = vec![0.0; n_classes];

        // Note: Real training would use optimization algorithm
        println!("   [Training logistic regression with L-BFGS...]");
    }

    fn predict(&self, X: &[Vec<f64>]) -> Vec<usize> {
        X.iter().map(|sample| {
            // Compute scores for each class
            let scores: Vec<f64> = self.coefficients
                .iter()
                .zip(&self.intercept)
                .map(|(coef, intercept)| {
                    coef.iter().zip(sample).map(|(c, x)| c * x).sum::<f64>() + intercept
                })
                .collect();

            // Return class with highest score
            scores
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                .map(|(idx, _)| idx)
                .unwrap()
        }).collect()
    }

    fn predict_proba(&self, X: &[Vec<f64>]) -> Vec<Vec<f64>> {
        X.iter().map(|sample| {
            // Compute raw scores
            let scores: Vec<f64> = self.coefficients
                .iter()
                .zip(&self.intercept)
                .map(|(coef, intercept)| {
                    coef.iter().zip(sample).map(|(c, x)| c * x).sum::<f64>() + intercept
                })
                .collect();

            // Apply softmax
            let max_score = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
            let exp_scores: Vec<f64> = scores.iter().map(|s| (s - max_score).exp()).collect();
            let sum_exp: f64 = exp_scores.iter().sum();

            exp_scores.iter().map(|e| e / sum_exp).collect()
        }).collect()
    }
}

/// Compute accuracy score
fn accuracy_score(y_true: &[usize], y_pred: &[usize]) -> f64 {
    let correct = y_true.iter().zip(y_pred).filter(|(a, b)| a == b).count();
    correct as f64 / y_true.len() as f64
}

/// Main ML classification pipeline
fn main() {
    println!("sklearn → Rust/Aprender Classification Pipeline");
    println!("{}", "=".repeat(50));

    // 1. Load data
    println!("\n1. Loading data...");
    let (X, y, feature_names, target_names) = load_data();
    println!("   Dataset shape: ({}, {})", X.len(), X[0].len());
    println!("   Classes: {:?}", target_names);
    println!("   Features: {:?}", feature_names);

    // 2. Split data
    println!("\n2. Splitting data (80% train, 20% test)...");
    let (X_train, X_test, y_train, y_test) = split_data(&X, &y, 0.2);
    println!("   Training samples: {}", X_train.len());
    println!("   Test samples: {}", X_test.len());

    // 3. Feature scaling
    println!("\n3. Scaling features...");
    let mut scaler = StandardScaler::new();
    scaler.fit(&X_train);
    let X_train_scaled = scaler.transform(&X_train);
    let X_test_scaled = scaler.transform(&X_test);
    println!("   Scaler mean: {:?}", scaler.mean);
    println!("   Scaler std: {:?}", scaler.std);

    // 4. Train model
    println!("\n4. Training logistic regression...");
    let mut model = LogisticRegression::new();
    model.fit(&X_train_scaled, &y_train);
    println!("   Model classes: {:?}", model.classes);
    println!("   Model coefficients shape: ({}, {})",
        model.coefficients.len(),
        model.coefficients[0].len()
    );

    // 5. Evaluate
    println!("\n5. Evaluating model...");
    let y_pred = model.predict(&X_test_scaled);
    let accuracy = accuracy_score(&y_test, &y_pred);
    println!("   Accuracy: {:.4}", accuracy);

    // 6. Predict on new sample
    println!("\n6. Predicting on new samples...");
    let new_sample = vec![vec![5.1, 3.5, 1.4, 0.2]]; // Likely setosa
    let new_sample_scaled = scaler.transform(&new_sample);

    let pred_class = model.predict(&new_sample_scaled)[0];
    let pred_proba = model.predict_proba(&new_sample_scaled)[0].clone();

    println!("   Sample: {:?}", new_sample[0]);
    println!("   Predicted class: {}", target_names[pred_class]);
    println!("   Probabilities:");
    for (i, class_name) in target_names.iter().enumerate() {
        println!("      {}: {:.4}", class_name, pred_proba[i]);
    }

    println!("\n✅ Pipeline complete!");

    println!("\n⚡ Performance Benefits:");
    println!("   • Type safety: Compile-time guarantees on data types");
    println!("   • Memory efficiency: No Python object overhead");
    println!("   • SIMD acceleration: Vectorized matrix operations");
    println!("   • Single binary: No runtime dependencies");
}
