//! Bench v0.9.0 — medições honestas para r-34 (críticas Qwen3.8 nº 4/D3 e
//! plano itens 1/3/5). Saída JSON em stdout.
//!
//! Cada bloco mede a MESMA computação em duas variantes (A/B) com
//! black_box e repetições suficientes; reporta mediana e razão.

use khtst_rust::rng::RngDeterministico;
use khtst_rust::simd::{self, BufAlinhado64};
use std::hint::black_box;
use std::time::Instant;

fn mediana(mut xs: Vec<f64>) -> f64 {
    xs.sort_by(|a, b| a.partial_cmp(b).unwrap());
    xs[xs.len() / 2]
}

fn bench<F: FnMut()>(mut f: F, reps: usize) -> f64 {
    let mut tempos = Vec::with_capacity(reps);
    for _ in 0..reps {
        let t = Instant::now();
        f();
        tempos.push(t.elapsed().as_secs_f64() * 1e3); // ms
    }
    mediana(tempos)
}

fn main() {
    let reps = 21usize;
    let mut resultado = serde_json::Map::new();
    resultado.insert("avx512_disponivel".into(), simd::tem_avx512().into());
    resultado.insert("cpus".into(), std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1).into());

    // ---------------- A) passo AdamW: escalar vs SIMD (despacho) ----------------
    {
        let n = 2_000_000usize;
        let mut rng = RngDeterministico::novo(1);
        let w: Vec<f32> = (0..n).map(|_| rng.normal(0.0, 0.5) as f32).collect();
        let g: Vec<f32> = (0..n).map(|_| rng.normal(0.0, 0.1) as f32).collect();
        let mut m = vec![0f32; n];
        let mut v = vec![0f32; n];
        let mut out = vec![0f32; n];
        let (lr, b1, b2, eps, wd, bc1, bc2) = (3e-4f64, 0.9, 0.999, 1e-8, 0.01, 0.91, 0.001);

        let t_esc = bench(|| {
            simd::passo_adamw_escalar(&w, &g, &mut m, &mut v, &mut out, lr, b1, b2, eps, wd, bc1, bc2);
            black_box(&out);
        }, reps);
        let t_simd = bench(|| {
            simd::passo_adamw(&w, &g, &mut m, &mut v, &mut out, lr, b1, b2, eps, wd, bc1, bc2);
            black_box(&out);
        }, reps);
        // bit-identidade em cena real (falsificação) — estados FRESCHOS iguais
        let (mut ma, mut va, mut oa) = (vec![0f32; n], vec![0f32; n], vec![0f32; n]);
        let (mut mb, mut vb, mut ob) = (vec![0f32; n], vec![0f32; n], vec![0f32; n]);
        simd::passo_adamw_escalar(&w, &g, &mut ma, &mut va, &mut oa, lr, b1, b2, eps, wd, bc1, bc2);
        simd::passo_adamw(&w, &g, &mut mb, &mut vb, &mut ob, lr, b1, b2, eps, wd, bc1, bc2);
        let identico = oa.iter().zip(ob.iter()).all(|(a, b)| a.to_bits() == b.to_bits());
        resultado.insert("adamw_escalar_ms".into(), t_esc.into());
        resultado.insert("adamw_simd_ms".into(), t_simd.into());
        resultado.insert("adamw_speedup".into(), (t_esc / t_simd).into());
        resultado.insert("adamw_bit_identico".into(), identico.into());
    }

    // ---------------- B) BMU 4D: Vec<Vec<f64>> vs SoA+AVX-512 ----------------
    {
        let (c, d, n_consultas) = (4096usize, 4usize, 2000usize);
        let mut rng = RngDeterministico::novo(2);
        let ps: Vec<Vec<f64>> = (0..c).map(|_| (0..d).map(|_| rng.normal(0.0, 2.0)).collect()).collect();
        let consultas: Vec<Vec<f64>> = (0..n_consultas).map(|_| (0..d).map(|_| rng.normal(0.0, 2.0)).collect()).collect();
        let dist = |a: &[f64], b: &[f64]| -> f64 { a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() };

        let t_vecvec = bench(|| {
            let mut soma = 0usize;
            for x in &consultas {
                let mut melhor = 0usize;
                let mut dm = f64::INFINITY;
                for (i, p) in ps.iter().enumerate() {
                    let dd = dist(x, p);
                    if dd < dm { dm = dd; melhor = i; }
                }
                soma += melhor;
            }
            black_box(soma);
        }, 7);

        let (soa, c2, d2) = simd::soa_de_prototipos(&ps);
        let mut buf = Vec::with_capacity(c2);
        let t_soa = bench(|| {
            let mut soma = 0usize;
            for x in &consultas {
                soma += simd::melhor_indice_soa_buf(&soa, c2, d2, x, &mut buf);
            }
            black_box(soma);
        }, 7);
        let (rm, _, _) = simd::rm_de_prototipos(&ps);
        let mut buf_rm = Vec::with_capacity(c2);
        let t_rm = bench(|| {
            let mut soma = 0usize;
            for x in &consultas {
                soma += simd::melhor_indice_rm_buf(&rm, c2, x, &mut buf_rm);
            }
            black_box(soma);
        }, 7);
        // equivalencia de indices
        let mut iguais = 0usize;
        for x in &consultas {
            let mut melhor = 0usize; let mut dm = f64::INFINITY;
            for (i, p) in ps.iter().enumerate() { let dd = dist(x, p); if dd < dm { dm = dd; melhor = i; } }
            if simd::melhor_indice_soa_buf(&soa, c2, d2, x, &mut buf) == melhor { iguais += 1; }
        }
        resultado.insert("bmu_vecvec_ms".into(), t_vecvec.into());
        resultado.insert("bmu_soa_ms".into(), t_soa.into());
        resultado.insert("bmu_rm_ms".into(), t_rm.into());
        resultado.insert("bmu_speedup_soa".into(), (t_vecvec / t_soa).into());
        resultado.insert("bmu_speedup_rm".into(), (t_vecvec / t_rm).into());
        resultado.insert("bmu_indices_iguais".into(), format!("{iguais}/{n_consultas}").into());
    }

    // ---------------- C) alinhamento 64B: alinhado vs malloc vs desalinhado ----------------
    {
        let n = 4_000_000usize;
        let mut alinhado = BufAlinhado64::novo_zeros(n);
        let mut comum: Vec<f64> = vec![0.42; n];
        // desalinhado proposital: pula 1 f64 (8 B) — pior caso de split de linha
        let mut base = BufAlinhado64::novo_zeros(n + 8);
        let desalinhado_ptr = unsafe { base.as_mut_slice().as_mut_ptr().add(1) };
        let fator = 1.0000001f64;
        let t_al = bench(|| {
            for x in alinhado.as_mut_slice() { *x = *x * fator + 0.1; }
            black_box(alinhado.as_slice().last());
        }, 7);
        let t_comum = bench(|| {
            for x in comum.iter_mut() { *x = *x * fator + 0.1; }
            black_box(comum.last());
        }, 7);
        let t_des = bench(|| {
            let fatia = unsafe { std::slice::from_raw_parts_mut(desalinhado_ptr, n) };
            for x in fatia.iter_mut() { *x = *x * fator + 0.1; }
            black_box(fatia.last());
        }, 7);
        resultado.insert("axpy_alinhado64_ms".into(), t_al.into());
        resultado.insert("axpy_malloc16_ms".into(), t_comum.into());
        resultado.insert("axpy_desalinhado8_ms".into(), t_des.into());
        resultado.insert("alinhado_endereco_mod64".into(), (alinhado.endereco() % 64 == 0).into());
    }

    // ---------------- D) f32 vs bf16 no candle CPU (crítica D3) ----------------
    {
        use candle_core::{DType, Device, Tensor};
        let dev = Device::Cpu;
        let n = 8_000_000usize;
        let dados: Vec<f32> = (0..n).map(|i| ((i as f32) * 0.001) % 1.0).collect();
        let a = Tensor::from_vec(dados.clone(), (n,), &dev).unwrap();
        let b = Tensor::from_vec(dados.clone(), (n,), &dev).unwrap();
        let t_f32 = bench(|| {
            let c = (&a + &b).unwrap().affine(1.5, 0.1).unwrap().sqr().unwrap();
            black_box(c.to_scalar::<f32>());
        }, 7);
        let ab = a.to_dtype(DType::BF16).unwrap();
        let bb = b.to_dtype(DType::BF16).unwrap();
        let t_bf16 = bench(|| {
            let c = (&ab + &bb).unwrap().affine(1.5, 0.1).unwrap().sqr().unwrap();
            black_box(c.to_dtype(DType::F32).unwrap().to_scalar::<f32>());
        }, 7);
        resultado.insert("candle_f32_ms".into(), t_f32.into());
        resultado.insert("candle_bf16_ms".into(), t_bf16.into());
        resultado.insert("bf16_vs_f32".into(), (t_f32 / t_bf16).into());
    }

    // ---------------- E) nosso AdamW vs candle_nn::optim::AdamW ----------------
    {
        use candle_core::{Device, Tensor, Var};
        use candle_nn::optim::{AdamW as AdamWC, Optimizer, ParamsAdamW};
        let dev = Device::Cpu;
        // 16 variáveis × 100k params ~ escala do tronco real
        let mut vars = Vec::new();
        for _ in 0..16 {
            let t = Tensor::randn(0f32, 0.1, (100_000,), &dev).unwrap();
            vars.push(Var::from_tensor(&t).unwrap());
        }
        let alvo = Tensor::randn(0f32, 1.0, (100_000,), &dev).unwrap();
        let roda = |nosso: bool| -> f64 {
            let mut opt_nosso = khtst_rust::treino::AdamWManual::novo(vars.clone(), 1e-3, 0.01);
            let params = ParamsAdamW { lr: 1e-3, beta1: 0.9, beta2: 0.999, eps: 1e-8, weight_decay: 0.01 };
            let mut opt_c = AdamWC::new(vars.clone(), params).unwrap();
            let t0 = Instant::now();
            for _ in 0..30 {
                let perda = {
                    let mut total = None::<Tensor>;
                    for v in &vars {
                        let d = v.as_tensor().sub(&alvo).unwrap().sqr().unwrap().mean_all().unwrap();
                        total = Some(match total { None => d, Some(t) => (&t + &d).unwrap() });
                    }
                    total.unwrap()
                };
                if nosso {
                    let grads = khtst_rust::treino::grads_achatados(&perda, &vars).unwrap();
                    opt_nosso.passo(grads);
                } else {
                    let gs = perda.backward().unwrap();
                    opt_c.step(&gs).unwrap();
                }
            }
            t0.elapsed().as_secs_f64() * 1e3 / 30.0
        };
        let t_nosso = mediana((0..5).map(|_| roda(true)).collect());
        let t_candle = mediana((0..5).map(|_| roda(false)).collect());
        resultado.insert("adamw_nosso_ms_por_passo".into(), t_nosso.into());
        resultado.insert("adamw_candle_nn_ms_por_passo".into(), t_candle.into());
        resultado.insert("nosso_vs_candle_nn".into(), (t_candle / t_nosso).into());
    }

    // ---------------- F) clip: escalar vs SIMD ----------------
    {
        let n = 2_000_000usize;
        let mut a: Vec<f32> = (0..n).map(|i| ((i as f32) * 0.001) % 0.5).collect();
        let mut b = a.clone();
        let t_esc = bench(|| {
            for x in a.iter_mut() { *x *= 0.37f32; }
            black_box(a.last());
        }, reps);
        let t_simd = bench(|| {
            simd::escalar_slice(&mut b, 0.37f32);
            black_box(b.last());
        }, reps);
        resultado.insert("clip_escalar_ms".into(), t_esc.into());
        resultado.insert("clip_simd_ms".into(), t_simd.into());
        resultado.insert("clip_speedup".into(), (t_esc / t_simd).into());
    }


    // ---------------- G) distâncias via MATMUL (gemm do candle: AVX-512+FMA) ----------------
    // D²[i,k] = ||q_i||² - 2·(q @ w^T)[i,k] + ||w_k||²  — muda a ordem de soma
    // (ε, não bit-a-bit); é a forma CANÔNICA de vetorizar distâncias em lote.
    {
        use candle_core::{Device, Tensor};
        let dev = Device::Cpu;
        let (c, d, nq) = (4096usize, 4usize, 2000usize);
        let mut rng = RngDeterministico::novo(4);
        let ps: Vec<Vec<f64>> = (0..c).map(|_| (0..d).map(|_| rng.normal(0.0, 2.0)).collect()).collect();
        let consultas: Vec<Vec<f64>> = (0..nq).map(|_| (0..d).map(|_| rng.normal(0.0, 2.0)).collect()).collect();
        let q_flat: Vec<f32> = consultas.iter().flat_map(|r| r.iter().map(|&v| v as f32)).collect();
        let w_flat: Vec<f32> = ps.iter().flat_map(|r| r.iter().map(|&v| v as f32)).collect();
        let q_t = Tensor::from_vec(q_flat, (nq, d), &dev).unwrap();
        let w_t = Tensor::from_vec(w_flat, (c, d), &dev).unwrap();
        let t_mm = bench(|| {
            // -2 q w^T
            let cross = q_t.matmul(&w_t.transpose(0, 1).unwrap()).unwrap().affine(-2.0, 0.0).unwrap();
            // + ||q||² por linha (broadcast) e + ||w||² por coluna
            let qn = q_t.sqr().unwrap().sum(1).unwrap().reshape((nq, 1)).unwrap().broadcast_as((nq, c)).unwrap();
            let wn = w_t.sqr().unwrap().sum(1).unwrap().reshape((1, c)).unwrap().broadcast_as((nq, c)).unwrap();
            let d2 = (&cross + &qn).unwrap().add(&wn).unwrap().clamp(0f32, f32::INFINITY).unwrap();
            black_box(d2.to_scalar::<f32>());
        }, 7);
        // referencia escalar p/ precisao
        let dist = |a: &[f64], b: &[f64]| -> f64 { a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() };
        let d2 = {
            let cross = q_t.matmul(&w_t.transpose(0, 1).unwrap()).unwrap().affine(-2.0, 0.0).unwrap();
            let qn = q_t.sqr().unwrap().sum(1).unwrap().reshape((nq, 1)).unwrap().broadcast_as((nq, c)).unwrap();
            let wn = w_t.sqr().unwrap().sum(1).unwrap().reshape((1, c)).unwrap().broadcast_as((nq, c)).unwrap();
            (&cross + &qn).unwrap().add(&wn).unwrap()
        };
        let mm_v = d2.to_vec2::<f32>().unwrap();
        let mut max_err = 0f64;
        for (i, qi) in consultas.iter().enumerate() {
            for (k, pk) in ps.iter().enumerate() {
                let exato = dist(qi, pk);
                let obt = mm_v[i][k] as f64;
                max_err = max_err.max((exato - obt).abs());
            }
        }
        resultado.insert("bmu_matmul_ms".into(), t_mm.into());
        resultado.insert("bmu_matmul_erro_max".into(), max_err.into());
        resultado.insert("bmu_speedup_matmul".into(), (15.313_f64 / t_mm).into());
    }
    let out = serde_json::json!({ "bench_v090": resultado });
    println!("{}", serde_json::to_string_pretty(&out).unwrap());
}