import time, random, os, sys
from queue import Queue
from shutil import copyfile
from subprocess import call
from threading import Thread
from model import num_of_GPU_cores, tryparameters
import datetime as dt

def bootstrap_the_experiment(resume=False, exp_id_toresume=None):
    def create_experiment_folder():
        while True:
            date = dt.datetime.now()
            exp_id = str(date.day) + "." + str(date.month) + "." + str(date.year) + "-" + str(date.hour) + "." + str(
                date.minute) + "." + str(date.second)
            path_to_exp_folder = os.path.join(os.getcwd(), 'Exp/Exp-{0}'.format(exp_id))
            if not os.path.exists(path_to_exp_folder):
                os.makedirs(path_to_exp_folder)
                return path_to_exp_folder, exp_id


    def copy_resume_code_to_root(path_to_code):
        print('To resume the training of older experiment, the original model code needs to be copied to root of the project.'
              'Any current "model.py" is backed up in "old_model.py"')
        path_to_root_code = os.path.join(os.getcwd(), '1000_model_code.py')
        os.rename(path_to_root_code, os.path.join(os.getcwd(), 'old_model.py'))
        copyfile(path_to_code, path_to_root_code)

        with open(path_to_root_code, 'r+') as file:
            lines = file.readlines()
            lines = lines[:-3]
            file.writelines(lines)

        with open(path_to_root_code, 'w') as file:
            [file.write(line) for line in lines]


    def model_runner(q, GPU_core):
        while True:
            param_set_id, start_from_epoch = q.get()
            q.task_done()

            if sys.platform == 'linux':
                os.environ["THEANO_FLAGS"] = "device=gpu{0}".format(GPU_core)

            elif sys.platform == 'win32':
                import theano.sandbox.cuda
                theano.sandbox.cuda.use('gpu{0}'.format(GPU_core))
            else:
                raise EnvironmentError('It seems your platform {0} is not yet supported. Please consider contributing.'.format(sys.platform))

            code = call(["python",
                         "experimentNew.py",
                         "{}".format(path_to_exp_folder),
                         "{}".format(exp_id),
                         "{}".format(param_set_id),
                         "{}".format(GPU_core),
                         "{}".format(start_from_epoch)])
            exitcodes.append(code)


    exp_start_time = time.time()
    num_threads = 1
    num_param_sets = len(tryparameters)
    q = Queue(maxsize=0)
    exitcodes = []

    if resume:
        exp_id = exp_id_toresume
        path_to_exp_folder = os.path.join(os.getcwd(), 'Exp', 'Exp-{0}'.format(exp_id))

        already_finished = [x for x in os.listdir(path_to_exp_folder) if '_' in x]
        already_finished_indices = [int(x[-4:])-1000 for x in already_finished]
        already_failed_indices = [int(x[-4:])-1000 for x in os.listdir(path_to_exp_folder) if len(os.listdir(os.path.join(path_to_exp_folder, x))) == 0]
        already_started = [x for x in os.listdir(path_to_exp_folder) if '_' not in x and len(os.listdir(os.path.join(path_to_exp_folder, x))) > 0]
        already_started_indices = [int(x[-4:])-1000 for x in already_started]

        to_restart = []
        for modelname in already_started:
            last_models = [x for x in os.listdir(os.path.join(path_to_exp_folder, modelname)) if 'h5' in x and 'best' not in x]
            idx = int(modelname[-4:])-1000
            if last_models:
                last_model_h5 = last_models[0]
                last_finished_epoch = int(last_model_h5[9:-3])
                to_restart.append((idx, last_finished_epoch + 1))

        if len(already_started) > 0:
            model_code_path = os.path.join(path_to_exp_folder, already_started[0],
                                           '1000_model_code.py'.format(already_started[0][-4:]))
        elif len(already_finished) > 0:
            model_code_path = os.path.join(path_to_exp_folder, already_finished[0],
                                           '1000_model_code.py'.format(already_finished[0][-4:]))
        else:
            raise EnvironmentError('It seems there is no model code from which to resume the experiment. '
                                   '\nCheck if you did not delete xxxx_exp_code.py from Model folders.')

        copy_resume_code_to_root(model_code_path)

    else:
        path_to_exp_folder, exp_id = create_experiment_folder()
        already_finished_indices, already_failed_indices, already_started_indices, to_restart = [], [], [], []


    for GPU_core in range(num_threads):
        worker = Thread(target=model_runner, args=(q, GPU_core))
        worker.setDaemon(True)
        worker.start()

    param_set_indices = [(x, 0) for x in range(num_param_sets)
                         if x not in already_finished_indices
                         and x not in already_failed_indices
                         and x not in already_started_indices]

    param_set_indices = to_restart + param_set_indices
    for param_set_id in param_set_indices:
        q.put(param_set_id)

    q.join()

    while (len(exitcodes) < len(param_set_indices)):
        time.sleep(1)

    num_successfull = exitcodes.count(0)
    num_failed = len(exitcodes) - num_successfull
    exp_elapsed_time_minutes = (time.time() - exp_start_time) / 60

    def brief_report():
        return '{5}\nALL {0} EXPERIMENTS DONE IN {1:.2f} MINUTES ON {2} GPU CORES\n{3} SUCCESSFUL | {4} FAILED\n{5}'\
            .format(num_param_sets, exp_elapsed_time_minutes, num_of_GPU_cores, num_successfull, num_failed, '#' * 80)

    print('\033[92m' + brief_report() + '\033[0m')

    with open(os.path.join(path_to_exp_folder, 'brief_report.txt'), 'a+') as file:
        file.write(brief_report())

if __name__ == "__main__":
    bootstrap_the_experiment()