onnx-ir 0.21.0

//! # Conv (1D)
//!
//! 1D convolution operation.
//!
//! **ONNX Spec**: <https://onnx.ai/onnx/operators/onnx__Conv.html>
//!
//! ## Opset Versions
//! - **Opset 1**: Initial version with basic convolution support
//! - **Opset 11**: No changes to Conv operator itself (broader ONNX updates)

use derive_new::new;
use onnx_ir_derive::NodeBuilder;

use crate::ir::{ArgType, Argument, Node, RawNode, TensorType};
use crate::processor::{
    InputSpec, NodeProcessor, NodeSpec, OutputPreferences, OutputSpec, ProcessError,
};

use super::padding::{AutoPad, PaddingConfig1d, padding_config_1d};

/// Node representation for Conv1d operation
#[derive(Debug, Clone, NodeBuilder)]
pub struct Conv1dNode {
    pub name: String,
    pub inputs: Vec<Argument>,
    pub outputs: Vec<Argument>,
    pub config: Conv1dConfig,
}

/// Configuration for Conv1d operations
#[derive(Debug, Clone, new)]
#[allow(clippy::too_many_arguments)]
pub struct Conv1dConfig {
    /// Kernel size
    pub kernel_size: usize,
    /// Stride
    pub stride: usize,
    /// Dilation
    pub dilation: usize,
    /// Number of groups
    pub groups: usize,
    /// Padding configuration (from explicit pads attribute)
    pub padding: PaddingConfig1d,
    /// Auto padding mode
    pub auto_pad: AutoPad,
}

/// Node processor for Conv1d operation
pub(crate) struct Conv1dProcessor;

impl NodeProcessor for Conv1dProcessor {
    type Config = Conv1dConfig;

    fn spec(&self) -> NodeSpec {
        NodeSpec {
            min_opset: 1,
            max_opset: None,
            inputs: InputSpec::Range(2, 3),
            outputs: OutputSpec::Exact(1),
        }
    }

    fn lift_constants(&self, node: &mut RawNode, _opset: usize) -> Result<(), ProcessError> {
        // Lift weight (input[1]) and optional bias (input[2])
        if node.inputs.len() > 1 && node.inputs[1].is_constant() {
            node.inputs[1].to_static()?;
        }
        if node.inputs.len() > 2 && node.inputs[2].is_constant() {
            node.inputs[2].to_static()?;
        }

        Ok(())
    }

