//! `pensamento::quebra_cabeca` — 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::*;


// ---------------------------------------------------------------------------
// QUEBRA-CABEÇA pós-treinamento (v0.7.0 — inovação com o conhecimento
// treinado; requisito do usuário; provas Q-1..Q-6 em r-31)
// ---------------------------------------------------------------------------

/// Configuração do quebra-cabeça (r-28 §4.3).
#[derive(Clone, Copy, Debug)]
pub struct ConfigQuebraCabeca {
    /// K peças em que o conhecimento é quebrado.
    pub n_pecas: usize,
    /// orçamento B de rodadas (terminação Q-4).
    pub n_rodadas_max: usize,
    /// limiar de validade (mesmo τ_c calibrado do verificador — Q-2).
    pub tau_c: f64,
    /// teto do ajuste por componente na árvore (regulação).
    pub delta_max: f64,
    /// remonta as peças com mistura ORTOGONAL por peça (rotação/ortogonalidade).
    pub rotacao_ortogonal: bool,
    /// para a execução após este nº de rodadas consecutivas sem aceite.
    pub paciencia: usize,
}

impl Default for ConfigQuebraCabeca {
    fn default() -> Self {
        ConfigQuebraCabeca {
            n_pecas: 4,
            n_rodadas_max: 12,
            tau_c: 0.55,
            delta_max: 0.05,
            rotacao_ortogonal: true,
            paciencia: 5,
        }
    }
}

impl ConfigQuebraCabeca {
    pub fn novo(n_pecas: usize, n_rodadas: usize, tau_c: f64) -> Self {
        ConfigQuebraCabeca {
            n_pecas,
            n_rodadas_max: n_rodadas,
            tau_c,
            delta_max: 0.05,
            rotacao_ortogonal: true,
            paciencia: (n_rodadas / 2).max(2),
        }
    }
}

/// Uma rodada do quebra-cabeça.
#[derive(Debug, Clone)]
pub struct RodadaQuebraCabeca {
    pub rodada: usize,
    /// permutação das peças aplicada na remontagem (reordenação).
    pub permutacao_pecas: Vec<usize>,
    /// confiança NCM do candidato remontado (reavaliação).
    pub confianca: f64,
    /// conf ≥ τ_c ⇒ VÁLIDO (ajusta a árvore); senão PUNIDO (descarta).
    pub aceito: bool,
    /// delta aplicado na árvore (0 quando descartado).
    pub delta_aplicado: f64,
    /// rotações ortogonais usadas (quando ativas) — por peça.
    pub usou_rotacao: bool,
}

/// Resultado consolidado do quebra-cabeça.
#[derive(Debug, Clone)]
pub struct ResultadoQuebraCabeca {
    pub rodadas: Vec<RodadaQuebraCabeca>,
    pub aceitos: usize,
    pub rejeitados: usize,
    pub confianca_base: f64,
    pub confianca_melhor: f64,
    pub telemetria: HashMap<String, f64>,
}

