import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Pre-generate a large set of uniform random data to be transformed
n_samples = 50000
uniform_data = np.random.uniform(0, 1, size=n_samples)

# Set up the figure and axis for plotting
fig, ax = plt.subplots(figsize=(8, 6))

# Set up the initial plot
line_cdf, = ax.plot([], [], label='CDF', color='blue')
line_quantile, = ax.plot([], [], label='Quantile function', color='orange')

# Labels and title
ax.legend()
ax.grid(True)

# Set axis limits for data and probability in [0, 1]
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)

# Animation function
def animate(beta_param):
    """Update the plot for each frame with different beta distribution shape parameters."""
    alpha = 2  # Keeping alpha constant for symmetry
    beta = beta_param  # Varying beta for the animation
    
    # Transform the pre-generated uniform data using the beta distribution's inverse CDF
    data = np.random.beta(alpha, beta, size=n_samples)
    
    # Sort the data for CDF and quantile function calculation
    sorted_data = np.sort(data)
    
    # Calculate the CDF values (empirical CDF)
    cdf_values = np.arange(1, n_samples + 1) / n_samples
    
    # Quantile function (inverse CDF)
    p_values = np.linspace(0, 1, n_samples)
    quantile_values = np.percentile(sorted_data, 100 * p_values)
    
    # Update CDF line (data on x-axis, probabilities on y-axis)
    line_cdf.set_data(sorted_data, cdf_values)
    
    # Update Quantile Function line (flipped: probabilities on x-axis, data on y-axis)
    line_quantile.set_data(p_values, quantile_values)
    
    return line_cdf, line_quantile

# List of beta parameters for animation
beta_params = np.linspace(0.5, 5, 100)

# Create animation with a longer interval for smoother transitions
ani = FuncAnimation(fig, animate, frames=beta_params, interval=200, blit=True)

if True:
    ani.save('cumulative-quantile.mp4', writer='ffmpeg', fps=10)

# Show the animation
plt.show()