//! Mixture-of-Experts components for the DQN policy network.
//!
//! K=8 learned-gate expert mixture per the spec at
//! `docs/superpowers/specs/2026-04-27-moe-regime-redesign-design.md`.
//! Phase 4 removed `RegimeConditionalDQN`; MoE is now the sole
//! regime-conditioning mechanism.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::cuda_autograd::GpuTensor;
use ml_core::MLError;
use rand::SeedableRng;
use rand::rngs::StdRng;
use rand::Rng;
use rand_distr::Normal;

/// K=8 experts (compile-time constant per spec §3).
pub const MOE_NUM_EXPERTS: usize = 8;

/// Expert bottleneck width per spec §4.2.
pub const MOE_EXPERT_BOTTLENECK: usize = 64;

/// Gate hidden width per spec §4.1.
pub const MOE_GATE_HIDDEN: usize = 64;

/// Gate subnetwork: state[42] → 64 → 8 → softmax over experts.
///
/// Parameters initialized to zero so g(s) = uniform 1/K initially —
/// no expert is favored by init noise; specialization emerges from data.
pub struct MoeGate {
    /// Linear 1 weights [state_dim, 64], zero-init.
    pub w1: GpuTensor,
    /// Linear 1 bias [64], zero-init.
    pub b1: GpuTensor,
    /// Linear 2 weights [64, 8], zero-init.
    pub w2: GpuTensor,
    /// Linear 2 bias [8], zero-init.
    pub b2: GpuTensor,
}

impl MoeGate {
    /// Construct with zero-initialized weights so initial gate = uniform 1/K.
    pub fn new(state_dim: usize, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let w1 = GpuTensor::zeros(&[state_dim, MOE_GATE_HIDDEN], stream)?;
        let b1 = GpuTensor::zeros(&[MOE_GATE_HIDDEN], stream)?;
        let w2 = GpuTensor::zeros(&[MOE_GATE_HIDDEN, MOE_NUM_EXPERTS], stream)?;
        let b2 = GpuTensor::zeros(&[MOE_NUM_EXPERTS], stream)?;
        Ok(Self { w1, b1, w2, b2 })
    }
}

/// Single expert MLP: h_s1[256] → 64 → h_s2[256] with bottleneck.
///
/// Each expert has independent Xavier init for initial differentiation.
pub struct MoeExpert {
    /// Linear 1 weights [256, 64], Xavier init.
    pub w1: GpuTensor,
    /// Linear 1 bias [64], zero-init.
    pub b1: GpuTensor,
    /// Linear 2 weights [64, 256], Xavier init.
    pub w2: GpuTensor,
    /// Linear 2 bias [256], zero-init.
    pub b2: GpuTensor,
}

impl MoeExpert {
    /// Construct with Xavier-initialized weights.
    ///
    /// `seed` provides per-expert distinct random init so experts have
    /// non-identical starting points.
    pub fn new(h_dim: usize, seed: u64, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        let w1 = init_xavier(
            &[h_dim, MOE_EXPERT_BOTTLENECK],
            seed,
            stream,
        )?;
        let b1 = GpuTensor::zeros(&[MOE_EXPERT_BOTTLENECK], stream)?;
        let w2 = init_xavier(
            &[MOE_EXPERT_BOTTLENECK, h_dim],
            seed.wrapping_add(1),
            stream,
        )?;
        let b2 = GpuTensor::zeros(&[h_dim], stream)?;
        Ok(Self { w1, b1, w2, b2 })
    }
}

/// Xavier-normal initialization helper.
///
/// One-time startup cost; not on the training hot-path.
fn init_xavier(shape: &[usize], seed: u64, stream: &Arc<CudaStream>) -> Result<GpuTensor, MLError> {
    let fan_in = shape.first().copied().unwrap_or(1) as f32;
    let fan_out = shape.get(1).copied().unwrap_or(1) as f32;
    let stddev = (2.0_f32 / (fan_in + fan_out)).sqrt();
    let normal = Normal::new(0.0_f32, stddev)
        .map_err(|e| MLError::ModelError(format!("MoE Xavier dist: {e}")))?;
    let numel: usize = shape.iter().product();
    let mut rng = StdRng::seed_from_u64(seed);
    let host: Vec<f32> = (0..numel).map(|_| rng.sample(normal)).collect();
    GpuTensor::from_host(&host, shape.to_vec(), stream)
        .map_err(|e| MLError::ModelError(format!("MoE init xavier: {e}")))
}
