//! 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; never the Hexagon, where int8 lost 6–9 dB.

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: &[f32], sigma: &[f32]) -> Result<Vec<f32>, DenoiseError> {
        let n = self.tile;
        let shape = ndarray::IxDyn(&[1, 1, n, n]);
        let m = ort::value::Tensor::from_array(
            ndarray::Array::from_shape_vec(shape.clone(), mosaic.to_vec())
                .map_err(|e| DenoiseError::Model(e.to_string()))?,
        )?;
        let s = ort::value::Tensor::from_array(
            ndarray::Array::from_shape_vec(shape, sigma.to_vec())
                .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}]"
            )));
        }
        Ok(data.to_vec())
    }
}