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

r = 3.56994567
alpha_min, alpha_max = 2.4, 2.8
highest_exponent = 8
fig, ax = plt.subplots(figsize=(30, 10))

frames = 12 * 3
alphas = alpha_max + (alpha_min - alpha_max) * np.linspace(0, 1, frames)

for i, alpha in tqdm(enumerate(alphas)):
    fig, ax = plt.subplots(figsize=(30, 10))

    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(0.5+(x_values-0.5)*(-alpha)**(n), 0.5+(y_values-0.5)*(-alpha)**(n), label=f'iteration = {2**n}', 
                alpha=0.1, color='black', linewidth=4)

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

        l_x = 5
        ax.set_xlim(0.5-l_x, 0.5+l_x)
        ax.set_ylim(-6, 4)

    fig.suptitle(f'Up to {2**(highest_exponent-1)} iterates of the Logistic map with r = {r}, each iterate scaled up by α = {alpha}',y=1.0)
    plt.tight_layout()

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

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