quantrs2_sim/tensor_network/

contraction.rs

//! Contraction strategies for tensor networks
//!
//! This module provides algorithms and interfaces for contracting
//! tensor networks efficiently.

use super::tensor::Tensor;
use quantrs2_core::error::{QuantRS2Error, QuantRS2Result};
use scirs2_core::ndarray::{ArrayD, Dimension, IxDyn};
use scirs2_core::Complex64;
use std::collections::{HashMap, HashSet};

/// Trait for a network of tensors that can be contracted
pub trait ContractableNetwork {
    /// Contract two tensors in the network, returning the ID of the resulting tensor
    fn contract_tensors(&mut self, tensor_id1: usize, tensor_id2: usize) -> QuantRS2Result<usize>;

    /// Optimize the contraction order of the network
    fn optimize_contraction_order(&mut self) -> QuantRS2Result<()>;
}

/// A contraction path for a tensor network
#[derive(Debug, Clone)]
pub struct ContractionPath {
    /// The sequence of tensor pairs to contract
    steps: Vec<(usize, usize)>,

    /// Estimated computational cost of this contraction path
    estimated_cost: f64,
}

impl ContractionPath {
    /// Create a new contraction path
    pub const fn new(steps: Vec<(usize, usize)>, estimated_cost: f64) -> Self {
        Self {
            steps,
            estimated_cost,
        }
    }

    /// Get the steps in this contraction path
    pub fn steps(&self) -> &[(usize, usize)] {
        &self.steps
    }

    /// Get the estimated cost of this contraction path
    pub const fn estimated_cost(&self) -> f64 {
        self.estimated_cost
    }
}

