import numpy as np
import matplotlib.pyplot as plt

def logistic_map(x, r):
    return r * x * (1 - x)

def function_iter(f, n, x, **kwargs):
    for _ in range(n):
        x = f(x, **kwargs)
    return x

import imageio.v3 as iio
import os
from natsort import natsorted
import moviepy.editor as mp

def export_to_video(dir_paths, fps=12):
    for dir_path in dir_paths:
        file_names = natsorted((fn for fn in os.listdir(dir_path) if fn.endswith('.png')))

        # Create a list of image files and set the frame rate
        images = []

        # Iterate over the file names and append the images to the list
        for file_name in file_names:
            file_path = os.path.join(dir_path, file_name)
            images.append(iio.imread(file_path))

        filename = dir_path[2:]
        iio.imwrite(f"{filename}.gif", images, duration=1000/fps, rewind=True)
        clip = mp.ImageSequenceClip(images, fps=fps)
        clip.write_videofile(f"{filename}.mp4")
    return

from tqdm import tqdm
import os
r = 1 + np.sqrt(8)
x_mid = 0.4
frames=24 * 6

x_maxs = x_mid + np.exp(np.linspace(0, -5, frames)) * 0.4
x_mins = x_mid - np.exp(np.linspace(0, -5, frames)) * 0.4
for i in range(frames):
    x_min = x_mins[i]
    x_max = x_maxs[i]
    fig, ax = plt.subplots(figsize=(30, 5))

    x_values = np.linspace(x_min, x_max, 10000)
    y_values = x_values

    y_values = function_iter(logistic_map,300, y_values, r=r)
    ax.plot(x_values, y_values, alpha=0.2, color='red')
    ax.set_xlim(x_min, x_max)

    dir_path = "./logistic_2"
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

    fig.savefig(f"{dir_path}/{i}.png")
    plt.close()

export_to_video(["./logistic_2"], fps=12)