diff --git a/NYT.py b/NYT.py
index 1d6f4c8..3f3c8f5 100644
--- a/NYT.py
+++ b/NYT.py
@@ -4,45 +4,40 @@ import json
 import numpy as np
 import time
 from pprint import pprint
-
-url = "https://www.nytimes.com/puzzles/sudoku/hard"
-
-page = requests.get(url)
-soup = BeautifulSoup(page.content, "html.parser") 
-table = soup.find_all("div", {"class":"pz-game-screen"})[0]
-script = table.find_all("script", {"type":"text/javascript"})
-text = script[0].get_text().split("=")[1]
-Dict = json.loads(text)
-grid = np.array( Dict["hard"]["puzzle_data"]["puzzle"] ).reshape(9,9).tolist()
-
-hard_puzzle_id = Dict["hard"]["puzzle_id"]
-hard_published = Dict["hard"]["published"]
-hard_print_date =  Dict["hard"]["print_date"]
-hard = "".join(str(x) for x in Dict["hard"]["puzzle_data"]["puzzle"])
-med = "".join(str(x) for x in Dict["medium"]["puzzle_data"]["puzzle"])
-easy = "".join(str(x) for x in Dict["easy"]["puzzle_data"]["puzzle"])
-
-
-# Import writer class from csv module
 from csv import writer
- 
-# List that we want to add as a new row
-HardList = ["Hard",Dict["hard"]["puzzle_id"],Dict["hard"]["published"],Dict["hard"]["print_date"],hard]
-MedList = ["Medium",Dict["medium"]["puzzle_id"],Dict["medium"]["published"],Dict["medium"]["print_date"],hard]
-EasyList = ["Easy",Dict["easy"]["puzzle_id"],Dict["easy"]["published"],Dict["easy"]["print_date"],hard]
 
+def get_puzzle(difficulty):
+     url = "https://www.nytimes.com/puzzles/sudoku/hard"
+     page = requests.get(url)
+     soup = BeautifulSoup(page.content, "html.parser") 
+     table = soup.find_all("div", {"class":"pz-game-screen"})[0]
+     script = table.find_all("script", {"type":"text/javascript"})
+     text = script[0].get_text().split("=")[1]
+     Dict = json.loads(text)
+     grid = np.array( Dict[difficulty]["puzzle_data"]["puzzle"] ).reshape(9,9).tolist()
 
-# Open our existing CSV file in append mode
-# Create a file object for this file
-with open('NYTSudoku.csv', 'a') as f_object:
-     # Pass this file object to csv.writer()
-    # and get a writer object
-    writer_object = writer(f_object)
-     # Pass the list as an argument into
-    # the writerow()
-    writer_object.writerow(HardList)
-    writer_object.writerow(MedList)
-    writer_object.writerow(EasyList)
- 
-    # Close the file object
-    f_object.close()
+     puzzle_id = Dict[difficulty]["puzzle_id"]
+     published = Dict[difficulty]["published"]
+     print_date =  Dict[difficulty]["print_date"]
+     puzzle_data = "".join(str(x) for x in Dict[difficulty]["puzzle_data"]["puzzle"])
+
+     # List that we want to add as a new row
+     csvrow = [difficulty,puzzle_id,published,print_date,puzzle_data]
+  
+     # Open our existing CSV file in append mode
+     # Create a file object for this file
+     with open('NYTSudoku.csv', 'a') as f_object:
+          # Pass this file object to csv.writer()
+     # and get a writer object
+     writer_object = writer(f_object)
+          # Pass the list as an argument into
+     # the writerow()
+     writer_object.writerow(csvrow)
+
+     # Close the file object
+     f_object.close()
+
+if __name__ == '__main__':
+     get_puzzle("hard")
+     get_puzzle("medium")
+     get_puzzle("easy")
\ No newline at end of file