Linear Algebra Utilities

def kak_decomposition(unitary: np.ndarray, *, atol: float = 1e-8) -> 'KakDecomposition':
    """KAK decomposition of two-qubit unitary matrices.
    
    Args:
        unitary: 4x4 unitary matrix to decompose
        atol: Absolute tolerance for comparisons
        
    Returns:
        KAK decomposition result
    """

class KakDecomposition:
    """Result of KAK decomposition: U = (A1 ⊗ A0) exp(i(x XX + y YY + z ZZ)) (A3 ⊗ A2)."""
    
    def __init__(self, global_phase: complex,
                 single_qubit_operations_before: Tuple[np.ndarray, np.ndarray],
                 interaction_coefficients: Tuple[float, float, float], 
                 single_qubit_operations_after: Tuple[np.ndarray, np.ndarray]) -> None:
        """Initialize KAK decomposition."""
        
    @property
    def global_phase(self) -> complex:
        """Global phase factor."""
        
    @property
    def single_qubit_operations_before(self) -> Tuple[np.ndarray, np.ndarray]:
        """Single-qubit operations before interaction."""
        
    @property
    def interaction_coefficients(self) -> Tuple[float, float, float]:
        """Interaction coefficients (x, y, z) for XX, YY, ZZ terms."""
        
    @property
    def single_qubit_operations_after(self) -> Tuple[np.ndarray, np.ndarray]:
        """Single-qubit operations after interaction."""
        
    def __mul__(self, other: 'KakDecomposition') -> 'KakDecomposition':
        """Multiply two KAK decompositions."""

def kak_vector(unitary: np.ndarray, *, atol: float = 1e-8) -> np.ndarray:
    """Extract KAK interaction coefficients from two-qubit unitary."""

def kak_canonicalize_vector(x: float, y: float, z: float, *, atol: float = 1e-8) -> Tuple[float, float, float]:
    """Canonicalize KAK vector to preferred coordinate system."""

def axis_angle(matrix: np.ndarray, *, atol: float = 1e-8) -> 'AxisAngleDecomposition':
    """Axis-angle decomposition of single-qubit unitary matrix.
    
    Args:
        matrix: 2x2 unitary matrix
        atol: Absolute tolerance
        
    Returns:
        Axis-angle decomposition
    """

class AxisAngleDecomposition:
    """Axis-angle decomposition of rotation: exp(-i θ/2 n⃗·σ⃗)."""
    
    def __init__(self, angle: float, axis: np.ndarray, global_phase: complex = 1) -> None:
        """Initialize axis-angle decomposition."""
        
    @property
    def angle(self) -> float:
        """Rotation angle in radians."""
        
    @property  
    def axis(self) -> np.ndarray:
        """Rotation axis as 3D unit vector."""
        
    @property
    def global_phase(self) -> complex:
        """Global phase factor."""

def deconstruct_single_qubit_matrix_into_angles(mat: np.ndarray) -> Tuple[float, float, float]:
    """Deconstruct single-qubit matrix into ZYZ Euler angles."""

def kron_factor_4x4_to_2x2s(matrix: np.ndarray, *, atol: float = 1e-8) -> Optional[Tuple[np.ndarray, np.ndarray]]:
    """Factor 4x4 matrix into Kronecker product of 2x2 matrices.
    
    Args:
        matrix: 4x4 matrix to factor
        atol: Tolerance for factorization
        
    Returns:
        Tuple of 2x2 matrices (A, B) such that matrix ≈ A ⊗ B, or None if not factorizable
    """

def unitary_eig(matrix: np.ndarray, *, atol: float = 1e-8) -> Tuple[np.ndarray, np.ndarray]:
    """Eigendecomposition of unitary matrix.
    
    Args:
        matrix: Unitary matrix
        atol: Tolerance for unitarity check
        
    Returns:
        Tuple of (eigenvalues, eigenvectors)
    """

