oxigdal-ml 0.1.7

//! Super-resolution model implementation

use ndarray::{Array2, Array3, Array4, Axis, s};
use oxionnx::{GraphOptimizationLevel, Session, SessionBuilder, Tensor};
use rayon::prelude::*;
use serde::{Deserialize, Serialize};
use std::path::Path;
use tracing::{debug, info};

use crate::error::{InferenceError, ModelError, Result};
use oxigdal_core::buffer::RasterBuffer;
use oxigdal_core::types::RasterDataType;

use super::UpscaleFactor;

/// Configuration for super-resolution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SuperResConfig {
    /// Upscale factor (2x or 4x)
    pub scale_factor: usize,
    /// Tile size (width and height in pixels)
    pub tile_size: usize,
    /// Overlap between tiles (in pixels)
    pub overlap: usize,
    /// Batch size for processing multiple tiles
    pub batch_size: usize,
}

impl SuperResConfig {
    /// Creates a new super-resolution configuration
    ///
    /// # Arguments
    ///
    /// * `scale_factor` - Upscale factor (2 or 4)
    /// * `tile_size` - Size of tiles for processing
    /// * `overlap` - Overlap between tiles in pixels
    ///
    /// # Example
    ///
    /// ```
    /// use oxigdal_ml::superres::SuperResConfig;
    ///
    /// let config = SuperResConfig::new(2, 256, 32);
    /// assert_eq!(config.scale_factor, 2);
    /// assert_eq!(config.tile_size, 256);
    /// assert_eq!(config.overlap, 32);
    /// ```
    #[must_use]
    pub fn new(scale_factor: usize, tile_size: usize, overlap: usize) -> Self {
        Self {
            scale_factor,
            tile_size,
            overlap,
            batch_size: 1,
        }
    }

    /// Set batch size for parallel processing
    #[must_use]
    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }
}

impl Default for SuperResConfig {
    fn default() -> Self {
        Self::new(2, 256, 32)
    }
}

/// Super-resolution model using ONNX Runtime
pub struct SuperResolution {
    session: Session,
    config: SuperResConfig,
}

impl SuperResolution {
    /// Load a super-resolution model from an ONNX file
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the ONNX model file
    /// * `config` - Super-resolution configuration
    ///
    /// # Errors
    ///
    /// Returns an error if the model file cannot be loaded or is invalid
    ///
    /// # Example
    ///
    /// ```no_run
    /// use oxigdal_ml::superres::{SuperResolution, SuperResConfig};
    ///
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let config = SuperResConfig::default();
    /// let model = SuperResolution::from_file("real_esrgan_2x.onnx", config)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn from_file<P: AsRef<Path>>(path: P, config: SuperResConfig) -> Result<Self> {
        let path = path.as_ref();

        if !path.exists() {
            return Err(ModelError::NotFound {
                path: path.display().to_string(),
            }
            .into());
        }

        // SessionBuilder methods return Self (no Result until commit_from_file)
        let session = SessionBuilder::new()
            .with_optimization_level(GraphOptimizationLevel::All)
            .commit_from_file(path)
            .map_err(|e: oxionnx::OnnxError| ModelError::LoadFailed {
                reason: e.to_string(),
            })?;

        info!("Loaded super-resolution model from {}", path.display());

