quantrs2-circuit 0.1.3

//! Tensor network compression for quantum circuits
//!
//! This module provides tensor network representations of quantum circuits
//! for efficient simulation and optimization.

use crate::builder::Circuit;
use crate::dag::{circuit_to_dag, CircuitDag, DagNode};
// SciRS2 POLICY compliant - using scirs2_core::Complex64
use quantrs2_core::{
    error::{QuantRS2Error, QuantRS2Result},
    gate::GateOp,
    qubit::QubitId,
};
use scirs2_core::ndarray::{Array2, ArrayView2};
use scirs2_core::Complex64;
use scirs2_linalg::svd;
use std::collections::{HashMap, HashSet};
use std::f64::consts::PI;

/// Complex number type
type C64 = Complex64;

/// Tensor representing a quantum gate or state
#[derive(Debug, Clone)]
pub struct Tensor {
    /// Tensor data in row-major order
    pub data: Vec<C64>,
    /// Shape of the tensor (dimensions)
    pub shape: Vec<usize>,
    /// Labels for each index
    pub indices: Vec<String>,
}

impl Tensor {
    /// Create a new tensor
    #[must_use]
    pub fn new(data: Vec<C64>, shape: Vec<usize>, indices: Vec<String>) -> Self {
        assert_eq!(shape.len(), indices.len());
        let total_size: usize = shape.iter().product();
        assert_eq!(data.len(), total_size);

        Self {
            data,
            shape,
            indices,
        }
    }

    /// Create an identity tensor
    #[must_use]
    pub fn identity(dim: usize, in_label: String, out_label: String) -> Self {
        let mut data = vec![C64::new(0.0, 0.0); dim * dim];
        for i in 0..dim {
            data[i * dim + i] = C64::new(1.0, 0.0);
        }

        Self::new(data, vec![dim, dim], vec![in_label, out_label])
    }

    /// Get the rank (number of indices)
    #[must_use]
    pub fn rank(&self) -> usize {
        self.shape.len()
    }

    /// Get the total number of elements
    #[must_use]
    pub fn size(&self) -> usize {
        self.data.len()
    }

    /// Contract two tensors along specified indices
    pub fn contract(&self, other: &Self, self_idx: &str, other_idx: &str) -> QuantRS2Result<Self> {
        // Find index positions
        let self_pos = self
            .indices
            .iter()
            .position(|s| s == self_idx)
            .ok_or_else(|| QuantRS2Error::InvalidInput(format!("Index {self_idx} not found")))?;
        let other_pos = other
            .indices
            .iter()
            .position(|s| s == other_idx)
            .ok_or_else(|| QuantRS2Error::InvalidInput(format!("Index {other_idx} not found")))?;

        // Check dimensions match
        if self.shape[self_pos] != other.shape[other_pos] {
            return Err(QuantRS2Error::InvalidInput(format!(
                "Dimension mismatch: {} vs {}",
                self.shape[self_pos], other.shape[other_pos]
            )));
        }

        // Compute new shape and indices
        let mut new_shape = Vec::new();
        let mut new_indices = Vec::new();

        for (i, (dim, idx)) in self.shape.iter().zip(&self.indices).enumerate() {
            if i != self_pos {
                new_shape.push(*dim);
                new_indices.push(idx.clone());
            }
        }

        for (i, (dim, idx)) in other.shape.iter().zip(&other.indices).enumerate() {
            if i != other_pos {
                new_shape.push(*dim);
                new_indices.push(idx.clone());
            }
        }

        // Perform contraction (simplified implementation)
        let new_size: usize = new_shape.iter().product();
        let mut new_data = vec![C64::new(0.0, 0.0); new_size];

        // This is a simplified contraction - in practice, would use optimized tensor libraries
        let contract_dim = self.shape[self_pos];

        // For now, return a placeholder
        Ok(Self::new(new_data, new_shape, new_indices))
    }

    /// Reshape the tensor
    pub fn reshape(&mut self, new_shape: Vec<usize>) -> QuantRS2Result<()> {
        let new_size: usize = new_shape.iter().product();
        if new_size != self.size() {
            return Err(QuantRS2Error::InvalidInput(format!(
                "Cannot reshape {} elements to shape {:?}",
                self.size(),
                new_shape
            )));
        }

        self.shape = new_shape;
        Ok(())
    }
}

/// Tensor network representation of a quantum circuit
#[derive(Debug)]
pub struct TensorNetwork {
    /// Tensors in the network
    tensors: Vec<Tensor>,
    /// Connections between tensors (`tensor_idx1`, idx1, `tensor_idx2`, idx2)
    bonds: Vec<(usize, String, usize, String)>,
    /// Open indices (external legs)
    open_indices: HashMap<String, (usize, usize)>, // index -> (tensor_idx, position)
}

impl Default for TensorNetwork {
    fn default() -> Self {
        Self::new()
    }
}

impl TensorNetwork {
    /// Create a new empty tensor network
    #[must_use]
    pub fn new() -> Self {
        Self {
            tensors: Vec::new(),
            bonds: Vec::new(),
            open_indices: HashMap::new(),
        }
    }

    /// Add a tensor to the network
    pub fn add_tensor(&mut self, tensor: Tensor) -> usize {
        let idx = self.tensors.len();

        // Track open indices
        for (pos, index) in tensor.indices.iter().enumerate() {
            self.open_indices.insert(index.clone(), (idx, pos));
        }

        self.tensors.push(tensor);
        idx
    }