/// Calculate the optimal contraction path for a tensor network
///
/// This function implements a greedy algorithm to determine a good
/// contraction order for a tensor network. It's not guaranteed to find
/// the optimal path, but it should produce reasonable results for
/// most practical cases.
pub fn calculate_greedy_contraction_path(
    tensors: &HashMap<usize, Tensor>,
    connections: &[(super::tensor::TensorIndex, super::tensor::TensorIndex)],
) -> QuantRS2Result<ContractionPath> {
    // Step 1: Build a graph of tensor connections
    let mut tensor_connections = HashMap::new();
    for (t1, t2) in connections {
        tensor_connections
            .entry(t1.tensor_id)
            .or_insert_with(HashSet::new)
            .insert(t2.tensor_id);
        tensor_connections
            .entry(t2.tensor_id)
            .or_insert_with(HashSet::new)
            .insert(t1.tensor_id);
    }

    // Step 2: Calculate dimensions of each tensor
    let mut tensor_dims = HashMap::new();
    for (&id, tensor) in tensors {
        tensor_dims.insert(id, tensor.dimensions.iter().product::<usize>());
    }

    // Step 3: Greedy algorithm: repeatedly find the pair of tensors that,
    // when contracted, minimizes the size of the resulting tensor
    let mut remaining_tensors: HashSet<usize> = tensors.keys().copied().collect();
    let mut steps = Vec::new();
    let mut total_cost = 0.0;

    while remaining_tensors.len() > 1 {
        let mut best_cost = f64::INFINITY;
        let mut best_pair = None;

        // Find the best pair to contract next
        for &t1 in &remaining_tensors {
            if let Some(connected) = tensor_connections.get(&t1) {
                for &t2 in connected {
                    if remaining_tensors.contains(&t2) {
                        // Calculate cost of contracting t1 and t2
                        let combined_dim = tensor_dims[&t1] * tensor_dims[&t2];
                        let cost = combined_dim as f64;

                        if cost < best_cost {
                            best_cost = cost;
                            best_pair = Some((t1, t2));
                        }
                    }
                }
            }
        }

        // If we found a pair to contract
        if let Some((t1, t2)) = best_pair {
            // Add to our contraction steps
            steps.push((t1, t2));
            total_cost += best_cost;

            // Remove contracted tensors
            remaining_tensors.remove(&t1);
            remaining_tensors.remove(&t2);

            // Add new contracted tensor
            let new_id = t1; // Reuse first tensor's ID
            remaining_tensors.insert(new_id);

            // Update connections for the new tensor
            let mut new_connections = HashSet::new();

            // Merge connections from t1
            // First collect all connected tensors from t1
            let mut t1_connected_tensors = Vec::new();
            if let Some(t1_connections) = tensor_connections.get(&t1) {
                for &connected_tensor in t1_connections {
                    if connected_tensor != t2 && remaining_tensors.contains(&connected_tensor) {
                        t1_connected_tensors.push(connected_tensor);
                        new_connections.insert(connected_tensor);
                    }
                }
            }

            // Now update their connections
            for connected_tensor in t1_connected_tensors {
                if let Some(other_connections) = tensor_connections.get_mut(&connected_tensor) {
                    other_connections.remove(&t1);
                    other_connections.remove(&t2);
                    other_connections.insert(new_id);
                }
            }

            // Merge connections from t2
            // First collect all connected tensors from t2
            let mut t2_connected_tensors = Vec::new();
            if let Some(t2_connections) = tensor_connections.get(&t2) {
                for &connected_tensor in t2_connections {
                    if connected_tensor != t1 && remaining_tensors.contains(&connected_tensor) {
                        t2_connected_tensors.push(connected_tensor);
                        new_connections.insert(connected_tensor);
                    }
                }
            }

            // Now update their connections
            for connected_tensor in t2_connected_tensors {
                if let Some(other_connections) = tensor_connections.get_mut(&connected_tensor) {
                    other_connections.remove(&t1);
                    other_connections.remove(&t2);
                    other_connections.insert(new_id);
                }
            }

            // Set the new tensor's connections
            tensor_connections.insert(new_id, new_connections);

            // Update the dimension of the new tensor (simplified)
            // In a real implementation, we'd calculate this based on the actual tensors
            tensor_dims.insert(new_id, (tensor_dims[&t1] * tensor_dims[&t2]) / 2);
        } else {
            // No connected tensors found, just contract the first two remaining
            let mut remaining_vec: Vec<_> = remaining_tensors.iter().copied().collect();
            remaining_vec.sort_unstable();

            if remaining_vec.len() >= 2 {
                let t1 = remaining_vec[0];
                let t2 = remaining_vec[1];

                steps.push((t1, t2));
                total_cost += (tensor_dims[&t1] * tensor_dims[&t2]) as f64;

                remaining_tensors.remove(&t1);
                remaining_tensors.remove(&t2);
                remaining_tensors.insert(t1);

                // Update dimensions
                tensor_dims.insert(t1, (tensor_dims[&t1] * tensor_dims[&t2]) / 2);
            } else {
                // Only one tensor left, we're done
                break;
            }
        }
    }

    Ok(ContractionPath::new(steps, total_cost))
}

/// Calculate the optimal contraction path using a more advanced algorithm
///
/// This function implements a more sophisticated algorithm that takes into account
/// the structure of the tensor network to find a better contraction path.
pub fn calculate_optimal_contraction_path(
    tensors: &HashMap<usize, Tensor>,
    connections: &[(super::tensor::TensorIndex, super::tensor::TensorIndex)],
) -> QuantRS2Result<ContractionPath> {
    // First, check if we can identify a specific circuit structure that has
    // a known optimal contraction pattern
    if let Some(path) = identify_circuit_structure(tensors, connections) {
        return Ok(path);
    }

    // If no special structure is identified, fall back to the greedy algorithm
    calculate_greedy_contraction_path(tensors, connections)
}