        Ok(Self { session, config })
    }

    /// Upscale a raster using the super-resolution model
    ///
    /// # Arguments
    ///
    /// * `input` - Input raster buffer
    ///
    /// # Errors
    ///
    /// Returns an error if inference fails or input validation fails
    ///
    /// # Example
    ///
    /// ```no_run
    /// use oxigdal_ml::superres::{SuperResolution, SuperResConfig};
    /// use oxigdal_core::buffer::RasterBuffer;
    /// use oxigdal_core::types::RasterDataType;
    ///
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let config = SuperResConfig::default();
    /// let mut model = SuperResolution::from_file("real_esrgan_2x.onnx", config)?;
    ///
    /// let input = RasterBuffer::zeros(512, 512, RasterDataType::Float32);
    /// let output = model.upscale(&input)?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn upscale(&mut self, input: &RasterBuffer) -> Result<RasterBuffer> {
        let width = input.width() as usize;
        let height = input.height() as usize;

        debug!(
            "Starting super-resolution upscaling: {}x{} -> {}x{}",
            width,
            height,
            width * self.config.scale_factor,
            height * self.config.scale_factor
        );

        // Extract tiles with overlap
        let tiles = self.extract_tiles(input)?;

        // Process tiles in batches (sequential for now due to &mut self requirement)
        let processed_tiles = self.process_batch(&tiles)?;

        // Merge tiles with blending
        let merged = self.merge_tiles(
            &processed_tiles,
            width * self.config.scale_factor,
            height * self.config.scale_factor,
        )?;

        // Create output buffer
        RasterBuffer::new(
            merged
                .as_slice()
                .ok_or_else(|| InferenceError::OutputParsingFailed {
                    reason: "Failed to convert array to slice".to_string(),
                })?
                .iter()
                .flat_map(|&v: &f32| v.to_le_bytes())
                .collect(),
            (width * self.config.scale_factor) as u64,
            (height * self.config.scale_factor) as u64,
            RasterDataType::Float32,
            input.nodata(),
        )
        .map_err(Into::into)
    }

    /// Extract tiles from input raster with overlap.
    ///
    /// Delegates tile layout to [`crate::tiling::compute_tile_grid`].
    fn extract_tiles(&self, input: &RasterBuffer) -> Result<Vec<TileInfo>> {
        let width = input.width() as usize;
        let height = input.height() as usize;
        let tile_size = self.config.tile_size;
        let overlap = self.config.overlap;

        // Compute tile positions via the shared tiling module
        let specs = crate::tiling::compute_tile_grid(width, height, tile_size, tile_size, overlap)?;

        let mut tiles = Vec::with_capacity(specs.len());

        for spec in &specs {
            tiles.push(TileInfo {
                x: spec.x_offset,
                y: spec.y_offset,
                width: spec.width,
                height: spec.height,
                data: self.extract_tile_data(
                    input,
                    spec.x_offset,
                    spec.y_offset,
                    spec.width,
                    spec.height,
                )?,
            });
        }

        debug!(
            "Extracted {} tiles from {}x{} image",
            tiles.len(),
            width,
            height
        );

        Ok(tiles)
    }

    /// Extract single tile data
    fn extract_tile_data(
        &self,
        input: &RasterBuffer,
        x: usize,
        y: usize,
        width: usize,
        height: usize,
    ) -> Result<Array3<f32>> {
        // For simplicity, assume single-band input (can be extended to multi-band)
        let mut tile = Array3::<f32>::zeros((1, height, width));

        // Extract data from buffer (simplified - assumes Float32)
        for ty in 0..height {
            for tx in 0..width {
                let pixel_idx = ((y + ty) * input.width() as usize + (x + tx)) * 4; // 4 bytes for Float32
                let bytes = input.as_bytes();

                if pixel_idx + 4 <= bytes.len() {
                    let value = f32::from_le_bytes([
                        bytes[pixel_idx],
                        bytes[pixel_idx + 1],
                        bytes[pixel_idx + 2],
                        bytes[pixel_idx + 3],
                    ]);
                    tile[[0, ty, tx]] = value;
                }
            }
        }

        Ok(tile)
    }

    /// Process a batch of tiles through the model
    fn process_batch(&mut self, tiles: &[TileInfo]) -> Result<Vec<ProcessedTile>> {
        let mut processed = Vec::with_capacity(tiles.len());

        // Get input/output names from session metadata
        let input_name = self
            .session
            .input_info()
            .first()
            .map(|i| i.name.clone())
            .unwrap_or_else(|| "input".to_string());

        let output_name = self
            .session
            .output_info()
            .first()
            .map(|o| o.name.clone())
            .unwrap_or_else(|| "output".to_string());

        for tile in tiles {
            // Create input tensor with batch dimension inserted
            let input_tensor_arr = tile.data.clone().insert_axis(Axis(0));

            // Create Tensor from ndarray (Tensor::from_ndarray returns Tensor directly)
            let input_value = Tensor::from_ndarray(input_tensor_arr);

            // Build inputs map using oxionnx::inputs! macro
            let inputs_map = oxionnx::inputs![input_name.as_str() => input_value].map_err(
                |e: oxionnx::OnnxError| InferenceError::Failed {
                    reason: format!("Failed to build inputs map: {}", e),
                },
            )?;

            // Run inference — session.run takes &HashMap<&str, Tensor>
            let outputs = self
                .session
                .run(&inputs_map)
                .map_err(|e: oxionnx::OnnxError| InferenceError::Failed {
                    reason: e.to_string(),
                })?;

            // Extract output tensor from HashMap<String, Tensor>
            let output_value = outputs.get(output_name.as_str()).ok_or_else(|| {
                InferenceError::OutputParsingFailed {
                    reason: format!("Output tensor '{}' not found", output_name),
                }
            })?;

            // try_extract_tensor returns (&[usize], &[f32]) — shape elements are usize
            let (shape, data) =
                output_value
                    .try_extract_tensor::<f32>()
                    .map_err(
                        |e: oxionnx::OnnxError| InferenceError::OutputParsingFailed {
                            reason: e.to_string(),
                        },
                    )?;

            // shape elements are already usize — no cast needed
            let shape_vec: Vec<usize> = shape.to_vec();

            // Convert to ndarray
            let output: Array4<f32> = Array4::from_shape_vec(
                (shape_vec[0], shape_vec[1], shape_vec[2], shape_vec[3]),
                data.to_vec(),
            )
            .map_err(|e| InferenceError::OutputParsingFailed {
                reason: format!("Failed to reshape output: {}", e),
            })?;

            processed.push(ProcessedTile {
                x: tile.x * self.config.scale_factor,
                y: tile.y * self.config.scale_factor,
                data: output.index_axis_move(Axis(0), 0),
            });
        }

        Ok(processed)
    }

    /// Merge processed tiles with overlap blending
    fn merge_tiles(
        &self,
        tiles: &[ProcessedTile],
        output_width: usize,
        output_height: usize,
    ) -> Result<Array3<f32>> {
        let mut output = Array3::<f32>::zeros((1, output_height, output_width));
        let mut weight_map = Array3::<f32>::zeros((1, output_height, output_width));

        let overlap = self.config.overlap * self.config.scale_factor;

        for tile in tiles {
            let tile_height = tile.data.shape()[1];
            let tile_width = tile.data.shape()[2];

            // Create weight matrix for alpha blending
            let weights = self.create_blend_weights(tile_width, tile_height, overlap);

            // Blend tile into output
            for c in 0..1 {
                for ty in 0..tile_height {
                    for tx in 0..tile_width {
                        let out_y = tile.y + ty;
                        let out_x = tile.x + tx;

                        if out_y < output_height && out_x < output_width {
                            let weight = weights[[ty, tx]];
                            output[[c, out_y, out_x]] += tile.data[[c, ty, tx]] * weight;
                            weight_map[[c, out_y, out_x]] += weight;
                        }
                    }
                }
            }
        }

        // Normalize by weight map
        output.zip_mut_with(&weight_map, |out, &w| {
            if w > 0.0 {
                *out /= w;
            }
        });

        Ok(output)
    }

    /// Create blend weights for smooth tile merging.
    ///
    /// Note: this uses a slightly different formula from the shared
    /// [`crate::tiling::compute_blend_weight`]: on the right and bottom
    /// edges it uses `(width - x)` rather than `(width - x - 1)`, which
    /// gives a non-zero weight on the very last pixel row/column.
    fn create_blend_weights(&self, width: usize, height: usize, overlap: usize) -> Array2<f32> {
        let mut weights = Array2::<f32>::ones((height, width));

        if overlap == 0 {
            return weights;
        }

        // Create linear blend in overlap regions
        for y in 0..height {
            for x in 0..width {
                let mut w = 1.0_f32;

                // Blend on edges
                if x < overlap {
                    w = w.min(x as f32 / overlap as f32);
                }
                if x >= width - overlap {
                    w = w.min((width - x) as f32 / overlap as f32);
                }
                if y < overlap {
                    w = w.min(y as f32 / overlap as f32);
                }
                if y >= height - overlap {
                    w = w.min((height - y) as f32 / overlap as f32);
                }

                weights[[y, x]] = w;
            }
        }

        weights
    }
}

/// Information about an extracted tile
struct TileInfo {
    x: usize,
    y: usize,
    width: usize,
    height: usize,
    data: Array3<f32>,
}

/// A processed (upscaled) tile
struct ProcessedTile {
    x: usize,
    y: usize,
    data: Array3<f32>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_config_creation() {
        let config = SuperResConfig::new(2, 256, 32);
        assert_eq!(config.scale_factor, 2);
        assert_eq!(config.tile_size, 256);
        assert_eq!(config.overlap, 32);
    }

    #[test]
    fn test_config_default() {
        let config = SuperResConfig::default();
        assert_eq!(config.scale_factor, 2);
        assert_eq!(config.batch_size, 1);
    }

    #[test]
    #[ignore = "Requires ONNX model to be installed"]
    fn test_blend_weights() {
        let _config = SuperResConfig::default();
        // We can't easily test this without a full model
        // Just verify the method signature compiles
    }
}