    /// Connect two tensor indices
    pub fn add_bond(
        &mut self,
        t1: usize,
        idx1: String,
        t2: usize,
        idx2: String,
    ) -> QuantRS2Result<()> {
        if t1 >= self.tensors.len() || t2 >= self.tensors.len() {
            return Err(QuantRS2Error::InvalidInput(
                "Tensor index out of range".to_string(),
            ));
        }

        // Remove from open indices
        self.open_indices.remove(&idx1);
        self.open_indices.remove(&idx2);

        self.bonds.push((t1, idx1, t2, idx2));
        Ok(())
    }

    /// Contract the entire network to a single tensor
    pub fn contract_all(&self) -> QuantRS2Result<Tensor> {
        if self.tensors.is_empty() {
            return Err(QuantRS2Error::InvalidInput(
                "Empty tensor network".to_string(),
            ));
        }

        // Simple contraction order: left to right
        // In practice, would use optimal contraction ordering
        let mut result = self.tensors[0].clone();

        for bond in &self.bonds {
            let (t1, idx1, t2, idx2) = bond;
            if *t1 == 0 {
                result = result.contract(&self.tensors[*t2], idx1, idx2)?;
            }
        }

        Ok(result)
    }

    /// Apply SVD-based bond compression to the tensor network.
    ///
    /// For each internal bond in the network:
    /// 1. Reshape the pair of connected tensors into a bipartite matrix M (rows = left legs, cols = right legs).
    /// 2. Compute the real-valued Schmidt decomposition via SVD on |M_ij|.
    /// 3. Truncate to `max_bond_dim` singular values (or drop those below `tolerance`).
    /// 4. Reconstruct the tensors: left ← U * diag(s), right ← Vt.
    pub fn compress(&mut self, max_bond_dim: usize, tolerance: f64) -> QuantRS2Result<()> {
        // Iterate over all bonds; for each bond compress the pair of adjacent tensors.
        // We collect bond info first to avoid borrow issues.
        let bond_indices: Vec<usize> = (0..self.bonds.len()).collect();

        for bond_idx in bond_indices {
            let (t1_idx, ref idx1, t2_idx, ref idx2) = self.bonds[bond_idx].clone();

            if t1_idx >= self.tensors.len() || t2_idx >= self.tensors.len() {
                continue;
            }

            // Build real-valued matrix from |amplitude|^2 of t1's data (left) × t2's data (right).
            // Rows = size of t1, cols = size of t2 (simplified: treat each tensor as a flattened vector).
            let rows = self.tensors[t1_idx].size();
            let cols = self.tensors[t2_idx].size();

            if rows == 0 || cols == 0 {
                continue;
            }

            // Build the coupling matrix: M[i,j] = Re(conj(t1[i]) * t2[j])
            let mut mat_data = Vec::with_capacity(rows * cols);
            for i in 0..rows {
                let a = self.tensors[t1_idx].data[i];
                for j in 0..cols {
                    let b = self.tensors[t2_idx].data[j];
                    // Real part of ⟨a|b⟩ coupling
                    mat_data.push(a.re * b.re + a.im * b.im);
                }
            }

            let mat = Array2::from_shape_vec((rows, cols), mat_data).map_err(|e| {
                QuantRS2Error::RuntimeError(format!("SVD matrix build failed: {e}"))
            })?;

            // Compute SVD: M = U * diag(s) * Vt
            let svd_result = svd(&mat.view(), false, None).map_err(|e| {
                QuantRS2Error::RuntimeError(format!("SVD failed on bond {bond_idx}: {e}"))
            });

            let (u_mat, s_vec, vt_mat) = match svd_result {
                Ok(result) => result,
                Err(_) => {
                    // If SVD fails (e.g., tiny matrix), leave bond unchanged
                    continue;
                }
            };

            // Determine truncation rank
            let s_total: f64 = s_vec.iter().copied().sum();
            let mut rank = s_vec.len();

            // Truncate by tolerance (keep singular values whose cumulative fraction > 1-tolerance)
            if s_total > 0.0 {
                let mut cumulative = 0.0;
                for (k, &sv) in s_vec.iter().enumerate() {
                    cumulative += sv / s_total;
                    if cumulative >= 1.0 - tolerance {
                        rank = k + 1;
                        break;
                    }
                }
            }

            // Apply max_bond_dim cap
            rank = rank.min(max_bond_dim).min(s_vec.len());

            if rank == 0 {
                rank = 1;
            }

            // Reconstruct left tensor data: new_left[i] = sum_k U[i,k] * s[k]  (for k < rank)
            // We store the result back as the left tensor's flat data (rows dimension preserved, compressed).
            let mut new_t1_data: Vec<C64> = self.tensors[t1_idx].data.clone();
            let mut new_t2_data: Vec<C64> = self.tensors[t2_idx].data.clone();

            // Project t1 onto the rank-truncated left singular vectors
            // new_t1[i] = sum_{k=0}^{rank-1} U[i,k] * s[k] * (old_t1[i] magnitude)
            for i in 0..rows {
                let mut proj = 0.0f64;
                for k in 0..rank {
                    proj += u_mat[[i, k]] * s_vec[k];
                }
                // Scale the complex amplitude by the projected singular value weight
                let original_norm = (new_t1_data[i].norm_sqr() + 1e-300_f64).sqrt();
                let scale = proj.abs() / (original_norm + 1e-300_f64);
                new_t1_data[i] = C64::new(new_t1_data[i].re * scale, new_t1_data[i].im * scale);
            }

            // Project t2 onto the rank-truncated right singular vectors
            for j in 0..cols {
                let mut proj = 0.0f64;
                for k in 0..rank {
                    proj += vt_mat[[k, j]];
                }
                let original_norm = (new_t2_data[j].norm_sqr() + 1e-300_f64).sqrt();
                let scale = proj.abs() / (original_norm + 1e-300_f64);
                new_t2_data[j] = C64::new(new_t2_data[j].re * scale, new_t2_data[j].im * scale);
            }

            self.tensors[t1_idx].data = new_t1_data;
            self.tensors[t2_idx].data = new_t2_data;
        }

        Ok(())
    }
}

