//! TRACES: S15 | FR-MRG-8
//! Load an ONNX file through the application's own runtime and run it once.
//!
//! ```sh
//! cargo run -p dr-segment --example onnx_probe --release -- model.onnx [1x1x768x1024]
//! ```
//!
//! The F6 check, as a tool. tract's operator coverage is the thing that can
//! sink a model choice — `segmentation.md` records a dynamic-shape export it
//! could not parse at all — and the only way to know is to load the file
//! under the backend the app ships and see. This does that for any model,
//! before any Rust is written against its outputs: it prints the declared
//! inputs and outputs, runs zeros through at the given shape, and times it.
//!
//! Written for S15.2 (XFeat), kept because the next model will need it too.

use std::time::Instant;

fn main() {
    let mut args = std::env::args().skip(1);
    let Some(path) = args.next() else {
        eprintln!("usage: onnx_probe <model.onnx> [NxCxHxW]");
        std::process::exit(2);
    };
    let shape: Vec<usize> = args
        .next()
        .map(|s| s.split('x').map(|d| d.parse().expect("dimension")).collect())
        .unwrap_or_else(|| vec![1, 1, 768, 1024]);

    let bytes = std::fs::read(&path).expect("read model");
    println!("{path}: {} bytes", bytes.len());

    let _ = ort::set_api(ort_tract::api());
    let t = Instant::now();
    let mut session = match ort::session::Session::builder().and_then(|mut b| b.commit_from_memory(&bytes)) {
        Ok(s) => s,
        Err(e) => {
            println!("FAIL  load: {e}");
            std::process::exit(1);
        }
    };
    println!("ok    loaded in {:?}", t.elapsed());
    for i in session.inputs().iter() {
        println!("      input  {} {:?}", i.name(), i.dtype());
    }
    for o in session.outputs().iter() {
        println!("      output {} {:?}", o.name(), o.dtype());
    }

    let n: usize = shape.iter().product();

    // Twice: the first run pays for tract's optimisation and plan, the second
    // is the number that matters. The tensor is built per run rather than
    // cloned — `Tensor::clone` under the tract backend panics.
    for pass in 1..=2 {
        let input =
            ndarray::Array::from_shape_vec(ndarray::IxDyn(&shape), vec![0.0f32; n]).expect("shape");
        let tensor = ort::value::Tensor::from_array(input).expect("tensor");
        let t = Instant::now();
        let outputs = match session.run(ort::inputs![tensor]) {
            Ok(o) => o,
            Err(e) => {
                println!("FAIL  run: {e}");
                std::process::exit(1);
            }
        };
        println!("ok    run {pass} in {:?}", t.elapsed());
        if pass == 2 {
            for i in 0..outputs.len() {
                match outputs[i].try_extract_tensor::<f32>() {
                    Ok((shape, data)) => {
                        let (lo, hi) = data
                            .iter()
                            .fold((f32::MAX, f32::MIN), |(lo, hi), &v| (lo.min(v), hi.max(v)));
                        println!("      output {i}: shape {shape:?}, range {lo:.4}..{hi:.4}");
                    }
                    Err(e) => println!("      output {i}: not f32 ({e})"),
                }
            }
        }
    }
}