from mountaincar import MountainCar
from gp_gym_info import info
import gym

# GP Parameters
info["env_name"] = "MountainCar-v0"
info["pop_size"] = 100
info["max_gens"] = 10
info["max_depth"] = 1
info["num_eps"] = 100

agent = MountainCar(info)
best_program = agent.train()
print(best_program)
f = agent.fit(best_program, 100, 200, render=False)
print(f)
info["env_name"] = "MountainCarContinuous-v0"
info["pop_size"] = 100
info["max_gens"] = 10
info["max_depth"] = 1
info["tournament_size"] = 5
info["num_eps"] = 10
agent = MountainCar(info)

solutions = {}
force_values = np.arange(0.0, 1.0, 0.1)
fitness_scores = []
counter = 1
for force in force_values:
    solution = "IFLTE(0.0, velocity, {}, {})".format(force, -force)
    f = agent.fit(solution, 100, 200, render=False)[0]
    fitness_scores.append(f)

    # Timing
    print(counter)
    counter += 1

    if f >= 90:
        solutions[solution] = f

# Print solutions and their fitness scores
print()
for s, f in solutions.items():
    print("{}: {}".format(s, f))

# Plot fitness scores