"""
Script to be used for tiny-imagenet formatting before launching experiments
"""

import os
import argparse


def create_val_folder(data_set_path):
    """
    Used for Tiny-imagenet dataset
    Copied from https://github.com/soumendukrg/BME595_DeepLearning/blob/master/Homework-06/train.py
    This method is responsible for separating validation images into separate sub folders,
    so that test and val data can be read by the pytorch dataloaders
    """
    path = os.path.join(data_set_path, 'val/images')  # path where validation data is present now
    filename = os.path.join(
        data_set_path, 'val/val_annotations.txt'
    )  # file where image2class mapping is present
    fp = open(filename, "r")
    data = fp.readlines()

    # Create a dictionary with image names as key and corresponding classes as values
    val_img_dict = {}
    for line in data:
        words = line.split("\t")
        val_img_dict[words[0]] = words[1]
    fp.close()

    # Create folder if not present, and move image into proper folder
    for img, folder in val_img_dict.items():
        newpath = (os.path.join(path, folder))
        if not os.path.exists(newpath):
            os.makedirs(newpath)

        if os.path.exists(os.path.join(path, img)):
            os.rename(os.path.join(path, img), os.path.join(newpath, img))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataplace", "-dp", type=str, default=None, help="Parent folder to data")
    args = parser.parse_args()
    create_val_folder(os.path.join(args.dataplace, 'tinyimagenet200-data'))

if __name__ == "__main__":
    main()
