Blame for src/diagnostico.rs

//! Diagnóstico de erros comuns durante refatoração (pedido do usuário:
//! "inicialmente analise também se os erros podem ser por digitação ou outra
//! causa comum durante refatoração de scripts").
//!
//! Fornece: (i) lint de config com sugestões por distância de Levenshtein
//! (detecta chaves digitadas errado / renomeadas durante a migração Python→Rust);
//! (ii) taxonomia de erros de execução com causa provável e correção sugerida;
//! (iii) classificador de falhas de NaN/Inf com origem provável.

use std::collections::BTreeMap;

/// Chaves canônicas conhecidas do config (espeelho do config.rs).
pub fn chaves_conhecidas() -> Vec<String> {
    let mut v = vec![];
    for secao in ["modelo", "ncm", "adba", "ponte", "treino", "telemetria", "dados", "memoria"] {
        v.push(secao.to_string());
        let campos: &[&str] = match secao {
            "modelo" => &["vocab", "d_modelo", "n_camadas", "n_cabecas", "d_ff",
                          "comprimento_ctx", "janela_sliding", "kv_cache_max", "usar_atencao_arvore",
                          "moe", "mtp", "roteador", "microunidades", "atencao_moe",
                          "moe_enc_dec", "autoajuste", "mamba3", "sonicmoe"],
            "ncm" => &["ativo", "niveis", "d_lat", "k_por_nivel", "ema_codebook", "beta_commit",
                       "corrupcao_cond", "dropout_cfg", "cfg_scale", "pesos_nivel", "escala_r",
                       "modo_quantizacao", "niveis_token",
                       // subseções (dicts aninhados — 2º nível)
                       "nested_dropout", "gsq", "lambdas", "temperaturas", "mbr"],
            "adba" => &["ativo", "gamma_ancora", "reponderacao_nivel", "temperatura_nivel"],
            "ponte" => &["ativo", "capacidade", "timeout_ms", "rodadas_demo"],
            "treino" => &["lr", "aquecimento", "passos_fase_a", "passos_fase_b", "lote", "clip",
                          "peso_decaimento", "telemetria_a_cada"],
            "telemetria" => &["jsonl", "deteccao_aprendizado", "ajuste_automatico",
                              "paciencia_plateau", "fator_lr_plateau"],
            "dados" => &["shard_treino", "shard_validacao", "max_artigos", "max_bytes"],
            "memoria" => &["som_lado", "lambda_jepa", "rm_dois_tempos"],
            _ => &[],
        };
        for c in campos {
            v.push(format!("{secao}.{c}"));
        }
    }
    v
}

/// Distância de Levenshtein clássica (com corte `max`).
pub fn levenshtein(a: &str, b: &str, max: usize) -> usize {
    let a: Vec<char> = a.chars().collect();
    let b: Vec<char> = b.chars().collect();
    if a.len().abs_diff(b.len()) > max {
        return max + 1;
    }
    let mut prev: Vec<usize> = (0..=b.len()).collect();
    let mut cur = vec![0usize; b.len() + 1];
    for i in 1..=a.len() {
        cur[0] = i;
        for j in 1..=b.len() {
            let custo = if a[i - 1] == b[j - 1] { 0 } else { 1 };
            cur[j] = (prev[j] + 1).min(cur[j - 1] + 1).min(prev[j - 1] + custo);
        }
        std::mem::swap(&mut prev, &mut cur);
    }
    prev[b.len()]
}

/// Resultado do lint de uma chave avulsa.
#[derive(Debug, Clone)]
pub struct AvisoLint {
    pub chave: String,
    pub sugestao: Option<String>,
    pub causa: String,
    pub correcao: String,
}

