import os
import json

import hpbandster.core.result as hpres

import ConfigSpace as CS
import ConfigSpace.hyperparameters as CSH
from ConfigSpace.read_and_write import json as config_space_json_r_w


import argparse
parser = argparse.ArgumentParser(description='Bobh_Eval')
parser.add_argument('--dir_log_files', type=str, default='bohb_logs/darts_1/')
parser.add_argument('--dir_configs', type=str, default='configs/darts_1_configspace.json')
parser.add_argument('--dir_configs_new', type=str, default='configs/darts_1_configspace.json')
parser.add_argument('--only_eval', action='store_true', help='Run only eval.')
parser.add_argument('--BSDS', action='store_true', help='Run on BSDS.')
args = parser.parse_args()


def main(args):
    # Load configspace
    config_path = args.dir_configs
    cs = config_space_json_r_w.read(open(config_path, 'r').read())

    # load the example run from the log files\n",
    result = hpres.logged_results_to_HBS_result(args.dir_log_files)

    # # get all executed runs\n",
    # all_runs = result.get_all_runs()

    # get the 'dict' that translates config ids to the actual configurations\n",
    id2conf = result.get_id2config_mapping()

    # Here is how you get he incumbent (best configuration)\n",
    inc_id = result.get_incumbent_id()

    # let's grab the run on the highest budget\n",
    inc_runs = result.get_runs_by_id(inc_id)
    inc_run = inc_runs[-1]

    # We have access to all information: the config, the loss observed during\n",
    # optimization, and all the additional information\n",
    if "args.BSDS":
        inc_config = id2conf[inc_id]['config']
        inc_val_metric = inc_run.info['loss'][0]
        inc_train_loss = inc_run.info['Train_Loss'][0]
        inc_val_loss = inc_run.info['Validation_Loss'][0]
        inc_trunc_loss = inc_run.info['Trunc_Loss'][0]
        inc_train_metric = inc_run.info['Train_Metric'][0]
        inc_trunc_metric = inc_run.info['Trunc_Metric'][0]

        print('Best found configuration:')
        print(inc_config)
        print(f'It achieved a validation PSNR of {inc_val_metric}, training PSNR of {inc_train_metric},'
              f'search training loss of {inc_train_loss}, search validation loss of {inc_val_loss}.')
        print('budget', inc_run.budget)

        config_dict = {
            'configs': inc_config,
            'budget': inc_run.budget,
            'val_metric': inc_val_metric,
            'train_metric': inc_train_metric,
            'train_loss': inc_train_loss,
            'val_loss': inc_val_loss,
            'trunc_loss': inc_trunc_loss,
            'trunc_metric': inc_trunc_metric
        }

    else:
        inc_config = id2conf[inc_id]['config']
        inc_eval_val_loss = inc_run.info['loss'][0]
        # inc_eval_train_loss = inc_run.info['Arch_Train'][0]
        inc_eval_train_loss = inc_run.info['Arch_Train'][0]
        inc_val_loss = inc_run.info['Validation_Loss'][0]
        inc_train_loss = inc_run.info['Train_Loss'][0]
        inc_test_trunc_loss = inc_run.info['Trunc_Loss'][0]

        print('Best found configuration:')
        print(inc_config)
        print(f'It achieved a retrained val loss of {inc_eval_val_loss}, retrained train loss of {inc_eval_train_loss}, '
              f'search training loss of {inc_train_loss}, search validation loss of {inc_val_loss}.')
        print('budget', inc_run.budget)

        config_dict = {
            'configs': inc_config,
            'budget': inc_run.budget,
            'eval_val_loss': inc_eval_val_loss,
            'eval_train_loss': inc_eval_train_loss,
            'train_loss': inc_train_loss,
            'val_loss': inc_val_loss,
            'trunc_loss': inc_test_trunc_loss
        }

    with open(os.path.join(args.dir_log_files, 'results.txt'), 'w') as file:
        json.dump(str(config_dict), file)
        file.write('\n')

    # Overwrite Configspace
    if not args.only_eval:

        config = CS.ConfigurationSpace()
        config.add_hyperparameter(CSH.CategoricalHyperparameter(name=cs.get_hyperparameter('alpha_scheduler').name,
                                                                choices=cs.get_hyperparameter(
                                                                    'alpha_scheduler').choices,
                                                                default_value=inc_config['alpha_scheduler']))
        config.add_hyperparameter(CSH.CategoricalHyperparameter(name=cs.get_hyperparameter('alpha_optimizer').name,
                                                                choices=cs.get_hyperparameter(
                                                                    'alpha_optimizer').choices,
                                                                default_value=inc_config['alpha_optimizer']))
        config.add_hyperparameter(CSH.CategoricalHyperparameter(name=cs.get_hyperparameter('alpha_warmup').name,
                                                                choices=cs.get_hyperparameter('alpha_warmup').choices,
                                                                default_value=inc_config['alpha_warmup']))
        config.add_hyperparameter(CSH.CategoricalHyperparameter(name=cs.get_hyperparameter('param_warmup').name,
                                                                choices=cs.get_hyperparameter('param_warmup').choices,
                                                                default_value=inc_config['param_warmup']))

        config.add_hyperparameter(CSH.UniformFloatHyperparameter(name=cs.get_hyperparameter('alpha_lr').name,
                                                                 lower=cs.get_hyperparameter('alpha_lr').lower,
                                                                 upper=cs.get_hyperparameter('alpha_lr').upper,
                                                                 default_value=inc_config['alpha_lr'],
                                                                 log=cs.get_hyperparameter('alpha_lr').log))
        config.add_hyperparameter(CSH.UniformFloatHyperparameter(name=cs.get_hyperparameter('param_lr').name,
                                                                 lower=cs.get_hyperparameter('param_lr').lower,
                                                                 upper=cs.get_hyperparameter('param_lr').upper,
                                                                 default_value=inc_config['param_lr'],
                                                                 log=cs.get_hyperparameter('param_lr').log))
        config.add_hyperparameter(CSH.UniformFloatHyperparameter(name=cs.get_hyperparameter('param_weight_decay').name,
                                                                 lower=cs.get_hyperparameter(
                                                                     'param_weight_decay').lower,
                                                                 upper=cs.get_hyperparameter(
                                                                     'param_weight_decay').upper,
                                                                 default_value=inc_config['param_weight_decay'],
                                                                 log=cs.get_hyperparameter('param_weight_decay').log))
        config.add_hyperparameter(CSH.UniformFloatHyperparameter(name=cs.get_hyperparameter('alpha_weight_decay').name,
                                                                 lower=cs.get_hyperparameter(
                                                                     'alpha_weight_decay').lower,
                                                                 upper=cs.get_hyperparameter(
                                                                     'alpha_weight_decay').upper,
                                                                 default_value=inc_config['alpha_weight_decay'],
                                                                 log=cs.get_hyperparameter('alpha_weight_decay').log))

        config.add_hyperparameter(CSH.Constant(name='epochs', value=cs.get_hyperparameter('epochs').value))

        with open(args.dir_configs_new, 'w') as fh:
            fh.write(config_space_json_r_w.write(config))

        print('Saved new configs to {}'.format(args.dir_configs_new))


if __name__ == '__main__':
    main(args)
