//! **MoEs de codificação/decodificação REDIRECIONÁVEIS com Mecanismo de
//! Atenção próprio e independente (paralelo)** (v0.6.0 — pedido do usuário).
//!
//! Cada `MoeRedirecionavel` possui:
//! - um roteador aprendido `d → E`;
//! - E especialistas `MlpJepa`, cada um com um PAPEL (Codificador ou
//!   Decodificador) — o conjunto de papéis é REDIRECIONÁVEL em tempo de
//!   execução (permutação), sem reconstrução do modelo;
//! - uma ATENÇÃO PRÓPRIA independente da atenção ADBA do tronco, com gate
//!   escalar nascido em ZERO (neutralidade no nascimento);
//! - capacidade por especialista e perda auxiliar de balanceamento.
//!
//! **Provas completas em `docs/memorias/r-25-*`; síntese:**
//! - **M-R1 (invariância sob redirecionamento):** o redirecionamento é uma
//!   permutação π dos papéis; o multiconjunto de atribuições
//!   (especialista, token) numa passada por papel é levado ao multiconjunto
//!   da passada π(papel) — a capacidade total e a perda auxiliar de
//!   balanceamento E·Σ f_e·p_e são invariantes (soma sobre todos os e).
//! - **M-R2 (neutralidade da atenção própria no nascimento):** com
//!   `g = 0`, `y = x` bit-a-bit; a perturbação é limitada por
//!   `‖Δy‖ ≤ L_soft·|g|·‖W_o·V‖` (Lipschitz do softmax com constante
//!   L_soft = 1 em probabilidade; estrutura idêntica ao T-P5).
//! - **M-R3 (paralelismo dos dois caminhos):** numa mesma entrada, as
//!   passadas Codificador e Decodificador tocam conjuntos DISJUNTOS de
//!   especialistas e compartilham apenas leitura — o DAG de computação tem
//!   duas sub-DAGs independentes ⇒ tempo paralelo (2 processadores) =
//!   max(T_enc, T_dec) + sobrecarga < T_enc + T_dec = tempo sequencial
//!   (modelo PRAM; medido empiricamente com threads).
//! - **M-R4 (fechamento da classe de funções):** redirecionar equivale a
//!   permutar colunas do roteador e re-rotular papéis; a classe de funções
//!   implementáveis é fechada sob π ⇒ expressividade inalterada.
//! - **M-R5 (water-filling por QE — Cauchy–Schwarz):** dadas capacidades
//!   totais C e erros por célula q_i, a alocação `c_i ∝ q_i` minimiza
//!   `Σ q_i²/c_i ≥ (Σq_i)²/C`, cota superior do erro refinado ponderado.

use candle_core::{Device, Result, Tensor, Var};
use crate::camadas::Linear;
use crate::jepa::MlpJepa;
use crate::rng::RngDeterministico;
use crate::treino::AdamWManual;

// ---------------------------------------------------------------------------
// Atenção própria (independente e paralela à ADBA do tronco)
// ---------------------------------------------------------------------------

pub struct AtencaoPropria {
    pub q: Linear,
    pub k: Linear,
    pub v: Linear,
    pub o: Linear,
    /// gate escalar — nasce ZERO (M-R2: neutralidade no nascimento).
    pub gate: Var,
    pub n_cabecas: usize,
}

impl AtencaoPropria {
    pub fn novo(d: usize, n_cabecas: usize, dev: &Device) -> Result<Self> {
        let desvio = 1.0 / ((d * 2) as f64).sqrt();
        let mk = || Linear::sem_vies(d, d, desvio, dev);
        let gate = Tensor::zeros((), candle_core::DType::F32, dev)?;
        Ok(AtencaoPropria {
            q: mk()?, k: mk()?, v: mk()?, o: mk()?,
            gate: Var::from_tensor(&gate)?,
            n_cabecas,
        })
    }

    /// x (n, d) → y = x + g·attn(x) — multi-cabeça compacta (sem máscara:
    /// atenção plena entre os tokens da entrada do MoE).
    pub fn fwd(&self, x: &Tensor) -> Result<Tensor> {
        let d = x.dim(candle_core::D::Minus1)?;
        let d_cab = d / self.n_cabecas;
        let n = x.dim(0)?;
        let forma = (self.n_cabecas, n, d_cab);
        let q = self.q.fwd(x)?.reshape(forma)?;
        let k = self.k.fwd(x)?.reshape(forma)?;
        let v = self.v.fwd(x)?.reshape(forma)?;
        let scores = q.matmul(&k.transpose(1, 2)?)?.affine(1.0 / (d_cab as f64).sqrt(), 0.0)?;
        let attn = candle_nn::ops::softmax(&scores, candle_core::D::Minus1)?;
        let ctx = attn.matmul(&v)?; // (h, n, d_cab)
        let ctx = ctx.reshape((n, d))?;
        let saida = self.o.fwd(&ctx)?;
        // y = x + g·saida (g=0 no nascimento ⇒ y = x — M-R2)
        let g = self.gate.as_tensor();
        let forma: Vec<usize> = saida.dims().to_vec();
        let delta = saida.broadcast_mul(&g.broadcast_as(forma)?)?;
        Ok((x + &delta)?)
    }

    pub fn vars(&self) -> Vec<Var> {
        let mut v = self.q.vars();
        v.extend(self.k.vars());
        v.extend(self.v.vars());
        v.extend(self.o.vars());
        v.push(self.gate.clone());
        v
    }
}

// ---------------------------------------------------------------------------
// MoE redirecionável
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Papel {
    Codificador,
    Decodificador,
}

impl Papel {
    pub fn nome(&self) -> &'static str {
        match self { Papel::Codificador => "codificador", Papel::Decodificador => "decodificador" }
    }
}

pub struct SaidaMoeR {
    pub saida: Tensor,
    pub perda_aux: Tensor,
    pub carga: Vec<usize>,
    pub descartados: usize,
    /// especialistas efetivamente usados nesta passada (índices globais).
    pub usados: Vec<usize>,
}

pub struct MoeRedirecionavel {
    pub roteador: Linear,             // d → E
    pub especialistas: Vec<MlpJepa>,  // E experts d → d
    pub papeis: Vec<Papel>,           // papel corrente de cada especialista
    pub atencao: AtencaoPropria,      // atenção PRÓPRIA (M-R2/M-R3)
    pub n_top: usize,
    pub capacidade: usize,
    pub lambda_aux: f64,
}