/// Lint: dado o conjunto de chaves usadas (achatadas "secao.campo"),
/// devolve avisos para chaves desconhecidas com sugestão por Levenshtein ≤ 2
/// e para renomeações conhecidas da migração Python→Rust.
pub fn lint_chaves(usadas: &BTreeMap<String, serde_json::Value>) -> Vec<AvisoLint> {
    let conhecidas = chaves_conhecidas();
    // Renomeações conhecidas da refatoração (v23 Python → KHTST_RUST).
    let renomeacoes: &[(&str, &str, &str)] = &[
        ("d_model", "modelo.d_modelo", "renomeação v23→Rust (digitação provável do nome antigo)"),
        ("n_layers", "modelo.n_camadas", "renomeação v23→Rust"),
        ("ctx", "modelo.comprimento_ctx", "renomeação v23→Rust"),
        ("niveis_tokenizador", "ncm.niveis", "renomeação do plano n-04 §3"),
        ("taxa_corrupcao", "ncm.corrupcao_cond", "renomeação v23→Rust"),
        ("beta_compromisso", "ncm.beta_commit", "tradução PT↔EN misturada na refatoração"),
        ("kv_cache", "modelo.kv_cache_max", "renomeação v23→Rust"),
        ("taxa_aprendizado", "treino.lr", "tradução PT↔EN misturada na refatoração"),
    ];
    let mut avisos = vec![];
    for (chave, _) in usadas {
        let plano = chave.split('.').next().unwrap_or(chave);
        let campo = chave.strip_prefix(&format!("{plano}.")).unwrap_or(chave);
        let secao_valida = conhecidas.iter().any(|k| k == plano);
        if !secao_valida {
            // seção inteira desconhecida
            let melhor = conhecidas
                .iter()
                .map(|k| (levenshtein(plano, k, 3), k))
                .filter(|(d, _)| *d <= 3)
                .min_by_key(|(d, _)| *d);
            match melhor {
                Some((_, k)) => avisos.push(AvisoLint {
                    chave: chave.clone(),
                    sugestao: Some(k.to_string()),
                    causa: "seção desconhecida — provável digitação ou nome da versão anterior".into(),
                    correcao: "usar a seção sugerida (ver r-02 §4)".into(),
                }),
                None => avisos.push(AvisoLint {
                    chave: chave.clone(),
                    sugestao: None,
                    causa: "seção desconhecida sem similar".into(),
                    correcao: "consultar chaves_conhecidas() e a memória r-02".into(),
                }),
            }
            continue;
        }
        let esperado = format!("{plano}.{campo}");
        if !conhecidas.contains(&esperado) {
            let melhor = conhecidas
                .iter()
                .filter(|k| k.starts_with(&format!("{plano}.")))
                .map(|k| (levenshtein(campo, k.split('.').nth(1).unwrap_or(k), 2), k))
                .filter(|(d, _)| *d <= 2)
                .min_by_key(|(d, _)| *d);
            match melhor {
                Some((_, k)) => avisos.push(AvisoLint {
                    chave: chave.clone(),
                    sugestao: Some(k.clone()),
                    causa: "campo desconhecido — digitação provável (Levenshtein ≤ 2)".into(),
                    correcao: format!("substituir por '{k}'"),
                }),
                None => {
                    // tenta renomeação conhecida
                    let ren = renomeacoes.iter().find(|(velho, _, _)| *velho == campo);
                    avisos.push(AvisoLint {
                        chave: chave.clone(),
                        sugestao: ren.map(|(_, novo, _)| novo.to_string()),
                        causa: ren
                            .map(|(_, _, motivo)| motivo.to_string())
                            .unwrap_or_else(|| "campo desconhecido na seção".into()),
                        correcao: "corrigir o nome do campo".into(),
                    });
                }
            }
        }
    }
    avisos
}