/// Convert a quantum circuit to tensor network representation
pub struct CircuitToTensorNetwork<const N: usize> {
    /// Maximum bond dimension for compression
    max_bond_dim: Option<usize>,
    /// Truncation tolerance
    tolerance: f64,
}

impl<const N: usize> Default for CircuitToTensorNetwork<N> {
    fn default() -> Self {
        Self::new()
    }
}

impl<const N: usize> CircuitToTensorNetwork<N> {
    /// Create a new converter
    #[must_use]
    pub const fn new() -> Self {
        Self {
            max_bond_dim: None,
            tolerance: 1e-10,
        }
    }

    /// Set maximum bond dimension
    #[must_use]
    pub const fn with_max_bond_dim(mut self, dim: usize) -> Self {
        self.max_bond_dim = Some(dim);
        self
    }

    /// Set truncation tolerance
    #[must_use]
    pub const fn with_tolerance(mut self, tol: f64) -> Self {
        self.tolerance = tol;
        self
    }

    /// Convert circuit to tensor network
    pub fn convert(&self, circuit: &Circuit<N>) -> QuantRS2Result<TensorNetwork> {
        let mut tn = TensorNetwork::new();
        let mut qubit_wires: HashMap<usize, String> = HashMap::new();

        // Initialize qubit wires
        for i in 0..N {
            qubit_wires.insert(i, format!("q{i}_in"));
        }

        // Convert each gate to a tensor
        for (gate_idx, gate) in circuit.gates().iter().enumerate() {
            let tensor = self.gate_to_tensor(gate.as_ref(), gate_idx)?;
            let tensor_idx = tn.add_tensor(tensor);

            // Connect to previous wires
            for qubit in gate.qubits() {
                let q = qubit.id() as usize;
                let prev_wire = qubit_wires
                    .get(&q)
                    .ok_or_else(|| {
                        QuantRS2Error::InvalidInput(format!("Qubit wire {q} not found"))
                    })?
                    .clone();
                let new_wire = format!("q{q}_g{gate_idx}");

                // Add bond from previous wire to this gate
                if gate_idx > 0 || prev_wire.contains("_g") {
                    tn.add_bond(
                        tensor_idx - 1,
                        prev_wire.clone(),
                        tensor_idx,
                        format!("in_{q}"),
                    )?;
                }

                // Update wire for next connection
                qubit_wires.insert(q, new_wire);
            }
        }

        Ok(tn)
    }

    /// Convert a gate to tensor representation
    fn gate_to_tensor(&self, gate: &dyn GateOp, gate_idx: usize) -> QuantRS2Result<Tensor> {
        let qubits = gate.qubits();
        let n_qubits = qubits.len();

        match n_qubits {
            1 => {
                // Single-qubit gate
                let matrix = self.get_single_qubit_matrix(gate)?;
                let q = qubits[0].id() as usize;

                Ok(Tensor::new(
                    matrix,
                    vec![2, 2],
                    vec![format!("in_{}", q), format!("out_{}", q)],
                ))
            }
            2 => {
                // Two-qubit gate
                let matrix = self.get_two_qubit_matrix(gate)?;
                let q0 = qubits[0].id() as usize;
                let q1 = qubits[1].id() as usize;

                Ok(Tensor::new(
                    matrix,
                    vec![2, 2, 2, 2],
                    vec![
                        format!("in_{}", q0),
                        format!("in_{}", q1),
                        format!("out_{}", q0),
                        format!("out_{}", q1),
                    ],
                ))
            }
            _ => Err(QuantRS2Error::UnsupportedOperation(format!(
                "{n_qubits}-qubit gates not yet supported for tensor networks"
            ))),
        }
    }

    /// Get matrix representation of single-qubit gate
    fn get_single_qubit_matrix(&self, gate: &dyn GateOp) -> QuantRS2Result<Vec<C64>> {
        // Simplified - would use actual gate matrices
        match gate.name() {
            "H" => Ok(vec![
                C64::new(1.0 / 2.0_f64.sqrt(), 0.0),
                C64::new(1.0 / 2.0_f64.sqrt(), 0.0),
                C64::new(1.0 / 2.0_f64.sqrt(), 0.0),
                C64::new(-1.0 / 2.0_f64.sqrt(), 0.0),
            ]),
            "X" => Ok(vec![
                C64::new(0.0, 0.0),
                C64::new(1.0, 0.0),
                C64::new(1.0, 0.0),
                C64::new(0.0, 0.0),
            ]),
            "Y" => Ok(vec![
                C64::new(0.0, 0.0),
                C64::new(0.0, -1.0),
                C64::new(0.0, 1.0),
                C64::new(0.0, 0.0),
            ]),
            "Z" => Ok(vec![
                C64::new(1.0, 0.0),
                C64::new(0.0, 0.0),
                C64::new(0.0, 0.0),
                C64::new(-1.0, 0.0),
            ]),
            _ => Ok(vec![
                C64::new(1.0, 0.0),
                C64::new(0.0, 0.0),
                C64::new(0.0, 0.0),
                C64::new(1.0, 0.0),
            ]),
        }
    }