/// Identify common quantum circuit structures and return their optimal contraction paths
///
/// This function analyzes the tensor network to identify if it corresponds to a
/// common quantum circuit structure (like a linear circuit, GHZ state preparation,
/// or a quantum Fourier transform). If identified, returns a pre-computed optimal
/// contraction path.
fn identify_circuit_structure(
    tensors: &HashMap<usize, Tensor>,
    connections: &[(super::tensor::TensorIndex, super::tensor::TensorIndex)],
) -> Option<ContractionPath> {
    // Build a graph of tensor connections for analysis
    let mut tensor_connections = HashMap::new();
    for (t1, t2) in connections {
        tensor_connections
            .entry(t1.tensor_id)
            .or_insert_with(HashSet::new)
            .insert(t2.tensor_id);
        tensor_connections
            .entry(t2.tensor_id)
            .or_insert_with(HashSet::new)
            .insert(t1.tensor_id);
    }

    // Get a sorted list of tensor IDs
    let mut tensor_ids: Vec<usize> = tensors.keys().copied().collect();
    tensor_ids.sort_unstable();

    // Pattern 1: Linear Circuit (CNOT chain)
    // In a linear circuit, most tensors connect to exactly 2 others,
    // forming a chain-like structure
    if is_linear_circuit(&tensor_connections, &tensor_ids) {
        // For linear circuits, we should contract from one end to the other
        let mut steps = Vec::new();
        let mut cost = 0.0;

        // Order tensors by their position in the chain
        let ordered_tensors = order_linear_circuit(&tensor_connections, &tensor_ids);

        // Contract tensors in sequence
        for ids in ordered_tensors.windows(2) {
            steps.push((ids[0], ids[1]));
            cost += 16.0; // Simplified cost model (2^2 * 2^2)
        }

        return Some(ContractionPath::new(steps, cost));
    }

    // Pattern 2: Star-shaped Circuit (like GHZ state preparation)
    // In a star circuit, one central tensor connects to many others,
    // and those others have few connections
    if is_star_circuit(&tensor_connections, &tensor_ids) {
        // For star circuits, we should contract the leaf nodes with the central node
        let mut steps = Vec::new();
        let mut cost = 0.0;

        // Find the central tensor (the one with most connections)
        let central = find_central_tensor(&tensor_connections);

        // Contract all leaf tensors with the central one
        let leaf_tensors: Vec<_> = tensor_ids
            .iter()
            .filter(|&&id| {
                id != central
                    && tensor_connections
                        .get(&id)
                        .is_some_and(|conns| conns.contains(&central))
            })
            .copied()
            .collect();

        for leaf in leaf_tensors {
            steps.push((central, leaf));
            cost += 16.0; // Simplified cost model
        }

        return Some(ContractionPath::new(steps, cost));
    }

    // Pattern 3: Quantum Fourier Transform (QFT) Circuit
    // QFT has a specific pattern of controlled-phase gates
    if is_qft_circuit(&tensor_connections, tensors) {
        return Some(optimize_qft_circuit(&tensor_connections, tensors));
    }

    // Pattern 4: QAOA Circuit
    // QAOA has alternating layers of problem and mixer Hamiltonians
    if is_qaoa_circuit(&tensor_connections, tensors) {
        return Some(optimize_qaoa_circuit(&tensor_connections, tensors));
    }

    // No special structure identified
    None
}

/// Check if the tensor network represents a Quantum Fourier Transform circuit
fn is_qft_circuit(
    tensor_connections: &HashMap<usize, HashSet<usize>>,
    tensors: &HashMap<usize, Tensor>,
) -> bool {
    // QFT typically has a triangular pattern of controlled-phase gates
    // followed by Hadamard gates and swaps

    // Count gate types and specific patterns that indicate a QFT structure
    let mut hadamard_count = 0;
    let mut controlled_phase_count = 0;
    let mut swap_count = 0;

    // This is a simplified check - a full check would inspect the actual tensor structure
    for tensor in tensors.values() {
        // Check dimensions to guess if it's a single-qubit gate (rank 2) or two-qubit gate (rank 4)
        if tensor.rank == 2 {
            hadamard_count += 1;
        } else if tensor.rank == 4 {
            // Try to classify the two-qubit gate
            if tensor.dimensions == vec![2, 2, 2, 2] {
                // Controlled-phase gates have entries at the (0,0), (1,1), (2,2), (3,3) positions
                // with specific phases - this is a simplified check
                controlled_phase_count += 1;
            }

            // Count potential swap gates
            if is_swap_like_tensor(tensor) {
                swap_count += 1;
            }
        }
    }

    // A QFT circuit typically has Hadamard gates on all qubits and controlled-phase gates
    // The specific pattern is a Hadamard gate on each qubit, followed by controlled-phase gates
    // with decreasing rotation angles, and finally SWAP gates to reverse the qubits

    // This is a simplified heuristic
    hadamard_count > 0 && controlled_phase_count > 0 && hadamard_count >= controlled_phase_count / 2
}

