from qiskit import QuantumCircuit, execute, Aer
from qiskit.quantum_info import DensityMatrix, entropy
from qiskit.visualization import plot_histogram
import numpy as np

# Helper function to compute von Neumann entropy
def von_neumann_entropy(rho):
    return entropy(DensityMatrix(rho))

# Step 1: Create the entangled state (Bell state)
qc = QuantumCircuit(2)
qc.h(0)
qc.cx(0, 1)

# Simulate to get the density matrix of the initial entangled state
backend = Aer.get_backend('statevector_simulator')
result = execute(qc, backend).result()
statevector = result.get_statevector(qc)
rho_AB = DensityMatrix(statevector).data

# Calculate the initial entropy
S_rho_AB = von_neumann_entropy(rho_AB)
print(f"Initial entropy S(ρ_AB): {S_rho_AB}")

# Step 2: Apply local operations and calculate the averaged density matrix
pauli_operators = [np.array([[1, 0], [0, 1]]),  # I
                   np.array([[0, 1], [1, 0]]),  # X
                   np.array([[0, -1j], [1j, 0]]),  # Y
                   np.array([[1, 0], [0, -1]])]  # Z

rho_bar_AB = np.zeros((4, 4), dtype=complex)
for sigma in pauli_operators:
    U = np.kron(sigma, np.eye(2))
    rho_bar_AB += U @ rho_AB @ U.conj().T

rho_bar_AB /= 4

# Calculate the entropy of the averaged state
S_rho_bar_AB = von_neumann_entropy(rho_bar_AB)
print(f"Entropy of averaged state S(ρ̄_AB): {S_rho_bar_AB}")

# Step 3: Compute the Holevo quantity
holevo_quantity = S_rho_bar_AB - S_rho_AB
print(f"Holevo quantity χ(ρ_AB): {holevo_quantity}")

# Optionally visualize the density matrices
import matplotlib.pyplot as plt
from qiskit.visualization import plot_state_city

# Plot the density matrix of the original state
plot_state_city(DensityMatrix(rho_AB), title="Original State ρ_AB")
plt.show()

# Plot the density matrix of the averaged state
plot_state_city(DensityMatrix(rho_bar_AB), title="Averaged State ρ̄_AB")
plt.show()
