//! Deep Gaussian Process demonstration
//!
//! This example shows how to use Deep GPs for regression with uncertainty quantification.

use rtx_ml_classic::bayesian::{
    AdvancedKernel, DeepGP, DeepGPConfig, InducingStrategy, PeriodicKernel, SVGP, SVGPConfig,
    SpectralMixtureKernel, select_inducing_points,
};
use rtx_tensor::{Device, Tensor};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    println!("=== Deep Gaussian Process Demonstration ===\n");

    let device = Device::cpu();

    // Example 1: Sparse Variational GP on synthetic data
    println!("1. Sparse Variational GP (SVGP)");
    println!("--------------------------------");

    // Generate synthetic data: y = sin(x) + noise
    let n_train = 100;
    let x_train_data: Vec<f32> = (0..n_train).map(|i| (i as f32) * 0.1).collect();
    let y_train_data: Vec<f32> = x_train_data
        .iter()
        .map(|&x| x.sin() + 0.1 * ((x * 10.0).sin()))
        .collect();

    let x_train = Tensor::from_data(
        x_train_data.iter().map(|&x| x).collect(),
        vec![n_train, 1],
        &device,
    )?;
    let y_train = Tensor::from_data(y_train_data, vec![n_train], &device)?;

    // Configure SVGP with 20 inducing points
    let svgp_config = SVGPConfig {
        num_inducing: 20,
        learn_inducing_locations: false,
        jitter: 1e-6,
        length_scale: 1.0,
        variance: 1.0,
        noise: 0.1,
    };

    let mut svgp = SVGP::new(svgp_config)?;
    svgp.initialize(&x_train, &y_train)?;

    // Compute ELBO
    let elbo = svgp.elbo(&x_train, &y_train)?;
    println!("SVGP ELBO: {:.4}", elbo);

    // Make predictions on test data
    let n_test = 50;
    let x_test_data: Vec<f32> = (0..n_test).map(|i| (i as f32) * 0.2).collect();
    let x_test = Tensor::from_data(
        x_test_data.iter().map(|&x| x).collect(),
        vec![n_test, 1],
        &device,
    )?;

    let (mean, variance) = svgp.predict(&x_test)?;
    let mean_data = mean.to_cpu()?;
    let var_data = variance.to_cpu()?;

    println!(
        "Predictions at x=0.0: mean={:.4}, std={:.4}",
        mean_data[0],
        var_data[0].sqrt()
    );
    println!(
        "Predictions at x=5.0: mean={:.4}, std={:.4}\n",
        mean_data[25],
        var_data[25].sqrt()
    );

    // Example 2: Deep Gaussian Process
    println!("2. Deep Gaussian Process (Deep GP)");
    println!("-----------------------------------");

    let dgp_config = DeepGPConfig {
        num_layers: 3,
        hidden_dims: vec![5, 3],
        num_inducing_per_layer: 15,
        length_scale: 1.0,
        variance: 1.0,
        noise: 0.1,
        jitter: 1e-6,
    };

    let mut dgp = DeepGP::new(1, 1, dgp_config)?;
    dgp.initialize(&x_train, &y_train)?;

    println!("Deep GP initialized with {} layers", dgp.num_layers());

    // Compute ELBO
    let dgp_elbo = dgp.elbo(&x_train, &y_train)?;
    println!("Deep GP ELBO: {:.4}", dgp_elbo);

    // Make predictions
    let (dgp_mean, dgp_var) = dgp.predict(&x_test)?;
    let dgp_mean_data = dgp_mean.to_cpu()?;
    let dgp_var_data = dgp_var.to_cpu()?;

    println!(
        "Deep GP predictions at x=0.0: mean={:.4}, std={:.4}",
        dgp_mean_data[0],
        dgp_var_data[0].sqrt()
    );
    println!(
        "Deep GP predictions at x=5.0: mean={:.4}, std={:.4}\n",
        dgp_mean_data[25],
        dgp_var_data[25].sqrt()
    );

    // Example 3: Advanced Kernels
    println!("3. Advanced Kernels");
    println!("-------------------");

    // Periodic kernel for periodic patterns
    let periodic = PeriodicKernel::new(1.0, 2.0, 0.5)?;
    let k_periodic = periodic.compute(&[0.0], &[1.0]);
    println!("Periodic kernel k(0.0, 1.0) = {:.4}", k_periodic);

    // Spectral Mixture kernel
    let weights = vec![0.7, 0.3];
    let means = vec![vec![1.0], vec![3.0]];
    let variances = vec![vec![0.5], vec![0.8]];
    let spectral = SpectralMixtureKernel::new(weights, means, variances)?;
    let k_spectral = spectral.compute(&[0.0], &[1.0]);
    println!("Spectral Mixture kernel k(0.0, 1.0) = {:.4}", k_spectral);

    // Composite kernel (sum of two kernels)
    let k1 = AdvancedKernel::Periodic(PeriodicKernel::new(1.0, 2.0, 0.5)?);
    let k2 = AdvancedKernel::SpectralMixture(spectral);
    let composite = rtx_ml_classic::bayesian::CompositeKernel::new(
        k1,
        k2,
        rtx_ml_classic::bayesian::KernelOp::Add,
    );
    let k_composite = composite.compute(&[0.0], &[1.0]);
    println!("Composite kernel k(0.0, 1.0) = {:.4}\n", k_composite);

    // Example 4: Inducing Point Selection Strategies
    println!("4. Inducing Point Selection");
    println!("----------------------------");

    // Random selection
    let inducing_random = select_inducing_points(&x_train, 10, InducingStrategy::Random, None)?;
    println!(
        "Random: selected {} inducing points",
        inducing_random.shape().dims()[0]
    );

    // K-means clustering
    let inducing_kmeans = select_inducing_points(&x_train, 10, InducingStrategy::KMeans, None)?;
    println!(
        "K-means: selected {} inducing points",
        inducing_kmeans.shape().dims()[0]
    );

    // Greedy selection
    let inducing_greedy = select_inducing_points(&x_train, 10, InducingStrategy::Greedy, None)?;
    println!(
        "Greedy: selected {} inducing points",
        inducing_greedy.shape().dims()[0]
    );

    println!("\n=== Demonstration Complete ===");
    println!("\nKey Features:");
    println!("- Sparse Variational GP scales to 10K+ data points");
    println!("- Deep GP models hierarchical patterns");
    println!("- Advanced kernels capture complex structures");
    println!("- Multiple inducing point selection strategies");
    println!("- Full uncertainty quantification");

    Ok(())
}