    /// Get matrix representation of two-qubit gate
    fn get_two_qubit_matrix(&self, gate: &dyn GateOp) -> QuantRS2Result<Vec<C64>> {
        // Simplified - would use actual gate matrices
        if gate.name() == "CNOT" {
            let mut matrix = vec![C64::new(0.0, 0.0); 16];
            matrix[0] = C64::new(1.0, 0.0); // |00⟩ -> |00⟩
            matrix[5] = C64::new(1.0, 0.0); // |01⟩ -> |01⟩
            matrix[15] = C64::new(1.0, 0.0); // |10⟩ -> |11⟩
            matrix[10] = C64::new(1.0, 0.0); // |11⟩ -> |10⟩
            Ok(matrix)
        } else {
            // Identity for unsupported gates
            let mut matrix = vec![C64::new(0.0, 0.0); 16];
            for i in 0..16 {
                matrix[i * 16 + i] = C64::new(1.0, 0.0);
            }
            Ok(matrix)
        }
    }
}

/// Matrix Product State representation of a circuit
#[derive(Debug)]
pub struct MatrixProductState {
    /// Site tensors
    tensors: Vec<Tensor>,
    /// Bond dimensions
    bond_dims: Vec<usize>,
    /// Number of qubits
    n_qubits: usize,
}

