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

def takacs_function(x, rho=0.5, max_bits=20):
    """
    Compute the Takacs-type singular function for x in [0, 1].
    """
    temp = x
    one_positions = []  # Store positions of '1' bits
    
    for i in range(max_bits):
        temp *= 2
        bit = int(temp)
        if bit == 1:
            one_positions.append(i)
        temp -= bit
    
    # Compute the summation
    one_positions = np.array(one_positions)
    powers_of_rho = rho ** np.arange(len(one_positions))
    terms = (1 + rho) ** -one_positions
    f_val = np.sum(powers_of_rho * terms)
    return f_val

# Parameters
rho = 0.5
num_points = 2000
xs = np.linspace(0, 1, num_points)

# Define the range of bits to animate over
min_bits = 1
max_bits_final = 30
bits_range = range(min_bits, max_bits_final+1)

# Precompute all frames for efficiency
all_frames = []
for mb in bits_range:
    ys = np.array([takacs_function(x, rho=rho, max_bits=mb) for x in xs])
    all_frames.append(ys)

# Set up the figure and line object
fig, ax = plt.subplots(figsize=(10, 6))
line, = ax.plot([], [], color='blue', linewidth=1.5)
ax.set_title("Takacs-type Singular Function Approximation")
ax.set_xlabel("x")
ax.set_ylabel("f(x)")
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)  # The function is bounded between 0 and 1
ax.grid(True, linestyle='--', alpha=0.5)

# Initialization function for FuncAnimation
def init():
    line.set_data([], [])
    return line,

# Update function for FuncAnimation
def update(frame):
    # frame is an index into bits_range
    ys = all_frames[frame]
    current_bits = bits_range[frame]
    line.set_data(xs, ys)
    ax.set_title(f"Takacs-type Singular Function (max_bits = {current_bits})")
    return line,

# Create animation
ani = FuncAnimation(fig, update, frames=len(bits_range), init_func=init, blit=True, interval=200)

ani.save('takacs_function.mp4', writer='ffmpeg', fps=5)
plt.show()