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 = 1+np.sqrt(8), 3.8494344
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 = 8
    fig, ax = plt.subplots(figsize=(30, 10))

    x_min, x_max = 0.46, 0.54
    x_values = np.linspace(x_min, x_max, 10000)
    y_values = x_values
    last_n_iter = 0
    for n in range(highest_exponent):
        n_iter = 3 * 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, x_min, x_max, linewidth=0.3, alpha=0.8, linestyle='--', color='black')
    # ax.legend()
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(0.45, 0.53)
    fig.suptitle(f'Up to {last_n_iter} iterates of the Logistic map with r = {r}',y=1.0)
    plt.tight_layout()

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

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