from cProfile import label
import os
import json
from PIL import Image
from tqdm import tqdm
import shutil
import numpy as np

import sys
sys.path.append("..")

from data_generate.split_generator import SplitGenerator

if __name__ == "__main__":
    # load the config file
    config_name = 'prompt.json'
    jsonfile = open(os.path.join('../config/5-shot', config_name))
    config = json.loads(jsonfile.read())

    dest_dir = os.path.join(os.path.join(config['data_dir'], '5-shot'), 'cifar_fs')
    label_dir = os.path.join(os.path.join(os.path.join(config['data_dir'], 'raw'), 'cifar100_raw'), "variant.txt")
    raw_data_dir = os.path.join(os.path.join(config['data_dir'], 'raw'), 'cifar100_raw')


    if os.path.exists(dest_dir):
        # clean the previous split
        shutil.rmtree(dest_dir)
    os.makedirs(dest_dir)

    metatrain = [os.path.join(os.path.join(raw_data_dir, "data"), line.rstrip('\n'))
                 for line in open(os.path.join(raw_data_dir, 'splits', 'bertinetto', 'train.txt'), 'r')]
    metaval = [os.path.join(os.path.join(raw_data_dir, "data"), line.rstrip('\n'))
               for line in open(os.path.join(raw_data_dir, 'splits', 'bertinetto', 'val.txt'), 'r')]
    metatest = [os.path.join(os.path.join(raw_data_dir, "data"), line.rstrip('\n'))
                for line in open(os.path.join(raw_data_dir, 'splits', 'bertinetto', 'test.txt'), 'r')]
    
    np.save(os.path.join(dest_dir, 'metatrain.npy'), metatrain)
    np.save(os.path.join(dest_dir, 'metaval.npy'), metaval)
    np.save(os.path.join(dest_dir, 'metatest.npy'), metatest)