import numpy as np
import matplotlib.pyplot as plt
from quantum_states import sample_and_calculate
from tqdm import tqdm

# Define range of dimensions to test
fixed_dim = 64
dimensions = np.arange(2, 64, 2)  # Test dimensions from 2 to 50 in steps of 2
expected_entropies = []
theoretical_entropies = []
predicted_entropies = []

# Calculate entropies for each dimension
for dim in tqdm(dimensions, desc="Calculating entropies"):
    # For each dimension, we'll keep one subsystem fixed at dim=2 
    # and vary the other dimension
    entropies = sample_and_calculate(dim, fixed_dim, n_samples=1000)
    expected_entropies.append(np.mean(entropies))
    theoretical_entropies.append(np.log2(min(dim, fixed_dim)))
    beta = min(dim, fixed_dim)/(2*np.log(2)*max(dim, fixed_dim))
    predicted_entropies.append(np.log2(min(dim, fixed_dim)) - beta)

# Create the plot
plt.figure(figsize=(10, 6))
plt.plot(dimensions, expected_entropies, 'b-', label='Expected Entropy')
plt.plot(dimensions, theoretical_entropies, 'r--', label='Theoretical Entropy')
plt.plot(dimensions, predicted_entropies, 'g--', label='Predicted Entropy')
plt.xlabel('Dimension of Subsystem B')
plt.ylabel('von Neumann Entropy (bits)')
plt.title(f'von Neumann Entropy vs. System Dimension, with Dimension of Subsystem A = {fixed_dim}')
plt.legend()
plt.grid(True)
plt.show()