def map_eigenvalues(matrix: np.ndarray, func: Callable[[complex], complex], *, atol: float = 1e-8) -> np.ndarray:
    """Apply function to eigenvalues of matrix.
    
    Args:
        matrix: Input matrix
        func: Function to apply to each eigenvalue
        atol: Tolerance for computations
        
    Returns:
        Matrix with transformed eigenvalues
    """

def so4_to_magic_su2s(so4_matrix: np.ndarray, *, atol: float = 1e-8) -> Tuple[np.ndarray, np.ndarray]:
    """Convert SO(4) matrix to pair of SU(2) matrices using magic basis."""

def dot(*values: np.ndarray) -> np.ndarray:
    """Enhanced matrix multiplication for multiple matrices.
    
    Args:
        *values: Matrices to multiply in sequence
        
    Returns:
        Product of all matrices
    """

def kron(*matrices: np.ndarray) -> np.ndarray:
    """Kronecker product of multiple matrices.
    
    Args:
        *matrices: Matrices to take Kronecker product of
        
    Returns:
        Kronecker product ⊗ᵢ matrices[i]
    """

def kron_with_controls(matrices: Sequence[np.ndarray], controls: Sequence[Sequence[int]]) -> np.ndarray:
    """Kronecker product with control structure.
    
    Args:
        matrices: Matrices for each subsystem
        controls: Control qubit specifications
        
    Returns:
        Controlled Kronecker product
    """

def block_diag(*blocks: np.ndarray) -> np.ndarray:
    """Block diagonal matrix from given blocks.
    
    Args:
        *blocks: Diagonal blocks
        
    Returns:
        Block diagonal matrix
    """

def apply_matrix_to_slices(target: np.ndarray,
                          matrix: np.ndarray, 
                          slices: Sequence[Union[slice, int, Ellipsis]],
                          *, out: Optional[np.ndarray] = None) -> np.ndarray:
    """Apply matrix to specified slices of target tensor.
    
    Args:
        target: Target tensor to modify
        matrix: Matrix to apply
        slices: Slices specifying which parts to operate on
        out: Output buffer (optional)
        
    Returns:
        Modified tensor
    """

def targeted_left_multiply(left_matrix: np.ndarray,
                          right_target: np.ndarray,
                          target_axes: Sequence[int],
                          *, out: Optional[np.ndarray] = None) -> np.ndarray:
    """Left-multiply target tensor by matrix on specified axes."""

def targeted_conjugate_about(tensor: np.ndarray,
                           matrix: np.ndarray,
                           indices: Sequence[int],
                           *, out: Optional[np.ndarray] = None) -> np.ndarray:
    """Conjugate tensor by matrix: M† @ tensor @ M on specified indices."""

def partial_trace(tensor: np.ndarray,
                 keep_indices: Sequence[int],
                 qid_shape: Tuple[int, ...]) -> np.ndarray:
    """Compute partial trace of tensor, keeping specified subsystems.
    
    Args:
        tensor: Input tensor (state vector or density matrix)
        keep_indices: Indices of subsystems to keep
        qid_shape: Shape of each subsystem
        
    Returns:
        Partial trace over discarded subsystems
    """

def partial_trace_of_state_vector_as_mixture(state_vector: np.ndarray,
                                           keep_indices: Sequence[int],
                                           qid_shape: Tuple[int, ...]) -> Tuple[Tuple[float, np.ndarray], ...]:
    """Partial trace of state vector as mixture representation."""

def sub_state_vector(state_vector: np.ndarray,
                    keep_indices: Sequence[int],
                    qid_shape: Tuple[int, ...],
                    *, default: Optional[np.ndarray] = None,
                    atol: float = 1e-8) -> Optional[np.ndarray]:
    """Extract sub-state vector for specified subsystems."""

def state_vector_kronecker_product(*factors: np.ndarray) -> np.ndarray:
    """Kronecker product of state vectors.
    
    Args:
        *factors: State vectors to tensor together
        
    Returns:
        Tensor product state vector
    """

def density_matrix_kronecker_product(*factors: np.ndarray) -> np.ndarray:
    """Kronecker product of density matrices.
    
    Args:
        *factors: Density matrices to tensor together
        
    Returns:
        Tensor product density matrix
    """

