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_min, r_max = 3, 3.56994567
delta = 4.6692
frames_per_delta = 24
total_deltas = 6
frames = frames_per_delta * total_deltas
scaling_factor = delta**(1/frames_per_delta)
rs = r_max + (r_min - r_max) / (scaling_factor ** np.array(range(frames)))

for i, r in tqdm(enumerate(rs)):
    highest_exponent = 9
    fig, ax = plt.subplots(figsize=(30, 6))

    x_values = np.linspace(0, 1, 100000)
    y_values = x_values
    last_n_iter = 0

    for n in range(highest_exponent):
        n_iter = 2**n
        y_values = function_iter(logistic_map, n_iter - last_n_iter, y_values, r=r)
        last_n_iter = n_iter

        ax.plot(x_values, y_values, label=f'iteration = {2**n}', alpha=0.2, color='blue')

    ax.hlines(0.5, 0, 1, linewidth=0.3, alpha=0.8, linestyle='--', color='black')

    fig.suptitle(f'Up to {2**(highest_exponent-1)} iterates of the Logistic map with r = {r}',y=1.0)
    plt.tight_layout()
    
    dir_path = "./logistic"
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

    fig.savefig(f"{dir_path}/{i}.png")
    plt.close()
export_to_video(["./logistic"], fps=12)