/// Check if a tensor might represent a SWAP-like operation
fn is_swap_like_tensor(tensor: &Tensor) -> bool {
    // SWAP gates have a pattern where the permutation of indices is non-trivial
    // This is a simplified check - a full check would inspect the actual tensor values
    tensor.rank == 4 && tensor.dimensions == vec![2, 2, 2, 2]
}

/// Generate an optimized contraction path for a QFT circuit
fn optimize_qft_circuit(
    tensor_connections: &HashMap<usize, HashSet<usize>>,
    tensors: &HashMap<usize, Tensor>,
) -> ContractionPath {
    // QFT circuits are best contracted starting from the least significant qubit (bottom)
    // and working upward. This follows the natural decomposition of the QFT.

    // Build the tensor IDs in the desired contraction order
    let mut ordered_tensors: Vec<usize> = Vec::new();
    let mut tensor_ids: Vec<usize> = tensors.keys().copied().collect();
    tensor_ids.sort_unstable();

    // Sort tensors by their connectivity pattern
    // In a QFT, we want to contract from bottom to top for optimal efficiency
    let mut steps = Vec::new();
    let mut cost = 0.0;

    // This is a simplified implementation - in a full implementation,
    // we'd analyze the QFT structure more carefully

    // First, try to identify layers of gates in the QFT
    let mut layers = identify_qft_layers(tensor_connections, &tensor_ids);

    // Contract each layer from bottom to top
    for layer in layers {
        // Contract tensors within the layer
        for i in 0..layer.len().saturating_sub(1) {
            steps.push((layer[i], layer[i + 1]));
            cost += 16.0; // Simplified cost model
        }
    }

    // If we couldn't identify layers properly, fall back to a basic contraction strategy
    if steps.is_empty() {
        for i in 0..tensor_ids.len().saturating_sub(1) {
            steps.push((tensor_ids[i], tensor_ids[i + 1]));
            cost += 16.0;
        }
    }

    ContractionPath::new(steps, cost)
}

/// Identify layers of a QFT circuit for optimal contraction
fn identify_qft_layers(
    tensor_connections: &HashMap<usize, HashSet<usize>>,
    tensor_ids: &[usize],
) -> Vec<Vec<usize>> {
    // Group tensors into layers based on their connections
    // In a QFT, we expect a specific pattern of connections between gates

    // This is a simplified implementation - in a real QFT optimizer,
    // we'd analyze the structure more carefully

    // For now, just group tensors by their degree (number of connections)
    let mut degree_groups: HashMap<usize, Vec<usize>> = HashMap::new();

    for &id in tensor_ids {
        let degree = tensor_connections.get(&id).map_or(0, |conns| conns.len());
        degree_groups.entry(degree).or_default().push(id);
    }

    // Order the groups by degree (descending)
    let mut degrees: Vec<usize> = degree_groups.keys().copied().collect();
    degrees.sort_by(|a, b| b.cmp(a));

    // Create layers based on degree groups
    let mut layers = Vec::new();
    for degree in degrees {
        if let Some(group) = degree_groups.get(&degree) {
            layers.push(group.clone());
        }
    }

    layers
}

