
import numpy as np

class Statevector:
    def __init__(self, data):
        self.data = np.array(data, dtype=complex)
        self.data /= np.linalg.norm(self.data)
        self.num_qubits = int(np.log2(len(self.data)))

    def probabilities(self): return np.abs(self.data)**2
    def measure(self): return np.random.choice(len(self.data), p=self.probabilities())
    def to_density_matrix(self): return np.outer(self.data, self.data.conj())
    def fidelity(self, other): return abs(np.dot(self.data.conj(), other.data))**2
    def entanglement_entropy(self, subsystem_qubits):
        n = self.num_qubits
        k = len(subsystem_qubits)
        rho = partial_trace(self.to_density_matrix(), n, subsystem_qubits)
        eigvals = np.linalg.eigvalsh(rho)
        eigvals = eigvals[eigvals > 1e-12]
        return float(-np.sum(eigvals * np.log2(eigvals)))
    def __repr__(self): return f"Statevector({self.num_qubits} qubits)"

class DensityMatrix:
    def __init__(self, data):
        self.data = np.array(data, dtype=complex)
        self.num_qubits = int(np.log2(self.data.shape[0]))
    def purity(self): return float(np.real(np.trace(self.data @ self.data)))
    def trace(self): return float(np.real(np.trace(self.data)))
    def fidelity(self, other):
        from scipy.linalg import sqrtm
        sq = sqrtm(self.data)
        F = np.real(np.trace(sqrtm(sq @ other.data @ sq)))
        return float(F**2)
    def __repr__(self): return f"DensityMatrix({self.num_qubits} qubits, purity={self.purity():.4f})"

class Pauli:
    MATRICES = {
        "I": np.eye(2, dtype=complex),
        "X": np.array([[0,1],[1,0]], dtype=complex),
        "Y": np.array([[0,-1j],[1j,0]]),
        "Z": np.array([[1,0],[0,-1]], dtype=complex),
    }
    def __init__(self, label):
        self.label = label
        mat = np.array([[1.0+0j]])
        for c in label:
            mat = np.kron(mat, self.MATRICES[c])
        self.matrix = mat
    def expectation_value(self, state):
        if isinstance(state, Statevector):
            return float(np.real(state.data.conj() @ self.matrix @ state.data))
        elif isinstance(state, DensityMatrix):
            return float(np.real(np.trace(self.matrix @ state.data)))
    def __repr__(self): return f"Pauli({self.label})"

def partial_trace(rho, n, keep_qubits):
    """Partial trace over qubits NOT in keep_qubits."""
    trace_qubits = [q for q in range(n) if q not in keep_qubits]
    dm = rho.reshape([2]*(2*n))
    for q in sorted(trace_qubits, reverse=True):
        dm = np.trace(dm, axis1=q, axis2=q+n)
        n -= 1
    k = len(keep_qubits)
    return dm.reshape(2**k, 2**k)

def state_fidelity(a, b):
    """Compute state fidelity between two Statevectors or DensityMatrices."""
    if isinstance(a, Statevector) and isinstance(b, Statevector):
        return a.fidelity(b)
    if isinstance(a, DensityMatrix) and isinstance(b, DensityMatrix):
        return a.fidelity(b)
    raise TypeError("Both must be Statevector or DensityMatrix")

def concurrence(rho):
    """Concurrence of a 2-qubit density matrix (entanglement measure)."""
    Y = np.array([[0,-1j],[1j,0]])
    YY = np.kron(Y, Y)
    rho_tilde = YY @ rho.conj() @ YY
    R = rho @ rho_tilde
    eigvals = np.sqrt(np.sort(np.abs(np.linalg.eigvals(R)))[::-1].astype(complex).real)
    return float(max(0, eigvals[0] - eigvals[1] - eigvals[2] - eigvals[3]))
