# This script creates and submits job files to apply the YOLOv5x models to their validation sets
import os
import time

applied_models = {'test_jpg_unaug_full_fold_1', 'test_jpg_aug_full_fold_1', 'test_jpg_aug_backbone_fold_1', 'test_jpg_unaug_backbone_fold_1', 'test_jpg_unaug_last_layer_fold_1', 'test_jpg_unaug_backbone_fold_2', 'test_jpg_aug_last_layer_fold_1', 'test_jpg_aug_backbone_fold_2', 'test_jpg_unaug_full_fold_2', 'test_jpg_aug_full_fold_2', 'test_jpg_unaug_last_layer_fold_2', 'test_jpg_aug_last_layer_fold_2', 'test_jpg_unaug_full_fold_3', 'test_jpg_aug_full_fold_3', 'test_jpg_aug_backbone_fold_3', 'test_jpg_unaug_backbone_fold_3', 'test_jpg_unaug_last_layer_fold_3', 'test_jpg_unaug_full_fold_4', 'test_jpg_aug_full_fold_4', 'test_jpg_aug_last_layer_fold_3', 'test_jpg_unaug_backbone_fold_4', 'test_jpg_aug_backbone_fold_4', 'test_jpg_unaug_last_layer_fold_4', 'test_jpg_aug_last_layer_fold_4', 'test_jpg_aug_backbone_fold_5', 'test_jpg_unaug_full_fold_5', 'test_jpg_aug_full_fold_5', 'test_jpg_unaug_backbone_fold_5', 'test_jpg_aug_backbone_fold_6', 'test_jpg_unaug_full_fold_6', 'test_jpg_aug_last_layer_fold_5', 'test_jpg_aug_full_fold_6', 'test_jpg_unaug_backbone_fold_6', 'test_jpg_unaug_last_layer_fold_5', 'test_jpg_unaug_last_layer_fold_6', 'test_jpg_aug_backbone_fold_7', 'test_jpg_aug_last_layer_fold_6', 'test_jpg_unaug_full_fold_7', 'test_jpg_unaug_backbone_fold_7', 'test_jpg_aug_full_fold_7', 'test_jpg_aug_backbone_fold_8', 'test_jpg_aug_last_layer_fold_7', 'test_jpg_unaug_full_fold_8', 'test_jpg_aug_full_fold_8', 'test_jpg_unaug_backbone_fold_8', 'test_jpg_unaug_last_layer_fold_7', 'test_jpg_aug_backbone_fold_9', 'test_jpg_aug_last_layer_fold_8', 'test_jpg_unaug_full_fold_9', 'test_jpg_unaug_backbone_fold_9', 'test_jpg_aug_full_fold_9', 'test_jpg_unaug_last_layer_fold_8', 'test_jpg_unaug_last_layer_fold_9', 'test_jpg_aug_full_fold_10', 'test_jpg_aug_last_layer_fold_9', 'test_jpg_unaug_full_fold_10', 'test_jpg_unaug_backbone_fold_10', 'test_jpg_aug_backbone_fold_10', 'test_jpg_unaug_last_layer_fold_10', 'test_jpg_aug_last_layer_fold_10', 'test_jpg_aug_full_fold_1', 'test_jpg_unaug_backbone_fold_1', 'test_jpg_unaug_full_fold_1', 'test_jpg_aug_backbone_fold_1', 'test_jpg_unaug_last_layer_fold_1', 'test_jpg_unaug_full_fold_2', 'test_jpg_unaug_backbone_fold_2', 'test_jpg_aug_backbone_fold_2', 'test_jpg_aug_full_fold_2', 'test_jpg_aug_last_layer_fold_1', 'test_jpg_unaug_last_layer_fold_2', 'test_jpg_unaug_full_fold_3', 'test_jpg_unaug_backbone_fold_3', 'test_jpg_aug_full_fold_3', 'test_jpg_aug_last_layer_fold_2', 'test_jpg_aug_backbone_fold_3', 'test_jpg_unaug_last_layer_fold_3', 'test_jpg_unaug_backbone_fold_4', 'test_jpg_aug_backbone_fold_4', 'test_jpg_aug_last_layer_fold_3', 'test_jpg_unaug_full_fold_4', 'test_jpg_aug_full_fold_4', 'test_jpg_unaug_last_layer_fold_4', 'test_jpg_unaug_full_fold_5', 'test_jpg_aug_last_layer_fold_4', 'test_jpg_unaug_backbone_fold_5', 'test_jpg_aug_full_fold_5', 'test_jpg_aug_backbone_fold_5', 'test_jpg_unaug_backbone_fold_6', 'test_jpg_aug_backbone_fold_6', 'test_jpg_unaug_last_layer_fold_5', 'test_jpg_aug_last_layer_fold_5', 'test_jpg_unaug_full_fold_6', 'test_jpg_aug_full_fold_6', 'test_jpg_unaug_full_fold_7', 'test_jpg_unaug_last_layer_fold_6', 'test_jpg_aug_full_fold_7', 'test_jpg_aug_last_layer_fold_6', 'test_jpg_unaug_backbone_fold_7', 'test_jpg_aug_backbone_fold_7', 'test_jpg_unaug_backbone_fold_8', 'test_jpg_unaug_full_fold_8', 'test_jpg_unaug_last_layer_fold_7', 'test_jpg_aug_backbone_fold_8', 'test_jpg_aug_last_layer_fold_7', 'test_jpg_aug_full_fold_8', 'test_jpg_unaug_full_fold_9', 'test_jpg_unaug_last_layer_fold_8', 'test_jpg_aug_last_layer_fold_8', 'test_jpg_aug_full_fold_9', 'test_jpg_unaug_backbone_fold_9', 'test_jpg_aug_backbone_fold_9', 'test_jpg_unaug_last_layer_fold_9', 'test_jpg_aug_last_layer_fold_9', 'test_jpg_unaug_backbone_fold_10', 'test_jpg_unaug_full_fold_10', 'test_jpg_aug_backbone_fold_10', 'test_jpg_aug_full_fold_10', 'test_jpg_unaug_last_layer_fold_10', 'test_jpg_aug_last_layer_fold_10', 'test_jpg_unaug_full_fold_1', 'test_jpg_aug_full_fold_1', 'test_jpg_aug_last_layer_fold_1', 'test_jpg_unaug_backbone_fold_1', 'test_jpg_unaug_full_fold_2', 'test_jpg_aug_backbone_fold_1', 'test_jpg_unaug_last_layer_fold_1', 'test_jpg_aug_backbone_fold_2', 'test_jpg_aug_full_fold_2', 'test_jpg_unaug_last_layer_fold_2', 'test_jpg_unaug_backbone_fold_2', 'test_jpg_aug_last_layer_fold_2', 'test_jpg_unaug_full_fold_3', 'test_jpg_aug_last_layer_fold_3', 'test_jpg_aug_full_fold_3', 'test_jpg_unaug_backbone_fold_3', 'test_jpg_aug_backbone_fold_3', 'test_jpg_unaug_last_layer_fold_3', 'test_jpg_unaug_full_fold_4', 'test_jpg_aug_backbone_fold_4', 'test_jpg_aug_full_fold_4', 'test_jpg_unaug_backbone_fold_4', 'test_jpg_unaug_last_layer_fold_4', 'test_jpg_aug_last_layer_fold_4', 'test_jpg_unaug_full_fold_5', 'test_jpg_aug_full_fold_5', 'test_jpg_aug_last_layer_fold_5', 'test_jpg_unaug_backbone_fold_5', 'test_jpg_unaug_full_fold_6', 'test_jpg_aug_backbone_fold_5', 'test_jpg_unaug_last_layer_fold_5', 'test_jpg_unaug_last_layer_fold_6', 'test_jpg_aug_full_fold_6', 'test_jpg_aug_backbone_fold_6', 'test_jpg_unaug_backbone_fold_6', 'test_jpg_aug_last_layer_fold_6', 'test_jpg_unaug_full_fold_7', 'test_jpg_aug_full_fold_7', 'test_jpg_aug_last_layer_fold_7', 'test_jpg_unaug_full_fold_8', 'test_jpg_unaug_backbone_fold_7', 'test_jpg_unaug_last_layer_fold_7', 'test_jpg_aug_backbone_fold_7', 'test_jpg_aug_full_fold_8', 'test_jpg_aug_backbone_fold_8', 'test_jpg_unaug_last_layer_fold_8', 'test_jpg_unaug_backbone_fold_8', 'test_jpg_aug_last_layer_fold_8', 'test_jpg_unaug_full_fold_9', 'test_jpg_aug_last_layer_fold_9', 'test_jpg_aug_full_fold_9', 'test_jpg_unaug_backbone_fold_9', 'test_jpg_aug_backbone_fold_9', 'test_jpg_unaug_full_fold_10', 'test_jpg_unaug_last_layer_fold_9'}