impl MoeRedirecionavel {
    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<_>>>()?;
        // papéis iniciais alternados (metade cod, metade dec)
        let papeis = (0..n_expertos)
            .map(|i| if i % 2 == 0 { Papel::Codificador } else { Papel::Decodificador })
            .collect();
        Ok(MoeRedirecionavel {
            roteador,
            especialistas,
            papeis,
            atencao: AtencaoPropria::novo(d, 2, dev)?,
            n_top: n_top.max(1),
            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.extend(self.atencao.vars());
        v
    }

    /// **M-R1/M-R4:** redireciona os papéis segundo a permutação `permuta`
    /// (novo_papel[i] = papel_anterior[permuta[i]]). `redirecionar_todos`
    /// troca Codificador ↔ Decodificador.
    pub fn redirecionar(&mut self, permuta: &[usize]) {
        let antigos = self.papeis.clone();
        for (i, &p) in permuta.iter().enumerate() {
            if p < antigos.len() {
                self.papeis[i] = antigos[p];
            }
        }
    }

    pub fn redirecionar_todos(&mut self) {
        for p in self.papeis.iter_mut() {
            *p = match *p {
                Papel::Codificador => Papel::Decodificador,
                Papel::Decodificador => Papel::Codificador,
            };
        }
    }

    fn indices_do_papel(&self, papel: Papel) -> Vec<usize> {
        self.papeis.iter().enumerate().filter(|(_, p)| **p == papel).map(|(i, _)| i).collect()
    }

    /// Passada por papel: atenção própria → roteamento top-k DENTRO do
    /// subconjunto de especialistas do papel → capacidade → soma ponderada.
    ///
    /// ⚠️ **v0.8.0 — bug v080-1 catalogado (r-33 §5):** esta versão
    /// materializa os valores com `to_vec2`/`from_vec` no acúmulo, o que
    /// SECCIONA o grafo de autograd — especialistas e roteador NÃO recebem
    /// gradiente por esta saída (só a atenção própria treina). Mantida como
    /// referência de VALORES (o teste M-R6 verifica a equivalência numérica
    /// com a versão vetorizada); para TREINO use `fwd_vetorizado`.
    pub fn fwd(&self, x: &Tensor, papel: Papel) -> Result<SaidaMoeR> {
        let y = self.atencao.fwd(x)?; // (n, d) — atenção própria independente
        let n = y.dim(0)?;
        let d = y.dim(1)?;
        let dev = y.device();
        let do_papel = self.indices_do_papel(papel);
        let e_local = do_papel.len();
        let mut descartados = 0usize;
        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)?;
        // probs de roteamento só sobre especialistas do papel (mascarado)
        let logits_todos = self.roteador.fwd(&y)?; // (n, E)
        let mut probs_v: Vec<Vec<f32>> = logits_todos.to_vec2::<f32>()?;
        for linha in probs_v.iter_mut() {
            // máscara: -inf fora do papel
            for (e, lv) in linha.iter_mut().enumerate() {
                if !do_papel.contains(&e) {
                    *lv = f32::NEG_INFINITY;
                }
            }
        }
        // softmax manual por linha (com -inf → 0)
        let probs: Vec<Vec<f32>> = probs_v
            .iter()
            .map(|linha| {
                let m = linha.iter().copied().fold(f32::NEG_INFINITY, f32::max);
                let exps: Vec<f32> = linha.iter().map(|v| (v - m).exp()).collect();
                let s: f32 = exps.iter().sum();
                if s <= 0.0 {
                    vec![0.0; linha.len()]
                } else {
                    exps.iter().map(|e| e / s).collect()
                }
            })
            .collect();
        // seleção top-k com capacidade (código extraído — MESMA semântica
        // de desempate da v0.6.0, compartilhada com fwd_vetorizado)
        let (designacao, carga) = selecionar_com_capacidade(
            &probs,
            &do_papel,
            n,
            self.n_top,
            self.capacidade,
            self.especialistas.len(),
        );
        let usados: Vec<usize> = designacao.iter().map(|(ei, _)| *ei).collect();
        for (ei, escolhidos) in designacao {
            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 = y.index_select(&idx_t, 0)?;
            let y_e = self.especialistas[ei].fwd(&tile)?;
            let w: Vec<f32> = escolhidos.iter().map(|&i| probs[i][ei]).collect();
            let w_t = Tensor::from_vec(w.clone(), (n_e, 1), &dev)?;
            let contrib = y_e.broadcast_mul(&w_t)?;
            let mut buf = saida.to_vec2::<f32>()?;
            let contrib_v = contrib.to_vec2::<f32>()?;
            for (kk, &i) in escolhidos.iter().enumerate() {
                for j in 0..d {
                    buf[i][j] += contrib_v[kk][j];
                }
            }
            saida = Tensor::from_vec(buf.into_iter().flatten().collect::<Vec<f32>>(), (n, d), &dev)?;
            let mut pt = peso_total.to_vec2::<f32>()?;
            for (kk, &i) in escolhidos.iter().enumerate() {
                pt[i][0] += w[kk];
            }
            peso_total = Tensor::from_vec(pt.into_iter().flatten().collect::<Vec<f32>>(), (n, 1), &dev)?;
        }
        let pt_v = peso_total.to_vec2::<f32>()?;
        for i in 0..n {
            if pt_v[i][0] == 0.0 {
                descartados += 1;
            }
        }
        let denom = peso_total.clamp(1e-6, f64::INFINITY)?;
        let saida = saida.broadcast_div(&denom)?;
        // saída residual: y + roteamento (atenção já embutida em y)
        let saida = (&y + &saida)?;
        // perda auxiliar sobre o SUBCONJUNTO do papel (invariante global — M-R1)
        let mut f = vec![0f32; e_local.max(1)];
        let mut m = 0usize;
        for i in 0..n {
            let (mut melhor, mut mp) = (0usize, f32::NEG_INFINITY);
            for (li, &g) in do_papel.iter().enumerate() {
                if probs[i][g] > mp {
                    mp = probs[i][g];
                    melhor = li;
                }
            }
            f[melhor] += 1.0;
            m += 1;
        }
        if m > 0 && e_local > 0 {
            for v in f.iter_mut() {
                *v /= m as f32;
            }
        }
        let f_t = Tensor::from_vec(f, (e_local.max(1),), &dev)?;
        let p_media_local: Vec<f32> = (0..e_local)
            .map(|li| {
                let g = do_papel[li];
                (0..n).map(|i| probs[i][g]).sum::<f32>() / n.max(1) as f32
            })
            .collect();
        let p_t = Tensor::from_vec(p_media_local, (e_local.max(1),), &dev)?;
        let perda_aux = ((&f_t * &p_t)?.sum_all()? * (e_local.max(1) as f64))?;
        Ok(SaidaMoeR { saida, perda_aux, carga, descartados, usados })
    }

