import os
import sys
import json
import argparse
import numpy as np
import ConfigSpace

from robo.fmin import bayesian_optimization

from tabular_benchmarks.wrn_cifar10_3h import WRNCIFAR103HBenchmark
from tabular_benchmarks.fcnet_year_prediction import FCNetYearPredictionBenchmark
from tabular_benchmarks.fcnet_benchmark import FCNetBenchmark

parser = argparse.ArgumentParser()
parser.add_argument('--run_id', default=0, type=int, nargs='?', help='unique number to identify this run')
parser.add_argument('--benchmark', default="wrn_cifar10", type=str, nargs='?', help='specifies the benchmark')
parser.add_argument('--n_iters', default=100, type=int, nargs='?', help='number of iterations for optimization method')
parser.add_argument('--output_path', default="./", type=str, nargs='?',
                    help='specifies the path where the results will be saved')
parser.add_argument('--data_dir', default="./", type=str, nargs='?', help='specifies the path to the tabular data')

args = parser.parse_args()

if args.benchmark == "wrn_cifar10":
    b = WRNCIFAR103HBenchmark(data_dir=args.data_dir)

elif args.benchmark == "fcnet_regression":
    b = FCNetYearPredictionBenchmark(data_dir=args.data_dir)

elif args.benchmark == "protein_structure":
    b = FCNetBenchmark(dataset=args.data_dir)

elif args.benchmark == "slice_localization":
    b = FCNetBenchmark(dataset=args.data_dir)

output_path = os.path.join(args.output_path, "rf_bo")
os.makedirs(os.path.join(output_path), exist_ok=True)

cs = b.get_configuration_space()


def obj_wrapper(x):
    config = dict()

    for i, h in enumerate(cs.get_hyperparameters()):

        if type(h) == ConfigSpace.hyperparameters.OrdinalHyperparameter:
            config[h.name] = h.sequence[int(round(x[i]))]
        else:
            config[h.name] = h.choices[int(round(x[i]))]

    y, c = b.objective_function(config)

    return y


lower = []
upper = []

for h in cs.get_hyperparameters():
    lower.append(0)
    if type(h) == ConfigSpace.hyperparameters.OrdinalHyperparameter:
        upper.append(len(h.sequence) - 1)
    else:
        upper.append(len(h.choices) - 1)

results = bayesian_optimization(obj_wrapper,
                                np.array(lower),
                                np.array(upper),
                                model_type="rf",
                                num_iterations=args.n_iters)

res = b.get_results()

fh = open(os.path.join(output_path, 'run_%d.json' % args.run_id), 'w')
json.dump(res, fh)
fh.close()

add rf based bo