/// Check if the tensor network represents a QAOA circuit
fn is_qaoa_circuit(
    tensor_connections: &HashMap<usize, HashSet<usize>>,
    tensors: &HashMap<usize, Tensor>,
) -> bool {
    // QAOA has alternating layers of problem Hamiltonian (typically ZZ interactions)
    // and mixer Hamiltonian (typically X rotations)

    // Count gate types associated with QAOA
    let mut x_rotation_count = 0;
    let mut zz_interaction_count = 0;

    // This is a simplified check - a full check would inspect the actual tensor structure
    for tensor in tensors.values() {
        // Single-qubit gate (possibly X rotation)
        if tensor.rank == 2 {
            x_rotation_count += 1; // Assume some are X rotations
        }
        // Two-qubit gate (possibly ZZ interaction)
        else if tensor.rank == 4 {
            zz_interaction_count += 1; // Assume some are ZZ interactions
        }
    }

    // QAOA typically has alternating layers of problem and mixer Hamiltonians,
    // so we expect to see both ZZ interactions and X rotations
    x_rotation_count > 0 && zz_interaction_count > 0
}

/// Generate an optimized contraction path for a QAOA circuit
fn optimize_qaoa_circuit(
    tensor_connections: &HashMap<usize, HashSet<usize>>,
    tensors: &HashMap<usize, Tensor>,
) -> ContractionPath {
    // For QAOA circuits, we want to prioritize contracting the problem Hamiltonian terms
    // (typically ZZ interactions) before the mixer Hamiltonian terms (X rotations)

    // First, sort tensors by rank (higher rank first)
    let mut tensor_ids: Vec<usize> = tensors.keys().copied().collect();
    tensor_ids.sort_by(|a, b| {
        if let (Some(tensor_a), Some(tensor_b)) = (tensors.get(a), tensors.get(b)) {
            tensor_b.rank.cmp(&tensor_a.rank) // Higher rank first
        } else {
            std::cmp::Ordering::Equal
        }
    });

    // Group tensors by rank (for QAOA, rank 4 = two-qubit gates, rank 2 = single-qubit gates)
    let mut rank_groups: HashMap<usize, Vec<usize>> = HashMap::new();

    for &id in &tensor_ids {
        if let Some(tensor) = tensors.get(&id) {
            rank_groups.entry(tensor.rank).or_default().push(id);
        }
    }

    // Create contraction steps prioritizing two-qubit gates (ZZ interactions)
    let mut steps = Vec::new();
    let mut cost = 0.0;

    // First, contract the two-qubit gates (problem Hamiltonian)
    if let Some(two_qubit_gates) = rank_groups.get(&4) {
        for (i, &id1) in two_qubit_gates.iter().enumerate() {
            for &id2 in two_qubit_gates.iter().skip(i + 1) {
                // Check if these tensors are connected
                if tensor_connections
                    .get(&id1)
                    .is_some_and(|conns| conns.contains(&id2))
                {
                    steps.push((id1, id2));
                    cost += 64.0; // Higher cost for two-qubit gate contraction (2^3 * 2^3)
                }
            }
        }
    }

    // Then, contract the single-qubit gates (mixer Hamiltonian)
    if let Some(single_qubit_gates) = rank_groups.get(&2) {
        for (i, &id1) in single_qubit_gates.iter().enumerate() {
            for &id2 in single_qubit_gates.iter().skip(i + 1) {
                // Check if these tensors are connected
                if tensor_connections
                    .get(&id1)
                    .is_some_and(|conns| conns.contains(&id2))
                {
                    steps.push((id1, id2));
                    cost += 16.0; // Lower cost for single-qubit gate contraction (2^2 * 2^2)
                }
            }
        }
    }

    // If no steps were created (no direct connections found),
    // fall back to a simple sequential contraction
    if steps.is_empty() {
        for i in 0..tensor_ids.len().saturating_sub(1) {
            steps.push((tensor_ids[i], tensor_ids[i + 1]));
            cost += 16.0; // Default cost
        }
    }

    ContractionPath::new(steps, cost)
}