impl MatrixProductState {
    /// Create MPS from a quantum circuit via explicit unitary tensor contraction.
    ///
    /// Algorithm:
    /// 1. Initialize the MPS as the |0...0⟩ product state: each site tensor is [1, 0] with
    ///    bond dimensions [1, ..., 1].
    /// 2. For each gate in the circuit:
    ///    - Single-qubit gate U on site `i`: contract the 2x2 unitary into the rank-3 site tensor
    ///      Γ\[i\] with shape \[χ_left, 2, χ_right\].
    ///    - Two-qubit gate U on sites (i, i+1): reshape the two adjacent site tensors into a
    ///      combined matrix of shape \[χ_left * 2, 2 * χ_right\], apply the 4x4 unitary, then
    ///      perform SVD to split back into two site tensors and update the bond dimension.
    pub fn from_circuit<const N: usize>(circuit: &Circuit<N>) -> QuantRS2Result<Self> {
        if N == 0 {
            return Ok(Self {
                tensors: Vec::new(),
                bond_dims: Vec::new(),
                n_qubits: 0,
            });
        }

        let converter = CircuitToTensorNetwork::<N>::new();
        // bond_dims[i] = bond dimension between site i and i+1 (length N-1).
        let mut bond_dims = vec![1usize; N.saturating_sub(1)];

        // Site tensors: Γ[i] has shape [χ_left, 2, χ_right] stored as flat Vec<C64>
        // For i=0: shape [1, 2, 1]; for the |0⟩ state: data = [1, 0] (physical index 0→1, 1→0)
        let mut site_tensors: Vec<Vec<C64>> = (0..N)
            .map(|_| {
                // [1, 2, 1] tensor for |0⟩: Γ[0,0,0]=1, Γ[0,1,0]=0
                vec![C64::new(1.0, 0.0), C64::new(0.0, 0.0)]
            })
            .collect();

        // Helper: retrieve 2×2 matrix for a single-qubit gate (reuse converter logic)
        let gate_to_single_mat = |g: &dyn GateOp| -> Option<[C64; 4]> {
            match g.name() {
                "H" => Some([
                    C64::new(1.0 / 2.0_f64.sqrt(), 0.0),
                    C64::new(1.0 / 2.0_f64.sqrt(), 0.0),
                    C64::new(1.0 / 2.0_f64.sqrt(), 0.0),
                    C64::new(-1.0 / 2.0_f64.sqrt(), 0.0),
                ]),
                "X" => Some([
                    C64::new(0.0, 0.0),
                    C64::new(1.0, 0.0),
                    C64::new(1.0, 0.0),
                    C64::new(0.0, 0.0),
                ]),
                "Y" => Some([
                    C64::new(0.0, 0.0),
                    C64::new(0.0, -1.0),
                    C64::new(0.0, 1.0),
                    C64::new(0.0, 0.0),
                ]),
                "Z" => Some([
                    C64::new(1.0, 0.0),
                    C64::new(0.0, 0.0),
                    C64::new(0.0, 0.0),
                    C64::new(-1.0, 0.0),
                ]),
                "RY" | "RZ" | "RX" | "S" | "T" | "SX" | "ID" | "I" => {
                    // Use identity as fallback for parameterized gates
                    Some([
                        C64::new(1.0, 0.0),
                        C64::new(0.0, 0.0),
                        C64::new(0.0, 0.0),
                        C64::new(1.0, 0.0),
                    ])
                }
                _ => None,
            }
        };

        // Helper: 4×4 CNOT unitary (control=row 0 of physical indices)
        let cnot_mat: [C64; 16] = {
            let mut m = [C64::new(0.0, 0.0); 16];
            m[0] = C64::new(1.0, 0.0); // |00⟩ → |00⟩
            m[5] = C64::new(1.0, 0.0); // |01⟩ → |01⟩
            m[14] = C64::new(1.0, 0.0); // |10⟩ → |11⟩
            m[11] = C64::new(1.0, 0.0); // |11⟩ → |10⟩
            m
        };

        // Default max_bond_dim during construction (no truncation limit)
        let max_bd = 32usize;

        for gate in circuit.gates() {
            let qubits = gate.qubits();
            match qubits.len() {
                1 => {
                    let qi = qubits[0].id() as usize;
                    if qi >= N {
                        continue;
                    }
                    if let Some(u) = gate_to_single_mat(gate.as_ref()) {
                        // Contract: new_Γ[α, σ', β] = Σ_σ U[σ', σ] * Γ[α, σ, β]
                        // Current shape: [χ_l, 2, χ_r] — chi_l=1, chi_r=1 for initial state
                        // Since we store flat [2] for the initial product state:
                        let old = site_tensors[qi].clone();
                        let phys = old.len(); // = 2 * χ_l * χ_r in general
                                              // Simple case: apply 2×2 unitary to physical index dimension 2
                        let half = phys / 2;
                        let mut new_site = vec![C64::new(0.0, 0.0); phys];
                        for alpha in 0..half {
                            let s0 = old[alpha]; // physical |0⟩
                            let s1 = old[alpha + half]; // physical |1⟩
                            new_site[alpha] = u[0] * s0 + u[1] * s1; // U[0,0]*|0⟩ + U[0,1]*|1⟩
                            new_site[alpha + half] = u[2] * s0 + u[3] * s1; // U[1,0]*|0⟩ + U[1,1]*|1⟩
                        }
                        site_tensors[qi] = new_site;
                    }
                }
                2 => {
                    let qi = qubits[0].id() as usize;
                    let qj = qubits[1].id() as usize;
                    // Only handle adjacent qubits (i, i+1)
                    if qi >= N || qj >= N || qj != qi + 1 {
                        continue;
                    }
                    let gate_name = gate.name();
                    let unitary_mat: [C64; 16] = if gate_name == "CNOT" || gate_name == "CX" {
                        cnot_mat
                    } else {
                        // Identity 4×4 for unsupported two-qubit gates
                        let mut id = [C64::new(0.0, 0.0); 16];
                        id[0] = C64::new(1.0, 0.0);
                        id[5] = C64::new(1.0, 0.0);
                        id[10] = C64::new(1.0, 0.0);
                        id[15] = C64::new(1.0, 0.0);
                        id
                    };

                    // Left site: [χ_l, 2, χ_m], right site: [χ_m, 2, χ_r]
                    // Merge into: Θ[χ_l * 2, 2 * χ_r] via Θ[α*2+σ, σ'*χ_r+β] = Σ_m Γ_i[α,σ,m] * Γ_j[m,σ',β]
                    let left = &site_tensors[qi];
                    let right = &site_tensors[qj];
                    let left_phys = left.len(); // χ_l * 2
                    let right_phys = right.len(); // 2 * χ_r
                    let chi_m = bond_dims.get(qi).copied().unwrap_or(1);
                    let chi_l = left_phys / 2; // should equal χ_l * χ_m / χ_m
                    let chi_r = right_phys / 2; // should equal χ_m * χ_r / χ_m

                    // Build merged tensor Θ: shape [chi_l * 2, chi_r * 2]
                    // Index convention: row = (chi_l_idx * 2 + sigma_i), col = (sigma_j * chi_r + chi_r_idx)
                    let nrows = chi_l * 2;
                    let ncols = chi_r * 2;
                    let mut theta = vec![C64::new(0.0, 0.0); nrows * ncols];

                    // Contract over χ_m (bond index between site qi and qj)
                    // left[alpha, sigma_i] = left_flat[sigma_i * chi_l + alpha]  (stored as [phys_0, phys_1])
                    // right[sigma_j, beta] = right_flat[sigma_j * chi_r + beta]
                    for sigma_i in 0..2usize {
                        for alpha in 0..chi_l {
                            let l_val = left
                                .get(sigma_i * chi_l + alpha)
                                .copied()
                                .unwrap_or(C64::new(0.0, 0.0));
                            for sigma_j in 0..2usize {
                                for beta in 0..chi_r {
                                    let r_val = right
                                        .get(sigma_j * chi_r + beta)
                                        .copied()
                                        .unwrap_or(C64::new(0.0, 0.0));
                                    let row = alpha * 2 + sigma_i;
                                    let col = sigma_j * chi_r + beta;
                                    if row < nrows && col < ncols {
                                        theta[row * ncols + col] += l_val * r_val;
                                    }
                                }
                            }
                        }
                    }

                    // Apply two-qubit unitary: Θ' = U * Θ (in the combined physical index space)
                    // U acts on (σ_i, σ_j) space (4×4), Θ rows ~ (α, σ_i), Θ cols ~ (σ_j, β)
                    // Θ'[α*2+σ'_i, σ'_j*χ_r+β] = Σ_{σ_i, σ_j} U[σ'_i*2+σ'_j, σ_i*2+σ_j] * Θ[α*2+σ_i, σ_j*χ_r+β]
                    let mut theta_prime = vec![C64::new(0.0, 0.0); nrows * ncols];
                    for alpha in 0..chi_l {
                        for sigma_i_out in 0..2usize {
                            for sigma_j_out in 0..2usize {
                                for beta in 0..chi_r {
                                    let row_out = alpha * 2 + sigma_i_out;
                                    let col_out = sigma_j_out * chi_r + beta;
                                    let mut val = C64::new(0.0, 0.0);
                                    for sigma_i_in in 0..2usize {
                                        for sigma_j_in in 0..2usize {
                                            let u_idx = (sigma_i_out * 2 + sigma_j_out) * 4
                                                + sigma_i_in * 2
                                                + sigma_j_in;
                                            let u_val = unitary_mat
                                                .get(u_idx)
                                                .copied()
                                                .unwrap_or(C64::new(0.0, 0.0));
                                            let row_in = alpha * 2 + sigma_i_in;
                                            let col_in = sigma_j_in * chi_r + beta;
                                            val += u_val
                                                * theta
                                                    .get(row_in * ncols + col_in)
                                                    .copied()
                                                    .unwrap_or(C64::new(0.0, 0.0));
                                        }
                                    }
                                    if row_out < nrows && col_out < ncols {
                                        theta_prime[row_out * ncols + col_out] = val;
                                    }
                                }
                            }
                        }
                    }

                    // SVD on the real part to get new bond dimension
                    let real_mat_data: Vec<f64> = theta_prime.iter().map(|c| c.re).collect();
                    let real_mat =
                        Array2::from_shape_vec((nrows, ncols), real_mat_data).map_err(|e| {
                            QuantRS2Error::RuntimeError(format!("MPS matrix reshape failed: {e}"))
                        })?;

                    let svd_res = svd(&real_mat.view(), false, None)
                        .map_err(|e| QuantRS2Error::RuntimeError(format!("MPS SVD failed: {e}")));

                    let (u_mat, s_vec, vt_mat) = match svd_res {
                        Ok(r) => r,
                        Err(_) => {
                            // Fallback: keep tensors unchanged
                            continue;
                        }
                    };

                    // Truncate bond dimension to max_bd
                    let new_chi_m = s_vec.len().min(max_bd);

                    // Reconstruct left site: shape [chi_l * 2, new_chi_m]
                    // new_left[row, k] = U[row, k] * sqrt(s[k])
                    let mut new_left = vec![C64::new(0.0, 0.0); chi_l * 2 * new_chi_m];
                    for row in 0..nrows {
                        for k in 0..new_chi_m {
                            let sv = s_vec[k].max(0.0).sqrt();
                            let idx = row * new_chi_m + k;
                            new_left[idx] = C64::new(u_mat[[row, k]] * sv, 0.0);
                        }
                    }

                    // Reconstruct right site: shape [new_chi_m, chi_r * 2]
                    // new_right[k, col] = sqrt(s[k]) * Vt[k, col]
                    let mut new_right = vec![C64::new(0.0, 0.0); new_chi_m * chi_r * 2];
                    for k in 0..new_chi_m {
                        let sv = s_vec[k].max(0.0).sqrt();
                        for col in 0..ncols {
                            let idx = k * ncols + col;
                            new_right[idx] = C64::new(vt_mat[[k, col]] * sv, 0.0);
                        }
                    }

                    site_tensors[qi] = new_left;
                    site_tensors[qj] = new_right;
                    if qi < bond_dims.len() {
                        bond_dims[qi] = new_chi_m;
                    }
                }
                _ => {
                    // Multi-qubit gates beyond 2-qubit: skip
                }
            }
        }

        // Build the site Tensors with correct shape annotations
        let tensors: Vec<Tensor> = site_tensors
            .into_iter()
            .enumerate()
            .map(|(i, data)| {
                let chi_l = if i == 0 { 1 } else { bond_dims[i - 1] };
                let chi_r = if i + 1 < N { bond_dims[i] } else { 1 };
                let shape = vec![chi_l, 2, chi_r];
                let indices = vec![
                    format!("bond_left_{i}"),
                    format!("phys_{i}"),
                    format!("bond_right_{i}"),
                ];
                // Ensure data length matches shape product
                let expected = chi_l * 2 * chi_r;
                let mut padded = data;
                padded.resize(expected, C64::new(0.0, 0.0));
                Tensor::new(padded, shape, indices)
            })
            .collect();

        Ok(Self {
            tensors,
            bond_dims,
            n_qubits: N,
        })
    }

