import numpy as np
import imageio
import torch

def convert_grayscale_to_rgb(images):
    return np.repeat(images, 3, axis=1)

def save_gif(imgs, output_path, args):
    imgs = imgs * 255.0
    imgs = imgs.detach().to('cpu', torch.uint8).numpy()

    if imgs.shape[1] == 1:
        imgs = convert_grayscale_to_rgb(imgs)

    # Reshape A to be suitable for image processing:
    # Change A from (100, 3, 64, 64) to (100, 64, 64, 3) for imageio compatibility
    imgs = imgs.transpose((0, 2, 3, 1))

    # Assuming 10 groups with 10 frames each
    num_groups = args.latent_dim
    frames_per_group = args.interval

    # Create the GIF, combining frames horizontally within each group

    # output_path = '/content/drive/MyDrive/CoLab/composite_output.gif'
    with imageio.get_writer(output_path, mode='I', duration=0.5) as writer:
        for i in range(frames_per_group):
            # Extract one frame from each group to form a composite image
            frames = [imgs[j * frames_per_group + i] for j in range(num_groups)]
            # Concatenate frames horizontally
            composite_frame = np.concatenate(frames, axis=1)  # Concatenate along the width
            writer.append_data(composite_frame)

    print(f"GIF saved to {output_path}")

