# -- Markdown Cell --
# Gigel created the validation dataset with help from an LLM. However, before that he tried doing it on his own (but he got bored). This was his progress:

# -- Code Cell --
from PIL import Image
import os
img = Image.open("training_data/train_img_01.jpg")
tiles_folder = "train_tiles"

tile = img.crop((0, 0, 128, 128)) #get a tile from image
tile_path = os.path.join(tiles_folder, "img1_tile0.jpg") 
tile.save(tile_path) #save to path

tile = img.crop((128, 0, 256, 128)) #once more
tile_path = os.path.join(tiles_folder, "img1_tile1.jpg")
tile.save(tile_path)

# -- Code Cell --
tile1 = [] #name of first tile
tile2 = [] #name of second tile
answer = [] #labels
tile1.append("img1_tile0.jpg") 
tile2.append("img1_tile1.jpg")
answer.append(1) #tile1 to left of tile2
tile1.append("img1_tile1.jpg")
tile2.append("img1_tile0.jpg")
answer.append(3) #tile1 to left of tile2


# -- Markdown Cell --
# Bam! now you have your brand new dataset with a whopping 2 pieces of data!
# 
# Also here is the incredible resnet Gigel heard about!

# -- Code Cell --
from torchvision import models
resnet = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) 
print(resnet)

# -- Markdown Cell --
# If for whatever reason, you don't want some parameters of the model to change you can do this

# -- Code Cell --
#freeze all parameters
for param in resnet.parameters():
    param.requires_grad = False 

#if for whatever reason you want the parameters of layer1, block 0 first conv  (this is purely for example purposes, it was chosen randomly)
#to be unfreezed (model can learn and change these parameters) this is how you would do it

for param in resnet.layer1[0].conv1.parameters():
  param.requires_grad = True

# -- Markdown Cell --
# Gigel also provided you with a way to generate the output DF!

# -- Code Cell --
import pandas as pd
test_df = pd.read_csv("test/test_pairs.csv")
#len(test_df)
subtaskID = [1] * len(test_df) + [2] * len(test_df)
datapointID = test_df['datapointID'].to_list() * 2

#REPLACE THIS WITH ACTUAL ANSWERS FOR SUBTASK 1
subtask1Answers = [0] * len(test_df)
#REPLACE THIS WITH ACTUAL ANSWERS FOR SUBTASK 2
subtask2Answers = [0] * len(test_df)


answer = subtask1Answers + subtask2Answers

output_df = pd.DataFrame({
    "subtaskID": subtaskID, 
    "datapointID": datapointID, 
    "answer": answer
})
output_df.to_csv("output.csv")