from evaluator.graphon_evaluator import DiscretizedGraphonEvaluatorFinite
from games.graphons import er_graphon, ranked_attachment_graphon, uniform_attachment_graphon
from games.finite.investment import InvestmentGraphon
from games.finite.sis import SISGraphon
from simulator.graphon_simulator import DiscretizedGraphonExactSimulatorFinite
from solver.graphon_solver import DiscretizedGraphonExactSolverFinite
from solver.policy.random_policy import RandomFinitePolicy

for graphon in [uniform_attachment_graphon, ranked_attachment_graphon, er_graphon]:
    for g in [SISGraphon, InvestmentGraphon]:
        game = g(graphon=graphon)

        simulator = DiscretizedGraphonExactSimulatorFinite()
        eval_solver = DiscretizedGraphonExactSolverFinite(eta=0, verbose=True)
        evaluator = DiscretizedGraphonEvaluatorFinite()

        policies = []
        mus = []
        logs = []

        """ Initialize with random policy """
        policy = RandomFinitePolicy(game.agent_observation_space, game.agent_action_space)
        mu, _ = simulator.simulate(game, policy)

        """ Evaluation of the policy under its induced mean field """
        best_response, info = eval_solver.solve(game, mu)
        eval_results_pi = evaluator.evaluate(game, mu, policy)
        eval_results_opt = evaluator.evaluate(game, mu, best_response)
        print("{} {} policy {} optimal {} expl {} "
              .format(g.__name__, graphon.__name__,
                      eval_results_pi["eval_mean_returns"], eval_results_opt["eval_mean_returns"],
                      eval_results_opt["eval_mean_returns"] - eval_results_pi["eval_mean_returns"]),
              flush=True)
