import os
from PIL import Image
import torch
import shutil

def resize_images(input_folder, output_folder, target_size=(32, 32)):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # Search all the files
    for root, dirs, files in os.walk(input_folder):
        for filename in files:
            if filename.endswith('.png'):
                file_path = os.path.join(root, filename)
                
                relative_path = os.path.relpath(file_path, input_folder)
                output_path = os.path.join(output_folder, relative_path)
                
                output_dir = os.path.dirname(output_path)
                if not os.path.exists(output_dir):
                    os.makedirs(output_dir)
                
                with Image.open(file_path) as img:
                    img_resized = img.resize(target_size, Image.LANCZOS)

                    # Save reseized images
                    img_resized.save(output_path)
                    print(f'Saved resized image to {output_path}')
                    
def select_and_copy_images(tensor_path, folder_path, output_folder):
    # Load the .pt file
    indices = torch.load(tensor_path)
    
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    
    # Iterate through all the files in the folder
    for filename in os.listdir(folder_path):
        # Check if the file is a .png file
        if filename.endswith('.png'):
            # Extract the number from the filename (assuming the number is in the fifth position)
            parts = filename.split('_')
            if len(parts) >= 5:
                file_number = int(parts[4].split('.')[0])  # Extract the number
                
                # Check if the number is in the tensor
                if file_number in indices.tolist():
                    # Copy the matching file to the output folder
                    shutil.copy(os.path.join(folder_path, filename), output_folder)
                    
                    # If you want to move the files instead of copying, use the following line:
                    # shutil.move(os.path.join(folder_path, filename), output_folder)
    
    print(f"Extraction complete, files have been copied to {output_folder}")
    
if __name__ == '__main__':
    import os
    import sys
    org_img_path = sys.argv[1]
    resized_img_path = sys.argv[2]
    target_size = sys.argv[3]
    tensor_path = sys.argv[4]
    mixed_path = sys.argv[5]

    resize_images(org_img_path, resized_img_path, target_size)
    select_and_copy_images(tensor_path, resized_img_path, mixed_path)