    fn infer_types(
        &self,
        node: &mut RawNode,
        _opset: usize,
        _output_preferences: &OutputPreferences,
    ) -> Result<(), ProcessError> {
        // TODO: Add test for zero or negative stride values - spec requires positive strides
        // TODO: Add test for zero or negative dilation values - spec requires positive dilations
        // TODO: Add test for zero or negative group values - spec requires positive groups
        // TODO: Validate channels_in divisible by groups - required by spec but not validated
        // TODO: Validate channels_out divisible by groups - required by spec but not validated
        // TODO: Add test for very large kernel_shape/stride/dilation - potential overflow/memory issues
        // TODO: Add test coverage for auto_pad values: SAME_UPPER, SAME_LOWER, VALID - currently unsupported

        // Validate attributes before extracting config
        for (key, value) in node.attrs.iter() {
            match key.as_str() {
                "kernel_shape" | "strides" | "pads" | "dilations" | "group" => {}
                "auto_pad" => {
                    // Validate the value is a known auto_pad string
                    AutoPad::parse(&value.clone().into_string())?;
                }
                _ => {
                    return Err(ProcessError::InvalidAttribute {
                        name: key.clone(),
                        reason: format!("Unexpected attribute for Conv1d: {key}"),
                    });
                }
            }
        }

        // Extract input tensor type
        let tensor = match &node.inputs[0].ty {
            ArgType::Tensor(tensor) => tensor,
            _ => {
                return Err(ProcessError::TypeMismatch {
                    expected: "Tensor".to_string(),
                    actual: format!("{:?}", node.inputs[0].ty),
                });
            }
        };

        // Validate input tensor rank - Conv1d expects rank 3 (N x C x L)
        if tensor.rank != 3 {
            return Err(ProcessError::Custom(format!(
                "Conv1d expects input tensor of rank 3 (N x C x L), got rank {}",
                tensor.rank
            )));
        }

        // Validate weight tensor type and rank
        let weight_tensor = match &node.inputs[1].ty {
            ArgType::Tensor(tensor) => tensor,
            _ => {
                return Err(ProcessError::TypeMismatch {
                    expected: "Tensor (weight)".to_string(),
                    actual: format!("{:?}", node.inputs[1].ty),
                });
            }
        };

        // Weight should be rank 3 (M x C/group x kW)
        if weight_tensor.rank != 3 {
            return Err(ProcessError::Custom(format!(
                "Conv1d expects weight tensor of rank 3 (M x C/group x kW), got rank {}",
                weight_tensor.rank
            )));
        }

        // Validate dtypes match
        if tensor.dtype != weight_tensor.dtype {
            return Err(ProcessError::TypeMismatch {
                expected: format!("Weight tensor with dtype {:?}", tensor.dtype),
                actual: format!("Weight tensor with dtype {:?}", weight_tensor.dtype),
            });
        }

        // Validate bias if present
        if node.inputs.len() > 2 {
            let bias_tensor = match &node.inputs[2].ty {
                ArgType::Tensor(tensor) => tensor,
                _ => {
                    return Err(ProcessError::TypeMismatch {
                        expected: "Tensor (bias)".to_string(),
                        actual: format!("{:?}", node.inputs[2].ty),
                    });
                }
            };

            // Bias should be rank 1 (M)
            if bias_tensor.rank != 1 {
                return Err(ProcessError::Custom(format!(
                    "Conv1d expects bias tensor of rank 1 (M), got rank {}",
                    bias_tensor.rank
                )));
            }

            // Validate bias dtype matches
            if tensor.dtype != bias_tensor.dtype {
                return Err(ProcessError::TypeMismatch {
                    expected: format!("Bias tensor with dtype {:?}", tensor.dtype),
                    actual: format!("Bias tensor with dtype {:?}", bias_tensor.dtype),
                });
            }
        }

        // Compute output static_shape: [batch, out_channels, L_out]
        let static_shape = {
            let batch = tensor
                .static_shape
                .as_ref()
                .and_then(|s| s.first().copied().flatten());
            let out_channels = node.inputs[1]
                .value()
                .and_then(|data| data.shape.first().copied())
                .or_else(|| {
                    weight_tensor
                        .static_shape
                        .as_ref()
                        .and_then(|s| s.first().copied().flatten())
                });
            let l_out = self.extract_config(node, _opset).ok().and_then(|config| {
                let l_in = tensor
                    .static_shape
                    .as_ref()
                    .and_then(|s| s.get(2).copied().flatten())?;
                let (pad_left, pad_right) = config.padding.as_tuple();
                let padding = pad_left + pad_right;
                let dilation = config.dilation;
                let kernel = config.kernel_size;
                let stride = config.stride;
                // L_out = floor((L_in + padding - dilation * (kernel - 1) - 1) / stride + 1)
                let numerator = l_in as isize + padding as isize
                    - dilation as isize * (kernel as isize - 1)
                    - 1;
                if numerator < 0 || stride == 0 {
                    return None;
                }
                Some(numerator as usize / stride + 1)
            });
            Some(vec![batch, out_channels, l_out])
        };

        // Conv1d preserves rank (same as input)
        node.outputs[0].ty = ArgType::Tensor(TensorType {
            dtype: tensor.dtype,
            rank: tensor.rank,
            static_shape,
        });

        Ok(())
    }

    fn extract_config(&self, node: &RawNode, _opset: usize) -> Result<Self::Config, ProcessError> {
        let mut kernel_shape = Vec::new();
        let mut strides = vec![1];
        let mut pads = vec![0, 0];
        let mut dilations = vec![1];
        let mut group: usize = 1;
        let mut auto_pad = AutoPad::NotSet;

        let weight_shape = node.inputs[1]
            .value()
            .ok_or_else(|| {
                ProcessError::Custom("Conv1d: weight tensor must be present".to_string())
            })?
            .shape
            .to_vec();

        for (key, value) in node.attrs.iter() {
            match key.as_str() {
                "kernel_shape" => kernel_shape = value.clone().into_i64s(),
                "strides" => strides = value.clone().into_i64s(),
                "pads" => pads = value.clone().into_i64s(),
                "dilations" => dilations = value.clone().into_i64s(),
                "group" => group = value.clone().into_i64() as usize,
                "auto_pad" => auto_pad = AutoPad::parse(&value.clone().into_string())?,
                _ => {}
            }
        }

        let padding = padding_config_1d(&pads);

        let kernel_size = if kernel_shape.is_empty() {
            if weight_shape.len() != 3 {
                return Err(ProcessError::Custom(format!(
                    "Conv1d: expected to infer kernel shape from a weight tensor of rank 3 but got shape {:?}",
                    weight_shape
                )));
            }
            weight_shape[2]
        } else {
            kernel_shape[0] as _
        };

        let config = Conv1dConfig::new(
            kernel_size,
            strides[0] as usize,
            dilations[0] as usize,
            group,
            padding,
            auto_pad,
        );

        Ok(config)
    }