    /// **M-R6 (v0.8.0):** passada VETORIZADA — roteamento, acúmulo e perda
    /// auxiliar em OPERAÇÕES DE TENSOR rastreadas (proposta Qwen-3.8 nº 4
    /// IMPLEMENTADA de verdade: o esboço original terminava em
    /// `unimplemented!()`; a fórmula da perda auxiliar `E·Σ f_e·p̄_e`
    /// estava CORRETA — Switch Transformer — e é mantida).
    ///
    /// **Correção do bug v080-1:** a acumulação usa `index_add` (op
    /// rastreada pelo autograd do candle — `IndexAdd` em op.rs) em vez do
    /// round-trip `to_vec2`/`from_vec`; os pesos de roteamento saem do
    /// PRÓPRIO tensor `probs` (vivo até o roteador); a perda auxiliar usa
    /// `p̄` do tensor vivo com `f` constante (padrão Switch: argmax é
    /// não-diferenciável, a média de probs carrega o gradiente). O
    /// gradiente flui por TODOS os caminhos (teste M-R7) com VALORES
    /// idênticos aos da `fwd` (mesma seleção — teste M-R6).
    pub fn fwd_vetorizado(&self, x: &Tensor, papel: Papel) -> Result<SaidaMoeR> {
        let y = self.atencao.fwd(x)?; // (n, d)
        let n = y.dim(0)?;
        let d = y.dim(1)?;
        let dev = y.device();
        let e_total = self.especialistas.len();
        let do_papel = self.indices_do_papel(papel);
        let e_local = do_papel.len();
        // 1) máscara do papel + softmax EM TENSOR — grafo VIVO até o roteador
        let mut mascara_v = vec![f32::NEG_INFINITY; e_total];
        for &e in &do_papel {
            mascara_v[e] = 0.0;
        }
        let mascara = Tensor::from_vec(mascara_v, (1, e_total), &dev)?;
        let logits = self.roteador.fwd(&y)?;
        let mascarados = logits.broadcast_add(&mascara)?;
        let probs = candle_nn::ops::softmax(&mascarados, candle_core::D::Minus1)?; // (n, E) VIVO
        // 2) seleção discreta (única materialização — o roteamento duro é
        // não-diferenciável por natureza; os PESOS ficam vivos no passo 3)
        let probs_v = probs.to_vec2::<f32>()?;
        let (designacao, carga) =
            selecionar_com_capacidade(&probs_v, &do_papel, n, self.n_top, self.capacidade, e_total);
        let usados: Vec<usize> = designacao.iter().map(|(ei, _)| *ei).collect();
        // 3) acúmulo VIVO: index_add é rastreada — especialistas recebem
        // gradiente pelas contribuições, e o roteador pelos pesos w
        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)?;
        for (ei, escolhidos) in &designacao {
            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 = y.index_select(&idx_t, 0)?; // VIVO
            let y_e = self.especialistas[*ei].fwd(&tile)?; // VIVO (gradiente!)
            // w = probs[i][ei] VIVO: gradiente flui ao roteador
            let linhas = probs.index_select(&idx_t, 0)?; // (n_e, E) VIVO
            let w_t = linhas.narrow(1, *ei, 1)?.contiguous()?; // (n_e, 1) VIVO
            let contrib = y_e.broadcast_mul(&w_t)?; // (n_e, d) VIVO
            saida = saida.index_add(&idx_t, &contrib, 0)?; // VIVO (M-R6)
            peso_total = peso_total.index_add(&idx_t, &w_t, 0)?; // VIVO
        }
        // 4) normalização idêntica à fwd (clamp 1e-6 nos descartados)
        let pt_v = peso_total.to_vec2::<f32>()?; // estatística apenas
        let mut descartados = 0usize;
        for i in 0..n {
            if pt_v[i][0] == 0.0 {
                descartados += 1;
            }
        }
        let denom = peso_total.clamp(1e-6, f64::INFINITY)?;
        let saida = saida.broadcast_div(&denom)?;
        let saida = (&y + &saida)?; // residual — como a fwd
        // 5) perda auxiliar EM TENSOR (fórmula M-R1/Switch: e_local·Σ f_e·p̄_e)
        let cols: Vec<u32> = do_papel.iter().map(|&e| e as u32).collect();
        let cols_t = Tensor::from_vec(cols, (e_local.max(1),), &dev)?;
        let probs_local = probs.index_select(&cols_t, 1)?; // (n, e_local) VIVO
        let argmax_v = probs_local.argmax(candle_core::D::Minus1)?.to_vec1::<u32>()?;
        let mut f = vec![0f32; e_local.max(1)];
        for &a in &argmax_v {
            f[a as usize] += 1.0;
        }
        if n > 0 {
            for v in f.iter_mut() {
                *v /= n as f32;
            }
        }
        let f_t = Tensor::from_vec(f, (e_local.max(1),), &dev)?;
        let p_barra = probs_local.mean_keepdim(0)?.flatten_all()?; // (e_local,) VIVO
        let perda_aux = ((&f_t * &p_barra)?.sum_all()? * (e_local.max(1) as f64))?;
        Ok(SaidaMoeR { saida, perda_aux, carga, descartados, usados })
    }

    /// **M-R3:** passadas Codificador e Decodificador em DUAS THREADS
    /// (DAGs independentes — especialistas disjuntos por papel). Devolve
    /// (enc, dec, tempo_paralelo_s, tempo_sequencial_s, speedup).
    pub fn fwd_ambos_par(&self, x: &Tensor) -> Result<(SaidaMoeR, SaidaMoeR, f64, f64, f64)> {
        use std::time::Instant;
        // sequencial (medição)
        let t0 = Instant::now();
        let enc_seq = self.fwd(x, Papel::Codificador)?;
        let dec_seq = self.fwd(x, Papel::Decodificador)?;
        let t_seq = t0.elapsed().as_secs_f64();
        // paralelo (2 threads com escopo — referências compartilhadas)
        let t0 = Instant::now();
        let (enc, dec) = {
            let xref = &x;
            let sref = &self;
            std::thread::scope(|sc| {
                let h1 = sc.spawn(move || sref.fwd(xref, Papel::Codificador));
                let h2 = sc.spawn(move || sref.fwd(xref, Papel::Decodificador));
                let r1 = h1.join().unwrap_or_else(|_| Err(candle_core::Error::Msg("thread enc".into())));
                let r2 = h2.join().unwrap_or_else(|_| Err(candle_core::Error::Msg("thread dec".into())));
                (r1, r2)
            })
        };
        let (enc, dec) = (enc?, dec?);
        let t_par = t0.elapsed().as_secs_f64();
        // determinismo: paralelo == sequencial (mesmas entradas, especialistas disjuntos)
        let dif = (&enc.saida - &enc_seq.saida)?.sqr()?.sum_all()?.to_scalar::<f32>()? as f64
            + (&dec.saida - &dec_seq.saida)?.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
        assert!(dif < 1e-9, "paralelo deve ser bit-a-bit igual ao sequencial (dif={dif})");
        let speedup = t_seq / t_par.max(1e-12);
        Ok((enc, dec, t_par, t_seq, speedup))
    }
}

// ---------------------------------------------------------------------------
// M-R5 — alocação water-filling por QE
// ---------------------------------------------------------------------------

