//! `jepa::sonic_moe` — extraído do monólito na v0.10.0
//! (divisão SRP, crítica Qwen3.8 lote 3 nº 1; r-35 §2).

use super::*;


// ---------------------------------------------------------------------------
// Parte 5 — SonicMoE (IO/tile-aware)
// ---------------------------------------------------------------------------

/// MoE top-k com capacidade por especialista e perda de balanceamento.
///
/// **Prova SM-1 (limite de capacidade):** o roteamento tile-aware mantém no
/// máximo `C` tokens por especialista por ladrilho — o drop é determinístico
/// (ordem de chegada) e o nº de tokens descartados é `Σ_e max(0, load_e − C)`.
///
/// **Prova SM-2 (balanceamento):** `L_aux = E·Σ_e f_e·p_e`; com roteamento
/// auto-consistente (`f = p`), `L_aux = E·Σ p_e² ≥ 1` com igualdade sse
/// `p_e = 1/E` (Cauchy–Schchwarz) — o mínimo força o roteador uniforme.
pub struct SonicMoe {
    pub roteador: Linear,            // d → E
    pub especialistas: Vec<MlpJepa>, // E experts d → d
    pub n_top: usize,
    pub capacidade: usize,
    pub lambda_aux: f64,
}

pub struct SaidaMoe {
    pub saida: Tensor,
    pub perda_aux: Tensor,
    pub carga: Vec<usize>,  // tokens efetivamente processados por especialista
    pub descartados: usize, // tokens com especialistas cortados pela capacidade
}

impl SonicMoe {
    pub fn novo(d: usize, n_expertos: usize, n_top: usize, capacidade: usize, dev: &Device) -> Result<Self> {
        let roteador = Linear::sem_vies(d, n_expertos, 0.1, dev)?;
        let especialistas = (0..n_expertos)
            .map(|_| MlpJepa::novo(d, 2 * d, d, dev))
            .collect::<Result<Vec<_>>>()?;
        Ok(Self { roteador, especialistas, n_top, capacidade, lambda_aux: 0.01 })
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v = self.roteador.vars();
        for e in &self.especialistas {
            v.extend(e.vars());
        }
        v
    }

    /// Forward tile-aware: (n, d) → (n, d).
    pub fn fwd(&self, x: &Tensor) -> Result<SaidaMoe> {
        let n = x.dim(0)?;
        let d = x.dim(1)?;
        let e = self.especialistas.len();
        let dev = x.device();
        let logits = self.roteador.fwd(x)?; // (n, E)
        let probs = candle_nn::ops::softmax(&logits, candle_core::D::Minus1)?;

        // top-k por token (não diferenciável — pesos do top-k são dif. via probs)
        let probs_v = probs.to_vec2::<f32>()?;
        let mut carga = vec![0usize; e];
        let mut descartados = 0usize;
        // acumulador de saída (n, d) — construído por soma de contribuições
        let mut saida = Tensor::zeros((n, d), candle_core::DType::F32, &dev)?;
        let mut peso_total = Tensor::zeros((n, 1), candle_core::DType::F32, &dev)?;
        // por especialista (tile): tokens cujo top-k contém e e com carga < C
        for ei in 0..e {
            // candidatos: tokens ordenados por prob do especialista ei
            let mut cand: Vec<(usize, f32)> = (0..n)
                .map(|i| (i, probs_v[i][ei]))
                .collect();
            cand.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
            let mut escolhidos: Vec<usize> = Vec::new();
            for (i, _) in cand {
                if carga[ei] >= self.capacidade {
                    break;
                }
                // verifica se ei está no top-k do token i
                let mut ordem: Vec<(f32, usize)> =
                    (0..e).map(|k| (probs_v[i][k], k)).collect();
                ordem.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap());
                let top: Vec<usize> = ordem.iter().rev().take(self.n_top).map(|(_, k)| *k).collect();
                if top.contains(&ei) {
                    escolhidos.push(i);
                    carga[ei] += 1;
                }
            }
            if escolhidos.is_empty() {
                continue;
            }
            // tile do especialista: uma única passagem em lote
            let idx: Vec<u32> = escolhidos.iter().map(|&i| i as u32).collect();
            let n_e = idx.len();
            let idx_t = Tensor::from_vec(idx, (n_e,), &dev)?;
            let tile = x.index_select(&idx_t, 0)?;
            let y = self.especialistas[ei].fwd(&tile)?; // (m, d)
            // pesos: prob do especialista ei para cada token escolhido
            let w: Vec<f32> = escolhidos.iter().map(|&i| probs_v[i][ei]).collect();
            let n_w = escolhidos.len();
            let w_t = Tensor::from_vec(w.clone(), (n_w, 1), &dev)?;
            let contrib = y.broadcast_mul(&w_t)?;
            // scatter-add: acumula nas posições
            let mut buf = saida.to_vec2::<f32>()?;
            let contrib_v = contrib.to_vec2::<f32>()?;
            for (k, &i) in escolhidos.iter().enumerate() {
                for j in 0..d {
                    buf[i][j] += contrib_v[k][j];
                }
            }
            saida = Tensor::from_vec(
                buf.into_iter().flatten().collect::<Vec<f32>>(),
                (n, d),
                &dev,
            )?;
            let mut pt = peso_total.to_vec2::<f32>()?;
            for (k, &i) in escolhidos.iter().enumerate() {
                pt[i][0] += w[k];
            }
            peso_total = Tensor::from_vec(
                pt.into_iter().flatten().collect::<Vec<f32>>(),
                (n, 1),
                &dev,
            )?;
        }
        // tokens sem especialista (todos cortados) → descartados; saída 0
        let pt_v = peso_total.to_vec2::<f32>()?;
        for i in 0..n {
            if pt_v[i][0] == 0.0 {
                descartados += 1;
            }
        }
        // renormaliza pelo peso efetivo (evita viés da capacidade)
        let denom = peso_total.clamp(1e-6, f64::INFINITY)?;
        let saida = saida.broadcast_div(&denom)?;

        // perda auxiliar: E·Σ f_e·p_e (f = fração de tokens cujo top-1 é e)
        let mut f = vec![0f32; e];
        for i in 0..n {
            let (mut melhor, mut mp) = (0usize, f32::NEG_INFINITY);
            for k in 0..e {
                if probs_v[i][k] > mp {
                    mp = probs_v[i][k];
                    melhor = k;
                }
            }
            f[melhor] += 1.0 / n as f32;
        }
        let f_t = Tensor::from_vec(f, (e,), &dev)?;
        let p_media = probs.mean(0)?; // (E,)
        let perda_aux = ((&f_t * &p_media)?.sum_all()? * (e as f64))?;
        Ok(SaidaMoe { saida, perda_aux, carga, descartados })
    }
}