    fn build_node(&self, builder: RawNode, opset: usize) -> Node {
        let config = self
            .extract_config(&builder, opset)
            .expect("Config extraction failed");

        Node::Conv1d(Conv1dNode {
            name: builder.name,
            inputs: builder.inputs,
            outputs: builder.outputs,
            config,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ir::NodeType;
    use crate::node::test_utils::TestNodeBuilder;

    fn create_test_node(
        kernel_shape: Vec<i64>,
        strides: Vec<i64>,
        pads: Vec<i64>,
        dilations: Vec<i64>,
        group: i64,
        has_bias: bool,
        auto_pad: Option<&str>,
    ) -> TestNodeBuilder {
        // Create weight tensor data
        let weight_data = vec![0.1; 16];

        let has_kernel_shape = !kernel_shape.is_empty();

        // Start building the node with input and weight
        let mut builder = TestNodeBuilder::new(NodeType::Conv1d, "test_conv1d")
            .input_tensor_f32("data", 3, None)
            .input_tensor_f32_data(
                "weight",
                weight_data,
                vec![2, 2, 4], // [out_channels, in_channels, kernel_size]
            )
            .output_tensor_f32("output", 3, None);

        // Add bias if needed
        if has_bias {
            builder = builder.input_tensor_f32_data("bias", vec![0.1, 0.2], vec![2]);
        }

        if let Some(auto_pad) = auto_pad {
            builder = builder.attr_string("auto_pad", auto_pad);
        }

        // Add attributes
        builder = builder
            .attr_ints("strides", strides)
            .attr_ints("pads", pads)
            .attr_ints("dilations", dilations)
            .attr_int("group", group);

        if has_kernel_shape {
            builder = builder.attr_ints("kernel_shape", kernel_shape);
        }

        builder
    }

    #[test]
    fn test_conv1d_config_basic() {
        let node = create_test_node(vec![4], vec![1], vec![0, 0], vec![1], 1, false, None)
            .build_with_graph_data(16);
        let mut node = node;
        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, 4);
        assert_eq!(config.stride, 1);
        assert_eq!(config.dilation, 1);
        assert_eq!(config.groups, 1);
        assert!(matches!(config.padding, PaddingConfig1d::Valid));
    }

    #[test]
    fn test_conv1d_config_with_padding() {
        let node = create_test_node(vec![4], vec![2], vec![2, 2], vec![1], 1, true, None)
            .build_with_graph_data(16);
        let mut node = node;
        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, 4);
        assert_eq!(config.stride, 2);
        assert_eq!(config.dilation, 1);
        assert_eq!(config.groups, 1);
        assert!(matches!(config.padding, PaddingConfig1d::Explicit(2, 2)));
    }

    #[test]
    fn test_conv1d_config_with_dilation() {
        let node = create_test_node(vec![4], vec![1], vec![0, 0], vec![2], 1, false, None)
            .build_with_graph_data(16);
        let mut node = node;
        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, 4);
        assert_eq!(config.stride, 1);
        assert_eq!(config.dilation, 2);
        assert_eq!(config.groups, 1);
        assert!(matches!(config.padding, PaddingConfig1d::Valid));
    }

