use crate::ir::{DataType, IRNode, IRNodeType, NodeId, Shape, SparseFormat};

/// Sparse tensor operations for the Unified IR
pub struct SparseOps;

impl SparseOps {
    pub fn create_sparse_matmul(
        id: NodeId,
        inputs: Vec<NodeId>,
        format: SparseFormat,
        output_shape: Shape,
    ) -> IRNode {
        IRNode::new(
            id,
            IRNodeType::SparseMatMul { format },
            inputs,
            vec![DataType::F32],
            vec![output_shape],
        )
    }

    pub fn create_csr_matmul(id: NodeId, inputs: Vec<NodeId>, output_shape: Shape) -> IRNode {
        Self::create_sparse_matmul(id, inputs, SparseFormat::CSR, output_shape)
    }

    pub fn create_coo_matmul(id: NodeId, inputs: Vec<NodeId>, output_shape: Shape) -> IRNode {
        Self::create_sparse_matmul(id, inputs, SparseFormat::COO, output_shape)
    }

    pub fn create_csc_matmul(id: NodeId, inputs: Vec<NodeId>, output_shape: Shape) -> IRNode {
        Self::create_sparse_matmul(id, inputs, SparseFormat::CSC, output_shape)
    }
}