/// Seleção top-k com capacidade por especialista (v0.8.0 — extraída
/// VERBATIM do loop da `fwd` v0.6.0: mesma ordem de candidatos, mesmos
/// desempates estáveis de `sort_by`). Compartilhada por `fwd` e
/// `fwd_vetorizado` ⇒ as duas atribuem EXATAMENTE os mesmos tokens aos
/// mesmos especialistas (base do teste M-R6 — equivalência de valores).
/// Devolve (designação por especialista, carga global por especialista).
fn selecionar_com_capacidade(
    probs: &[Vec<f32>],
    do_papel: &[usize],
    n: usize,
    n_top: usize,
    capacidade: usize,
    n_especialistas: usize,
) -> (Vec<(usize, Vec<usize>)>, Vec<usize>) {
    let mut carga = vec![0usize; n_especialistas];
    let mut designacao: Vec<(usize, Vec<usize>)> = Vec::new();
    for &ei in do_papel.iter() {
        // candidatos ordenados por prob do especialista ei (desc)
        let mut cand: Vec<(usize, f32)> = (0..n).map(|i| (i, probs[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] >= capacidade {
                break;
            }
            // ei está no top-k(papel) do token i?
            let mut ordem: Vec<(f32, usize)> =
                do_papel.iter().map(|&g| (probs[i][g], g)).collect();
            ordem.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
            let top: Vec<usize> = ordem.iter().rev().take(n_top).map(|(_, g)| *g).collect();
            if top.contains(&ei) {
                escolhidos.push(i);
                carga[ei] += 1;
            }
        }
        if escolhidos.is_empty() {
            continue;
        }
        designacao.push((ei, escolhidos));
    }
    (designacao, carga)
}

/// Aloca `capacidade_total` especialistas entre células com QE `qes`
/// proporcionalmente ao QE (Cauchy–Schwarz: minimiza Σ q_i²/c_i).
/// Devolve c_i ≥ 1 com Σ c_i = capacidade_total (se capacidade ≥ células).
pub fn alocar_por_qe(qes: &[f64], capacidade_total: usize) -> Vec<usize> {
    let m = qes.len();
    if m == 0 {
        return Vec::new();
    }
    let soma: f64 = qes.iter().sum::<f64>().max(1e-12);
    let mut c: Vec<usize> = qes
        .iter()
        .map(|q| ((q / soma) * capacidade_total as f64).round() as usize)
        .collect();
    // garante mínimo 1 e corrige a soma por maior resíduo
    for ci in c.iter_mut() {
        *ci = (*ci).max(1);
    }
    let mut total: usize = c.iter().sum();
    while total > capacidade_total.max(m) {
        // remove da maior alocação
        let (mut imax, mut vmax) = (0usize, 0usize);
        for (i, &v) in c.iter().enumerate() {
            if v > vmax {
                vmax = v;
                imax = i;
            }
        }
        if c[imax] > 1 {
            c[imax] -= 1;
        }
        total -= 1;
    }
    while total < capacidade_total {
        let (mut imax, mut vmax) = (0usize, 0.0f64);
        for (i, (&v, &q)) in c.iter().zip(qes.iter()).enumerate() {
            let res = q / (v as f64);
            if res > vmax {
                vmax = res;
                imax = i;
            }
        }
        c[imax] += 1;
        total += 1;
    }
    c
}

// ---------------------------------------------------------------------------
// Aplicação JEPA — codificador/preditor redirecionáveis com MoE
// ---------------------------------------------------------------------------

pub struct RelatorioJepaMoe {
    pub modo: &'static str,
    pub perda_pred_inicial: f64,
    pub perda_pred_final: f64,
    pub perda_vc_final: f64,
    pub std_emb: f64,
    pub carga_max: usize,
    pub descartados: usize,
}

/// Treino JEPA compacto com três braços (ablação honesta):
/// - `"pequena"`: codificador/preditor MLP simples (baseline);
/// - `"densa"`: MLP com o TOTAL de parâmetros dos especialistas (capacidade
///   equiparada, SEM roteamento — isola o ganho do roteamento);
/// - `"moe"`: `MoeRedirecionavel` nos dois caminhos (Codificador no
///   codificador, Decodificador no preditor) com atenção própria.
pub fn treinar_jepa_com_moe(
    patches: &Tensor,
    modo: &'static str,
    passos: usize,
    frac_mascara: f64,
    semente: u64,
) -> Result<RelatorioJepaMoe> {
    let dev = patches.device();
    let n = patches.dim(0)?;
    let d = patches.dim(1)?;
    let mut rng = RngDeterministico::novo(semente);
    // alvo = patches originais (detach); contexto = patches mascarados
    let alvo = patches.detach();
    let contexto = {
        let mut v = patches.flatten_all()?.to_vec1::<f32>()?;
        for x in v.iter_mut() {
            if rng.uniforme() < frac_mascara {
                *x = 0.0;
            }
        }
        Tensor::from_vec(v, (n, d), &dev)?
    };
    let e = 4usize; // especialistas por papel no braço moe
    let n_top = 2;
    let capacidade = n;

    // braços
    enum Braco {
        Pequena(MlpJepa, MlpJepa),
        Densa(MlpJepa, MlpJepa),
        Moe(MoeRedirecionavel),
    }
    let braco = match modo {
        "pequena" => Braco::Pequena(
            MlpJepa::novo(d, 2 * d, d, &dev)?,
            MlpJepa::novo(d, 2 * d, d, &dev)?,
        ),
        "densa" => {
            // capacidade equiparada: oculta = soma das ocultas dos especialistas
            let d_oculta = 2 * d * e;
            Braco::Densa(
                MlpJepa::novo(d, d_oculta, d, &dev)?,
                MlpJepa::novo(d, d_oculta, d, &dev)?,
            )
        }
        _ => Braco::Moe(MoeRedirecionavel::novo(d, 2 * e, n_top, capacidade, &dev)?),
    };

    let perda_de = |b: &Braco, x: &Tensor| -> Result<(Tensor, Tensor, usize)> {
        match b {
            Braco::Pequena(e1, p1) => {
                let z = e1.fwd(x)?;
                let yh = p1.fwd(&z)?;
                let dif = (&yh - &alvo)?;
                Ok(((dif.sqr()?.sum_all()? / (n as f64 * d as f64))?,
                    z.sqr()?.mean_all()?, 0))
            }
            Braco::Densa(e1, p1) => {
                let z = e1.fwd(x)?;
                let yh = p1.fwd(&z)?;
                let dif = (&yh - &alvo)?;
                Ok(((dif.sqr()?.sum_all()? / (n as f64 * d as f64))?,
                    z.sqr()?.mean_all()?, 0))
            }
            Braco::Moe(m) => {
                // v0.8.0 (bug v080-1): fwd_vetorizado — gradiente flui aos
                // especialistas e ao roteador (a fwd seccionava o grafo)
                let s_enc = m.fwd_vetorizado(x, Papel::Codificador)?;
                let s_dec = m.fwd_vetorizado(&s_enc.saida, Papel::Decodificador)?;
                let dif = (&s_dec.saida - &alvo)?;
                let lp = (dif.sqr()?.sum_all()? / (n as f64 * d as f64))?;
                // perda total = preditiva + aux balanceamento (ambos papéis)
                let lt = (lp.affine(1.0, 0.0)? + s_enc.perda_aux.affine(m.lambda_aux, 0.0)?
                    + s_dec.perda_aux.affine(m.lambda_aux, 0.0)?)?;
                let carga_max = s_enc.carga.iter().copied().fold(0, usize::max)
                    .max(s_dec.carga.iter().copied().fold(0, usize::max));
                Ok((lt, lp, carga_max))
            }
        }
    };

    let (mut opt, vars): (AdamWManual, Vec<Var>) = match &braco {
        Braco::Pequena(a, b) => {
            let v = [a.vars(), b.vars()].concat();
            (AdamWManual::novo(v.clone(), 3e-3, 0.0), v)
        }
        Braco::Densa(a, b) => {
            let v = [a.vars(), b.vars()].concat();
            (AdamWManual::novo(v.clone(), 3e-3, 0.0), v)
        }
        Braco::Moe(m) => (AdamWManual::novo(m.vars(), 3e-3, 0.0), m.vars()),
    };
    // atenção própria: gate nasce 0; sobe 1 passo em 8 (aquecimento honesto
    // do mecanismo — T-P5 análogo: neutralidade no nascimento, gradiente flui)
    if let Braco::Moe(m) = &braco {
        let g0 = m.atencao.gate.as_tensor().to_scalar::<f32>()? as f64;
        let _ = g0;
    }

    let (l0, _, _) = perda_de(&braco, &contexto)?;
    let perda_inicial = l0.to_scalar::<f32>()? as f64;
    let mut perda_final = perda_inicial;
    let mut vc_final = 0.0;
    let mut carga_max = 0usize;
    let mut descartados = 0usize;
    let mut std_emb = 0.0;
    for passo in 0..passos {
        let (lt, lp, cm) = perda_de(&braco, &contexto)?;
        if let Braco::Moe(m) = &braco {
            // abre o gate gradualmente (0 → 1 em 25% dos passos)
            let alvo_g = ((passo as f64 / (passos as f64 / 4.0).max(1.0)).min(1.0)) as f32;
            let g_atual = m.atencao.gate.as_tensor().to_scalar::<f32>()? as f64;
            if (alvo_g as f64 - g_atual).abs() > 1e-9 {
                m.atencao.gate.set(&Tensor::new(alvo_g, &dev)?)?;
            }
        }
        let grads = crate::treino::grads_achatados(&lt, &vars)?;
        AdamWManual::clipar(&mut grads.clone(), 1.0);
        let mut g2 = grads;
        AdamWManual::clipar(&mut g2, 5.0);
        opt.passo(g2);
        perda_final = lp.to_scalar::<f32>()? as f64;
        carga_max = cm;
        // VC / std do embedding (para o relatório)
        let (_, vc, _) = perda_de(&braco, &contexto)?;
        vc_final = vc.to_scalar::<f32>()? as f64;
        std_emb = {
            let z = match &braco {
                Braco::Pequena(a, _) => a.fwd(&contexto)?,
                Braco::Densa(a, _) => a.fwd(&contexto)?,
                Braco::Moe(m) => m.fwd(&contexto, Papel::Codificador)?.saida,
            };
            z.flatten_all()?.to_vec1::<f32>()?.iter().map(|x| *x as f64).collect::<Vec<f64>>()
                .chunks(n).map(|c| {
                    let m = c.iter().sum::<f64>() / c.len() as f64;
                    (c.iter().map(|x| (x - m).powi(2)).sum::<f64>() / c.len() as f64).sqrt()
                }).sum::<f64>() / n as f64
        };
    }
    if let Braco::Moe(m) = &braco {
        let s = m.fwd_vetorizado(&contexto, Papel::Codificador)?;
        descartados = s.descartados;
    }
    Ok(RelatorioJepaMoe {
        modo,
        perda_pred_inicial: perda_inicial,
        perda_pred_final: perda_final,
        perda_vc_final: vc_final,
        std_emb,
        carga_max,
        descartados,
    })
}

// ---------------------------------------------------------------------------
// Aplicação SOM — refinamento residual com alocação por QE (M-R5)
// ---------------------------------------------------------------------------

pub struct RelatorioSomMoe {
    pub qe_antes: f64,
    pub qe_depois: f64,
    pub qe_alocacao_uniforme: f64,
    pub carga_por_celula: Vec<usize>,
    pub tempo_treino_s: f64,
    /// qualidade da DECODIFICAÇÃO após redirecionamento (mse w_bmu → x).
    pub mse_decodificador: f64,
}

/// Refina um SOM 4D treinado com um MoE redirecionável:
/// 1. aloca especialistas por célula ∝ QE (water-filling — M-R5);
/// 2. treina o papel Codificador como corretor residual
///    `x̂ = w_bmu(x) + E(x)`;
/// 3. mede QE′; repete com alocação uniforme (ablação do water-filling);
/// 4. REDIRECIONA os especialistas para Decodificador e treina
///    `x̂ = E(w_bmu)` (reconstrução a partir do código) — mede o MSE.
pub fn refinar_som_com_moe(
    dados: &Tensor,
    centroides: &[Vec<f64>],
    passos: usize,
    semente: u64,
) -> Result<RelatorioSomMoe> {
    use std::time::Instant;
    let t0 = Instant::now();
    let dev = dados.device();
    let n = dados.dim(0)?;
    let d = dados.dim(1)?;
    let m = centroides.len();
    // BMU + QE por célula (tolera F32/F64)
    let dados_v: Vec<Vec<f64>> = match dados.to_vec2::<f64>() {
        Ok(v) => v,
        Err(_) => dados
            .to_vec2::<f32>()?
            .iter()
            .map(|l| l.iter().map(|x| *x as f64).collect())
            .collect(),
    };
    let bmu: Vec<usize> = dados_v
        .iter()
        .map(|x| {
            centroides
                .iter()
                .enumerate()
                .min_by(|(i, a), (j, b)| {
                    let da: f64 = a.iter().zip(x).map(|(w, v)| (w - v).powi(2)).sum();
                    let db: f64 = b.iter().zip(x).map(|(w, v)| (w - v).powi(2)).sum();
                    da.partial_cmp(&db).unwrap_or(std::cmp::Ordering::Equal).then(i.cmp(j))
                })
                .map(|(i, _)| i)
                .unwrap_or(0)
        })
        .collect();
    let mut qe_cel = vec![0f64; m];
    let mut cont_cel = vec![0usize; m];
    for (i, b) in bmu.iter().enumerate() {
        let dist: f64 = centroides[*b].iter().zip(&dados_v[i]).map(|(w, v)| (w - v).powi(2)).sum();
        qe_cel[*b] += dist.sqrt();
        cont_cel[*b] += 1;
    }
    for i in 0..m {
        if cont_cel[i] > 0 {
            qe_cel[i] /= cont_cel[i] as f64;
        }
    }
    let qe_antes: f64 = dados_v
        .iter()
        .enumerate()
        .map(|(i, b)| {
            let w = &centroides[bmu[i]];
            (w.iter().zip(b).map(|(a, v)| (a - v).powi(2)).sum::<f64>()).sqrt()
        })
        .sum::<f64>()
        / n as f64;

    // ---------- braço water-filling ----------
    // capacidade = m células + 8 especialistas EXTRA distribuídos ∝ QE
    // (com C < m o water-filling degenera em uniforme — mínimo 1 por célula)
    let capacidade_total = m + 8;
    let aloca_wf = alocar_por_qe(&qe_cel, capacidade_total);
    let qe_wf = treinar_corretor(dados, centroides, &bmu, &aloca_wf, passos, semente, &dev)?;
    // ---------- braço uniforme (ablação: MESMO total e MESMA semente —
    // única diferença é a alocação; round-robin soma exatamente C) ----------
    let aloca_un = {
        let mut a = vec![1usize; m];
        let mut i = 0usize;
        while a.iter().sum::<usize>() < capacidade_total {
            a[i % m] += 1;
            i += 1;
        }
        a
    };
    let qe_un = treinar_corretor(dados, centroides, &bmu, &aloca_un, passos, semente, &dev)?;

    // ---------- redirecionamento → decodificador ----------
    // mesma pool de especialistas, agora aprendendo w_bmu → x
    let mse_dec = treinar_decodificador(dados, centroides, &bmu, &aloca_wf, passos, semente + 2, &dev)?;

    Ok(RelatorioSomMoe {
        qe_antes,
        qe_depois: qe_wf,
        qe_alocacao_uniforme: qe_un,
        carga_por_celula: aloca_wf,
        tempo_treino_s: t0.elapsed().as_secs_f64(),
        mse_decodificador: mse_dec,
    })
}

/// Corretor residual: x̂ = w_bmu + E_{célula}(x). Treina com AdamW.
fn treinar_corretor(
    dados: &Tensor,
    centroides: &[Vec<f64>],
    bmu: &[usize],
    alocacao: &[usize],
    passos: usize,
    semente: u64,
    dev: &Device,
) -> Result<f64> {
    let n = dados.dim(0)?;
    let d = dados.dim(1)?;
    // um "especialista" = MlpJepa por CÉLULA (capacidade alocada = nº de
    // especialistas extra treinados nessa célula, somados em ensemble)
    let n_esp_tot: usize = alocacao.iter().sum();
    let mut ensembles: Vec<Vec<MlpJepa>> = Vec::with_capacity(centroides.len());
    for c in alocacao {
        let mut v = Vec::new();
        for _ in 0..(*c).max(1) {
            v.push(MlpJepa::novo(d, 4 * d, d, dev)?);
        }
        ensembles.push(v);
    }
    let mut vars: Vec<Var> = Vec::new();
    for ens in &ensembles {
        for e in ens {
            vars.extend(e.vars());
        }
    }
    let mut opt = AdamWManual::novo(vars.clone(), 3e-3, 0.0);
    // tensores de apoio
    let w_bmu = {
        let v: Vec<f32> = (0..n).flat_map(|i| centroides[bmu[i]].iter().map(|x| *x as f32).collect::<Vec<f32>>()).collect();
        Tensor::from_vec(v, (n, d), dev)?
    };
    let alvo_res = (dados - &w_bmu)?;
    for _ in 0..passos {
        let mut soma = Tensor::zeros((), candle_core::DType::F32, dev)?;
        for (i, ens) in ensembles.iter().enumerate() {
            // tokens da célula i
            let idx: Vec<u32> = (0..n).filter(|&k| bmu[k] == i).map(|k| k as u32).collect();
            if idx.is_empty() {
                continue;
            }
            let m_idx = idx.len();
            let it = Tensor::from_vec(idx, (m_idx,), dev)?;
            let x_c = dados.index_select(&it, 0)?;
            let r_c = alvo_res.index_select(&it, 0)?;
            // ensemble: média das saídas
            let mut y = Tensor::zeros((m_idx, d), candle_core::DType::F32, dev)?;
            for e in ens {
                y = (y.affine(1.0, 0.0)? + e.fwd(&x_c)?.affine(1.0, 0.0)?)?;
            }
            let k = ens.len() as f64;
            let yh = y.affine(1.0 / k, 0.0)?;
            let dif = (&yh - &r_c)?;
            soma = (soma.affine(1.0, 0.0)? + dif.sqr()?.sum_all()?.affine(1.0 / (n as f64 * d as f64), 0.0)?)?;
        }
        let _ = n_esp_tot;
        let grads = crate::treino::grads_achatados(&soma, &vars)?;
        let mut g = grads;
        AdamWManual::clipar(&mut g, 5.0);
        opt.passo(g);
    }
    // QE refinado
    let mut qe = 0f64;
    for (i, ens) in ensembles.iter().enumerate() {
        let idx: Vec<u32> = (0..n).filter(|&k| bmu[k] == i).map(|k| k as u32).collect();
        if idx.is_empty() {
            continue;
        }
        let m_idx = idx.len();
        let it = Tensor::from_vec(idx, (m_idx,), dev)?;
        let x_c = dados.index_select(&it, 0)?;
        let mut y = Tensor::zeros((m_idx, d), candle_core::DType::F32, dev)?;
        for e in ens {
            y = (y.affine(1.0, 0.0)? + e.fwd(&x_c)?.affine(1.0, 0.0)?)?;
        }
        let k = ens.len() as f64;
        let yh = y.affine(1.0 / k, 0.0)?;
        let xhat = (w_bmu.index_select(&it, 0)? + &yh)?;
        let dif = (&x_c - &xhat)?;
        qe += dif.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
    }
    Ok((qe / (n as f64 * d as f64)).sqrt() * (d as f64).sqrt()) // QE médio por ponto (norma L2)
}

/// Decodificador (pós-redirecionamento): x̂ = E(w_bmu) — reconstrução
/// a partir apenas do código (centróide). Mede o MSE final.
fn treinar_decodificador(
    dados: &Tensor,
    centroides: &[Vec<f64>],
    bmu: &[usize],
    alocacao: &[usize],
    passos: usize,
    semente: u64,
    dev: &Device,
) -> Result<f64> {
    let _ = semente;
    let n = dados.dim(0)?;
    let d = dados.dim(1)?;
    let mut ensembles: Vec<Vec<MlpJepa>> = Vec::with_capacity(centroides.len());
    for c in alocacao {
        let mut v = Vec::new();
        for _ in 0..(*c).max(1) {
            v.push(MlpJepa::novo(d, 4 * d, d, dev)?);
        }
        ensembles.push(v);
    }
    let mut vars: Vec<Var> = Vec::new();
    for ens in &ensembles {
        for e in ens {
            vars.extend(e.vars());
        }
    }
    let mut opt = AdamWManual::novo(vars.clone(), 3e-3, 0.0);
    let w_bmu = {
        let v: Vec<f32> = (0..n).flat_map(|i| centroides[bmu[i]].iter().map(|x| *x as f32).collect::<Vec<f32>>()).collect();
        Tensor::from_vec(v, (n, d), dev)?
    };
    for _ in 0..passos {
        let mut soma = Tensor::zeros((), candle_core::DType::F32, dev)?;
        for (i, ens) in ensembles.iter().enumerate() {
            let idx: Vec<u32> = (0..n).filter(|&k| bmu[k] == i).map(|k| k as u32).collect();
            if idx.is_empty() {
                continue;
            }
            let m_idx = idx.len();
            let it = Tensor::from_vec(idx, (m_idx,), dev)?;
            let w_c = w_bmu.index_select(&it, 0)?;
            let x_c = dados.index_select(&it, 0)?;
            let mut y = Tensor::zeros((m_idx, d), candle_core::DType::F32, dev)?;
            for e in ens {
                y = (y.affine(1.0, 0.0)? + e.fwd(&w_c)?.affine(1.0, 0.0)?)?;
            }
            let k = ens.len() as f64;
            let yh = y.affine(1.0 / k, 0.0)?;
            let dif = (&yh - &x_c)?;
            soma = (soma.affine(1.0, 0.0)? + dif.sqr()?.sum_all()?.affine(1.0 / (n as f64 * d as f64), 0.0)?)?;
        }
        let grads = crate::treino::grads_achatados(&soma, &vars)?;
        let mut g = grads;
        AdamWManual::clipar(&mut g, 5.0);
        opt.passo(g);
    }
    // MSE final do decodificador
    let mut mse = 0f64;
    for (i, ens) in ensembles.iter().enumerate() {
        let idx: Vec<u32> = (0..n).filter(|&k| bmu[k] == i).map(|k| k as u32).collect();
        if idx.is_empty() {
            continue;
        }
        let m_idx = idx.len();
        let it = Tensor::from_vec(idx, (m_idx,), dev)?;
        let w_c = w_bmu.index_select(&it, 0)?;
        let x_c = dados.index_select(&it, 0)?;
        let mut y = Tensor::zeros((m_idx, d), candle_core::DType::F32, dev)?;
        for e in ens {
            y = (y.affine(1.0, 0.0)? + e.fwd(&w_c)?.affine(1.0, 0.0)?)?;
        }
        let k = ens.len() as f64;
        let yh = y.affine(1.0 / k, 0.0)?;
        let dif = (&yh - &x_c)?;
        mse += dif.sqr()?.sum_all()?.to_scalar::<f32>()? as f64;
    }
    Ok(mse / (n as f64 * d as f64))
}

pub fn versao() -> &'static str {
    "moe-redirecionavel-0.6.0"
}

