torsh-graph 0.2.0

//! 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)]
/// Crate-local result alias: the error type defaults to [`TorshError`],
/// so both `Result<T>` and `Result<T, OtherError>` stay valid.
type Result<T, E = torsh_core::error::TorshError> = std::result::Result<T, E>;

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) -> Result<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)?;

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

        Ok(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) -> Result<GraphData> {
        let incidence_data = self.incidence_matrix.to_vec()?;
        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])?
        } else {
            let num_edges = edges.len() / 2;
            from_vec(edges, &[2, num_edges], torsh_core::device::DeviceType::Cpu)?
        };

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

    /// Convert to regular graph using star expansion
    pub fn to_graph_star_expansion(&self) -> Result<GraphData> {
        let incidence_data = self.incidence_matrix.to_vec()?;
        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])?
        } else {
            let num_edges = edges.len() / 2;
            from_vec(edges, &[2, num_edges], torsh_core::device::DeviceType::Cpu)?
        };

        // Extend node features with virtual nodes
        let virtual_features: Tensor = randn(&[self.num_hyperedges, self.x.shape().dims()[1]])?;
        // Concatenate original and virtual node features
        let node_data = self.x.to_vec()?;
        let virtual_data = virtual_features.to_vec()?;
        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,
        )?;

        Ok(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,
    ) -> Result<Self> {
        let weight = Parameter::new(randn(&[in_features, out_features])?);
        let bias = if bias {
            Some(Parameter::new(zeros(&[out_features])?))
        } else {
            None
        };

        let attention_weight = if use_attention {
            Some(Parameter::new(randn(&[out_features])?))
        } else {
            None
        };

        Ok(Self {
            in_features,
            out_features,
            weight,
            bias,
            use_attention,
            attention_weight,
            dropout,
        })
    }

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

        // 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())?
        } else {
            node_features_transformed
        };

        // Create output hypergraph with updated node features
        Ok(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,
    ) -> Result<Tensor> {
        if let Some(ref attention_weight) = self.attention_weight {
            // Compute attention scores
            let attention_scores = hyperedge_features.matmul(&attention_weight.clone_data())?;
            let attention_probs = attention_scores.softmax(-1)?;

            // Apply attention to features
            let attention_expanded = attention_probs.unsqueeze(-1)?;
            Ok(hyperedge_features.mul(&attention_expanded)?)
        } else {
            Ok(hyperedge_features.clone())
        }
    }

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

        // Add epsilon to prevent division by zero
        let safe_degrees = degrees.add_scalar(epsilon)?;
        let inv_degrees = safe_degrees.reciprocal()?;

        // 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)?
        } else {
            inv_degrees
        };
        let inv_degrees_expanded = inv_degrees_squeezed.unsqueeze(-1)?;
        Ok(features.mul(&inv_degrees_expanded)?)
    }
}

impl GraphLayer for HGCNConv {
    fn forward(&self, graph: &GraphData) -> Result<GraphData> {
        // Convert regular graph to hypergraph and back for compatibility
        let hypergraph = graph_to_hypergraph(graph)?;
        let output_hypergraph = HGCNConv::forward(self, &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,
    ) -> Result<Self> {
        let head_dim = out_features / heads;

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

        let bias = if bias {
            Some(Parameter::new(zeros(&[out_features])?))
        } else {
            None
        };

        Ok(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) -> Result<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())?;
        let keys = hypergraph.x.matmul(&self.key_weight.clone_data())?;
        let values = hypergraph.x.matmul(&self.value_weight.clone_data())?;

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

        // 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])?;
        let mut output = concatenated.matmul(&self.output_weight.clone_data())?;

        // Add bias if present
        if let Some(ref bias) = self.bias {
            output = output.add(&bias.clone_data())?;
        }

        // Create output hypergraph
        Ok(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,
    ) -> Result<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])?;

        let incidence_data = hypergraph.incidence_matrix.to_vec()?;

        // 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)?;
            }
        }

        Ok(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,
    ) -> Result<()> {
        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])?;
            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)?
                        .slice_tensor(1, head, head + 1)?
                        .squeeze_tensor(0)?
                        .squeeze_tensor(0)?;

                    let k_j = k
                        .slice_tensor(0, node_j, node_j + 1)?
                        .slice_tensor(1, head, head + 1)?
                        .squeeze_tensor(0)?
                        .squeeze_tensor(0)?;

                    let v_j = v
                        .slice_tensor(0, node_j, node_j + 1)?
                        .slice_tensor(1, head, head + 1)?
                        .squeeze_tensor(0)?
                        .squeeze_tensor(0)?;

                    // Compute attention weight (simplified)
                    let attention_score = q_i.dot(&k_j)?.mul_scalar(scale)?;
                    let weight = attention_score.exp()?.item()?;

                    // Aggregate values
                    let weighted_value = v_j.mul_scalar(weight)?;
                    aggregated = aggregated.add(&weighted_value)?;
                    total_weight += weight;
                }
            }

            // Normalize and update output
            if total_weight > 0.0 {
                aggregated = aggregated.div_scalar(total_weight)?;

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

        Ok(())
    }
}

