# -- Code Cell --
import os
import pandas as pd
from PIL import Image

# This should point to the folder that contains the "train" and "test" folders and "xrays.png"
PATH_TO_STARTER_KIT = '.' # Change this if you saved your starter kit somewhere else

# Starting code for the first subtask
xrays_task1 = Image.open(os.path.join(PATH_TO_STARTER_KIT, "xrays.png"))

"""
Solve the first subtask
"""

subtask1_sol= [{
    'subtaskID':1,
    'datapointID':1,
    'answer': 0 # Replace this with the calculated value for subtask 1
}]
df_sub1 = pd.DataFrame.from_dict(subtask1_sol)

# -- Code Cell --
# Starting code for the second subtask
df_train = pd.read_csv(os.path.join(PATH_TO_STARTER_KIT, 'train_data.csv'))
df_test = pd.read_csv(os.path.join(PATH_TO_STARTER_KIT, 'test_data.csv'))

train_folder_path = os.path.join(PATH_TO_STARTER_KIT, "train")
test_folder_path = os.path.join(PATH_TO_STARTER_KIT, "test")

# Reading images. Processing a single image at a time to not overload the RAM
for train_image_name in df_train['file_name']:
    train_image_path = os.path.join(train_folder_path, train_image_name)
    train_image = Image.open(train_image_path)
    
    # Do what you need with an image from train data...
    
    train_image.close()
    
for test_image_name in df_test['file_name']:
    test_image_path = os.path.join(test_folder_path, test_image_name)
    test_image = Image.open(test_image_path)
    
    # Do what you need with an image from test data...
    
    test_image.close()

# -- Code Cell --
"""
Solve the second subtask
"""

# Replace this with your predictions on the test data 
# Make sure to have them in the same order as in test_data.csv
y_test = ['unknown'] * 300 
df_ans = df_test.copy()
df_ans['subtaskID']=2
df_ans['answer']=y_test
df_ans = df_ans.drop(columns=['file_name', 'label'])
df_final = pd.concat([df_sub1, df_ans], axis=0)
df_final.to_csv('submission.csv', index=False)