#[cfg(test)]
mod testes {
    use super::*;

    fn dados_sinteticos(n: usize, d: usize, semente: u64) -> Tensor {
        let mut rng = crate::rng::RngDeterministico::novo(semente);
        let v: Vec<f32> = (0..n * d).map(|_| rng.normal(0.0, 1.0) as f32).collect();
        Tensor::from_vec(v, (n, d), &Device::Cpu).unwrap()
    }

    #[test]
    fn mr1_invariancia_sob_redirecionamento() {
        let dev = Device::Cpu;
        let x = dados_sinteticos(64, 16, 3);
        let mut moe = MoeRedirecionavel::novo(16, 6, 2, 64, &dev).unwrap();
        let s1 = moe.fwd(&x, Papel::Codificador).unwrap();
        let perda_antes = s1.perda_aux.to_scalar::<f32>().unwrap() as f64;
        let carga_antes: usize = s1.carga.iter().sum();
        moe.redirecionar_todos();
        // após trocar TODOS os papéis, a passada Decodificador usa os
        // especialistas que eram do codificador — a perda auxiliar do
        // SUBCONJUNTO é comparável (mesma distribuição de roteamento)
        let s2 = moe.fwd(&x, Papel::Decodificador).unwrap();
        let perda_depois = s2.perda_aux.to_scalar::<f32>().unwrap() as f64;
        let carga_depois: usize = s2.carga.iter().sum();
        assert_eq!(carga_antes, carga_depois, "capacidade total conservada (M-R1)");
        assert!(
            (perda_antes - perda_depois).abs() < 2.0 + 1e-6,
            "perda auxiliar da família conservada: {perda_antes:.4} vs {perda_depois:.4}"
        );
    }

