"""
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 alpha values

IGNORE THE CONSTANT C

NOTE there is bug in the program, You should fix it if you want to use the visualization, it relates to the alpha range and you should not plot the prob of 0
"""

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

# Set dimensions
db = 16
da_values = [8, 16, 32]
alpha_range = np.linspace(0, 2, 100)  # Range of alpha values to plot
n_samples = 100000

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

for da in tqdm(da_values, desc="Processing d_A values"):
    # Calculate beta according to the formula
    beta = da / (np.log(2) * db)
    
    # Calculate probability for each alpha
    predicted_probabilities = []
    actual_probabilities = []
    for alpha in tqdm(alpha_range, desc=f"Calculating probabilities for d_A={da}", leave=False):
        # 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)
        actual_probabilities.append(np.sum(entropies > np.log2(da) - alpha - beta) / n_samples)
    
    # plt.plot(alpha_range, predicted_probabilities, label=f'$d_A={da}$', linestyle='--')
    plt.plot(alpha_range, actual_probabilities, label=f'$d_A={da}$', linestyle='-')

plt.xlabel(r'$\alpha$')
plt.ylabel('Probability')
plt.title(r'$\operatorname{Pr}[H(\psi_A) <\log_2(d_A)-\alpha-\beta]$ vs $\alpha$ for different $d_A$')
plt.legend()
plt.grid(True)
plt.yscale('log')  # Use log scale for better visualization
plt.show()