#room4v3
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from smac.env.multiagentenv import MultiAgentEnv

import atexit
from operator import attrgetter
from copy import deepcopy
import numpy as np
import enum
import math
from absl import logging
import random

def generate(d):
    ret=np.random.random(d)
    l=(ret**2).sum()**0.5
    while l>1:
        ret=np.random.random(d)
        l=(ret**2).sum()**0.5
    return ret/l
def generate_apart(vec,d):
    ret=np.random.random(d)
    l=(ret**2).sum()**0.5
    while l>1 or (ret/l).dot(vec)>0.6:
        ret=np.random.random(d)
        l=(ret**2).sum()**0.5
    return ret/l

class Toy1Env(MultiAgentEnv):
    """The StarCraft II environment for decentralised multi-agent
    micromanagement scenarios.
    """
    def __init__(
            self,
            delay=5,
            reward_win=10,
            obs_last_action=False,
            state_last_action=True,
            is_print=False,
            print_rew=False,
            print_steps=1000,
            seed=None
    ):
        # Map arguments
        self.print_rew = print_rew
        self.is_print = is_print
        self.print_steps = print_steps
        self._seed = random.randint(0, 9999)
        np.random.seed(self._seed)
        self.n_agents = 2

        # Observations and state
        self.obs_last_action = obs_last_action
        self.state_last_action = state_last_action

        # Rewards args
        self.reward_win = reward_win

        # Other
        self._seed = seed

        # Actions
        self.n_actions = 3

        # Statistics
        self._episode_count = 0
        self._episode_steps = 0
        self._total_steps = 0
        self.battles_won = 0
        self.battles_game = 0

        self.p_step = 0
        self.rew_gather = []
        self.is_print_once = False

        self.last_action = np.zeros((self.n_agents, self.n_actions))

        self.episode_limit = delay

        self.dim=3

        self.true_embed=generate(self.dim)

        # initialize agents
        if random.random()<0.5:
            self.state_n = [generate_apart(self.true_embed,self.dim)]
            self.ans=1
        else:
            self.state_n =[self.true_embed.copy()]
            self.ans=0
        self.state_n += [generate_apart(self.true_embed,self.dim)]

    def step(self, actions):
        """Returns reward, terminated, info."""
        self._total_steps += 1
        self._episode_steps += 1
        info = {}

        '''if self.is_print:
            if self._episode_steps==1:
                print('goal',self.goal_loc_1,self.goal_loc_2)
            print(self._episode_steps,self.state_n.reshape(-1),actions,self.reach_1,self.reach_2)'''

        reward=0
        terminated = False
        info['battle_won'] = False

        self.state_n = [generate_apart(self.true_embed,self.dim) for _ in range(2)]

        if self._episode_steps >= self.episode_limit:
            terminated = True
            if actions[0]==self.ans:
                reward=self.reward_win
                self.battles_won+=1
                info['battle_won']=True

        if terminated:
            self._episode_count += 1
            self.battles_game += 1

        if self.print_rew:
            self.p_step += 1
            if terminated:
                self.rew_gather.append(reward)
            if self.p_step % self.print_steps == 0:
                print('steps: %d, average rew: %.3lf' % (self.p_step,
                                                         float(np.mean(self.rew_gather)) / self.reward_win))
                self.is_print_once = True

        return reward, terminated, info

    def get_obs(self):
        """Returns all agent observations in a list."""
        return [self.get_obs_agent(i) for i in range(self.n_agents)]

    def get_obs_agent(self, agent_id):
        """Returns observation for agent_id."""
        return np.array(self.state_n[agent_id])

    def get_obs_size(self):
        """Returns the size of the observation."""
        return self.dim

    def get_state(self):
        """Returns the global state."""
        return np.concatenate(self.state_n)#np.concatenate([self.state_n.reshape(-1),self.goal_loc.reshape(-1)])

    def get_state_size(self):
        """Returns the size of the global state."""
        return 2*self.dim

    def get_avail_actions(self):
        """Returns the available actions of all agents in a list."""
        return [self.get_avail_agent_actions(i) for i in range(self.n_agents)]

    def get_avail_agent_actions(self, agent_id):
        """Returns the available actions for agent_id."""
        return [1] * 2

    def get_total_actions(self):
        """Returns the total number of actions an agent could ever take."""
        return 2

    def reset(self):
        """Returns initial observations and states."""
        self._episode_steps = 0
        self.last_action = np.zeros((self.n_agents, self.n_actions))

        if random.random()<0.5:
            self.state_n = [generate_apart(self.true_embed,self.dim)]
            self.ans=1
        else:
            self.state_n =[self.true_embed.copy()]
            self.ans=0
        self.state_n += [generate_apart(self.true_embed,self.dim)]

        self.is_print=False
        if random.random()<0.001:
            self.is_print=True

        return self.get_obs(), self.get_state()

    def render(self):
        pass

    def close(self):
        pass

    def seed(self):
        pass

    def save_replay(self):
        """Save a replay."""
        pass

    def get_env_info(self):
        env_info = {"state_shape": self.get_state_size(),
                    "obs_shape": self.get_obs_size(),
                    "n_actions": self.get_total_actions(),
                    "n_agents": self.n_agents,
                    "episode_limit": self.episode_limit}
        return env_info

    def get_stats(self):
        stats = {
            "battles_won": self.battles_won,
            "battles_game": self.battles_game,
            "win_rate": self.battles_won / self.battles_game
        }
        return stats

    def clean(self):
        self.p_step = 0
        self.rew_gather = []
        self.is_print_once = False
