torsh-graph 0.1.3

//! Hypergraph Neural Networks
//!
//! Advanced implementation of hypergraph neural networks for multi-relational learning.
//! Hypergraphs generalize graphs by allowing edges (hyperedges) to connect any number of nodes,
//! enabling modeling of complex multi-way relationships in data.
//!
//! # Features:
//! - Hypergraph data structures with efficient storage
//! - Multiple hypergraph convolution layers (HGCN, HyperGAT, HGNN)
//! - Hypergraph attention mechanisms
//! - Advanced pooling and coarsening operations
//! - Spectral hypergraph methods
//! - Dynamic hypergraph construction

// Framework infrastructure - components designed for future use
#![allow(dead_code)]
use crate::parameter::Parameter;
use crate::{GraphData, GraphLayer};
use torsh_tensor::{
    creation::{from_vec, randn, zeros},
    Tensor,
};

/// Hypergraph data structure representing multi-way relationships
#[derive(Debug, Clone)]
pub struct HypergraphData {
    /// Node feature matrix (num_nodes x num_features)
    pub x: Tensor,
    /// Hyperedge incidence matrix (num_nodes x num_hyperedges)
    pub incidence_matrix: Tensor,
    /// Hyperedge weights (optional)
    pub hyperedge_weights: Option<Tensor>,
    /// Hyperedge features (optional)
    pub hyperedge_features: Option<Tensor>,
    /// Node degrees (sum of incident hyperedge weights)
    pub node_degrees: Tensor,
    /// Hyperedge cardinalities (number of nodes per hyperedge)
    pub hyperedge_cardinalities: Tensor,
    /// Number of nodes
    pub num_nodes: usize,
    /// Number of hyperedges
    pub num_hyperedges: usize,
}

impl HypergraphData {
    /// Create a new hypergraph from node features and incidence matrix
    pub fn new(x: Tensor, incidence_matrix: Tensor) -> Self {
        let num_nodes = x.shape().dims()[0];
        let num_hyperedges = incidence_matrix.shape().dims()[1];

        // Compute node degrees (sum over hyperedges - axis 1)
        let node_degrees = incidence_matrix
            .sum_dim(&[1], false)
            .expect("sum_dim node_degrees should succeed");

        // Compute hyperedge cardinalities (sum over nodes - axis 0)
        let hyperedge_cardinalities = incidence_matrix
            .sum_dim(&[0], false)
            .expect("sum_dim hyperedge_cardinalities should succeed");

        Self {
            x,
            incidence_matrix,
            hyperedge_weights: None,
            hyperedge_features: None,
            node_degrees,
            hyperedge_cardinalities,
            num_nodes,
            num_hyperedges,
        }
    }

    /// Add hyperedge weights
    pub fn with_hyperedge_weights(mut self, weights: Tensor) -> Self {
        self.hyperedge_weights = Some(weights);
        self
    }

    /// Add hyperedge features
    pub fn with_hyperedge_features(mut self, features: Tensor) -> Self {
        self.hyperedge_features = Some(features);
        self
    }

    /// Convert to regular graph using clique expansion
    pub fn to_graph_clique_expansion(&self) -> GraphData {
        let incidence_data = self
            .incidence_matrix
            .to_vec()
            .expect("conversion should succeed");
        let mut edges = Vec::new();

        // For each hyperedge, create clique (all pairs of nodes)
        for e in 0..self.num_hyperedges {
            let mut nodes_in_hyperedge = Vec::new();

            // Find nodes in this hyperedge
            for v in 0..self.num_nodes {
                let idx = v * self.num_hyperedges + e;
                if incidence_data[idx] > 0.0 {
                    nodes_in_hyperedge.push(v as f32);
                }
            }

            // Create all pairs within the hyperedge
            for i in 0..nodes_in_hyperedge.len() {
                for j in (i + 1)..nodes_in_hyperedge.len() {
                    edges.extend_from_slice(&[nodes_in_hyperedge[i], nodes_in_hyperedge[j]]);
                    edges.extend_from_slice(&[nodes_in_hyperedge[j], nodes_in_hyperedge[i]]);
                }
            }
        }

        let edge_index = if edges.is_empty() {
            zeros(&[2, 0]).expect("zeros empty edge_index should succeed")
        } else {
            let num_edges = edges.len() / 2;
            from_vec(edges, &[2, num_edges], torsh_core::device::DeviceType::Cpu)
                .expect("from_vec edge_index should succeed")
        };

        GraphData::new(self.x.clone(), edge_index)
    }

