Blame for src/config.rs

//! Config do KHTST_RUST — espelha `configs/base.json` do KHTST v23 + seção `ncm`
//! (n-04 §3) + seções novas (rust/adba/ponte/telemetria). Defaults neutros
//! (nascimento neutro, T10): `ncm.ativo=false` E `lambdas.conceito=0.0` (dupla trava).

use serde::{Deserialize, Serialize};

fn padrao_niveis() -> Vec<usize> {
    vec![4, 8, 16, 32, 64, 128, 256]
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct MoE {
    #[serde(default = "d_true")]
    pub ativo: bool,
    #[serde(default = "d_usize4")]
    pub n_experts_fusao: usize,
    #[serde(default = "d_usize2")]
    pub top_k: usize,
}
fn d_true() -> bool { true }
fn d_usize2() -> usize { 2 }
fn d_usize4() -> usize { 4 }

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MTP {
    #[serde(default = "d_true")]
    pub ativo: bool,
    #[serde(default = "d_usize2")]
    pub k_cabecas: usize,
    #[serde(default = "d_f005")]
    pub beta_entropia: f64,
    #[serde(default = "d_true")]
    pub k_adaptativo: bool,
}
impl Default for MTP {
    fn default() -> Self {
        MTP { ativo: true, k_cabecas: 2, beta_entropia: 0.05, k_adaptativo: true }
    }
}
fn d_f005() -> f64 { 0.05 }

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Modelo {
    pub vocab: usize,
    pub d_modelo: usize,
    pub n_camadas: usize,
    pub n_cabecas: usize,
    pub d_ff: usize,
    pub comprimento_ctx: usize,
    pub janela_sliding: usize,
    pub kv_cache_max: usize,
    #[serde(default)]
    pub moe: MoE,
    #[serde(default)]
    pub mtp: MTP,
    #[serde(default = "d_true")]
    pub usar_atencao_arvore: bool,
}
impl Default for Modelo {
    fn default() -> Self {
        Modelo {
            vocab: 16384,
            d_modelo: 192,
            n_camadas: 6,
            n_cabecas: 4,
            d_ff: 512,
            comprimento_ctx: 256,
            janela_sliding: 192,
            kv_cache_max: 2048,
            moe: MoE::default(),
            mtp: MTP::default(),
            usar_atencao_arvore: true,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Gsq {
    pub tau0: f64,
    pub tau_min: f64,
    pub t_alvo: usize,
}
impl Default for Gsq {
    fn default() -> Self { Gsq { tau0: 1.0, tau_min: 0.5, t_alvo: 400 } }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NestedDropout {
    #[serde(default = "d_tipo_nd")]
    pub tipo: String,
}
fn d_tipo_nd() -> String { "uniforme".into() }

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Lambdas {
    pub conceito: f64,
    pub rec: f64,
    pub rho_ancora: f64,
}
impl Default for Lambdas {
    fn default() -> Self { Lambdas { conceito: 0.0, rec: 1.0, rho_ancora: 0.05 } }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Temperaturas {
    pub auto: bool,
    pub entropia_alvo: Option<f64>,
    pub limites: [f64; 2],
}
impl Default for Temperaturas {
    fn default() -> Self { Temperaturas { auto: true, entropia_alvo: None, limites: [0.5, 1.5] } }
}

/// Seção `ncm` — schema aditivo do n-04 §3 (defaults 100% neutros).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Ncm {
    pub ativo: bool,
    #[serde(default = "padrao_niveis")]
    pub niveis: Vec<usize>,
    pub d_lat: usize,
    pub k_por_nivel: usize,
    pub ema_codebook: f64,
    pub beta_commit: f64,
    pub nested_dropout: NestedDropout,
    pub corrupcao_cond: f64,
    pub dropout_cfg: f64,
    pub cfg_scale: f64,
    pub gsq: Gsq,
    #[serde(default = "d_pesos_nivel")]
    pub pesos_nivel: String,
    #[serde(default)]
    pub lambdas: Lambdas,
    #[serde(default)]
    pub temperaturas: Temperaturas,
    pub mbr: Mbr,
    /// Modo da escala R_ℓ dos codebooks: "ema_projecao" | "otima_teo" | "aprendida" (r-01 §4.4).
    #[serde(default = "d_escala_r")]
    pub escala_r: String,
    /// Modo de quantização: "gsq_rco" (padrão, Teo D.4) | "ste_ema" (Teo A.0).
    #[serde(default = "d_modo_quant")]
    pub modo_quantizacao: String,
    /// Níveis como tokens (modo D.6 — MTP equivalente).
    #[serde(default)]
    pub niveis_token: bool,
}
fn d_pesos_nivel() -> String { "uniforme".into() }
fn d_escala_r() -> String { "ema_projecao".into() }
fn d_modo_quant() -> String { "gsq_rco".into() }

impl Default for Ncm {
    fn default() -> Self {
        Ncm {
            ativo: false,
            niveis: padrao_niveis(),
            d_lat: 64,
            k_por_nivel: 64,
            ema_codebook: 0.99,
            beta_commit: 0.25,
            nested_dropout: NestedDropout { tipo: "uniforme".into() },
            corrupcao_cond: 0.15,
            dropout_cfg: 0.10,
            cfg_scale: 1.0,
            gsq: Gsq::default(),
            pesos_nivel: "uniforme".into(),
            lambdas: Lambdas::default(),
            temperaturas: Temperaturas::default(),
            mbr: Mbr { n_cand: 1 },
            escala_r: "ema_projecao".into(),
            modo_quantizacao: "gsq_rco".into(),
            niveis_token: false,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Mbr {
    pub n_cand: usize,
}
impl Default for Mbr {
    fn default() -> Self { Mbr { n_cand: 1 } }
}

/// ADBA — Atenção Dual-Bloco Adaptativa (aprimoramento 1; prova em r-04).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Adba {
    pub ativo: bool,
    /// Peso FiLM do resumo-âncora por nível (0 = NCM puro).
    pub gamma_ancora: f64,
    /// Reponderação aprendível dos logits por nível-fonte (true = ativa).
    pub reponderacao_nivel: bool,
    /// Temperatura aprendível por nível (true = ativa).
    pub temperatura_nivel: bool,
}
impl Default for Adba {
    fn default() -> Self {
        Adba { ativo: true, gamma_ancora: 1.0, reponderacao_nivel: true, temperatura_nivel: true }
    }
}

/// Ponte Dupla — comunicação paralela bidirecional entre modelos (aprimoramento 2; prova em r-05).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Ponte {
    pub ativo: bool,
    pub capacidade: usize,
    pub timeout_ms: u64,
    pub rodadas_demo: usize,
}
impl Default for Ponte {
    fn default() -> Self { Ponte { ativo: true, capacidade: 64, timeout_ms: 2000, rodadas_demo: 24 } }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Treino {
    pub lr: f64,
    pub aquecimento: usize,
    pub passos_fase_a: usize,
    pub passos_fase_b: usize,
    pub lote: usize,
    pub clip: f64,
    pub peso_decaimento: f64,
    #[serde(default = "d_usize25")]
    pub telemetria_a_cada: usize,
}
fn d_usize25() -> usize { 25 }
impl Default for Treino {
    fn default() -> Self {
        Treino {
            lr: 3.0e-4,
            aquecimento: 100,
            passos_fase_a: 1500,
            passos_fase_b: 2000,
            lote: 8,
            clip: 1.0,
            peso_decaimento: 0.01,
            telemetria_a_cada: 25,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TelemetriaCfg {
    pub jsonl: bool,
    pub deteccao_aprendizado: bool,
    pub ajuste_automatico: bool,
    pub paciencia_plateau: usize,
    pub fator_lr_plateau: f64,
}
impl Default for TelemetriaCfg {
    fn default() -> Self {
        TelemetriaCfg {
            jsonl: true,
            deteccao_aprendizado: true,
            ajuste_automatico: true,
            paciencia_plateau: 200,
            fator_lr_plateau: 0.6,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Dados {
    pub shard_treino: String,
    pub shard_validacao: String,
    pub max_artigos: usize,
    pub max_bytes: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MemoriaCfg {
    pub som_lado: usize,
    pub lambda_jepa: f64,
    pub rm_dois_tempos: bool,
}
impl Default for MemoriaCfg {
    fn default() -> Self { MemoriaCfg { som_lado: 16, lambda_jepa: 0.1, rm_dois_tempos: true } }
}

#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ConfigKhtst {
    #[serde(default)]
    pub modelo: Modelo,
    #[serde(default)]
    pub ncm: Ncm,
    #[serde(default)]
    pub adba: Adba,
    #[serde(default)]
    pub ponte: Ponte,
    #[serde(default)]
    pub treino: Treino,
    #[serde(default)]
    pub telemetria: TelemetriaCfg,
    #[serde(default)]
    pub dados: Dados,
    #[serde(default)]
    pub memoria: MemoriaCfg,
}

impl ConfigKhtst {
    pub fn validar(&self) -> Result<(), String> {
        let m = &self.modelo;
        if m.d_modelo % m.n_cabecas != 0 {
            return Err(format!(
                "modelo.d_modelo ({}) não é divisível por modelo.n_cabecas ({})",
                m.d_modelo, m.n_cabecas
            ));
        }
        let n = &self.ncm;
        if n.niveis.is_empty() || n.niveis.last() != Some(&m.comprimento_ctx) {
            return Err(format!(
                "ncm.niveis deve terminar em modelo.comprimento_ctx ({}); atual: {:?}",
                m.comprimento_ctx, n.niveis
            ));
        }
        for w in n.niveis.windows(2) {
            if w[1] <= w[0] {
                return Err(format!("ncm.niveis deve ser estritamente crescente: {:?}", n.niveis));
            }
        }
        if n.k_por_nivel > n.d_lat {
            // Cor. A.1.2: K > d sai do Stiefel estrito — exige nota (permitido, com aviso).
            log::warn!("ncm.k_por_nivel ({}) > ncm.d_lat ({}) — frame apertado (Cor. A.1.2), penalidade Welch recomendada", n.k_por_nivel, n.d_lat);
        }
        if !(0.0..=1.0).contains(&n.corrupcao_cond) {
            return Err("ncm.corrupcao_cond deve estar em [0,1]".into());
        }
        if self.treino.lote == 0 || self.treino.lr <= 0.0 {
            return Err("treino.lote >= 1 e treino.lr > 0 exigidos".into());
        }
        Ok(())
    }

    /// L (número de níveis) — Teo B.2(ii).
    pub fn l_niveis(&self) -> usize { self.ncm.niveis.len() }

    /// Verificação de orçamento do KV-cache (T2): |cond| + Σ_{j<L} n_j ≤ kv_cache_max.
    pub fn checar_orcamento_cache(&self) -> Result<usize, String> {
        let soma: usize = self.ncm.niveis.iter().copied().sum();
        let usado = self.modelo.comprimento_ctx + soma - self.modelo.comprimento_ctx; // cond + níveis < L
        if usado > self.modelo.kv_cache_max {
            return Err(format!(
                "orçamento do KV-cache violado (T2): {} > kv_cache_max = {}",
                usado, self.modelo.kv_cache_max
            ));
        }
        Ok(usado)
    }
}