/// Check if the tensor network represents a linear circuit
fn is_linear_circuit(
    tensor_connections: &HashMap<usize, HashSet<usize>>,
    tensor_ids: &[usize],
) -> bool {
    // Check that most tensors have exactly 2 connections (except the endpoints)
    let mut num_endpoints = 0;

    for &id in tensor_ids {
        let degree = tensor_connections.get(&id).map_or(0, |conns| conns.len());

        if degree > 2 {
            // If any tensor has more than 2 connections, it's not linear
            return false;
        } else if degree == 1 {
            // Count tensors with only one connection (should be exactly 2 for a chain)
            num_endpoints += 1;
        }
    }

    // A linear circuit should have exactly 2 endpoints
    num_endpoints == 2
}

/// Order tensors in a linear circuit from one end to the other
fn order_linear_circuit(
    tensor_connections: &HashMap<usize, HashSet<usize>>,
    tensor_ids: &[usize],
) -> Vec<usize> {
    let mut result = Vec::new();

    // Find one endpoint
    let mut current = tensor_ids
        .iter()
        .find(|&&id| {
            tensor_connections
                .get(&id)
                .is_some_and(|conns| conns.len() == 1)
        })
        .copied();

    if let Some(start) = current {
        // Start from this endpoint
        result.push(start);
        let mut visited = HashSet::new();
        visited.insert(start);

        // Keep adding the next unvisited neighbor
        while let Some(id) = current {
            if let Some(connections) = tensor_connections.get(&id) {
                let next = connections
                    .iter()
                    .find(|&&next_id| !visited.contains(&next_id))
                    .copied();

                if let Some(next_id) = next {
                    result.push(next_id);
                    visited.insert(next_id);
                    current = Some(next_id);
                } else {
                    // No more unvisited neighbors
                    current = None;
                }
            } else {
                current = None;
            }
        }
    }

    // If we couldn't order it (not actually linear), just return original order
    if result.len() != tensor_ids.len() {
        return tensor_ids.to_vec();
    }

    result
}

/// Check if the tensor network represents a star-shaped circuit
fn is_star_circuit(
    tensor_connections: &HashMap<usize, HashSet<usize>>,
    tensor_ids: &[usize],
) -> bool {
    // Count degrees of each tensor
    let mut degree_counts = HashMap::new();

    for &id in tensor_ids {
        let degree = tensor_connections.get(&id).map_or(0, |conns| conns.len());
        *degree_counts.entry(degree).or_insert(0) += 1;
    }

    // A star circuit has one central node with high degree,
    // and many leaf nodes with degree 1
    let high_degree = degree_counts.keys().filter(|&&d| d > 2).count();
    let degree_one = degree_counts.get(&1).copied().unwrap_or(0);

    // One high-degree node and multiple degree-1 nodes
    high_degree == 1 && degree_one > 2
}

/// Find the central tensor in a star-shaped circuit
fn find_central_tensor(tensor_connections: &HashMap<usize, HashSet<usize>>) -> usize {
    let mut max_degree = 0;
    let mut central = 0;

    for (&id, connections) in tensor_connections {
        let degree = connections.len();
        if degree > max_degree {
            max_degree = degree;
            central = id;
        }
    }

    central
}

/// Extract the shared axis pairs between two tensors from a connection list.
///
/// Each returned `(axis_in_id1, axis_in_id2)` corresponds to a bond linking the
/// two tensors, oriented so the first element indexes `id1` and the second
/// `id2`.
pub(crate) fn shared_axis_pairs(
    connections: &[(super::tensor::TensorIndex, super::tensor::TensorIndex)],
    id1: usize,
    id2: usize,
) -> Vec<(usize, usize)> {
    let mut pairs = Vec::new();
    for (a, b) in connections {
        if a.tensor_id == id1 && b.tensor_id == id2 {
            pairs.push((a.index, b.index));
        } else if a.tensor_id == id2 && b.tensor_id == id1 {
            pairs.push((b.index, a.index));
        }
    }
    pairs
}

