//! The requantisation probe: `R(t·m, s) + post` (§2, §5.1) through the kernel's
//! epilogue, as every operator that requantises forms it, for each rounding of §2 and
//! every shift from 0 to 63, over cases whose result fits 32 bits (the values an `acc32`
//! graph forms). Each line holds one multiplier and post bias, as one channel does:
//! powers of two (so that a quarter of its cases are exact ties and some are one unit
//! off a tie), random values at every width, and the extremes.

#![allow(dead_code)]

use wai_ops_npu::Npu;
use wai_ops_npu::lines::{Epi, MAX_N, RQ_RESCALE, round_code, rq_line};
use wai_ops_npu::probe::{ROUNDS, Rng, round_name, spec};

/// One line's cases: the inputs, the multiplier, the shift, the rounding and the results
/// §2 gives.
type Want = (Vec<i32>, i32, u32, &'static str, Vec<i64>);

/// What one probe run gave: the cases run, those that differ, and the first that did.
pub fn run(npu: &Npu, reps: usize, seed: u64) -> (u64, String) {
    let (_, bad, first) = run_counted(npu, reps, seed);
    (bad, first)
}

pub fn run_counted(npu: &Npu, reps: usize, seed: u64) -> (u64, u64, String) {
    let mut rng = Rng::new(seed);
    let mut lines = Vec::new();
    let mut want: Vec<Want> = Vec::new();
    for s in 0..=63u32 {
        for r in ROUNDS {
            for _ in 0..reps {
                let class = rng.next_u64() % 6;
                let (m, pow_b) = match class {
                    0 | 1 => {
                        let b = if s == 0 { 0 } else { (rng.next_u64() % s.min(31) as u64) as u32 };
                        let v = 1i64 << b;
                        ((if class == 0 { v } else { -v }).clamp(i32::MIN as i64, i32::MAX as i64) as i32, Some(b))
                    }
                    2 => (rng.i32(), None),
                    3 => (rng.bits(16), None),
                    4 => ([i32::MIN, i32::MAX, -1, 1][(rng.next_u64() % 4) as usize], None),
                    _ => (rng.bits(8), None),
                };
                let post = if rng.next_u64().is_multiple_of(3) { 0 } else { rng.bits(20) };
                let mut ts = Vec::with_capacity(MAX_N);
                let mut ys = Vec::with_capacity(MAX_N);
                let mut tries = 0;
                while ts.len() < MAX_N && tries < 200 * MAX_N {
                    tries += 1;
                    let k = ts.len();
                    let t = match pow_b {
                        Some(b) if s >= 1 && k % 4 < 2 && s - 1 - b <= 29 => {
                            let a = s - 1 - b;
                            let u = rng.bits(31 - a) as i64 >> 1;
                            let t = ((2 * u + 1) << a) as i32;
                            if k % 4 == 1 { t.saturating_add(if rng.next_u64() & 1 == 0 { 1 } else { -1 }) } else { t }
                        }
                        _ => match rng.next_u64() % 4 {
                            0 => rng.i32(),
                            1 => {
                                let bt = 1 + (rng.next_u64() % 32) as u32;
                                rng.bits(bt)
                            }
                            2 => [i32::MIN, i32::MAX, -1, 0, 1][(rng.next_u64() % 5) as usize],
                            _ => rng.bits(12),
                        },
                    };
                    let q = spec(t, m, s, r);
                    let y = q + post as i128;
                    if i32::try_from(q).is_ok() && i32::try_from(y).is_ok() {
                        ts.push(t);
                        ys.push(y as i64);
                    }
                }
                if ts.is_empty() {
                    continue;
                }
                let x: Vec<i64> = ts.iter().map(|&t| t as i64).collect();
                lines.push(rq_line(RQ_RESCALE, &x, &[], 0, Epi { round: round_code(r), shift: s as i32, mult: m, post }).unwrap());
                want.push((ts, m, s, round_name(r), ys));
            }
        }
    }
    let outs = npu.dispatch(&lines).expect("the probe's dispatches");
    let (mut cases, mut bad, mut first) = (0u64, 0u64, String::new());
    for ((ts, m, s, r, ys), o) in want.iter().zip(&outs) {
        for (i, (&t, &y)) in ts.iter().zip(ys).enumerate() {
            cases += 1;
            if o[i] as i64 != y {
                if bad == 0 {
                    first = format!("{r} s={s} t={t} m={m}: want {y}, got {}", o[i]);
                }
                bad += 1;
            }
        }
    }
    (cases, bad, first)
}