    #[test]
    fn test_conv1d_config_with_groups() {
        let node = create_test_node(vec![4], vec![1], vec![0, 0], vec![1], 2, false, None)
            .build_with_graph_data(16);
        let mut node = node;
        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, 4);
        assert_eq!(config.stride, 1);
        assert_eq!(config.dilation, 1);
        assert_eq!(config.groups, 2);
        assert!(matches!(config.padding, PaddingConfig1d::Valid));
    }

    #[test]
    fn test_conv1d_config_asymmetric_padding() {
        let node = create_test_node(vec![4], vec![1], vec![1, 2], vec![1], 1, false, None)
            .build_with_graph_data(16);
        let processor = Conv1dProcessor;
        let config = processor.extract_config(&node, 16).unwrap();
        // Asymmetric padding should now be captured instead of panicking
        assert!(matches!(config.padding, PaddingConfig1d::Explicit(1, 2)));
        assert!(config.padding.is_asymmetric());
    }

    #[test]
    #[should_panic(expected = "Negative pad values are not supported")]
    fn test_conv1d_config_negative_padding() {
        let node = create_test_node(vec![4], vec![1], vec![-1, -1], vec![1], 1, false, None)
            .build_with_graph_data(16);
        let processor = Conv1dProcessor;
        let _ = processor.extract_config(&node, 16);
    }

    #[test]
    fn test_conv1d_config_autopad_not_set() {
        let node = create_test_node(
            vec![4],
            vec![1],
            vec![0, 0],
            vec![1],
            1,
            false,
            Some("NOTSET"),
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, 4);
        assert_eq!(config.stride, 1);
        assert_eq!(config.dilation, 1);
        assert_eq!(config.groups, 1);
        assert!(matches!(config.padding, PaddingConfig1d::Valid));
    }

    #[test]
    fn test_conv1d_config_autopad_same_upper() {
        let node = create_test_node(
            vec![4],
            vec![1],
            vec![0, 0],
            vec![1],
            1,
            false,
            Some("SAME_UPPER"),
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        processor.infer_types(&mut node, 16, &prefs).unwrap();
        let config = processor.extract_config(&node, 16).unwrap();
        assert_eq!(config.auto_pad, AutoPad::SameUpper);
    }

    #[test]
    fn test_conv1d_config_kernel_shape_not_set() {
        let node = create_test_node(
            vec![],
            vec![1, 1],
            vec![0, 0, 0, 0],
            vec![1, 1],
            1,
            false,
            None,
        )
        .build_with_graph_data(16);
        let mut node = node;
        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        let config = processor.extract_config(&node, 16).unwrap();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        assert_eq!(config.kernel_size, 4); // Inferred via weight tensor shape
    }

    #[test]
    fn test_conv1d_static_shape_known() {
        // Input [1, 2, 10], weight [2, 2, 4], stride=1, pad=0, dilation=1
        // L_out = (10 + 0 - 1*(4-1) - 1) / 1 + 1 = (10 - 3 - 1) / 1 + 1 = 7
        let mut node = TestNodeBuilder::new(NodeType::Conv1d, "test")
            .input_tensor_f32("data", 3, Some(vec![1, 2, 10]))
            .input_tensor_f32_data("weight", vec![0.1; 16], vec![2, 2, 4])
            .output_tensor_f32("output", 3, None)
            .attr_ints("kernel_shape", vec![4])
            .attr_ints("strides", vec![1])
            .attr_ints("pads", vec![0, 0])
            .attr_ints("dilations", vec![1])
            .attr_int("group", 1)
            .build_with_graph_data(16);

        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        match &node.outputs[0].ty {
            ArgType::Tensor(t) => {
                assert_eq!(t.rank, 3);
                assert_eq!(t.static_shape, Some(vec![Some(1), Some(2), Some(7)]));
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_conv1d_static_shape_with_padding() {
        // Input [1, 2, 10], weight [2, 2, 4], stride=2, pad=[2, 2], dilation=1
        // L_out = (10 + 4 - 1*(4-1) - 1) / 2 + 1 = (10 + 4 - 3 - 1) / 2 + 1 = 10/2 + 1 = 6
        let mut node = TestNodeBuilder::new(NodeType::Conv1d, "test")
            .input_tensor_f32("data", 3, Some(vec![1, 2, 10]))
            .input_tensor_f32_data("weight", vec![0.1; 16], vec![2, 2, 4])
            .output_tensor_f32("output", 3, None)
            .attr_ints("kernel_shape", vec![4])
            .attr_ints("strides", vec![2])
            .attr_ints("pads", vec![2, 2])
            .attr_ints("dilations", vec![1])
            .attr_int("group", 1)
            .build_with_graph_data(16);

        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        match &node.outputs[0].ty {
            ArgType::Tensor(t) => {
                assert_eq!(t.rank, 3);
                assert_eq!(t.static_shape, Some(vec![Some(1), Some(2), Some(6)]));
            }
            _ => panic!("Expected tensor output"),
        }
    }

    #[test]
    fn test_conv1d_static_shape_no_input_shape() {
        // No input static_shape -> batch and spatial are None, out_channels is known
        let node = create_test_node(vec![4], vec![1], vec![0, 0], vec![1], 1, false, None)
            .build_with_graph_data(16);
        let mut node = node;
        let processor = Conv1dProcessor;
        let prefs = OutputPreferences::new();
        processor.infer_types(&mut node, 16, &prefs).unwrap();

        match &node.outputs[0].ty {
            ArgType::Tensor(t) => {
                assert_eq!(t.rank, 3);
                assert_eq!(t.static_shape, Some(vec![None, Some(2), None]));
            }
            _ => panic!("Expected tensor output"),
        }
    }
}