    /// Convert to regular graph using star expansion
    pub fn to_graph_star_expansion(&self) -> GraphData {
        let incidence_data = self
            .incidence_matrix
            .to_vec()
            .expect("conversion should succeed");
        let mut edges = Vec::new();

        // For each hyperedge, create a star with center at virtual node
        let virtual_node_offset = self.num_nodes;

        for e in 0..self.num_hyperedges {
            let virtual_node = (virtual_node_offset + e) as f32;

            // Connect all nodes in hyperedge to virtual center
            for v in 0..self.num_nodes {
                let idx = v * self.num_hyperedges + e;
                if incidence_data[idx] > 0.0 {
                    let node = v as f32;
                    edges.extend_from_slice(&[node, virtual_node]);
                    edges.extend_from_slice(&[virtual_node, node]);
                }
            }
        }

        let edge_index = if edges.is_empty() {
            zeros(&[2, 0]).expect("zeros empty edge_index should succeed")
        } else {
            let num_edges = edges.len() / 2;
            from_vec(edges, &[2, num_edges], torsh_core::device::DeviceType::Cpu)
                .expect("from_vec edge_index should succeed")
        };

        // Extend node features with virtual nodes
        let virtual_features: Tensor = randn(&[self.num_hyperedges, self.x.shape().dims()[1]])
            .expect("randn virtual_features should succeed");
        // Concatenate original and virtual node features
        let node_data = self.x.to_vec().expect("conversion should succeed");
        let virtual_data = virtual_features
            .to_vec()
            .expect("conversion should succeed");
        let mut extended_data = node_data;
        extended_data.extend(virtual_data);

        let total_nodes = self.num_nodes + self.num_hyperedges;
        let features_dim = self.x.shape().dims()[1];
        let extended_x = from_vec(
            extended_data,
            &[total_nodes, features_dim],
            torsh_core::device::DeviceType::Cpu,
        )
        .expect("from_vec extended_x should succeed");

        GraphData::new(extended_x, edge_index)
    }
}

/// Hypergraph Convolutional Network (HGCN) layer
#[derive(Debug)]
pub struct HGCNConv {
    in_features: usize,
    out_features: usize,
    weight: Parameter,
    bias: Option<Parameter>,
    use_attention: bool,
    attention_weight: Option<Parameter>,
    dropout: f32,
}

impl HGCNConv {
    /// Create a new HGCN layer
    pub fn new(
        in_features: usize,
        out_features: usize,
        bias: bool,
        use_attention: bool,
        dropout: f32,
    ) -> Self {
        let weight = Parameter::new(
            randn(&[in_features, out_features]).expect("randn weight should succeed"),
        );
        let bias = if bias {
            Some(Parameter::new(
                zeros(&[out_features]).expect("zeros bias should succeed"),
            ))
        } else {
            None
        };

        let attention_weight = if use_attention {
            Some(Parameter::new(
                randn(&[out_features]).expect("randn attention_weight should succeed"),
            ))
        } else {
            None
        };

        Self {
            in_features,
            out_features,
            weight,
            bias,
            use_attention,
            attention_weight,
            dropout,
        }
    }

    /// Forward pass through HGCN layer
    pub fn forward(&self, hypergraph: &HypergraphData) -> HypergraphData {
        // Simplified implementation for API compatibility
        // Step 1: Transform node features
        let node_features_transformed = hypergraph
            .x
            .matmul(&self.weight.clone_data())
            .expect("operation should succeed");

        // Step 2: Simplified hypergraph convolution (skip complex aggregation for now)
        let output_features = if let Some(ref bias) = self.bias {
            node_features_transformed
                .add(&bias.clone_data())
                .expect("operation should succeed")
        } else {
            node_features_transformed
        };

        // Create output hypergraph with updated node features
        HypergraphData {
            x: output_features,
            incidence_matrix: hypergraph.incidence_matrix.clone(),
            hyperedge_weights: hypergraph.hyperedge_weights.clone(),
            hyperedge_features: hypergraph.hyperedge_features.clone(),
            node_degrees: hypergraph.node_degrees.clone(),
            hyperedge_cardinalities: hypergraph.hyperedge_cardinalities.clone(),
            num_nodes: hypergraph.num_nodes,
            num_hyperedges: hypergraph.num_hyperedges,
        }
    }

    /// Apply attention mechanism to hyperedge features
    fn apply_attention(&self, hyperedge_features: &Tensor, _hypergraph: &HypergraphData) -> Tensor {
        if let Some(ref attention_weight) = self.attention_weight {
            // Compute attention scores
            let attention_scores = hyperedge_features
                .matmul(&attention_weight.clone_data())
                .expect("operation should succeed");
            let attention_probs = attention_scores
                .softmax(-1)
                .expect("softmax should succeed");

            // Apply attention to features
            let attention_expanded = attention_probs
                .unsqueeze(-1)
                .expect("unsqueeze should succeed");
            hyperedge_features
                .mul(&attention_expanded)
                .expect("operation should succeed")
        } else {
            hyperedge_features.clone()
        }
    }

    /// Normalize aggregated features by node degrees
    fn normalize_by_degrees(&self, features: &Tensor, hypergraph: &HypergraphData) -> Tensor {
        let degrees = &hypergraph.node_degrees;
        let epsilon = 1e-8;

        // Add epsilon to prevent division by zero
        let safe_degrees = degrees
            .add_scalar(epsilon)
            .expect("add_scalar should succeed");
        let inv_degrees = safe_degrees
            .reciprocal()
            .expect("reciprocal should succeed");

        // Expand inverse degrees to match feature dimensions
        // First squeeze to ensure we have shape [num_nodes] rather than [num_nodes, 1]
        let inv_degrees_squeezed = if inv_degrees.shape().dims().len() > 1 {
            inv_degrees
                .squeeze_tensor(1)
                .expect("squeeze_tensor should succeed")
        } else {
            inv_degrees
        };
        let inv_degrees_expanded = inv_degrees_squeezed
            .unsqueeze(-1)
            .expect("unsqueeze should succeed");
        features
            .mul(&inv_degrees_expanded)
            .expect("operation should succeed")
    }
}

