import torch
import pickle


def test_agent(env, Q, device, n_test_runs, test_log_file, n_steps,
               visualize=False):
    """
        env: Environment
        Q: Q-network
        device: cpu or gpu
        n_test_runs: Number of episodes for which performance is tested
        test_log_file: Path to log file for test results
        n_steps: Number of steps for which the DQN has currently been trained
        visualize: Flag for rendering environment
    """
    ep_returns = []
    for _ in range(n_test_runs):
        s = env.reset()
        s = torch.from_numpy(s).float().to(device)
        done = False
        ep_return = 0

        while not done:
            q = Q(s)
            a = torch.argmax(q)
            s, r, done, _ = env.step(a)
            ep_return += r
            s = torch.from_numpy(s).float().to(device)

        ep_returns.append(ep_return)

    logfile = open(test_log_file, 'a+b')
    pickle.dump({'steps': n_steps, 'returns': ep_returns}, logfile)
    logfile.close()

    print('Steps:', n_steps,
          'Avg. return:', sum(ep_returns) / n_test_runs,
          'Episodic return:', ep_returns)

    return ep_returns