use aprender::ensemble::{MixtureOfExperts, GatingNetwork, SoftmaxGating};

// Define experts
let experts = vec![
    RandomForest::new(100, 10),  // scope expert
    RandomForest::new(100, 10),  // type expert
    RandomForest::new(100, 10),  // method expert
];

// Gating network (routes inputs to experts)
let gating = SoftmaxGating::new(n_features, n_experts);

// MoE ensemble
let moe = MixtureOfExperts::builder()
    .experts(experts)
    .gating(gating)
    .top_k(2)  // sparse: only top 2 experts per input
    .build();

// Train end-to-end
moe.fit(&X_train, &y_train)?;

// Predict (weighted combination of expert outputs)
let predictions = moe.predict(&X_test)?;

pub trait GatingNetwork: Send + Sync {
    /// Compute expert weights for input
    fn forward(&self, x: &[f32]) -> Vec<f32>;
    
    /// Train gating network
    fn fit(&mut self, X: &[Vec<f32>], expert_losses: &[Vec<f32>]) -> Result<()>;
}

pub struct SoftmaxGating {
    weights: Matrix<f32>,  // [n_features, n_experts]
    temperature: f32,
}

pub struct MixtureOfExperts<E: Estimator, G: GatingNetwork> {
    experts: Vec<E>,
    gating: G,
    top_k: usize,
    load_balance_weight: f32,  // optional: encourage even expert usage
}