impl GraphLayer for HGCNConv {
    fn forward(&self, graph: &GraphData) -> GraphData {
        // Convert regular graph to hypergraph and back for compatibility
        let hypergraph = graph_to_hypergraph(graph);
        let output_hypergraph = self.forward(&hypergraph);
        output_hypergraph.to_graph_clique_expansion()
    }

    fn parameters(&self) -> Vec<Tensor> {
        let mut params = vec![self.weight.clone_data()];
        if let Some(ref bias) = self.bias {
            params.push(bias.clone_data());
        }
        if let Some(ref attention_weight) = self.attention_weight {
            params.push(attention_weight.clone_data());
        }
        params
    }
}

/// Hypergraph Attention Network (HyperGAT) layer
#[derive(Debug)]
pub struct HyperGATConv {
    in_features: usize,
    out_features: usize,
    heads: usize,
    query_weight: Parameter,
    key_weight: Parameter,
    value_weight: Parameter,
    hyperedge_attention: Parameter,
    output_weight: Parameter,
    bias: Option<Parameter>,
    dropout: f32,
}

impl HyperGATConv {
    /// Create a new HyperGAT layer
    pub fn new(
        in_features: usize,
        out_features: usize,
        heads: usize,
        dropout: f32,
        bias: bool,
    ) -> Self {
        let head_dim = out_features / heads;

        let query_weight = Parameter::new(
            randn(&[in_features, out_features]).expect("randn query_weight should succeed"),
        );
        let key_weight = Parameter::new(
            randn(&[in_features, out_features]).expect("randn key_weight should succeed"),
        );
        let value_weight = Parameter::new(
            randn(&[in_features, out_features]).expect("randn value_weight should succeed"),
        );
        let hyperedge_attention = Parameter::new(
            randn(&[heads, 2 * head_dim]).expect("randn hyperedge_attention should succeed"),
        );
        let output_weight = Parameter::new(
            randn(&[out_features, out_features]).expect("randn output_weight should succeed"),
        );

        let bias = if bias {
            Some(Parameter::new(
                zeros(&[out_features]).expect("zeros bias should succeed"),
            ))
        } else {
            None
        };

        Self {
            in_features,
            out_features,
            heads,
            query_weight,
            key_weight,
            value_weight,
            hyperedge_attention,
            output_weight,
            bias,
            dropout,
        }
    }

    /// Forward pass through HyperGAT layer
    pub fn forward(&self, hypergraph: &HypergraphData) -> HypergraphData {
        let num_nodes = hypergraph.num_nodes;
        let head_dim = self.out_features / self.heads;

        // Linear transformations
        let queries = hypergraph
            .x
            .matmul(&self.query_weight.clone_data())
            .expect("operation should succeed");
        let keys = hypergraph
            .x
            .matmul(&self.key_weight.clone_data())
            .expect("operation should succeed");
        let values = hypergraph
            .x
            .matmul(&self.value_weight.clone_data())
            .expect("operation should succeed");

        // Reshape for multi-head attention
        let q = queries
            .view(&[num_nodes as i32, self.heads as i32, head_dim as i32])
            .expect("view should succeed");
        let k = keys
            .view(&[num_nodes as i32, self.heads as i32, head_dim as i32])
            .expect("view should succeed");
        let v = values
            .view(&[num_nodes as i32, self.heads as i32, head_dim as i32])
            .expect("view should succeed");

        // Perform hyperedge-based attention
        let attended_features = self.hyperedge_attention_mechanism(&q, &k, &v, hypergraph);

        // Reshape back and apply output transformation
        let concatenated = attended_features
            .view(&[num_nodes as i32, self.out_features as i32])
            .expect("view should succeed");
        let mut output = concatenated
            .matmul(&self.output_weight.clone_data())
            .expect("operation should succeed");

        // Add bias if present
        if let Some(ref bias) = self.bias {
            output = output
                .add(&bias.clone_data())
                .expect("operation should succeed");
        }

        // Create output hypergraph
        HypergraphData {
            x: output,
            incidence_matrix: hypergraph.incidence_matrix.clone(),
            hyperedge_weights: hypergraph.hyperedge_weights.clone(),
            hyperedge_features: hypergraph.hyperedge_features.clone(),
            node_degrees: hypergraph.node_degrees.clone(),
            hyperedge_cardinalities: hypergraph.hyperedge_cardinalities.clone(),
            num_nodes: hypergraph.num_nodes,
            num_hyperedges: hypergraph.num_hyperedges,
        }
    }