    /// Compress the MPS via a left-to-right SVD sweep with bond truncation.
    ///
    /// For each bond between site i and i+1:
    /// 1. Reshape tensors\[i\] (shape \[χ_l, 2, χ_m\]) and tensors\[i+1\] (shape \[χ_m, 2, χ_r\])
    ///    into a combined matrix Θ of shape \[χ_l\*2, χ_r\*2\].
    /// 2. Compute SVD Θ = U Σ Vt.
    /// 3. Truncate to min(max_bond_dim, rank where σ_k / σ_0 > tolerance).
    /// 4. Set tensors\[i\] = U\[:, :new_χ\] \* diag(Σ\[:new_χ\])^(1/2),
    ///    tensors\[i+1\] = diag(Σ\[:new_χ\])^(1/2) \* Vt\[:new_χ, :\].
    pub fn compress(&mut self, max_bond_dim: usize, tolerance: f64) -> QuantRS2Result<()> {
        let n = self.n_qubits;
        if n <= 1 {
            return Ok(());
        }

        for i in 0..(n - 1) {
            if i + 1 >= self.tensors.len() {
                break;
            }

            let chi_l_i = self.tensors[i].shape.first().copied().unwrap_or(1);
            let chi_r_i = self.tensors[i].shape.get(2).copied().unwrap_or(1); // = chi_m
            let chi_r_j = self.tensors[i + 1].shape.get(2).copied().unwrap_or(1);

            let nrows = chi_l_i * 2;
            let ncols = chi_r_j * 2;

            // Build combined real-valued matrix from amplitudes
            // Θ[alpha*2+sigma_i, sigma_j*chi_r_j+beta] = Σ_m Γ_i[alpha,sigma_i,m] * Γ_{i+1}[m,sigma_j,beta]
            let left = &self.tensors[i].data;
            let right = &self.tensors[i + 1].data;
            let mut theta_real = vec![0.0f64; nrows * ncols];

            for alpha in 0..chi_l_i {
                for sigma_i in 0..2usize {
                    for m in 0..chi_r_i {
                        let l_idx = (alpha * 2 + sigma_i) * chi_r_i + m;
                        let l_val = left.get(l_idx).map(|c| c.re).unwrap_or(0.0);
                        if l_val == 0.0 {
                            continue;
                        }
                        for sigma_j in 0..2usize {
                            for beta in 0..chi_r_j {
                                let r_idx = (m * 2 + sigma_j) * chi_r_j + beta;
                                let r_val = right.get(r_idx).map(|c| c.re).unwrap_or(0.0);
                                let row = alpha * 2 + sigma_i;
                                let col = sigma_j * chi_r_j + beta;
                                if row < nrows && col < ncols {
                                    theta_real[row * ncols + col] += l_val * r_val;
                                }
                            }
                        }
                    }
                }
            }

            let mat = Array2::from_shape_vec((nrows, ncols), theta_real).map_err(|e| {
                QuantRS2Error::RuntimeError(format!("MPS compress reshape failed: {e}"))
            })?;

            let svd_res = svd(&mat.view(), false, None).map_err(|e| {
                QuantRS2Error::RuntimeError(format!("MPS compress SVD failed at bond {i}: {e}"))
            });

            let (u_mat, s_vec, vt_mat) = match svd_res {
                Ok(r) => r,
                Err(_) => continue,
            };

            // Determine truncation rank
            let sigma_max = s_vec.first().copied().unwrap_or(0.0);
            let rank = if sigma_max > 0.0 {
                s_vec
                    .iter()
                    .take_while(|&&sv| sv / sigma_max > tolerance)
                    .count()
            } else {
                1
            };
            let new_chi_m = rank.min(max_bond_dim).min(s_vec.len()).max(1);

            // Rebuild left tensor: shape [chi_l_i, 2, new_chi_m]
            let new_left_size = chi_l_i * 2 * new_chi_m;
            let mut new_left = vec![C64::new(0.0, 0.0); new_left_size];
            for row in 0..(chi_l_i * 2) {
                for k in 0..new_chi_m {
                    let sv = s_vec[k].max(0.0).sqrt();
                    let flat_idx = row * new_chi_m + k;
                    new_left[flat_idx] = C64::new(u_mat[[row, k]] * sv, 0.0);
                }
            }

            // Rebuild right tensor: shape [new_chi_m, 2, chi_r_j]
            let new_right_size = new_chi_m * 2 * chi_r_j;
            let mut new_right = vec![C64::new(0.0, 0.0); new_right_size];
            for k in 0..new_chi_m {
                let sv = s_vec[k].max(0.0).sqrt();
                for col in 0..(2 * chi_r_j) {
                    let flat_idx = k * 2 * chi_r_j + col;
                    new_right[flat_idx] = C64::new(vt_mat[[k, col]] * sv, 0.0);
                }
            }

            // Update tensors in place
            self.tensors[i].data = new_left;
            self.tensors[i].shape = vec![chi_l_i, 2, new_chi_m];

            self.tensors[i + 1].data = new_right;
            self.tensors[i + 1].shape = vec![new_chi_m, 2, chi_r_j];

            if i < self.bond_dims.len() {
                self.bond_dims[i] = new_chi_m;
            }
        }

        Ok(())
    }

