//! RNG determinístico com normal padrão correta (Box–Muller).
//!
//! **Lição v0.2.0→v0.4.0 (r-17 §3.4, do outro agente):** substituir
//! `Tensor::randn` (semeado pela entropia do SO) por um RNG próprio exige
//! implementar a DISTRIBUIÇÃO certa — um `randn` que gera uniforme falha em
//! silêncio e muda o regime de aprendizado inteiro (perda inicial 16,21 vs
//! 9,73). Aqui: Box–Muller exato + testes de média/desvio/simetria/
//! determinismo.
//!
//! **Prova de corretude (Box–Muller).** Sejam `u1 ~ U(0,1]`, `u2 ~ U[0,1)`,
//! `r = sqrt(-2 ln u1)`, `θ = 2π u2`. Então `(r cos θ, r sin θ)` são
//! independentes N(0,1):
//! a densidade conjunta de `(r, θ)` em coordenadas polares é
//! `f(r,θ) = (1/2π)·r·e^{-r²/2}`; o jacobiano da transformação
//! `(z1,z2) → (r,θ)` tem determinante `r`, logo
//! `f(z1,z2) = f(r,θ)/r = (2π)^{-1} e^{-(z1²+z2²)/2}` — produto de duas
//! N(0,1). ∎ (Teste `teste_randn_e_normal` verifica média≈0, desvio≈1 e
//! simetria; `teste_rng_deterministico` verifica bit-a-bit.)

use rand::prelude::*;
use rand::SeedableRng;

/// RNG determinístico (StdRng semeado) com amostragem normal exata.
pub struct RngDeterministico {
    interno: StdRng,
    sobra_normal: Option<f64>,
}

impl RngDeterministico {
    pub fn novo(semente: u64) -> Self {
        Self { interno: StdRng::seed_from_u64(semente), sobra_normal: None }
    }

    /// Uniforme em [0,1).
    pub fn uniforme(&mut self) -> f64 {
        self.interno.gen::<f64>()
    }

    /// Uniforme em [a,b).
    pub fn uniforme_intervalo(&mut self, a: f64, b: f64) -> f64 {
        a + (b - a) * self.uniforme()
    }

    /// Normal padrão N(0,1) por Box–Muller (aproveita as duas amostras).
    pub fn normal_padrao(&mut self) -> f64 {
        if let Some(z) = self.sobra_normal.take() {
            return z;
        }
        let u1 = self.uniforme().max(f64::MIN_POSITIVE); // evita log(0)
        let u2 = self.uniforme();
        let r = (-2.0 * u1.ln()).sqrt();
        let teta = 2.0 * std::f64::consts::PI * u2;
        self.sobra_normal = Some(r * teta.sin());
        r * teta.cos()
    }

    /// Normal com média e desvio: `mean + stdev * Z`.
    pub fn normal(&mut self, mean: f64, stdev: f64) -> f64 {
        mean + stdev * self.normal_padrao()
    }

    /// Inteiro uniforme em [0, n).
    pub fn inteiro(&mut self, n: usize) -> usize {
        self.interno.gen_range(0..n)
    }

    /// Amostra sem reposição de `k` índices de `0..n` (Fisher–Yates parcial).
    pub fn amostra_k_de_n(&mut self, k: usize, n: usize) -> Vec<usize> {
        let k = k.min(n);
        let mut idx: Vec<usize> = (0..n).collect();
        for i in 0..k {
            let j = i + self.inteiro(n - i);
            idx.swap(i, j);
        }
        idx.truncate(k);
        idx
    }

    /// Gama(a, 1) por Marsaglia–Tsang (a > 0) — v0.9.0 (Thompson Sampling).
    ///
    /// **Prova de corretude (esboço; teo T-TS0 em r-34 §7).** Para a ≥ 1:
    /// com d = a − 1/3, c = 1/√(9d), a aceitação
    /// `u < 1 − 0,0331·z⁴` (squeeze) ou `ln u < z²/2 + d(1 − v + ln v)`,
    /// `v = (1+cz)³ > 0`, devolve `x = d·v` exatamente com densidade
    /// ∝ x^{a−1}e^{−x} (Marsaglia & Tsang 2000, §3–4: a razão de aceitação
    /// é a razão entre a densidade alvo e a densidade proposta deformada
    /// por v(z) — o método da rejeição clássico). Para a < 1: reforço
    /// `G(a) = G(a+1)·U^{1/a}` ( Stuart, teoria da ordem estatística:
    /// se X ~ G(a+1) e U ~ U(0,1) independentes, X·U^{1/a} ~ G(a)).
    /// Requer normal + uniforme — ambos exatos neste RNG.
    pub fn gamma(&mut self, a: f64) -> f64 {
        debug_assert!(a > 0.0, "gamma exige a > 0");
        if a < 1.0 {
            let u = self.uniforme().max(f64::MIN_POSITIVE);
            return self.gamma(a + 1.0) * u.powf(1.0 / a);
        }
        let d = a - 1.0 / 3.0;
        let c = 1.0 / (9.0 * d).sqrt();
        loop {
            let z = self.normal_padrao();
            let v = (1.0 + c * z).powi(3);
            if v <= 0.0 {
                continue;
            }
            let u = self.uniforme();
            if u < 1.0 - 0.0331 * z * z * z * z {
                return d * v;
            }
            if u.ln() < 0.5 * z * z + d * (1.0 - v + v.ln()) {
                return d * v;
            }
        }
    }

    /// Beta(α, β) via duas gamas independentes: X/(X+Y) com
    /// X ~ G(α), Y ~ G(β) (identidade clássica da construção da Beta).
    pub fn beta(&mut self, alfa: f64, beta: f64) -> f64 {
        debug_assert!(alfa > 0.0 && beta > 0.0, "beta exige parâmetros > 0");
        let x = self.gamma(alfa);
        let y = self.gamma(beta);
        x / (x + y)
    }
}