    /// Hyperedge-based attention mechanism
    fn hyperedge_attention_mechanism(
        &self,
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        hypergraph: &HypergraphData,
    ) -> Tensor {
        let num_nodes = hypergraph.num_nodes;
        let head_dim = self.out_features / self.heads;

        // Initialize output
        let mut output =
            zeros(&[num_nodes, self.heads, head_dim]).expect("zeros output should succeed");

        let incidence_data = hypergraph
            .incidence_matrix
            .to_vec()
            .expect("conversion should succeed");

        // Process each hyperedge separately
        for e in 0..hypergraph.num_hyperedges {
            let mut nodes_in_hyperedge = Vec::new();

            // Find nodes in this hyperedge
            for v in 0..num_nodes {
                let idx = v * hypergraph.num_hyperedges + e;
                if incidence_data[idx] > 0.0 {
                    nodes_in_hyperedge.push(v);
                }
            }

            if nodes_in_hyperedge.len() < 2 {
                continue; // Skip hyperedges with less than 2 nodes
            }

            // Compute attention within hyperedge for each head
            for head in 0..self.heads {
                self.compute_hyperedge_attention(head, &nodes_in_hyperedge, q, k, v, &mut output);
            }
        }

        output
    }

    /// Compute attention for a specific hyperedge and head
    fn compute_hyperedge_attention(
        &self,
        head: usize,
        nodes: &[usize],
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        output: &mut Tensor,
    ) {
        let head_dim = self.out_features / self.heads;
        let scale = 1.0 / (head_dim as f32).sqrt();

        // For simplicity, use mean pooling within hyperedge
        // In practice, this would use more sophisticated attention
        for &node_i in nodes {
            let mut aggregated = zeros(&[head_dim]).expect("zeros aggregated should succeed");
            let mut total_weight = 0.0;

            for &node_j in nodes {
                if node_i != node_j {
                    // Get query and key for these nodes and head
                    let q_i = q
                        .slice_tensor(0, node_i, node_i + 1)
                        .expect("slice_tensor q_i should succeed")
                        .slice_tensor(1, head, head + 1)
                        .expect("slice_tensor q_i head should succeed")
                        .squeeze_tensor(0)
                        .expect("squeeze_tensor should succeed")
                        .squeeze_tensor(0)
                        .expect("squeeze_tensor should succeed");

                    let k_j = k
                        .slice_tensor(0, node_j, node_j + 1)
                        .expect("slice_tensor k_j should succeed")
                        .slice_tensor(1, head, head + 1)
                        .expect("slice_tensor k_j head should succeed")
                        .squeeze_tensor(0)
                        .expect("squeeze_tensor should succeed")
                        .squeeze_tensor(0)
                        .expect("squeeze_tensor should succeed");

                    let v_j = v
                        .slice_tensor(0, node_j, node_j + 1)
                        .expect("slice_tensor v_j should succeed")
                        .slice_tensor(1, head, head + 1)
                        .expect("slice_tensor v_j head should succeed")
                        .squeeze_tensor(0)
                        .expect("squeeze_tensor should succeed")
                        .squeeze_tensor(0)
                        .expect("squeeze_tensor should succeed");

                    // Compute attention weight (simplified)
                    let attention_score = q_i
                        .dot(&k_j)
                        .expect("dot should succeed")
                        .mul_scalar(scale)
                        .expect("mul_scalar should succeed");
                    let weight = attention_score
                        .exp()
                        .expect("exp should succeed")
                        .item()
                        .expect("tensor should have single item");

                    // Aggregate values
                    let weighted_value = v_j.mul_scalar(weight).expect("mul_scalar should succeed");
                    aggregated = aggregated
                        .add(&weighted_value)
                        .expect("operation should succeed");
                    total_weight += weight;
                }
            }

            // Normalize and update output
            if total_weight > 0.0 {
                aggregated = aggregated
                    .div_scalar(total_weight)
                    .expect("div_scalar should succeed");

                // Update output tensor (simplified assignment)
                let aggregated_data = aggregated.to_vec().expect("conversion should succeed");
                for (j, &val) in aggregated_data.iter().enumerate() {
                    output
                        .set_item(&[node_i, head, j], val)
                        .expect("set_item should succeed");
                }
            }
        }
    }
}

impl GraphLayer for HyperGATConv {
    fn forward(&self, graph: &GraphData) -> GraphData {
        let hypergraph = graph_to_hypergraph(graph);
        let output_hypergraph = self.forward(&hypergraph);
        output_hypergraph.to_graph_clique_expansion()
    }

    fn parameters(&self) -> Vec<Tensor> {
        let mut params = vec![
            self.query_weight.clone_data(),
            self.key_weight.clone_data(),
            self.value_weight.clone_data(),
            self.hyperedge_attention.clone_data(),
            self.output_weight.clone_data(),
        ];

        if let Some(ref bias) = self.bias {
            params.push(bias.clone_data());
        }

        params
    }
}

/// Hypergraph Neural Network (HGNN) layer based on spectral methods
#[derive(Debug)]
pub struct HGNNConv {
    in_features: usize,
    out_features: usize,
    weight: Parameter,
    bias: Option<Parameter>,
    use_spectral: bool,
}

impl HGNNConv {
    /// Create a new HGNN layer
    pub fn new(in_features: usize, out_features: usize, bias: bool, use_spectral: bool) -> Self {
        let weight = Parameter::new(
            randn(&[in_features, out_features]).expect("randn weight should succeed"),
        );
        let bias = if bias {
            Some(Parameter::new(
                zeros(&[out_features]).expect("zeros bias should succeed"),
            ))
        } else {
            None
        };

        Self {
            in_features,
            out_features,
            weight,
            bias,
            use_spectral,
        }
    }