/// Achata um JSON de config em "secao.campo" → valor (1 nível de aninhamento;
/// dicionários internos são representados como "<json>").
pub fn achatar(v: &serde_json::Value) -> BTreeMap<String, serde_json::Value> {
    let mut out = BTreeMap::new();
    if let serde_json::Value::Object(map) = v {
        for (secao, val) in map {
            match val {
                serde_json::Value::Object(sub) => {
                    for (campo, vv) in sub {
                        out.insert(format!("{secao}.{campo}"), vv.clone());
                    }
                }
                _ => {
                    out.insert(secao.clone(), val.clone());
                }
            }
        }
    }
    out
}

/// Taxonomia de erros comuns de refatoração (causa → detecção → correção).
pub fn classificar_erro(msg: &str) -> Vec<(&'static str, &'static str, &'static str)> {
    let m = msg.to_lowercase();
    let mut out = vec![];
    let regras: &[(&str, &str, &str, &str)] = &[
        ("shape", "incompat", "shape incompatível",
         "causa comum de refatoração: ordem de eixos trocada (linha/coluna) ou batch ausente — \
          verificar se o tensor está (lote, seq, d) e se a matriz de peso é (saída, entrada)"),
        ("dtype", "", "dtype divergente",
         "verificar to_dtype f32 antes de ops customizadas (CustomOp exige f32 contíguo)"),
        ("nan", "", "NaN detectado",
         "verificar: lr alto demais (reduzir 10×), log de probabilidade −inf (clamp), \
          τ do GSQ pequeno demais com logits grandes, ou resíduo divergente no nível fino"),
        ("inf", "", "Inf detectado",
         "verificar escala R_ℓ do codebook e o decaimento do SSD (a fora de (0,1))"),
        ("index", "out of bounds", "índice fora do intervalo",
         "causa comum: id de código ≥ K_ℓ ou id de token ≥ vocab — digitação em config de \
          k_por_nivel/vocab ou interpolação de nível com tamanho errado"),
        ("not found", "", "arquivo/chave não encontrado",
         "digitação em caminho de shard ou nome de campo renomeado na migração — rodar lint_chaves"),
        ("contiguous", "", "tensor não contíguo",
         "chamar .contiguous() antes de apply_op (layout de transposta/matmul)"),
        ("cuda", "", "backend CUDA ausente",
         "KHTST_RUST é CPU-first — remover referência a device cuda"),
    ];
    for (a, b, causa, corr) in regras {
        if m.contains(a) && (b.is_empty() || m.contains(b)) {
            out.push((*causa, *corr, "padrão da taxonomia de refatoração"));
        }
    }
    out
}

/// Varredura de NaN/Inf por posição com origem provável (nível do tokenizador).
pub fn origem_nan(vals: &[f32], nome: &str) -> Option<String> {
    let (mut n_nan, mut n_inf, mut pior) = (0usize, 0usize, 0f32);
    for (i, v) in vals.iter().enumerate() {
        if v.is_nan() {
            n_nan += 1;
            if pior == 0.0 { pior = i as f32; }
        } else if v.is_infinite() {
            n_inf += 1;
        }
    }
    if n_nan + n_inf == 0 {
        return None;
    }
    Some(format!(
        "{nome}: {n_nan} NaN / {n_inf} Inf (primeiro na posição ~{pior:.0}); \
         provável: gradiente explodido (reduzir lr), τ_min pequeno, ou R_ℓ desajustada"
    ))
}

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

    #[test]
    fn levenshtein_basico() {
        assert_eq!(levenshtein("beta_commit", "beta_comit", 2), 1);
        assert_eq!(levenshtein("corrupcao_cond", "corrupcao_cond", 2), 0);
        assert!(levenshtein("abc", "xyz", 2) > 2);
    }

    #[test]
    fn lint_detecta_digitacao() {
        let mut m = BTreeMap::new();
        m.insert("ncm.beta_comit".to_string(), serde_json::json!(0.25));
        let avisos = lint_chaves(&m);
        assert_eq!(avisos.len(), 1);
        assert_eq!(avisos[0].sugestao.as_deref(), Some("ncm.beta_commit"));
    }
}