"""
plot the probability of the entropy of the reduced density matrix of the pure state being greater than log2(d_A) - alpha - beta

for different d_A values, with fixed alpha and d_B Note, d_B>d_A
"""

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

# Set dimensions
db = 32
alpha = 0
da_range = np.arange(2, 10, 1)  # Range of d_A values to plot
n_samples = 1000000

plt.figure(figsize=(10, 6))

predicted_probabilities = []
actual_probabilities = []

for da in tqdm(da_range, desc="Processing d_A values"):
    # Calculate beta according to the formula
    beta = da / (np.log(2) * db)
    
    # Calculate probability according to the formula
    # Ignoring constant C as requested
    prob = np.exp(-((da * db - 1) * alpha**2 / (np.log2(da)**2)))
    predicted_probabilities.append(prob)
    # Calculate actual probability
    entropies = sample_and_calculate(da, db, n_samples=n_samples)
    count = np.sum(entropies < np.log2(da) - alpha - beta)
    # early stop if count is 0
    if count != 0:
        actual_probabilities.append(count / n_samples)
    else:
        actual_probabilities.extend([np.nan] * (len(da_range) - len(actual_probabilities)))
        break
    # debug
    print(f'da={da}, theoretical_prob={prob}, threshold={np.log2(da) - alpha - beta}, actual_prob={actual_probabilities[-1]}, entropy_heads={entropies[:10]}')
    
# plt.plot(da_range, predicted_probabilities, label=f'$d_A={da}$', linestyle='--')
plt.plot(da_range, actual_probabilities, label=f'$d_A={da}$', linestyle='-')

plt.xlabel(r'$d_A$')
plt.ylabel('Probability')
plt.title(r'$\operatorname{Pr}[H(\psi_A) < \log_2(d_A)-\alpha-\beta]$ vs $d_A$ for fixed $\alpha=$'+str(alpha)+r' and $d_B=$' +str(db)+ r' with $n=$' +str(n_samples))
# plt.legend()
plt.grid(True)
plt.yscale('log')  # Use log scale for better visualization
plt.show()