import numpy as np
import torch
from itertools import chain


class Evaluator(object):
    def __init__(self, environment, model, logger, agents, max_steps):
        self.env = environment
        self.model = model
        self.logger = logger
        self.agents = agents
        self.max_steps = max_steps

    def play_n_episodes(self, render=False, fixed_spawn=None, silent=False):
        """
        wraps play_one_episode, playing a single episode at a time and logs
        results used when playing demos.
        """
        if fixed_spawn is None:
            num_runs = 1 
            fixed_spawn = [None]
        else:
            # fixed_spawn should be, for example, [0.5 , 0.5 , 0.5, 0, 0, 0] for 2 runs
            # In the first run agents spawn in the middle and in the second they will spawn from the corner
            fixed_spawn = np.array(fixed_spawn).reshape((-1, 3)) # 3 dimensions
            num_runs = fixed_spawn.shape[0]
            # Set all the agents to the same spawn point
            fixed_spawn = np.stack([fixed_spawn for _ in range(self.agents)], axis=-1)

        num_files = self.env.files.num_files
        self.model.train(False)
        headers = ["number"] + list(chain.from_iterable(zip(
            [f"Filename {i}" for i in range(self.agents)],
            [f"Agent {i} pos x" for i in range(self.agents)],
            [f"Agent {i} pos y" for i in range(self.agents)],
            [f"Agent {i} pos z" for i in range(self.agents)],
            [f"Landmark {i} pos x" for i in range(self.agents)],
            [f"Landmark {i} pos y" for i in range(self.agents)],
            [f"Landmark {i} pos z" for i in range(self.agents)],
            [f"Distance {i}" for i in range(self.agents)])))
        self.logger.write_locations(headers)
        distances = []
        for j in range(num_runs):
            for k in range(num_files):
                score, start_dists, q_values, info = self.play_one_episode(render, fixed_spawn=fixed_spawn[j])
                row = [j * num_files + k + 1] + list(chain.from_iterable(zip(
                    [info[f"filename_{i}"] for i in range(self.agents)],
                    [info[f"agent_xpos_{i}"] for i in range(self.agents)],
                    [info[f"agent_ypos_{i}"] for i in range(self.agents)],
                    [info[f"agent_zpos_{i}"] for i in range(self.agents)],
                    [info.get(f"landmark_xpos_{i}", "N/A") for i in range(self.agents)],
                    [info.get(f"landmark_ypos_{i}", "N/A") for i in range(self.agents)],
                    [info.get(f"landmark_zpos_{i}", "N/A") for i in range(self.agents)],
                    [info.get(f"distError_{i}", "N/A") for i in range(self.agents)])))
                for i in range(self.agents):
                    key = f"distError_{i}"
                    if key in info:
                        distances.append(info[key])
                self.logger.write_locations(row)
        if len(distances) == 0:
            return None, None # No distance mean and std for task "play" as there is no ground truth
        mean = np.mean(distances, 0)
        std = np.std(distances, 0, ddof=1)
        if not silent:
            self.logger.log(f"mean distances {mean}")
            self.logger.log(f"Std distances {std}")
        return mean, std

    def play_one_episode(self, render=False, frame_history=4, fixed_spawn=None):
        device = next(self.model.parameters()).device
        def predict(obs_stack):
            """
            Run a full episode, mapping observation to action,
            using greedy policy.
            """
            inputs = torch.from_numpy(obs_stack).float().permute(
                0, 4, 1, 2, 3).unsqueeze(0).to(device)
            with torch.no_grad(): 
                q_vals = self.model(inputs)
            idx = torch.max(q_vals, -1)[1]
            greedy_steps = np.array(idx, dtype=np.int32).flatten()
            return greedy_steps, q_vals.detach().cpu().numpy().squeeze(0)

        obs_stack = self.env.reset(fixed_spawn)
        # Here obs have shape (agent, *image_size, frame_history)
        sum_r = np.zeros((self.agents))
        isOver = [False] * self.agents
        start_dists = None
        steps = 0
        while steps < self.max_steps and not np.all(isOver):
            acts, q_values = predict(obs_stack)
            obs_stack, r, isOver, info = self.env.step(acts, q_values, isOver)
            steps += 1
            # Machine-parseable trajectory log line (agent 0 only). Consumed
            # by app.py to render the agent's search path as an animation.
            print("STEP_LOC: {} {} {} {}".format(
                steps,
                info.get('agent_xpos_0', 'NA'),
                info.get('agent_ypos_0', 'NA'),
                info.get('agent_zpos_0', 'NA'),
            ))
            if start_dists is None:
                start_dists = [
                    info.get('distError_' + str(i), "N/A") for i in range(self.agents)]
            if render:
                self.env.render()
            for i in range(self.agents):
                if not isOver[i]:
                    sum_r[i] += r[i]
        return sum_r, start_dists, q_values, info