import os
import random
import shutil
from plantcv.plantcv import fatal_error


def sample_images(source_path, dest_path, num=100):
    if not os.path.exists(source_path):
        raise IOError(f"Directory does not exist: {source_path}")

    if not os.path.exists(dest_path):
        os.makedirs(dest_path)  # exist_ok argument does not exist in python 2

    img_element_array = []
    sample_array = []
    num_images = []
    img_extensions = ['.png', '.jpg', '.jpeg', '.tif', '.tiff', '.gif']

    # If SnapshotInfo exists then need to make a new csv for the random image sample
    if os.path.exists(os.path.join(source_path, 'SnapshotInfo.csv')):

        line_array = []
        input_csv = open(os.path.join(source_path, 'SnapshotInfo.csv'))
        header = input_csv.readline()
        for line in input_csv:
            line = line.rstrip("\n")
            element_arr = line.split(',')
            if element_arr[-1]:
                line_array.append(element_arr)
        input_csv.close()

        # Check to make sure number of imgs to select is less than number of images found
        if num > len(line_array):
            fatal_error(f"Number of images found ({len(line_array)}) less than 'num'.")

        for i in range(0, num):
            r = random.randint(0, len(line_array) - 1)
            while r in num_images:
                r = random.randint(0, len(line_array) - 1)
            sample_array.append(line_array[r])
            num_images.append(r)

        out_file = open(os.path.join(dest_path, 'SnapshotInfo.csv'), 'w')
        out_file.write(header)
        for element in sample_array:
            out_file.write(','.join(element))
            snap_path = os.path.join(source_path, "snapshot" + element[1])
            folder_path = os.path.join(dest_path, "snapshot" + element[1])
            if not os.path.exists(folder_path):
                os.mkdir(folder_path)  # the beginning of folder_path (dest_path) already exists from above
            for root, dirs, files in os.walk(snap_path):
                for file in files:
                    shutil.copy(os.path.join(root, file), folder_path)
    else:
        for root, dirs, files in os.walk(source_path):
            for file in files:
                # Check file type so that only images get copied over
                name, ext = os.path.splitext(file)
                if ext.lower() in img_extensions:
                    img_element_array.append(os.path.join(root, file))

        # Check to make sure number of imgs to select is less than number of images found
        if num > len(img_element_array):
            fatal_error(f"Number of images found ({len(img_element_array)}) less than 'num'.")

        # Get random images
        for i in range(0, num):
            r = random.randint(0, len(img_element_array) - 1)
            while r in num_images:
                r = random.randint(0, len(img_element_array) - 1)
            sample_array.append(img_element_array[r])
            num_images.append(r)

        # Copy images over to destination
        for element in sample_array:
            shutil.copy(element, dest_path)

@@ -381,50 +381,8 @@
  %5B%5D%0A
-    sample_array = %5B%5D%0A    num_images = %5B%5D%0A
     
@@ -1181,261 +1181,34 @@
     
-for i in range(0, num):%0A            r = random.randint(0, len(line_array) - 1)%0A            while r in num_images:%0A                r = random.randint(0, len(line_array) - 1)%0A            sample_array.append(line_array%5Br%5D)%0A            num_images.append(r)%0A
+# Create SnapshotInfo file
 %0A   
@@ -1309,24 +1309,25 @@
 header)%0A
+%0A
         
 for elem
@@ -1322,35 +1322,129 @@
     
-for element in sample_array
+# Get random snapshots%0A        random_index = random.sample(range(0, len(line_array)), num)%0A        for i in random_index
 :%0A  
@@ -1477,23 +1477,29 @@
 ,'.join(
-element
+line_array%5Bi%5D
 ))%0A     
@@ -1548,39 +1548,45 @@
 h, %22snapshot%22 + 
-element
+line_array%5Bi%5D
 %5B1%5D)%0A           
@@ -1637,23 +1637,29 @@
 shot%22 + 
-element
+line_array%5Bi%5D
 %5B1%5D)%0A   
@@ -2580,282 +2580,79 @@
     
-for i in range(0, num):%0A            r = random.randint(0, len(img_element_array) - 1)%0A            while r in num_images:%0A                r = random.randint(0, len(img_element_array) - 1)%0A            sample_array.append(img_element_array%5Br%5D)%0A            num_images.append(r)%0A
+random_index = random.sample(range(0, len(img_element_array) - 1), num)
 %0A   
@@ -2706,31 +2706,25 @@
 for 
-element in sample_array
+i in random_index
 :%0A  
@@ -2745,23 +2745,36 @@
 il.copy(
+img_
 element
+_array%5Bi%5D
 , dest_p

Use random sample method