impl GraphLayer for HyperGATConv {
    fn forward(&self, graph: &GraphData) -> Result<GraphData> {
        let hypergraph = graph_to_hypergraph(graph)?;
        let output_hypergraph = HyperGATConv::forward(self, &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,
    ) -> Result<Self> {
        let weight = Parameter::new(randn(&[in_features, out_features])?);
        let bias = if bias {
            Some(Parameter::new(zeros(&[out_features])?))
        } else {
            None
        };

        Ok(Self {
            in_features,
            out_features,
            weight,
            bias,
            use_spectral,
        })
    }

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

        // 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 output_features = output_features?;
        let final_features = if let Some(ref bias) = self.bias {
            output_features.add(&bias.clone_data())?
        } else {
            output_features
        };

        Ok(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,
    ) -> Result<Tensor> {
        // Compute normalized hypergraph Laplacian
        let laplacian = self.compute_hypergraph_laplacian(hypergraph);

        // Apply Laplacian: L @ X
        Ok(laplacian?.matmul(features)?)
    }

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

        // Hyperedge-to-node aggregation
        let aggregated = hypergraph.incidence_matrix.matmul(&hyperedge_features)?;

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

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

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

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

        let node_deg_data = node_degrees.to_vec()?;
        let hyperedge_deg_data = hyperedge_degrees.to_vec()?;

        // 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))?;
        }

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

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

        let identity = eye(num_nodes);
        Ok(identity?.sub(&intermediate)?)
    }

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

        let safe_degrees = degrees.add_scalar(epsilon)?;
        let inv_sqrt_degrees = safe_degrees.pow_scalar(-0.5)?;

        // 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)?
        } else {
            inv_sqrt_degrees
        };
        let inv_degrees_expanded = inv_degrees_squeezed.unsqueeze(-1)?;

        Ok(features.mul(&inv_degrees_expanded)?)
    }
}

impl GraphLayer for HGNNConv {
    fn forward(&self, graph: &GraphData) -> Result<GraphData> {
        let hypergraph = graph_to_hypergraph(graph)?;
        let output_hypergraph = HGNNConv::forward(self, &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,
    ) -> Result<Tensor> {
        match method {
            PoolingMethod::Mean => Ok(hypergraph.x.mean(Some(&[0]), false)?),
            PoolingMethod::Max => Ok(hypergraph.x.max(Some(0), false)?),
            PoolingMethod::Sum => Ok(hypergraph.x.sum_dim(&[0], false)?),
            PoolingMethod::Attention => attention_pool(hypergraph),
        }
    }

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

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

    /// Hierarchical hypergraph pooling
    pub fn hierarchical_hypergraph_pool(
        hypergraph: &HypergraphData,
        num_clusters: usize,
    ) -> Result<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) -> Result<Tensor> {
        // Simplified attention pooling
        let attention_scores = hypergraph.x.sum_dim(&[1], false)?;
        let attention_weights = attention_scores.softmax(0)?;
        let attention_expanded = attention_weights.unsqueeze(-1)?;

        let weighted_features = hypergraph.x.mul(&attention_expanded)?;
        Ok(weighted_features.sum_dim(&[0], false)?)
    }

    /// 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],
    ) -> Result<HypergraphData> {
        let num_clusters = cluster_assignments.iter().max().copied().unwrap_or(0) + 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()?;

        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,
        )?;

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

        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,
    ) -> Result<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])?; // Default features
        let incidence_matrix = from_vec(
            incidence_data,
            &[num_nodes, num_hyperedges],
            torsh_core::device::DeviceType::Cpu,
        )?;

        let hyperedge_weights = from_vec(
            weights,
            &[num_hyperedges],
            torsh_core::device::DeviceType::Cpu,
        )?;

        Ok(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,
    ) -> Result<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])?;
        let incidence_matrix = from_vec(
            incidence_data,
            &[num_nodes, num_hyperedges],
            torsh_core::device::DeviceType::Cpu,
        )?;

        HypergraphData::new(features, incidence_matrix)
    }

    /// Hypergraph metrics
    pub fn hypergraph_metrics(hypergraph: &HypergraphData) -> Result<HypergraphMetrics> {
        let node_degrees = hypergraph.node_degrees.to_vec()?;
        let hyperedge_cardinalities = hypergraph.hyperedge_cardinalities.to_vec()?;

        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;

        Ok(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) -> Result<HypergraphData> {
    let edge_data = crate::utils::tensor_to_vec2::<f32>(&graph.edge_index)?;
    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,
    )?;

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

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

#[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).expect("operation should succeed");

        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).expect("operation should succeed");

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

        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).expect("operation should succeed");

        let graph = hypergraph
            .to_graph_clique_expansion()
            .expect("operation should succeed");
        assert_eq!(graph.num_nodes, 3);

        let star_graph = hypergraph
            .to_graph_star_expansion()
            .expect("operation should succeed");
        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).expect("operation should succeed");

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

        let pooled_max = pooling::global_hypergraph_pool(&hypergraph, pooling::PoolingMethod::Max);
        assert_eq!(
            pooled_max.expect("operation should succeed").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).expect("operation should succeed");
        assert_eq!(hypergraph.num_nodes, 4);
        assert_eq!(hypergraph.num_hyperedges, 3);

        let metrics = utils::hypergraph_metrics(&hypergraph).expect("operation should succeed");
        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).expect("operation should succeed");
        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).expect("operation should succeed");

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

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