    /// Forward pass through HGNN layer
    pub fn forward(&self, hypergraph: &HypergraphData) -> HypergraphData {
        // Transform node features
        let x_transformed = hypergraph
            .x
            .matmul(&self.weight.clone_data())
            .expect("operation should succeed");

        // Compute hypergraph Laplacian and apply convolution
        let output_features = if self.use_spectral {
            self.spectral_convolution(&x_transformed, hypergraph)
        } else {
            self.spatial_convolution(&x_transformed, hypergraph)
        };

        // Add bias if present
        let final_features = if let Some(ref bias) = self.bias {
            output_features
                .add(&bias.clone_data())
                .expect("operation should succeed")
        } else {
            output_features
        };

        HypergraphData {
            x: final_features,
            incidence_matrix: hypergraph.incidence_matrix.clone(),
            hyperedge_weights: hypergraph.hyperedge_weights.clone(),
            hyperedge_features: hypergraph.hyperedge_features.clone(),
            node_degrees: hypergraph.node_degrees.clone(),
            hyperedge_cardinalities: hypergraph.hyperedge_cardinalities.clone(),
            num_nodes: hypergraph.num_nodes,
            num_hyperedges: hypergraph.num_hyperedges,
        }
    }

    /// Spectral convolution using hypergraph Laplacian
    fn spectral_convolution(&self, features: &Tensor, hypergraph: &HypergraphData) -> Tensor {
        // Compute normalized hypergraph Laplacian
        let laplacian = self.compute_hypergraph_laplacian(hypergraph);

        // Apply Laplacian: L @ X
        laplacian
            .matmul(features)
            .expect("operation should succeed")
    }

    /// Spatial convolution using incidence matrix
    fn spatial_convolution(&self, features: &Tensor, hypergraph: &HypergraphData) -> Tensor {
        // Node-to-hyperedge aggregation
        let incidence_t = hypergraph
            .incidence_matrix
            .transpose(0, 1)
            .expect("transpose should succeed");
        let hyperedge_features = incidence_t
            .matmul(features)
            .expect("operation should succeed");

        // Hyperedge-to-node aggregation
        let aggregated = hypergraph
            .incidence_matrix
            .matmul(&hyperedge_features)
            .expect("operation should succeed");

        // Normalize by node degrees
        self.normalize_by_degrees(&aggregated, hypergraph)
    }

    /// Compute normalized hypergraph Laplacian
    fn compute_hypergraph_laplacian(&self, hypergraph: &HypergraphData) -> Tensor {
        let h = &hypergraph.incidence_matrix;
        let num_nodes = hypergraph.num_nodes;

        // Compute degree matrices
        let node_degrees = h
            .sum_dim(&[1], false)
            .expect("sum_dim node_degrees should succeed");
        let hyperedge_degrees = h
            .sum_dim(&[0], false)
            .expect("sum_dim hyperedge_degrees should succeed");

        // Create diagonal degree matrices (simplified)
        let mut d_v = zeros(&[num_nodes, num_nodes]).expect("zeros d_v should succeed");
        let mut d_e = zeros(&[hypergraph.num_hyperedges, hypergraph.num_hyperedges])
            .expect("zeros d_e should succeed");

        let node_deg_data = node_degrees.to_vec().expect("conversion should succeed");
        let hyperedge_deg_data = hyperedge_degrees
            .to_vec()
            .expect("conversion should succeed");

        // Fill diagonal matrices
        for i in 0..num_nodes {
            let degree = node_deg_data[i].max(1e-8); // Avoid division by zero
            d_v.set_item(&[i, i], degree.powf(-0.5))
                .expect("set_item d_v should succeed");
        }

        for i in 0..hypergraph.num_hyperedges {
            let degree = hyperedge_deg_data[i].max(1e-8);
            d_e.set_item(&[i, i], degree.recip())
                .expect("set_item d_e should succeed");
        }

        // Compute normalized Laplacian: I - D_v^{-1/2} H D_e H^T D_v^{-1/2}
        let h_t = h.transpose(0, 1).expect("transpose should succeed");
        let intermediate = d_v
            .matmul(h)
            .expect("operation should succeed")
            .matmul(&d_e)
            .expect("operation should succeed")
            .matmul(&h_t)
            .expect("operation should succeed")
            .matmul(&d_v)
            .expect("operation should succeed");

        let identity = eye(num_nodes);
        identity
            .sub(&intermediate)
            .expect("operation should succeed")
    }

    /// Normalize features by node degrees
    fn normalize_by_degrees(&self, features: &Tensor, hypergraph: &HypergraphData) -> Tensor {
        let degrees = &hypergraph.node_degrees;
        let epsilon = 1e-8;

        let safe_degrees = degrees
            .add_scalar(epsilon)
            .expect("add_scalar should succeed");
        let inv_sqrt_degrees = safe_degrees
            .pow_scalar(-0.5)
            .expect("pow_scalar should succeed");

        // First squeeze to ensure we have shape [num_nodes] rather than [num_nodes, 1]
        let inv_degrees_squeezed = if inv_sqrt_degrees.shape().dims().len() > 1 {
            inv_sqrt_degrees
                .squeeze_tensor(1)
                .expect("squeeze_tensor should succeed")
        } else {
            inv_sqrt_degrees
        };
        let inv_degrees_expanded = inv_degrees_squeezed
            .unsqueeze(-1)
            .expect("unsqueeze should succeed");

        features
            .mul(&inv_degrees_expanded)
            .expect("operation should succeed")
    }
}