/// Semente corrente do treino (global, com padrão seguro).
use std::sync::atomic::{AtomicU64, Ordering};
static SEMENTE_GLOBAL: AtomicU64 = AtomicU64::new(42);

pub fn definir_semente(s: u64) {
    SEMENTE_GLOBAL.store(s, Ordering::SeqCst);
}

pub fn semente_atual() -> u64 {
    SEMENTE_GLOBAL.load(Ordering::SeqCst)
}

pub fn rng_global() -> RngDeterministico {
    RngDeterministico::novo(semente_atual())
}

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

    #[test]
    fn teste_randn_e_normal() {
        // Critério de falsificação: média ≈ 0, desvio ≈ 1, SIMETRIA (existem
        // valores negativos — o bug da v0.2.0 gerava uniforme com mean=0 e
        // TODOS os valores ficavam positivos).
        let mut rng = RngDeterministico::novo(7);
        let n = 200_000u32;
        let mut soma = 0.0;
        let mut soma_q = 0.0;
        let mut negativos = 0u32;
        for _ in 0..n {
            let z = rng.normal_padrao();
            soma += z;
            soma_q += z * z;
            if z < 0.0 {
                negativos += 1;
            }
        }
        let media = soma / n as f64;
        let desvio = (soma_q / n as f64 - media * media).sqrt();
        assert!((media - 0.0).abs() < 0.01, "media = {media} (esperado 0)");
        assert!((desvio - 1.0).abs() < 0.02, "desvio = {desvio} (esperado 1)");
        // simetria: proporção de negativos ∈ [0,47, 0,53]
        let frac = negativos as f64 / n as f64;
        assert!((frac - 0.5).abs() < 0.01, "fracao negativos = {frac}");
        // cauda: P(|Z|>3) ≈ 0,0027 — uniforme U(0,1) teria 0
        let mut cauda = 0u32;
        let mut rng2 = RngDeterministico::novo(7);
        for _ in 0..n {
            if rng2.normal_padrao().abs() > 3.0 {
                cauda += 1;
            }
        }
        let fc = cauda as f64 / n as f64;
        assert!(fc > 0.0015 && fc < 0.0045, "P(|Z|>3) = {fc} (esperado ~0,0027)");
    }

    #[test]
    fn teste_rng_deterministico() {
        // mesma semente ⇒ sequência bit-a-bit idêntica; sementes diferentes ⇒
        // sequências diferentes (não-degenerado).
        let mut a = RngDeterministico::novo(123);
        let mut b = RngDeterministico::novo(123);
        let mut c = RngDeterministico::novo(124);
        let va: Vec<f64> = (0..1000).map(|_| a.normal_padrao()).collect();
        let vb: Vec<f64> = (0..1000).map(|_| b.normal_padrao()).collect();
        let vc: Vec<f64> = (0..1000).map(|_| c.normal_padrao()).collect();
        assert_eq!(va, vb, "mesma semente deve dar bit-a-bit igual");
        assert_ne!(va, vc, "sementes diferentes devem divergir");
    }

    #[test]
    fn teste_amostra_k_de_n() {
        let mut rng = RngDeterministico::novo(9);
        let s = rng.amostra_k_de_n(5, 10);
        assert_eq!(s.len(), 5);
        let mut ordenado = s.clone();
        ordenado.sort();
        ordenado.dedup();
        assert_eq!(ordenado.len(), 5, "sem reposição: sem duplicatas");
        assert!(s.iter().all(|&i| i < 10));
        // k > n ⇒ devolve n
        let s2 = rng.amostra_k_de_n(15, 10);
        assert_eq!(s2.len(), 10);
    }

    /// T-TS0 (falsificação): média ≈ a e variância ≈ a da Gama(a,1)
    /// (os dois momentos identificam a família gama de escala 1); e
    /// média da Beta(α,β) ≈ α/(α+β).
    #[test]
    fn teste_gamma_beta() {
        for a in [0.5f64, 1.0, 2.5, 8.0] {
            let mut rng = RngDeterministico::novo(21);
            let n = 200_000u32;
            let (mut soma, mut soma_q) = (0.0, 0.0);
            for _ in 0..n {
                let x = rng.gamma(a);
                assert!(x > 0.0, "gamma deve ser > 0");
                soma += x;
                soma_q += x * x;
            }
            let media = soma / n as f64;
            let var = soma_q / n as f64 - media * media;
            let tol_m = 0.05 * a.max(1.0) + 0.02;
            let tol_v = 0.08 * a.max(1.0) + 0.05;
            assert!((media - a).abs() < tol_m, "G({a}): média {media} ≠ {a}");
            assert!((var - a).abs() < tol_v, "G({a}): var {var} ≠ {a}");
        }
        // Beta(2, 5): média = 2/7
        let mut rng = RngDeterministico::novo(33);
        let n = 200_000u32;
        let soma: f64 = (0..n).map(|_| rng.beta(2.0, 5.0)).sum();
        let media = soma / n as f64;
        assert!((media - 2.0 / 7.0).abs() < 0.005, "Beta(2,5): média {media} ≠ {}", 2.0 / 7.0);
        // determinismo
        let mut r1 = RngDeterministico::novo(77);
        let mut r2 = RngDeterministico::novo(77);
        let v1: Vec<f64> = (0..500).map(|_| r1.beta(1.3, 2.7)).collect();
        let v2: Vec<f64> = (0..500).map(|_| r2.beta(1.3, 2.7)).collect();
        assert_eq!(v1, v2, "mesma semente ⇒ mesma sequência");
    }
}