//! TRACES: FR-DEV-3g
//! The denoise network under the inference engine.
//!
//! The shipped export takes `mosaic` and `sigma`, `1×1×1408×1408`, and
//! returns `rgb`, `1×3×1408×1408` (darkroom-denoise `denoise/export.py`,
//! fixed shape because every model the engine runs is). The engine picks the
//! rung: fp16 on TensorRT and MIGraphX, which measured 0.00 dB from f32; f32
//! on CUDA and the CPU; on the Hexagon the `.a16w16.onnx` sibling, 16-bit
//! activations and weights, 0.00 dB from f32 on the tablet itself where int8
//! lost 5–9 dB (docs/dev/inference.md §1.5). That sibling is the same network
//! with the Bayer packing spelled `SpaceToDepth`, which QNN can hold and the
//! 6-D reshape it replaces it cannot.

use crate::tile::TileNet;
use crate::DenoiseError;
use dr_inference_engine::{Model, Role};

/// The edge of the tile the shipped export takes.
pub const TILE: usize = 1408;

pub struct OnnxNet {
    model: Model,
    tile: usize,
}

impl OnnxNet {
    pub fn from_path(path: &std::path::Path) -> Result<Self, DenoiseError> {
        let (path, form) = dr_inference_engine::resolve_model(Role::Denoiser, path);
        let bytes = std::fs::read(&path)?;
        Ok(OnnxNet {
            model: dr_inference_engine::open(Role::Denoiser, form, &bytes)?,
            tile: TILE,
        })
    }

    /// Where it runs, for a status line.
    pub fn rung(&self) -> Result<dr_inference_engine::Rung, DenoiseError> {
        Ok(self.model.acquire()?.rung())
    }
}

impl TileNet for OnnxNet {
    fn tile(&self) -> usize {
        self.tile
    }

    fn run(
        &mut self,
        mosaic: Vec<f32>,
        sigma: Vec<f32>,
        write: &mut dyn FnMut(&[f32]),
    ) -> Result<(), DenoiseError> {
        let n = self.tile;
        let shape = ndarray::IxDyn(&[1, 1, n, n]);
        // The vectors become the tensors: no copy on the way in.
        let m = ort::value::Tensor::from_array(
            ndarray::Array::from_shape_vec(shape.clone(), mosaic)
                .map_err(|e| DenoiseError::Model(e.to_string()))?,
        )?;
        let s = ort::value::Tensor::from_array(
            ndarray::Array::from_shape_vec(shape, sigma)
                .map_err(|e| DenoiseError::Model(e.to_string()))?,
        )?;
        let acquired = self.model.acquire()?;
        let mut session = acquired.lock();
        let outputs = session.run(ort::inputs!["mosaic" => m, "sigma" => s])?;
        let (shape, data) = outputs[0].try_extract_tensor::<f32>()?;
        let dims: Vec<i64> = shape.iter().copied().collect();
        if dims != [1, 3, n as i64, n as i64] {
            return Err(DenoiseError::Model(format!(
                "output is {dims:?}, expected [1, 3, {n}, {n}]"
            )));
        }
        // And none on the way out: the frame is written from the runtime's buffer.
        write(data);
        Ok(())
    }
}