impl GraphLayer for HGNNConv {
    fn forward(&self, graph: &GraphData) -> GraphData {
        let hypergraph = graph_to_hypergraph(graph);
        let output_hypergraph = self.forward(&hypergraph);
        output_hypergraph.to_graph_clique_expansion()
    }

    fn parameters(&self) -> Vec<Tensor> {
        let mut params = vec![self.weight.clone_data()];
        if let Some(ref bias) = self.bias {
            params.push(bias.clone_data());
        }
        params
    }
}

/// Hypergraph pooling operations
pub mod pooling {
    use super::*;

    /// Global hypergraph pooling
    pub fn global_hypergraph_pool(hypergraph: &HypergraphData, method: PoolingMethod) -> Tensor {
        match method {
            PoolingMethod::Mean => hypergraph
                .x
                .mean(Some(&[0]), false)
                .expect("mean pooling should succeed"),
            PoolingMethod::Max => hypergraph
                .x
                .max(Some(0), false)
                .expect("max pooling should succeed"),
            PoolingMethod::Sum => hypergraph
                .x
                .sum_dim(&[0], false)
                .expect("sum pooling should succeed"),
            PoolingMethod::Attention => attention_pool(hypergraph),
        }
    }

    /// Hyperedge-aware pooling
    pub fn hyperedge_pool(hypergraph: &HypergraphData, method: PoolingMethod) -> Tensor {
        let incidence_t = hypergraph
            .incidence_matrix
            .transpose(0, 1)
            .expect("transpose should succeed");

        match method {
            PoolingMethod::Mean => {
                // Average pooling over hyperedges
                let hyperedge_features = incidence_t
                    .matmul(&hypergraph.x)
                    .expect("operation should succeed");
                hyperedge_features
                    .mean(Some(&[0]), false)
                    .expect("mean pooling should succeed")
            }
            PoolingMethod::Max => {
                let hyperedge_features = incidence_t
                    .matmul(&hypergraph.x)
                    .expect("operation should succeed");
                hyperedge_features
                    .max(Some(0), false)
                    .expect("max pooling should succeed")
            }
            PoolingMethod::Sum => {
                let hyperedge_features = incidence_t
                    .matmul(&hypergraph.x)
                    .expect("operation should succeed");
                hyperedge_features
                    .sum_dim(&[0], false)
                    .expect("sum pooling should succeed")
            }
            PoolingMethod::Attention => {
                // Attention over hyperedges
                attention_pool(hypergraph)
            }
        }
    }

    /// Hierarchical hypergraph pooling
    pub fn hierarchical_hypergraph_pool(
        hypergraph: &HypergraphData,
        num_clusters: usize,
    ) -> HypergraphData {
        // Simplified clustering-based pooling
        let cluster_assignments = cluster_nodes(hypergraph, num_clusters);
        coarsen_hypergraph(hypergraph, &cluster_assignments)
    }

    /// Attention-based pooling
    fn attention_pool(hypergraph: &HypergraphData) -> Tensor {
        // Simplified attention pooling
        let attention_scores = hypergraph
            .x
            .sum_dim(&[1], false)
            .expect("sum_dim should succeed");
        let attention_weights = attention_scores.softmax(0).expect("softmax should succeed");
        let attention_expanded = attention_weights
            .unsqueeze(-1)
            .expect("unsqueeze should succeed");

        let weighted_features = hypergraph
            .x
            .mul(&attention_expanded)
            .expect("operation should succeed");
        weighted_features
            .sum_dim(&[0], false)
            .expect("sum_dim should succeed")
    }

    /// Simple node clustering for hierarchical pooling
    fn cluster_nodes(hypergraph: &HypergraphData, num_clusters: usize) -> Vec<usize> {
        let num_nodes = hypergraph.num_nodes;
        let mut assignments = vec![0; num_nodes];

        // Simple clustering by node index (for demonstration)
        for i in 0..num_nodes {
            assignments[i] = i % num_clusters;
        }

        assignments
    }

    /// Coarsen hypergraph based on cluster assignments
    fn coarsen_hypergraph(
        hypergraph: &HypergraphData,
        cluster_assignments: &[usize],
    ) -> HypergraphData {
        let num_clusters = cluster_assignments
            .iter()
            .max()
            .expect("reduction should succeed")
            + 1;
        let original_features = hypergraph.x.shape().dims()[1];

        // Average node features within clusters (simplified implementation)
        let mut coarse_features_data = vec![0.0; num_clusters * original_features];
        let mut cluster_counts = vec![0; num_clusters];

        let node_data = hypergraph.x.to_vec().expect("conversion should succeed");

        for (node, &cluster) in cluster_assignments.iter().enumerate() {
            cluster_counts[cluster] += 1;
            for feat in 0..original_features {
                let node_feat_idx = node * original_features + feat;
                let cluster_feat_idx = cluster * original_features + feat;
                coarse_features_data[cluster_feat_idx] += node_data[node_feat_idx];
            }
        }

        // Normalize by cluster size
        for cluster in 0..num_clusters {
            if cluster_counts[cluster] > 0 {
                for feat in 0..original_features {
                    let cluster_feat_idx = cluster * original_features + feat;
                    coarse_features_data[cluster_feat_idx] /= cluster_counts[cluster] as f32;
                }
            }
        }

        let coarse_features = from_vec(
            coarse_features_data,
            &[num_clusters, original_features],
            torsh_core::device::DeviceType::Cpu,
        )
        .expect("from_vec coarse_features should succeed");

        // Create coarse incidence matrix (simplified)
        let coarse_incidence = zeros(&[num_clusters, hypergraph.num_hyperedges])
            .expect("zeros coarse_incidence should succeed");

        HypergraphData::new(coarse_features, coarse_incidence)
    }