    /// Calculate overlap with another MPS
    pub fn overlap(&self, other: &Self) -> QuantRS2Result<C64> {
        if self.n_qubits != other.n_qubits {
            return Err(QuantRS2Error::InvalidInput(
                "MPS have different number of qubits".to_string(),
            ));
        }

        // Calculate ⟨ψ|φ⟩
        Ok(C64::new(1.0, 0.0)) // Placeholder
    }

    /// Calculate expectation value of observable
    pub const fn expectation_value(&self, observable: &TensorNetwork) -> QuantRS2Result<f64> {
        // Calculate ⟨ψ|O|ψ⟩
        Ok(0.0) // Placeholder
    }
}

/// Circuit compression using tensor networks
pub struct TensorNetworkCompressor {
    /// Maximum bond dimension
    max_bond_dim: usize,
    /// Truncation tolerance
    tolerance: f64,
    /// Compression method
    method: CompressionMethod,
}

#[derive(Debug, Clone)]
pub enum CompressionMethod {
    /// Singular Value Decomposition
    SVD,
    /// Density Matrix Renormalization Group
    DMRG,
    /// Time-Evolving Block Decimation
    TEBD,
}

impl TensorNetworkCompressor {
    /// Create a new compressor
    #[must_use]
    pub const fn new(max_bond_dim: usize) -> Self {
        Self {
            max_bond_dim,
            tolerance: 1e-10,
            method: CompressionMethod::SVD,
        }
    }

    /// Set compression method
    #[must_use]
    pub const fn with_method(mut self, method: CompressionMethod) -> Self {
        self.method = method;
        self
    }

    /// Compress a circuit
    pub fn compress<const N: usize>(
        &self,
        circuit: &Circuit<N>,
    ) -> QuantRS2Result<CompressedCircuit<N>> {
        let mps = MatrixProductState::from_circuit(circuit)?;

        Ok(CompressedCircuit {
            mps,
            original_gates: circuit.num_gates(),
            compression_ratio: 1.0, // Placeholder
        })
    }
}

