import random
import os
from pathlib import Path
import re

DEBUG = False

PERCENT_OF_SEARCH_SPACE_TO_SEARCH = 1.0
MAX_NUM_COMBINATIONS = 50

CV = 2
hyperparams = {'data': ['har_large'], # 'sim_easy', 'sim_hard'
               'ds_factor': [50],
               'epochs': [200],
               'lr': [0.001, 0.002, 0.0006],
               'z_dim':[1, 2, 4, 8],
               'rnn_dim': [16, 32],
               'transition_dim': [16],
               'n_cats': [20]}

# Best cpap
hyperparams = {'data': ['cpap'],
               'ds_factor': [50],
               'epochs': [200],
               'lr': [0.001],
               'z_dim':[2],
               'rnn_dim': [32],
               'transition_dim': [16],
               'n_cats': [20]}

combos_tried = [] # Hyper paramter combinations tried
possible_num_combos = 1


for key in hyperparams:
    if type(hyperparams[key]) == list:
        possible_num_combos *= len(hyperparams[key])


while True:
    hyper_param_dict = {}
    for key in hyperparams:
        if type(hyperparams[key]) == list:
            choice = random.choice(hyperparams[key])
            if type(choice) == tuple:
                hyper_param_dict[key.split(',')[0]] = choice[0]
                hyper_param_dict[key.split(',')[1]] = choice[1]
            else:
                hyper_param_dict[key] = choice
        else:
            hyper_param_dict[key] = hyperparams[key]
    

    if hyper_param_dict not in combos_tried:
        job_file = 'hyper_param_opt.sh'
        if job_file in os.listdir():
            os.remove(job_file) # Remove it so we start with an empty file
        
        dataset_chosen = hyper_param_dict['data']
        long_name = ''
        for hp in hyper_param_dict:
            long_name += hp + '_' + str(hyper_param_dict[hp]) + '_'
        long_name = long_name[:-1]
        
        if not os.path.exists(dataset_chosen):
            os.makedirs(dataset_chosen)

        file = Path(dataset_chosen + '/id_mapper.txt')
        file.touch(exist_ok=True)
        combo_exists = False
        short_name = ''
        max_short_name = 0
        with open(dataset_chosen + '/id_mapper.txt', 'r') as f:
            for line in f.readlines():
                if long_name in line:
                    combo_exists=True
                    short_name = line.split(':')[0]
                try:
                    max_short_name = max(max_short_name, int(line.split(':')[0]))
                except ValueError:
                    pass
        new_cv = False
        if combo_exists:
            cvs_used = [int(re.search('output_(\d+).txt', f).group(1)) for f in os.listdir(dataset_chosen + '/' + short_name) if re.match('output_\d+.txt', f)]
            if CV not in cvs_used:
                new_cv = True

        if not combo_exists or new_cv:
            if not combo_exists:
                short_name = '0'*(4-len(str(max_short_name+1))) + str(max_short_name+1)
                with open(dataset_chosen + '/id_mapper.txt', 'a') as f:
                    f.write('\n')
                    f.write(short_name + ':' + long_name)
        
        
            if not os.path.exists(dataset_chosen + '/' + short_name):
                os.makedirs(dataset_chosen + '/' + short_name)
            text_file = './' +  dataset_chosen + '/' + short_name + '/output_%d.txt'%CV
            
            
            with open(text_file, 'w') as f:
                pass # Simply create the file so that the text file is there even if the slurm job hasn't been queued yet.
        
            with open(job_file, 'w') as fh:
                fh.writelines("#!/bin/bash\n")
                fh.writelines("#SBATCH --job-name=%s\n" % short_name)
                fh.writelines("#SBATCH --output=%s\n" % (text_file))
                fh.writelines('#SBATCH --qos=nopreemption\n')
                fh.writelines("#SBATCH --partition='cpu'\n")
                fh.writelines("#SBATCH --mem=32G\n")
                fh.writelines("#SBATCH --open-mode=append\n")

                python_call = "python run_DMM.py "
                for key in hyper_param_dict:
                    s = '--'
                    s += key
                    s += ' '
                    s += str(hyper_param_dict[key])
                    s += ' '
                    python_call += s
                python_call += ' --short_name ' + short_name
                python_call += ' --cv ' + str(CV)

                fh.writelines('%s\n'%python_call)
            os.system("sbatch %s" %job_file)

            combos_tried.append(hyper_param_dict)
    if len(combos_tried) >= int(PERCENT_OF_SEARCH_SPACE_TO_SEARCH*possible_num_combos) or len(combos_tried) >= MAX_NUM_COMBINATIONS:
        break