    /// Pooling methods
    #[derive(Debug, Clone, Copy)]
    pub enum PoolingMethod {
        Mean,
        Max,
        Sum,
        Attention,
    }
}

/// Utility functions for hypergraph operations
pub mod utils {
    use super::*;

    /// Convert edge list to hypergraph
    pub fn edge_list_to_hypergraph(
        edges: &[(Vec<usize>, f32)],
        num_nodes: usize,
    ) -> HypergraphData {
        let num_hyperedges = edges.len();
        let mut incidence_data = vec![0.0; num_nodes * num_hyperedges];
        let mut weights = Vec::new();

        for (e, (edge_nodes, weight)) in edges.iter().enumerate() {
            weights.push(*weight);
            for &node in edge_nodes {
                if node < num_nodes {
                    incidence_data[node * num_hyperedges + e] = 1.0;
                }
            }
        }

        let features = randn(&[num_nodes, 16]).expect("randn features should succeed"); // Default features
        let incidence_matrix = from_vec(
            incidence_data,
            &[num_nodes, num_hyperedges],
            torsh_core::device::DeviceType::Cpu,
        )
        .expect("from_vec incidence_matrix should succeed");

        let hyperedge_weights = from_vec(
            weights,
            &[num_hyperedges],
            torsh_core::device::DeviceType::Cpu,
        )
        .expect("from_vec hyperedge_weights should succeed");

        HypergraphData::new(features, incidence_matrix).with_hyperedge_weights(hyperedge_weights)
    }

    /// Generate random hypergraph
    pub fn random_hypergraph(
        num_nodes: usize,
        num_hyperedges: usize,
        edge_prob: f32,
        features_dim: usize,
    ) -> HypergraphData {
        let mut rng = scirs2_core::random::thread_rng();
        let mut incidence_data = vec![0.0; num_nodes * num_hyperedges];

        // Generate random hyperedges
        for e in 0..num_hyperedges {
            for v in 0..num_nodes {
                if rng.gen_range(0.0..1.0) < edge_prob {
                    incidence_data[v * num_hyperedges + e] = 1.0;
                }
            }
        }

        let features = randn(&[num_nodes, features_dim]).expect("randn features should succeed");
        let incidence_matrix = from_vec(
            incidence_data,
            &[num_nodes, num_hyperedges],
            torsh_core::device::DeviceType::Cpu,
        )
        .expect("from_vec incidence_matrix should succeed");

        HypergraphData::new(features, incidence_matrix)
    }

    /// Hypergraph metrics
    pub fn hypergraph_metrics(hypergraph: &HypergraphData) -> HypergraphMetrics {
        let node_degrees = hypergraph
            .node_degrees
            .to_vec()
            .expect("conversion should succeed");
        let hyperedge_cardinalities = hypergraph
            .hyperedge_cardinalities
            .to_vec()
            .expect("conversion should succeed");

        let avg_node_degree = node_degrees.iter().sum::<f32>() / node_degrees.len() as f32;
        let avg_hyperedge_size =
            hyperedge_cardinalities.iter().sum::<f32>() / hyperedge_cardinalities.len() as f32;

        let density = node_degrees.iter().sum::<f32>()
            / (hypergraph.num_nodes * hypergraph.num_hyperedges) as f32;

        HypergraphMetrics {
            avg_node_degree,
            avg_hyperedge_size,
            density,
            num_nodes: hypergraph.num_nodes,
            num_hyperedges: hypergraph.num_hyperedges,
        }
    }

    /// Hypergraph statistics
    #[derive(Debug, Clone)]
    pub struct HypergraphMetrics {
        pub avg_node_degree: f32,
        pub avg_hyperedge_size: f32,
        pub density: f32,
        pub num_nodes: usize,
        pub num_hyperedges: usize,
    }
}

/// Convert regular graph to hypergraph (each edge becomes a hyperedge)
pub fn graph_to_hypergraph(graph: &GraphData) -> HypergraphData {
    let edge_data = crate::utils::tensor_to_vec2::<f32>(&graph.edge_index)
        .expect("tensor_to_vec2 should succeed");
    let num_edges = edge_data[0].len();
    let num_nodes = graph.num_nodes;

    // Each edge becomes a hyperedge connecting two nodes
    let mut incidence_data = vec![0.0; num_nodes * num_edges];

    for e in 0..num_edges {
        let src = edge_data[0][e] as usize;
        let dst = edge_data[1][e] as usize;

        if src < num_nodes && dst < num_nodes {
            incidence_data[src * num_edges + e] = 1.0;
            incidence_data[dst * num_edges + e] = 1.0;
        }
    }

    let incidence_matrix = from_vec(
        incidence_data,
        &[num_nodes, num_edges],
        torsh_core::device::DeviceType::Cpu,
    )
    .expect("from_vec incidence_matrix should succeed");

    HypergraphData::new(graph.x.clone(), incidence_matrix)
}