    #[test]
    fn mr2_neutralidade_no_nascimento() {
        let dev = Device::Cpu;
        let x = dados_sinteticos(32, 8, 5);
        let moe = MoeRedirecionavel::novo(8, 4, 2, 32, &dev).unwrap();
        let g = moe.atencao.gate.as_tensor().to_scalar::<f32>().unwrap();
        assert!(g == 0.0, "gate nasce zero");
        let y = moe.atencao.fwd(&x).unwrap();
        let dif = (&y - &x).unwrap().sqr().unwrap().sum_all().unwrap().to_scalar::<f32>().unwrap();
        assert!(dif == 0.0, "y = x bit-a-bit no nascimento (dif={dif})");
    }

    #[test]
    fn mr3_paralelismo_e_determinismo() {
        let dev = Device::Cpu;
        let x = dados_sinteticos(96, 24, 7);
        let moe = MoeRedirecionavel::novo(24, 8, 2, 96, &dev).unwrap();
        let (_enc, _dec, t_par, t_seq, speedup) = moe.fwd_ambos_par(&x).unwrap();
        // determinismo já garantido por assert interno; speedup registrado
        assert!(t_par > 0.0 && t_seq > 0.0);
        let _ = speedup;
    }

    #[test]
    fn mr5_water_filling_minimiza_cota() {
        // cota Σ q_i²/c_i: water-filling ≤ uniforme (Cauchy–Schwarz)
        let qes = vec![0.2, 0.5, 1.0, 0.3];
        let c_wf = alocar_por_qe(&qes, 8);
        let c_un = vec![2usize; 4];
        let f = |c: &[usize]| qes.iter().zip(c).map(|(q, ci)| q * q / *ci as f64).sum::<f64>();
        assert!(
            f(&c_wf) <= f(&c_un) + 1e-9,
            "water-filling {c_wf:?} (cota {:.4}) deve ≤ uniforme (cota {:.4})",
            f(&c_wf), f(&c_un)
        );
        assert_eq!(c_wf.iter().sum::<usize>(), 8);
    }

