import os
import pickle
import imageio
import matplotlib.pyplot as plt
from qdhf_things import run_qdhf, many_pictures

EXAMPLE_PROMPTS = [
    'an image of a cat on the sofa',
    'an image of a bear in a national park',
    'a photo of an astronaut riding a horse on mars',
    'a drawing of a tree behind a fence',
    'a painting of a sunset over the ocean',
    'a sketch of a racoon sitting on a mushroom',
    'a picture of a dragon flying over a castle',
    'a photo of a robot playing the guitar',
]

if __name__ == '__main__':
    print('Hello! I am a script!')

    for i, example_prompt in enumerate(EXAMPLE_PROMPTS):
        # Initialize list to store images for GIF
        images = []

        # Run QDHF
        for archive, plt in run_qdhf(example_prompt):
            # Save current plot to a temporary file
            temp_filename = f'./examples/temp_plot_{i}.png'
            plt.savefig(temp_filename)
            plt.close()

            # Read the saved image and append to images list
            images.append(imageio.imread(temp_filename))
            os.remove(temp_filename)

        # Create a GIF from the images
        gif_filename = f'./examples/archive_{i}.gif'
        imageio.mimsave(gif_filename, images, duration=0.5)  # Adjust duration as needed

        # Save archive with pickle
        pickle.dump(archive, open(f'./examples/archive_{i}.pkl', 'wb'))

        # Save the final archive plot
        plt = many_pictures(archive, example_prompt)
        plt.savefig(f'./examples/archive_pics_{i}.png')
        plt.close()