/// **Quebra-cabeça de inovação pós-treinamento** (r-28 §4.3):
/// 1. quebra o conhecimento treinado `conceitos` (m, d) em K peças por
///    permutação aleatória das linhas (Fisher–Yates com RNG determinístico);
/// 2. remonta as peças em nova ordem (e com mistura ortogonal por peça
///    quando `rotacao_ortogonal`) — "descobrir novas possibilidades";
/// 3. reavalia cada remontagem pelo verificador NCM (calibrado — V-2);
///    VÁLIDO (conf ≥ τ_c) ⇒ ajuste regulado na árvore de pesos indexada
///    (nó "pensamento", crédito fundo-topo); PUNIDO ⇒ descarta e o incremento
///    de punição sobe no nó (alimenta o gate do FNH — integração v0.7.0);
/// 4. repete até o orçamento B ou a paciência esgotar (Q-4).
///
/// **Provas (r-31):** Q-1 cobertura (permutações uniformes ⇒ S_K no suporte;
/// rotações varrem O(d)); Q-2 soundness herdada (mesmo limiar calibrado do
/// verificador — P(aceite inválido) ≤ 1−τ_c+ECE); Q-3 potencial monotônico
/// (Φ_r = Σ conf dos aceitos é não-decrescente; rejeições não tocam Φ);
/// Q-4 terminação (B finito, rodadas O(m·d)); Q-5 invariantes da árvore
/// preservadas (ajustes só tocam métricas, nunca a estrutura — T-A1 intacto;
/// agregação pós-ajuste mantém T-A2); Q-6 causalidade (permutação e confiança
/// medidas antes do ajuste — filtração F_r).
pub fn quebra_cabeca_inovar(
    conceitos: &Tensor,
    verificador: &crate::verificador::VerificadorNcm,
    logits_fn: &dyn Fn(&Tensor) -> Result<Tensor>,
    arvore: &mut crate::arvore_pesos::ArvorePesos,
    cfg: &ConfigQuebraCabeca,
    rng: &mut RngDeterministico,
) -> Result<ResultadoQuebraCabeca> {
    let dev = conceitos.device().clone();
    let (m, d) = (conceitos.dims()[0], conceitos.dims()[1]);
    let k = cfg.n_pecas.clamp(1, m);
    let no_pens = arvore.no("pensamento").unwrap_or(0);

    // linha de base: confiança do conhecimento original (agregado médio)
    let base_agg = conceitos.mean(0)?;
    let conf_base = {
        let lg = logits_fn(&base_agg)?;
        verificador.confianca_1d(&lg)?
    };

    let mut rodadas = Vec::new();
    let mut aceitos = 0usize;
    let mut rejeitados = 0usize;
    let mut conf_melhor = conf_base;
    let mut sem_aceite = 0usize;

    for r in 1..=cfg.n_rodadas_max {
        if sem_aceite >= cfg.paciencia {
            break;
        }
        // (1) QUEBRA: permutação de linhas Fisher–Yates (Q-6: amostrada ANTES)
        let mut perm: Vec<usize> = (0..m).collect();
        for i in (1..m).rev() {
            let j = rng.inteiro(i + 1);
            perm.swap(i, j);
        }
        // fatia em K peças contíguas (tamanhos quase iguais)
        let tamanho = (m + k - 1) / k;
        let mut pecas = Vec::with_capacity(k);
        let mut ini = 0usize;
        while ini < m {
            let fim = (ini + tamanho).min(m);
            let idx: Vec<u32> = perm[ini..fim].iter().map(|&x| x as u32).collect();
            let pedaco = conceitos.index_select(&Tensor::from_vec(idx, (fim - ini,), &dev)?, 0)?;
            pecas.push(pedaco);
            ini = fim;
        }
        // (2) REMONTA: nova ordem aleatória das peças (permutação de peças)
        let mut perm_pecas: Vec<usize> = (0..pecas.len()).collect();
        for i in (1..pecas.len()).rev() {
            let j = rng.inteiro(i + 1);
            perm_pecas.swap(i, j);
        }
        let mut remontada = Vec::with_capacity(pecas.len());
        for &pi in &perm_pecas {
            let mut pedaco = pecas[pi].clone();
            if cfg.rotacao_ortogonal && d >= 2 {
                // mistura ortogonal da peça: Q ∈ O(d) via polar de gaussiana
                // do RNG determinístico (preserva normas — ‖xQ‖_F = ‖x‖_F)
                let q = q_ortogonal_deterministica(d, rng, &dev)?;
                pedaco = pedaco.matmul(&q)?;
            }
            remontada.push(pedaco);
        }
        let candidata = Tensor::cat(&remontada, 0)?; // (m, d) remontado
        // (3) REAVALIAÇÃO pelo verificador NCM (Q-2: mesmo limiar calibrado)
        let agg = candidata.mean(0)?;
        let lg = logits_fn(&agg)?;
        let conf = verificador.confianca_1d(&lg)?;
        let aceito = conf >= cfg.tau_c;
        let delta = if aceito {
            // VÁLIDO: ajusta a árvore bidirecional (indexada, com teto)
            let escalado = (conf - cfg.tau_c).min(cfg.delta_max);
            arvore.aplicar_ajuste_indexado(no_pens, "confianca", escalado, cfg.delta_max)
        } else {
            // PUNIDO: descarta; a punição sobe no nó (sinal para o gate FNH)
            rejeitados += 1;
            sem_aceite += 1;
            arvore.aplicar_ajuste_indexado(no_pens, "punicao", cfg.delta_max * 0.5, cfg.delta_max);
            0.0
        };
        if aceito {
            aceitos += 1;
            sem_aceite = 0;
            conf_melhor = conf_melhor.max(conf);
        }
        rodadas.push(RodadaQuebraCabeca {
            rodada: r,
            permutacao_pecas: perm_pecas,
            confianca: conf,
            aceito,
            delta_aplicado: delta,
            usou_rotacao: cfg.rotacao_ortogonal,
        });
    }
    // crédito fundo-topo após todas as rodadas (T-A2)
    arvore.agregar_fundo_topo("confianca");
    arvore.agregar_fundo_topo("punicao");

    let n_rod = rodadas.len();
    let mut telemetria = HashMap::new();
    telemetria.insert("rodadas".into(), n_rod as f64);
    telemetria.insert("aceitos".into(), aceitos as f64);
    telemetria.insert("rejeitados".into(), rejeitados as f64);
    telemetria.insert("confianca_base".into(), conf_base);
    telemetria.insert("confianca_melhor".into(), conf_melhor);
    telemetria.insert(
        "taxa_aceitacao".into(),
        if n_rod > 0 { aceitos as f64 / n_rod as f64 } else { 0.0 },
    );
    telemetria.insert("rotacao_ativa".into(), if cfg.rotacao_ortogonal { 1.0 } else { 0.0 });
    Ok(ResultadoQuebraCabeca {
        rodadas,
        aceitos,
        rejeitados,
        confianca_base: conf_base,
        confianca_melhor: conf_melhor,
        telemetria,
    })
}

/// Q ∈ O(d) determinística: gaussiana do RNG próprio → retração polar.
/// Ortogonal EXATA (δ_ort ≤ 1e-6) e reprodutível (mesma semente ⇒ mesmo Q).
pub fn q_ortogonal_deterministica(
    d: usize,
    rng: &mut RngDeterministico,
    dev: &Device,
) -> Result<Tensor> {
    let n = d * d;
    let v: Vec<f32> = (0..n).map(|_| rng.normal_padrao() as f32).collect();
    let g = Tensor::from_vec(v, (d, d), dev)?;
    crate::stiefel::retrair_polar_frobenius(&g, 90)
}