    #[test]
    fn jepa_moe_aprende() {
        let x = dados_sinteticos(64, 12, 11);
        let r = treinar_jepa_com_moe(&x, "moe", 60, 0.5, 42).unwrap();
        assert!(r.perda_pred_final < r.perda_pred_inicial, "perda preditiva cai");
    }

    /// **M-R6 (v0.8.0):** equivalência de VALORES entre `fwd` e
    /// `fwd_vetorizado` — mesma seleção (auxiliar compartilhado), mesma
    /// saída, mesma perda auxiliar, mesma carga/descartados.
    #[test]
    fn mr6_equivalencia_de_valores() {
        let dev = Device::Cpu;
        for (n, d, e, k, cap, seed) in [
            (64usize, 16usize, 6usize, 2usize, 64usize, 3u64),
            (96, 24, 8, 2, 96, 5),
            (48, 12, 4, 1, 24, 9), // capacidade apertada: força descartados
        ] {
            let x = dados_sinteticos(n, d, seed);
            let moe = MoeRedirecionavel::novo(d, e, k, cap, &dev).unwrap();
            let s_ref = moe.fwd(&x, Papel::Codificador).unwrap();
            let s_vet = moe.fwd_vetorizado(&x, Papel::Codificador).unwrap();
            assert_eq!(s_ref.carga, s_vet.carga, "M-R6: mesma carga (n={n})");
            assert_eq!(s_ref.usados, s_vet.usados, "M-R6: mesmos usados");
            assert_eq!(s_ref.descartados, s_vet.descartados, "M-R6: mesmos descartados");
            let dif = (&s_ref.saida - &s_vet.saida).unwrap()
                .abs().unwrap().max_all().unwrap().to_scalar::<f32>().unwrap();
            assert!(dif < 1e-4, "M-R6: max|Δsaída| = {dif} (n={n}, cap={cap})");
            let pa = s_ref.perda_aux.to_scalar::<f32>().unwrap();
            let pb = s_vet.perda_aux.to_scalar::<f32>().unwrap();
            assert!(
                (pa - pb).abs() < 1e-4,
                "M-R6: perda auxiliar {pa:.6} vs {pb:.6} (n={n})"
            );
        }
    }