/// Compressed circuit representation
#[derive(Debug)]
pub struct CompressedCircuit<const N: usize> {
    /// MPS representation
    mps: MatrixProductState,
    /// Original number of gates
    original_gates: usize,
    /// Compression ratio
    compression_ratio: f64,
}

impl<const N: usize> CompressedCircuit<N> {
    /// Get compression ratio
    #[must_use]
    pub const fn compression_ratio(&self) -> f64 {
        self.compression_ratio
    }

    /// Decompress back to circuit
    pub fn decompress(&self) -> QuantRS2Result<Circuit<N>> {
        // Convert MPS back to circuit representation
        // This is non-trivial and would require gate synthesis
        Ok(Circuit::<N>::new())
    }

    /// Get fidelity with original circuit
    pub const fn fidelity(&self, original: &Circuit<N>) -> QuantRS2Result<f64> {
        // Calculate |⟨ψ_compressed|ψ_original⟩|²
        Ok(0.99) // Placeholder
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use quantrs2_core::gate::single::Hadamard;

    #[test]
    fn test_tensor_creation() {
        let data = vec![
            C64::new(1.0, 0.0),
            C64::new(0.0, 0.0),
            C64::new(0.0, 0.0),
            C64::new(1.0, 0.0),
        ];
        let tensor = Tensor::new(data, vec![2, 2], vec!["in".to_string(), "out".to_string()]);

        assert_eq!(tensor.rank(), 2);
        assert_eq!(tensor.size(), 4);
    }

    #[test]
    fn test_tensor_network() {
        let mut tn = TensorNetwork::new();

        let t1 = Tensor::identity(2, "a".to_string(), "b".to_string());
        let t2 = Tensor::identity(2, "c".to_string(), "d".to_string());

        let idx1 = tn.add_tensor(t1);
        let idx2 = tn.add_tensor(t2);

        tn.add_bond(idx1, "b".to_string(), idx2, "c".to_string())
            .expect("Failed to add bond between tensors");

        assert_eq!(tn.tensors.len(), 2);
        assert_eq!(tn.bonds.len(), 1);
    }

    #[test]
    fn test_circuit_to_tensor_network() {
        let mut circuit = Circuit::<2>::new();
        circuit
            .add_gate(Hadamard { target: QubitId(0) })
            .expect("Failed to add Hadamard gate");

        let converter = CircuitToTensorNetwork::<2>::new();
        let tn = converter
            .convert(&circuit)
            .expect("Failed to convert circuit to tensor network");

        assert!(!tn.tensors.is_empty());
    }

    #[test]
    fn test_compression() {
        let circuit = Circuit::<2>::new();
        let compressor = TensorNetworkCompressor::new(32);

        let compressed = compressor
            .compress(&circuit)
            .expect("Failed to compress circuit");
        assert!(compressed.compression_ratio() <= 1.0);
    }

    #[test]
    fn test_tensor_network_svd_compress() {
        use quantrs2_core::gate::multi::CNOT;

        // Build a circuit with a Hadamard + CNOT (Bell state preparation)
        let mut circuit = Circuit::<2>::new();
        circuit
            .add_gate(Hadamard { target: QubitId(0) })
            .expect("H gate");
        circuit
            .add_gate(CNOT {
                control: QubitId(0),
                target: QubitId(1),
            })
            .expect("CNOT gate");

        let converter = CircuitToTensorNetwork::<2>::new();
        let mut tn = converter.convert(&circuit).expect("Convert to TN");

        // Compress with max bond dim 4 and tolerance 1e-6
        tn.compress(4, 1e-6).expect("TN compress");
        // If we got here without panic, the test passes; verify structure
        assert_eq!(tn.tensors.len(), 2);
    }

    #[test]
    fn test_mps_from_circuit_trivial() {
        // Empty circuit → valid MPS
        let circuit = Circuit::<2>::new();
        let mps = MatrixProductState::from_circuit(&circuit).expect("MPS from empty circuit");
        assert_eq!(mps.n_qubits, 2);
        assert_eq!(mps.tensors.len(), 2);
    }

    #[test]
    fn test_mps_from_circuit_with_hadamard() {
        use quantrs2_core::gate::single::Hadamard;

        let mut circuit = Circuit::<3>::new();
        circuit
            .add_gate(Hadamard { target: QubitId(0) })
            .expect("H gate");

        let mps = MatrixProductState::from_circuit(&circuit).expect("MPS from H circuit");
        assert_eq!(mps.n_qubits, 3);
        assert_eq!(mps.tensors.len(), 3);
    }

    #[test]
    fn test_mps_compress_reduces_bond_dim() {
        use quantrs2_core::gate::multi::CNOT;

        // Bell state: H + CNOT should create a non-trivial entangled MPS
        let mut circuit = Circuit::<2>::new();
        circuit
            .add_gate(Hadamard { target: QubitId(0) })
            .expect("H gate");
        circuit
            .add_gate(CNOT {
                control: QubitId(0),
                target: QubitId(1),
            })
            .expect("CNOT gate");

        let mut mps = MatrixProductState::from_circuit(&circuit).expect("MPS from Bell circuit");

        // Compress with max bond dim 1 (strong truncation)
        mps.compress(1, 1e-10).expect("MPS compress");
        // Bond dims should be ≤ max_bond_dim
        for &bd in &mps.bond_dims {
            assert!(bd <= 1, "Bond dim {} exceeds max", bd);
        }
    }
}