/// Contract two tensors over a set of shared index pairs, summing over all of
/// them simultaneously (a multi-index Einstein summation).
///
/// `shared[i] = (axis_in_a, axis_in_b)` lists the axis pairs to contract; the
/// two axes in each pair must have equal dimension. The result's axes are `a`'s
/// non-contracted axes (in order) followed by `b`'s non-contracted axes (in
/// order). With an empty `shared` slice this is the outer (tensor) product.
///
/// This is the genuine contraction primitive used to execute contraction paths;
/// it performs real arithmetic over the tensor data rather than returning an
/// existing tensor.
pub(crate) fn contract_pair_multi(
    a: &Tensor,
    b: &Tensor,
    shared: &[(usize, usize)],
) -> QuantRS2Result<Tensor> {
    for &(ax_a, ax_b) in shared {
        if ax_a >= a.rank || ax_b >= b.rank {
            return Err(QuantRS2Error::CircuitValidationFailed(format!(
                "contract_pair_multi: axis out of range ({ax_a}, {ax_b})"
            )));
        }
        if a.dimensions[ax_a] != b.dimensions[ax_b] {
            return Err(QuantRS2Error::CircuitValidationFailed(format!(
                "contract_pair_multi: dimension mismatch {} vs {}",
                a.dimensions[ax_a], b.dimensions[ax_b]
            )));
        }
    }

    let a_contracted: HashSet<usize> = shared.iter().map(|&(ax, _)| ax).collect();
    let b_contracted: HashSet<usize> = shared.iter().map(|&(_, ax)| ax).collect();

    let a_free: Vec<usize> = (0..a.rank)
        .filter(|ax| !a_contracted.contains(ax))
        .collect();
    let b_free: Vec<usize> = (0..b.rank)
        .filter(|ax| !b_contracted.contains(ax))
        .collect();

    let mut result_dims: Vec<usize> = a_free.iter().map(|&ax| a.dimensions[ax]).collect();
    result_dims.extend(b_free.iter().map(|&ax| b.dimensions[ax]));

    let result_is_scalar = result_dims.is_empty();
    let result_shape = if result_is_scalar {
        IxDyn(&[1usize])
    } else {
        IxDyn(result_dims.as_slice())
    };
    let mut result_data = ArrayD::<Complex64>::zeros(result_shape);

    for (a_idx, a_val) in a.data.indexed_iter() {
        let a_raw = a_idx.slice();
        for (b_idx, b_val) in b.data.indexed_iter() {
            let b_raw = b_idx.slice();

            // Keep only element pairs whose shared indices agree.
            if shared
                .iter()
                .any(|&(ax_a, ax_b)| a_raw[ax_a] != b_raw[ax_b])
            {
                continue;
            }

            let mut res_idx: Vec<usize> = a_free.iter().map(|&ax| a_raw[ax]).collect();
            res_idx.extend(b_free.iter().map(|&ax| b_raw[ax]));

            let target = if result_is_scalar {
                &mut result_data[IxDyn(&[0usize])]
            } else {
                &mut result_data[IxDyn(res_idx.as_slice())]
            };
            *target += *a_val * *b_val;
        }
    }

    let final_data = if result_is_scalar {
        let scalar = result_data[IxDyn(&[0usize])];
        ArrayD::from_elem(IxDyn(&[]), scalar)
    } else {
        result_data
    };

    Ok(Tensor::new(final_data))
}

/// A tensor together with a global label for each of its axes. Two axes that
/// carry the same label are bonded and get summed over when their tensors meet.
struct LabeledTensor {
    tensor: Tensor,
    labels: Vec<usize>,
}