    /// **M-R7 (v0.8.0):** restauração do fluxo de gradiente — bug v080-1.
    /// Com a `fwd` (referência de valores) os especialistas e o roteador
    /// NÃO recebem gradiente (grafo seccionado por to_vec2/from_vec); com
    /// a `fwd_vetorizado` TODOS recebem (index_add/index_select vivos).
    #[test]
    fn mr7_restauracao_do_gradiente() {
        let dev = Device::Cpu;
        let x = dados_sinteticos(48, 12, 13);
        let moe = MoeRedirecionavel::novo(12, 6, 2, 48, &dev).unwrap();
        let norma_grads = |vetorizada: bool| -> (f64, f64) {
            let s = if vetorizada {
                moe.fwd_vetorizado(&x, Papel::Codificador).unwrap()
            } else {
                moe.fwd(&x, Papel::Codificador).unwrap()
            };
            let perda = s.saida.sqr().unwrap().sum_all().unwrap();
            let grads = perda.backward().unwrap();
            let mut n_roteador = 0.0f64;
            for v in moe.roteador.vars() {
                if let Some(g) = grads.get(&v) {
                    n_roteador += g.sqr().unwrap().sum_all().unwrap()
                        .to_scalar::<f32>().unwrap() as f64;
                }
            }
            let mut n_experts = 0.0f64;
            for e in &moe.especialistas {
                for v in e.vars() {
                    if let Some(g) = grads.get(&v) {
                        n_experts += g.sqr().unwrap().sum_all().unwrap()
                            .to_scalar::<f32>().unwrap() as f64;
                    }
                }
            }
            (n_roteador, n_experts)
        };
        // (a) o bug v080-1 é REAL: na fwd, roteador e especialistas têm
        // gradiente ZERO (grafo seccionado) — documentado, não corrigido na
        // fwd (referência de valores preservada por decisão de projeto)
        let (r_ref, e_ref) = norma_grads(false);
        assert!(
            r_ref == 0.0 && e_ref == 0.0,
            "M-R7: bug v080-1 documentado — fwd secciona o grafo (roteador {r_ref}, experts {e_ref})"
        );
        // (b) a correção: na vetorizada, TODOS recebem gradiente
        let (r_vet, e_vet) = norma_grads(true);
        assert!(
            r_vet > 0.0,
            "M-R7: roteador recebe gradiente na vetorizada ({r_vet})"
        );
        assert!(
            e_vet > 0.0,
            "M-R7: especialistas recebem gradiente na vetorizada ({e_vet})"
        );
    }

    /// **M-R8 (v0.8.0):** fim-a-fim — com a vetorizada, o treino REALMENTE
    /// atualiza os pesos dos especialistas (regressão do bug v080-1: antes
    /// só a atenção própria treinava no braço MoE).
    #[test]
    fn mr8_treino_atualiza_especialistas() {
        let dev = Device::Cpu;
        let x = dados_sinteticos(48, 12, 17);
        let moe = MoeRedirecionavel::novo(12, 6, 2, 48, &dev).unwrap();
        // snapshot com dados PRÓPRIOS (to_vec1) — detach() compartilha o
        // storage e o set() do otimizador sobrescreveria a "cópia"
        let snapshot: Vec<Vec<f32>> = moe
            .especialistas
            .iter()
            .map(|e| {
                e.vars()[0].as_tensor().flatten_all().unwrap().to_vec1::<f32>().unwrap()
            })
            .collect();
        // treino curto com gradiente fluindo
        let vars = moe.vars();
        let mut opt = crate::treino::AdamWManual::novo(vars.clone(), 3e-3, 0.0);
        for _ in 0..5 {
            let s = moe.fwd_vetorizado(&x, Papel::Codificador).unwrap();
            let perda = s.saida.sqr().unwrap().sum_all().unwrap();
            let grads = crate::treino::grads_achatados(&perda, &vars).unwrap();
            opt.passo(grads);
        }
        let mut delta_max = 0.0f64;
        for (e, antes) in moe.especialistas.iter().zip(snapshot.iter()) {
            let agora = e.vars()[0].as_tensor().flatten_all().unwrap()
                .to_vec1::<f32>().unwrap();
            let dif = agora.iter().zip(antes.iter()).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max);
            delta_max = delta_max.max(dif as f64);
        }
        let _ = &vars;
        assert!(
            delta_max > 0.0,
            "M-R8: especialistas devem mudar no treino (delta_max={delta_max})"
        );
    }
}

#[cfg(test)]
mod teste_som_moe_integrado {
    use super::*;

    #[test]
    fn debug_som_com_moe_caminho() {
        let dev = Device::Cpu;
        let dados_v: Vec<Vec<f64>> = (0..32).map(|i| (0..16).map(|j| (i * 7 + j * 3) as f64 / 100.0).collect()).collect();
        let (n, d) = (dados_v.len(), 16);
        let t = Tensor::from_vec(dados_v.into_iter().flatten().collect::<Vec<f64>>(), (n, d), &dev).unwrap()
            .to_dtype(candle_core::DType::F32).unwrap();
        let cfg = crate::som4d::CfgSom4D { nx: 2, ny: 2, nz: 2, nt: 2, ..Default::default() };
        let som = crate::som4d::treinar_som4d(&t, &cfg, &dev).unwrap();
        let centroides = som.prototipos.clone();
        let r = refinar_som_com_moe(&t, &centroides, 30, 1).unwrap();
        println!("qe {} -> {} (wf) | uniforme {}", r.qe_antes, r.qe_depois, r.qe_alocacao_uniforme);
        assert!(r.qe_depois < r.qe_antes, "corretor MoE reduz o QE");
        assert!(r.mse_decodificador < r.qe_antes.powi(2), "redirect decodifica");
    }
}