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

# Define range of dimensions to test
dimensionsA = np.arange(2, 64, 2)  # Test dimensions from 2 to 50 in steps of 2
dimensionsB = np.arange(2, 64, 2)  # Test dimensions from 2 to 50 in steps of 2

# Create meshgrid for 3D plot
X, Y = np.meshgrid(dimensionsA, dimensionsB)
Z = np.zeros_like(X, dtype=float)

# Calculate entropies for each dimension combination
total_iterations = len(dimensionsA) * len(dimensionsB)
pbar = tqdm(total=total_iterations, desc="Calculating entropies")

for i, dim_a in enumerate(dimensionsA):
    for j, dim_b in enumerate(dimensionsB):
        entropies = sample_and_calculate(dim_a, dim_b, n_samples=100)
        Z[j,i] = np.mean(entropies)
        pbar.update(1)
pbar.close()

# Create the 3D plot
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')

# Plot the surface
surf = ax.plot_surface(X, Y, Z, cmap='viridis')

# Add labels and title with larger font sizes
ax.set_xlabel('Dimension of Subsystem A', fontsize=12, labelpad=10)
ax.set_ylabel('Dimension of Subsystem B', fontsize=12, labelpad=10)
ax.set_zlabel('von Neumann Entropy (bits)', fontsize=12, labelpad=10)
ax.set_title('von Neumann Entropy vs. System Dimensions', fontsize=14, pad=20)

# Add colorbar
cbar = fig.colorbar(surf, ax=ax, label='Entropy')
cbar.ax.set_ylabel('Entropy', fontsize=12)

# Add tick labels with larger font size
ax.tick_params(axis='x', labelsize=10)
ax.tick_params(axis='y', labelsize=10)
ax.tick_params(axis='z', labelsize=10)

# Rotate the plot for better visibility
ax.view_init(elev=30, azim=45)

plt.show()