/// Contract a tensor network to a single tensor by executing the given
/// contraction `path`.
///
/// The `connections` list defines the bonds of the network; each bond is turned
/// into a shared axis label. The `path` gives the order of pairwise
/// contractions — for each step `(id1, id2)` the two tensors are contracted over
/// all axes they share (matching labels), and the merged tensor takes the place
/// of `id1` (the same "merge into the first ID" convention produced by
/// [`calculate_greedy_contraction_path`]). Any tensors left uncontracted by the
/// path are folded together via outer products.
///
/// This performs the real contraction arithmetic via [`contract_pair_multi`];
/// it no longer returns an arbitrary existing tensor.
pub fn contract_network_along_path(
    tensors: &mut HashMap<usize, Tensor>,
    connections: &mut Vec<(super::tensor::TensorIndex, super::tensor::TensorIndex)>,
    path: &ContractionPath,
    next_id: &mut usize,
) -> QuantRS2Result<Tensor> {
    let _ = next_id; // IDs follow the path's merge-into-first convention.

    if tensors.is_empty() {
        return Err(QuantRS2Error::CircuitValidationFailed(
            "contract_network_along_path: empty tensor network".to_string(),
        ));
    }

    // Assign a fresh unique label to every axis of every tensor.
    let mut next_label = 0usize;
    let mut working: HashMap<usize, LabeledTensor> = HashMap::new();
    for (&id, tensor) in tensors.iter() {
        let labels: Vec<usize> = (0..tensor.rank)
            .map(|_| {
                let label = next_label;
                next_label += 1;
                label
            })
            .collect();
        working.insert(
            id,
            LabeledTensor {
                tensor: tensor.clone(),
                labels,
            },
        );
    }

    // Bond the endpoints of each connection by giving them a shared label.
    for (a, b) in connections.iter() {
        let bond_label = next_label;
        next_label += 1;
        if let Some(lt) = working.get_mut(&a.tensor_id) {
            if a.index < lt.labels.len() {
                lt.labels[a.index] = bond_label;
            }
        }
        if let Some(lt) = working.get_mut(&b.tensor_id) {
            if b.index < lt.labels.len() {
                lt.labels[b.index] = bond_label;
            }
        }
    }

    // Execute the path, contracting each pair over their shared (matching-label)
    // axes and storing the merged tensor under the first ID.
    for &(id1, id2) in path.steps() {
        if id1 == id2 {
            continue;
        }
        let lt1 = working.remove(&id1).ok_or_else(|| {
            QuantRS2Error::CircuitValidationFailed(format!(
                "contract_network_along_path: tensor {id1} missing from path step"
            ))
        })?;
        let lt2 = working.remove(&id2).ok_or_else(|| {
            QuantRS2Error::CircuitValidationFailed(format!(
                "contract_network_along_path: tensor {id2} missing from path step"
            ))
        })?;

        let merged = contract_labeled(&lt1, &lt2)?;
        working.insert(id1, merged);
    }

    // Fold any remaining tensors (a well-formed path leaves exactly one).
    let mut remaining: Vec<LabeledTensor> = working.into_values().collect();
    let mut acc = remaining.remove(0);
    for lt in remaining {
        acc = contract_labeled(&acc, &lt)?;
    }

    Ok(acc.tensor)
}

/// Contract two labeled tensors over every axis pair whose labels match,
/// returning the merged tensor with its surviving axis labels.
fn contract_labeled(a: &LabeledTensor, b: &LabeledTensor) -> QuantRS2Result<LabeledTensor> {
    let mut shared_pairs = Vec::new();
    for (ai, &la) in a.labels.iter().enumerate() {
        for (bi, &lb) in b.labels.iter().enumerate() {
            if la == lb {
                shared_pairs.push((ai, bi));
            }
        }
    }

    let merged = contract_pair_multi(&a.tensor, &b.tensor, &shared_pairs)?;

    let a_contracted: HashSet<usize> = shared_pairs.iter().map(|&(ai, _)| ai).collect();
    let b_contracted: HashSet<usize> = shared_pairs.iter().map(|&(_, bi)| bi).collect();

    let mut labels: Vec<usize> = a
        .labels
        .iter()
        .enumerate()
        .filter(|&(i, _)| !a_contracted.contains(&i))
        .map(|(_, &l)| l)
        .collect();
    labels.extend(
        b.labels
            .iter()
            .enumerate()
            .filter(|&(i, _)| !b_contracted.contains(&i))
            .map(|(_, &l)| l),
    );

    Ok(LabeledTensor {
        tensor: merged,
        labels,
    })
}