# For each of the 10 folds
for fold in range(1, 11):
    # Determine device (i.e., CPU) to run on. Each node has two CPUs, so we alternate between 0 and 1 to get the most resources possible
    device = 0 if fold % 2 == 0 else 1

    # Freeze none of the model, backbone, and all but last layer
    for freeze in [0, 10, 24]:
        # Use augmented and unaugmented models
        for mode in ['unaug', 'aug']:
            # Create fold name 
            fold_name = f'augmented_fold_{fold}' if mode == 'aug' else f'fold_{fold}'

            # Create filenames and fold names
            if freeze == 0:
                file_name = f"{mode}_full_fold_{fold}.job"
            elif freeze == 10:
                file_name = f"{mode}_backbone_fold_{fold}.job"
            elif freeze == 24:
                file_name = f"{mode}_last_layer_fold_{fold}.job"

            wall_time = 1

            # Skip already trained models
            if f'test_jpg_{file_name[:-4]}' in applied_models:
                continue

            # Create job file
            with open(file_name, 'w') as output:
                # Specify resources and other arguments for cluster
                output.write(f"#!/bin/bash\n#SBATCH --time={wall_time}:00:00\n#SBATCH --nodes=1 --ntasks-per-node=1\n#SBATCH --partition=gpu\n#SBATCH --mail-type=BEGIN,END\n#SBATCH --job-name=inference\n\n")

                # Load python
                output.write("module load anaconda-python3\n")
                # Switch to model directory
                output.write("cd $HOME/yolov5\n") 
                # Activate environment
                output.write('source /software/python/anaconda3/etc/profile.d/conda.sh\n')

                # Apply model
                output.write(f'python3 detect.py --exist-ok --device 0 --nosave --save-txt --save-conf --img 640 --weights /home/username/yolov5/runs/train/{file_name[:-4]}/weights/best.pt --source /home/username/obj_Test_data/test_images_jpg --name test_jpg_{file_name[:-4]}')

            # Submit job file
            os.system(f"sbatch {file_name}")
            time.sleep(5)