/// Create identity matrix
fn eye(n: usize) -> Tensor {
    let mut data = vec![0.0; n * n];
    for i in 0..n {
        data[i * n + i] = 1.0;
    }
    from_vec(data, &[n, n], torsh_core::device::DeviceType::Cpu)
        .expect("from_vec eye should succeed")
}

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

    #[test]
    fn test_hypergraph_creation() {
        let features = randn(&[4, 3]).unwrap();
        let incidence_data = vec![
            1.0, 0.0, 1.0, // Node 0 in hyperedges 0 and 2
            1.0, 1.0, 0.0, // Node 1 in hyperedges 0 and 1
            0.0, 1.0, 1.0, // Node 2 in hyperedges 1 and 2
            0.0, 0.0, 1.0, // Node 3 in hyperedge 2
        ];
        let incidence_matrix = from_vec(incidence_data, &[4, 3], DeviceType::Cpu).unwrap();

        let hypergraph = HypergraphData::new(features, incidence_matrix);

        assert_eq!(hypergraph.num_nodes, 4);
        assert_eq!(hypergraph.num_hyperedges, 3);
        assert_eq!(hypergraph.x.shape().dims(), &[4, 3]);
        assert_eq!(hypergraph.incidence_matrix.shape().dims(), &[4, 3]);
    }

    #[test]
    fn test_hgcn_layer() {
        let features = randn(&[3, 4]).unwrap();
        let incidence_matrix =
            from_vec(vec![1.0, 0.0, 1.0, 1.0, 0.0, 1.0], &[3, 2], DeviceType::Cpu).unwrap();
        let hypergraph = HypergraphData::new(features, incidence_matrix);

        let hgcn = HGCNConv::new(4, 8, true, false, 0.1);
        let output = hgcn.forward(&hypergraph);

        assert_eq!(output.x.shape().dims(), &[3, 8]);
        assert_eq!(output.num_nodes, 3);
        assert_eq!(output.num_hyperedges, 2);
    }

    #[test]
    fn test_hypergraph_to_graph_conversion() {
        let features = randn(&[3, 4]).unwrap();
        let incidence_matrix =
            from_vec(vec![1.0, 0.0, 1.0, 1.0, 0.0, 1.0], &[3, 2], DeviceType::Cpu).unwrap();
        let hypergraph = HypergraphData::new(features, incidence_matrix);

        let graph = hypergraph.to_graph_clique_expansion();
        assert_eq!(graph.num_nodes, 3);

        let star_graph = hypergraph.to_graph_star_expansion();
        assert_eq!(star_graph.num_nodes, 5); // 3 original + 2 virtual nodes
    }

    #[test]
    fn test_hypergraph_pooling() {
        let features = randn(&[4, 6]).unwrap();
        let incidence_matrix = from_vec(
            vec![1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0],
            &[4, 2],
            DeviceType::Cpu,
        )
        .unwrap();
        let hypergraph = HypergraphData::new(features, incidence_matrix);

        let pooled_mean =
            pooling::global_hypergraph_pool(&hypergraph, pooling::PoolingMethod::Mean);
        assert_eq!(pooled_mean.shape().dims(), &[6]);

        let pooled_max = pooling::global_hypergraph_pool(&hypergraph, pooling::PoolingMethod::Max);
        assert_eq!(pooled_max.shape().dims(), &[6]);
    }

    #[test]
    fn test_hypergraph_utils() {
        let edges = vec![
            (vec![0, 1, 2], 1.0),
            (vec![1, 3], 0.8),
            (vec![0, 2, 3], 1.2),
        ];

        let hypergraph = utils::edge_list_to_hypergraph(&edges, 4);
        assert_eq!(hypergraph.num_nodes, 4);
        assert_eq!(hypergraph.num_hyperedges, 3);

        let metrics = utils::hypergraph_metrics(&hypergraph);
        assert!(metrics.avg_node_degree > 0.0);
        assert!(metrics.avg_hyperedge_size > 0.0);
    }

    #[test]
    fn test_random_hypergraph_generation() {
        let hypergraph = utils::random_hypergraph(5, 3, 0.6, 8);
        assert_eq!(hypergraph.num_nodes, 5);
        assert_eq!(hypergraph.num_hyperedges, 3);
        assert_eq!(hypergraph.x.shape().dims(), &[5, 8]);
    }

    #[test]
    fn test_hypergat_layer() {
        let features = randn(&[4, 6]).unwrap();
        let incidence_matrix = from_vec(
            vec![1.0, 0.0, 1.0, 1.0, 0.0, 1.0, 0.0, 1.0],
            &[4, 2],
            DeviceType::Cpu,
        )
        .unwrap();
        let hypergraph = HypergraphData::new(features, incidence_matrix);

        let hypergat = HyperGATConv::new(6, 12, 3, 0.1, true);
        let output = hypergat.forward(&hypergraph);

        assert_eq!(output.x.shape().dims(), &[4, 12]);
        assert_eq!(output.num_nodes, 4);
    }
}