//! Writes the limb-product model in the portable format (fp16 and int32), and a
//! batch of inputs with the products they must give, for a neural-processor
//! service that runs such models:
//!
//! ```text
//! cargo run --example npu_onnx --features quantum_fabric -- <dir>
//! ```
//!
//! In `<dir>`: `limbs_f16.onnx`, `limbs_i32.onnx`; `x_f16.bin`, `y_f16.bin`
//! (fp16 bits, little-endian), `x_i32.bin`, `y_i32.bin` (int32, little-endian),
//! each `[batch, 15, 8]`; `pairs.bin` (the operands, int32 pairs) and
//! `products.bin` (each exact `a·b`, int64). A service's output `z`
//! (`[batch, 15]`) is right when `Σ_s z[p, s]·16^s` is `products[p]` for every
//! `p`.
use std::io::Write;
use wai_quantum::quantum_fabric::lanes;
use wai_quantum::quantum_fabric::npu::{BATCH, Elem, LIMBS, SHIFTS, expand, expand_i32, onnx_limb_model};

fn write(path: &std::path::Path, bytes: &[u8]) {
    std::fs::File::create(path).and_then(|mut f| f.write_all(bytes)).unwrap_or_else(|e| panic!("{}: {e}", path.display()));
}

fn main() {
    let dir = std::path::PathBuf::from(std::env::args().nth(1).expect("an output directory"));
    std::fs::create_dir_all(&dir).expect("the directory");
    write(&dir.join("limbs_f16.onnx"), &onnx_limb_model(BATCH as u64, SHIFTS as u64, LIMBS as u64, Elem::F16));
    write(&dir.join("limbs_i32.onnx"), &onnx_limb_model(BATCH as u64, SHIFTS as u64, LIMBS as u64, Elem::I32));
    let mut pairs = lanes::probe_pairs(BATCH, 0x0a1b);
    pairs.truncate(BATCH);
    let n = BATCH * SHIFTS * LIMBS;
    let (mut x, mut y) = (vec![0u16; n], vec![0u16; n]);
    expand(&pairs, &mut x, &mut y);
    write(&dir.join("x_f16.bin"), &x.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<u8>>());
    write(&dir.join("y_f16.bin"), &y.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<u8>>());
    let (mut xi, mut yi) = (vec![0i32; n], vec![0i32; n]);
    expand_i32(&pairs, &mut xi, &mut yi);
    write(&dir.join("x_i32.bin"), &xi.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<u8>>());
    write(&dir.join("y_i32.bin"), &yi.iter().flat_map(|v| v.to_le_bytes()).collect::<Vec<u8>>());
    write(&dir.join("pairs.bin"), &pairs.iter().flat_map(|&(a, b)| [a.to_le_bytes(), b.to_le_bytes()].concat()).collect::<Vec<u8>>());
    write(&dir.join("products.bin"), &pairs.iter().flat_map(|&(a, b)| (i64::from(a) * i64::from(b)).to_le_bytes()).collect::<Vec<u8>>());
    println!("wrote {} pairs of a batch of {BATCH} to {}", pairs.len(), dir.display());
}