from bayes_opt import BayesianOptimization, JSONLogger, Events

from libs.cell_graph_dispatcher import CellGraphDispatcher
from libs.cell_graph_validator import CellGraphValidator


import numpy as np



N = 15
seed = 42



np.random.seed(seed)
width = np.random.randint(20, 150, (N,))
height = np.random.randint(20, 150, (N,))
nr_train = np.random.randint(50, 200, (N,))
n_cities = np.random.randint(2, 35, (N,))
grid_distribution_of_cities = False
max_rails_between_cities = np.random.randint(2, 4, (N,))
max_rail_in_city = np.random.randint(3, 6, (N,))
malfunction_rate = np.random.randint(500, 4000, (N,))
#???
prop_malfunction = np.random.uniform(0.01, 0.01, (N,))
min_duration = np.random.randint(20, 80, (N,))
max_duration = np.random.randint(20, 80, (N,))
max_duration = np.maximum(min_duration, max_duration)

speed_ration_map = {1.: 0.25,  # Fast passenger train
                    1. / 2.: 0.25,  # Fast freight train
                    1. / 3.: 0.25,  # Slow commuter train
                    1. / 4.: 0.25}  # Slow freight train


#make half with square sizes + last one is 150x150x200 trains
width[N-1] = 150
height[N//2:] = width[N//2:]
nr_train[N-1] = 200

test = {
    "width": width,
    "height": height,
    "trains": nr_train,
    "seed": seed,
    "cities": n_cities,
    "rails_between_cities": max_rails_between_cities,
    "rails_in_city": max_rail_in_city,
    "malfunction_rate": malfunction_rate,
    "prop_malfunction": prop_malfunction,
    "min_prop": min_duration,
    "max_prop": max_duration
}

#Run


def flatland_function(speed_coef, time_coef):

    def get_dispatcher(env):
        def get_sort_function(dispatcher: CellGraphDispatcher):
            def sort(idx):
                time = dispatcher.controllers[idx].dist_to_target[
                    dispatcher.graph._vertex_idx_from_point(env.agents[idx].initial_position), env.agents[
                        idx].initial_direction]
                speed = env.agents[idx].speed_data['speed']

                return speed * speed_coef + time * time_coef

            return sort

        return CellGraphDispatcher(env, sort_function=get_sort_function)

    res = CellGraphValidator.multiple_tests(get_dispatcher, **test)
    return res["finished"]


pbounds = {'speed_coef' : (-10000, 10000), 'time_coef': (-1, 1)}

optimizer = BayesianOptimization(
    f=flatland_function,
    pbounds=pbounds,
    random_state=seed,
)

logger = JSONLogger(path="./opt_log.json")
optimizer.subscribe(Events.OPTMIZATION_STEP, logger)

optimizer.probe({'speed_coef' : -10000, 'time_coef': 1})
optimizer.probe({'speed_coef' : -10000, 'time_coef': -1})
optimizer.probe({'speed_coef' : +10000, 'time_coef': 1})
optimizer.maximize(init_points=10, n_iter=100)


print(optimizer.max)