def is_unitary(matrix: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if matrix is unitary (U† U = I).
    
    Args:
        matrix: Matrix to test
        atol: Absolute tolerance
        
    Returns:
        True if matrix is unitary
    """

def is_orthogonal(matrix: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if matrix is orthogonal (O^T O = I)."""

def is_special_unitary(matrix: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if matrix is special unitary (unitary with determinant 1)."""

def is_special_orthogonal(matrix: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if matrix is special orthogonal (orthogonal with determinant 1)."""

def is_hermitian(matrix: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if matrix is Hermitian (M† = M).
    
    Args:
        matrix: Matrix to test
        atol: Absolute tolerance
        
    Returns:
        True if matrix is Hermitian
    """

def is_normal(matrix: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if matrix is normal (M M† = M† M)."""

def is_diagonal(matrix: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if matrix is diagonal."""

def is_cptp(choi: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if Choi matrix represents completely positive trace-preserving map.
    
    Args:
        choi: Choi matrix representation
        atol: Tolerance for tests
        
    Returns:
        True if map is CPTP
    """

def allclose_up_to_global_phase(a: np.ndarray, b: np.ndarray, *, 
                               rtol: float = 1e-5, atol: float = 1e-8) -> bool:
    """Check if matrices are equal up to global phase.
    
    Args:
        a, b: Matrices to compare
        rtol: Relative tolerance
        atol: Absolute tolerance
        
    Returns:
        True if matrices are equal up to global phase
    """

def matrix_commutes(m1: np.ndarray, m2: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if two matrices commute (AB = BA).
    
    Args:
        m1, m2: Matrices to test
        atol: Absolute tolerance
        
    Returns:
        True if matrices commute
    """

def diagonalize_real_symmetric_matrix(matrix: np.ndarray, *,
                                     atol: float = 1e-8) -> Tuple[np.ndarray, np.ndarray]:
    """Diagonalize real symmetric matrix.
    
    Args:
        matrix: Real symmetric matrix
        atol: Tolerance for symmetry check
        
    Returns:
        Tuple of (eigenvalues, eigenvectors)
    """

def diagonalize_real_symmetric_and_sorted_diagonal_matrices(symmetric_matrix: np.ndarray,
                                                          diagonal_matrix: np.ndarray, *,
                                                          atol: float = 1e-8) -> Tuple[np.ndarray, np.ndarray]:
    """Simultaneously diagonalize symmetric and diagonal matrices with sorted eigenvalues."""

def bidiagonalize_real_matrix_pair_with_symmetric_products(mat1: np.ndarray, mat2: np.ndarray, *,
                                                         atol: float = 1e-8) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Bidiagonalize pair of real matrices with symmetric products."""

def bidiagonalize_unitary_with_special_orthogonals(mat: np.ndarray, *,
                                                  atol: float = 1e-8) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Bidiagonalize unitary matrix using special orthogonal matrices."""

PAULI_BASIS: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]
"""Standard Pauli basis matrices [I, X, Y, Z]."""

def expand_matrix_in_orthogonal_basis(matrix: np.ndarray,
                                    basis: Sequence[np.ndarray]) -> np.ndarray:
    """Expand matrix in orthogonal basis.
    
    Args:
        matrix: Matrix to expand
        basis: Orthogonal basis matrices
        
    Returns:
        Coefficients in the basis expansion
    """

def matrix_from_basis_coefficients(coefficients: Sequence[complex],
                                 basis: Sequence[np.ndarray]) -> np.ndarray:
    """Construct matrix from basis coefficients.
    
    Args:
        coefficients: Expansion coefficients
        basis: Basis matrices
        
    Returns:
        Matrix = Σᵢ coefficients[i] * basis[i]
    """

def kron_bases(*bases: Sequence[np.ndarray]) -> List[np.ndarray]:
    """Kronecker product of multiple bases.
    
    Args:
        *bases: Bases to take Kronecker product of
        
    Returns:
        List of all Kronecker products
    """

def hilbert_schmidt_inner_product(mat1: np.ndarray, mat2: np.ndarray) -> complex:
    """Hilbert-Schmidt inner product: ⟨A, B⟩ = tr(A† B).
    
    Args:
        mat1, mat2: Matrices for inner product
        
    Returns:
        Inner product value
    """

def pow_pauli_combination(pauli_map: Dict[str, float], exponent: float) -> Dict[str, complex]:
    """Raise Pauli combination to a power."""

def all_near_zero(a: np.ndarray, *, atol: float = 1e-8) -> bool:
    """Check if all elements are near zero.
    
    Args:
        a: Array to check
        atol: Absolute tolerance
        
    Returns:
        True if all elements have magnitude < atol
    """

def all_near_zero_mod(a: np.ndarray, period: float, *, atol: float = 1e-8) -> bool:
    """Check if all elements are near zero modulo period."""

def match_global_phase(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Multiply b by phase to match global phase of a.
    
    Args:
        a: Reference matrix
        b: Matrix to adjust
        
    Returns:
        Phase-adjusted version of b
    """

def slice_for_qubits_equal_to(target_qubit_axes: Sequence[int],
                             little_endian_qureg_value: int = 0,
                             qid_shape: Optional[Sequence[int]] = None) -> Tuple[Union[slice, int], ...]:
    """Create slice for computational basis state."""

def phase_delta(mat: np.ndarray, mat2: np.ndarray) -> float:
    """Phase difference between matrices."""

def reflection_matrix_pow(reflection_matrix: np.ndarray, exponent: float) -> np.ndarray:
    """Raise reflection matrix to a power."""

def to_special(matrix: np.ndarray) -> np.ndarray:
    """Convert to special unitary/orthogonal (determinant 1)."""

def extract_right_diag(left: np.ndarray, middle: np.ndarray, right: np.ndarray) -> np.ndarray:
    """Extract right diagonal matrix from decomposition."""

def num_cnots_required(mat: np.ndarray) -> int:
    """Minimum number of CNOTs required to implement two-qubit unitary.
    
    Args:
        mat: 4x4 unitary matrix
        
    Returns:
        Minimum CNOT count (0, 1, 2, or 3)
    """

def can_numpy_support_shape(shape: Sequence[int]) -> bool:
    """Check if numpy can support given tensor shape."""

def transpose_flattened_array(array: np.ndarray,
                            shape: Tuple[int, ...],
                            permutation: Sequence[int]) -> np.ndarray:
    """Transpose flattened array according to permutation."""

def scatter_plot_normalized_kak_interaction_coefficients(mat: np.ndarray) -> None:
    """Create scatter plot of normalized KAK interaction coefficients."""

CONTROL_TAG: str
"""Tag for control operations in linear algebra contexts."""

import cirq
import numpy as np

# KAK decomposition of two-qubit gate
print("=== KAK Decomposition ===")
# Create a two-qubit unitary (CNOT gate)
cnot_matrix = np.array([[1, 0, 0, 0],
                       [0, 1, 0, 0], 
                       [0, 0, 0, 1],
                       [0, 0, 1, 0]], dtype=complex)

kak = cirq.kak_decomposition(cnot_matrix)
print(f"Global phase: {kak.global_phase}")
print(f"Interaction coefficients: {kak.interaction_coefficients}")
print(f"Single-qubit ops before shape: {[m.shape for m in kak.single_qubit_operations_before]}")

# Verify decomposition
reconstructed = kak.global_phase * np.kron(kak.single_qubit_operations_before[1], 
                                          kak.single_qubit_operations_before[0])
# Apply interaction term...
print(f"Decomposition valid: {np.allclose(cnot_matrix, reconstructed, atol=1e-10)}")

# Axis-angle decomposition for single qubit
print("\n=== Axis-Angle Decomposition ===")
h_matrix = np.array([[1, 1], [1, -1]]) / np.sqrt(2)  # Hadamard
axis_angle = cirq.axis_angle(h_matrix)
print(f"Rotation angle: {axis_angle.angle:.4f} radians")
print(f"Rotation axis: {axis_angle.axis}")

import cirq
import numpy as np

print("=== Matrix Properties ===")

# Test various matrix properties
matrices = {
    'Identity': np.eye(2),
    'Hadamard': np.array([[1, 1], [1, -1]]) / np.sqrt(2),
    'Pauli-X': np.array([[0, 1], [1, 0]]),
    'Random': np.random.randn(2, 2) + 1j * np.random.randn(2, 2)
}

for name, matrix in matrices.items():
    print(f"\n{name} matrix:")
    print(f"  Unitary: {cirq.is_unitary(matrix)}")
    print(f"  Hermitian: {cirq.is_hermitian(matrix)}")
    print(f"  Normal: {cirq.is_normal(matrix)}")
    
    if cirq.is_unitary(matrix):
        print(f"  Special unitary: {cirq.is_special_unitary(matrix)}")

# Test quantum channel properties
print("\n=== Quantum Channel Properties ===")
# Depolarizing channel Choi matrix
p = 0.1
depol_choi = (1-p) * np.kron([[1, 0], [0, 0]], np.eye(2)) + p/3 * sum(
    np.kron(pauli, pauli) for pauli in [[[0, 1], [1, 0]], [[0, -1j], [1j, 0]], [[1, 0], [0, -1]]]
)
print(f"Depolarizing channel is CPTP: {cirq.is_cptp(depol_choi)}")

import cirq
import numpy as np

print("=== Quantum State Operations ===")

# Create entangled state
bell_state = np.array([1, 0, 0, 1]) / np.sqrt(2)  # (|00⟩ + |11⟩)/√2
density_matrix = np.outer(bell_state, bell_state)

# Partial trace to get reduced density matrix of first qubit  
reduced_dm = cirq.partial_trace(density_matrix, keep_indices=[0], qid_shape=(2, 2))
print(f"Reduced density matrix of first qubit:")
print(reduced_dm)

# Check if reduced state is pure
eigenvals = np.linalg.eigvals(reduced_dm)
purity = np.sum(eigenvals**2)
print(f"Purity of reduced state: {purity.real:.4f}")
print(f"Is pure state: {np.isclose(purity, 1.0)}")

# Kronecker products
state1 = np.array([1, 0])  # |0⟩
state2 = np.array([0, 1])  # |1⟩
product_state = cirq.state_vector_kronecker_product(state1, state2)
print(f"Product state |01⟩: {product_state}")

import cirq
import numpy as np

print("=== Linear Algebra Utilities ===")

# Matrix operations
A = np.array([[1, 2], [3, 4]])
B = np.array([[5, 6], [7, 8]])
C = np.array([[1, 0], [0, -1]])

# Enhanced dot product
product = cirq.dot(A, B, C)
print(f"A @ B @ C shape: {product.shape}")

# Kronecker product
kron_product = cirq.kron(A, B)
print(f"A ⊗ B shape: {kron_product.shape}")

# Try to factor back
factors = cirq.kron_factor_4x4_to_2x2s(kron_product)
if factors is not None:
    A_recovered, B_recovered = factors
    print(f"Successfully factored: {np.allclose(A, A_recovered) and np.allclose(B, B_recovered)}")

# Pauli basis expansion
pauli_x = np.array([[0, 1], [1, 0]])
coeffs = cirq.expand_matrix_in_orthogonal_basis(pauli_x, cirq.PAULI_BASIS)
print(f"Pauli-X coefficients in Pauli basis: {coeffs}")

# Reconstruct from coefficients
reconstructed = cirq.matrix_from_basis_coefficients(coeffs, cirq.PAULI_BASIS)
print(f"Reconstruction successful: {np.allclose(pauli_x, reconstructed)}")

# Near-zero checks
small_matrix = np.array([[1e-10, 1e-9], [1e-11, 1e-12]])
print(f"All elements near zero (atol=1e-8): {cirq.all_near_zero(small_matrix, atol=1e-8)}")