import numpy as np
import math

from pysc2.lib import actions, features, units

from environment.starcraft_env import StarCraftEnv # environment

class DefeatZealots2v2(StarCraftEnv): 
    _PLAYER_RELATIVE = features.SCREEN_FEATURES.player_relative.index
    _PLAYER_SELF = 1
    _PLAYER_NEUTRAL = 3

    _NO_OP = actions.FUNCTIONS.no_op.id
    _MOVE_SCREEN = actions.FUNCTIONS.Move_screen.id
    _ATTACK_SCREEN = actions.FUNCTIONS.Attack_screen.id
    _SELECT_ARMY = actions.FUNCTIONS.select_army.id

    _SELECT_ALL = [0]
    _NOT_QUEUED = [0]
    _QUEUED = [1]

    _MOVE_VAL = 5.5
    _RADIO_VAL = 20
    _RANGE_VAL = 5

    _UP = 0
    _UP_LEFT = 1
    _LEFT = 2
    _DOWN_LEFT = 3
    _DOWN = 4
    _DOWN_RIGHT = 5
    _RIGHT = 6
    _UP_RIGHT = 7
    _ATTACK_CLOSEST = 8
    _ATTACK_LOWEST = 9

    possible_actions = [
        _UP,
        _UP_LEFT,
        _LEFT,
        _DOWN_LEFT,
        _DOWN,
        _DOWN_RIGHT,
        _RIGHT,
        _UP_RIGHT,
        _ATTACK_CLOSEST,
        _ATTACK_LOWEST
    ]

    UNIT_AGENT_1 = units.Protoss.Stalker
    UNIT_AGENT_2 = units.Zerg.Roach
    UNIT_ENEMY = units.Protoss.Zealot

    def __init__(self):
        super().__init__(args=['DefeatZealotswithBlink_2vs2'])

        self.current_agent = 1

    '''
        Switch flag of current agent used
    '''
    def switch(self):
        if self.current_agent == 1:
            self.current_agent = 2
        else:
            self.current_agent = 1

    '''
        Prepare basic parameters. This is called before start the episode.
    '''
    def prepare(self):
        # shared values
        self.enemy_totalHP = self._get_group_totalHP(group_type=self.UNIT_ENEMY)
        self.enemy_originalHP = self.enemy_totalHP
        self.enemy_onlyHP = self._get_group_totalHP(group_type=self.UNIT_ENEMY)
        self.current_enemies = 2

        # agent 1
        self.ally_totalHP_1 = self._get_group_totalHP(group_type=self.UNIT_AGENT_1)
        self.ally_originalHP_1 = self.ally_totalHP_1

        self.last_can_shoot_1 = False
        self.end_1 = False


        unit = self._get_group(group_type=self.UNIT_AGENT_1)[0]
        self.select_1 = actions.FUNCTIONS.select_point("select",(unit.x, unit.y))
        self.action_1 = self.select_1

        # agent 2
        self.ally_totalHP_2 = self._get_group_totalHP(group_type=self.UNIT_AGENT_2)
        self.ally_originalHP_2 = self.ally_totalHP_2

        self.last_can_shoot_2 = False
        self.end_2 = False


        unit = self._get_group(group_type=self.UNIT_AGENT_2)[0]
        self.select_2 = actions.FUNCTIONS.select_point("select",(unit.x, unit.y))
        self.action_2 = self.select_2

    '''
        Update values
    '''
    def update(self, deltaTime):
        if self.current_agent == 1:
            # agent 1
            if not self.end_1:
                self.last_can_shoot_1 = self.current_can_shoot_1

                unit = self._get_group(group_type=self.UNIT_AGENT_1)[0]
                self.select_1 = actions.FUNCTIONS.select_point("select",(unit.x, unit.y))
        else:
            # agent 2
            if not self.end_2:
                self.last_can_shoot_2 = self.current_can_shoot_2

                unit = self._get_group(group_type=self.UNIT_AGENT_2)[0]
                self.select_2 = actions.FUNCTIONS.select_point("select",(unit.x, unit.y))
        
    '''
        Do step of the environment
    '''
    def step(self):
        if not self.end_1:
            self.obs = self.env.step(actions=[self.select_1])[0]
            self.get_end()

            if not self.end_1:
                self.__check_action_available()
                self.obs = self.env.step(actions=[self.action_1])[0]
                self.get_end()

        self.switch()

        if not self.end_2:
            self.obs = self.env.step(actions=[self.select_2])[0]
            self.get_end()

            if not self.end_2:
                self.__check_action_available()
                self.obs = self.env.step(actions=[self.action_2])[0]
                
        self.switch()

    '''
        Return action of environment
    '''
    def get_action(self, action):
        if self.current_agent == 1: 
            if self.end_1: 
                return
            group_type = self.UNIT_AGENT_1
        elif self.current_agent == 2: 
            if self.end_2: 
                return
            group_type = self.UNIT_AGENT_2
        
        func = actions.FunctionCall(self._NO_OP, [])

        if self.possible_actions[action] == self._ATTACK_CLOSEST:
            # ATTACK ACTION
            if self._ATTACK_SCREEN in self.obs.observation.available_actions:
                unitx, unity = self._get_meangroup_position(group_type=group_type)
                zealot = self._get_zealot_closest(unit=[unitx,unity])
                func = actions.FunctionCall(self._ATTACK_SCREEN, [self._NOT_QUEUED, [zealot.x, zealot.y]])

        elif self.possible_actions[action] == self._ATTACK_LOWEST:
            # ATTACK ACTION
            if self._ATTACK_SCREEN in self.obs.observation.available_actions:
                unitx, unity = self._get_meangroup_position(group_type=group_type)
                zealot = self._get_zealot_lowest()
                func = actions.FunctionCall(self._ATTACK_SCREEN, [self._NOT_QUEUED, [zealot.x, zealot.y]])

        else:
            # MOVING ACTION
            if self._MOVE_SCREEN in self.obs.observation.available_actions:
                unitx, unity = self._get_meangroup_position(group_type=group_type)

                if  self.possible_actions[action] == self._UP:
                    if(unity - self._MOVE_VAL < 3.5):
                        unity += self._MOVE_VAL
                    nextPosition = [unitx, unity - self._MOVE_VAL]

                elif self.possible_actions[action] == self._UP_LEFT:
                    if(unity - self._MOVE_VAL < 3.5):
                        unity += self._MOVE_VAL
                    if(unitx - self._MOVE_VAL < 3.5):
                        unitx += self._MOVE_VAL
                    nextPosition = [unitx - (self._MOVE_VAL/2), unity - (self._MOVE_VAL/2)]

                elif self.possible_actions[action] == self._LEFT:
                    if(unitx - self._MOVE_VAL < 3.5):
                        unitx += self._MOVE_VAL
                    nextPosition = [unitx - self._MOVE_VAL, unity]

                elif self.possible_actions[action] == self._DOWN_LEFT:
                    if(unity + self._MOVE_VAL > 44.5):
                        unity -= self._MOVE_VAL
                    if(unitx - self._MOVE_VAL < 3.5):
                        unitx += self._MOVE_VAL
                    nextPosition = [unitx - (self._MOVE_VAL/2), unity + (self._MOVE_VAL/2)]

                elif self.possible_actions[action] == self._DOWN:
                    if(unity + self._MOVE_VAL > 44.5):
                        unity -= self._MOVE_VAL
                    nextPosition = [unitx, unity + self._MOVE_VAL]

                elif self.possible_actions[action] == self._DOWN_RIGHT:
                    if(unity + self._MOVE_VAL > 44.5):
                        unity -= self._MOVE_VAL
                    if(unitx + self._MOVE_VAL > 60.5):
                        unitx -= self._MOVE_VAL
                    nextPosition = [unitx + (self._MOVE_VAL/2), unity + (self._MOVE_VAL/2)]

                elif self.possible_actions[action] == self._RIGHT:
                    if(unitx + self._MOVE_VAL > 60.5):
                        unitx -= self._MOVE_VAL
                    nextPosition = [unitx + self._MOVE_VAL, unity]

                elif self.possible_actions[action] == self._UP_RIGHT:
                    if(unity - self._MOVE_VAL < 3.5):
                        unity += self._MOVE_VAL
                    if(unitx + self._MOVE_VAL > 60.5):
                        unitx -= self._MOVE_VAL
                    nextPosition = [unitx + (self._MOVE_VAL/2), unity - (self._MOVE_VAL/2)]

                func = actions.FunctionCall(self._MOVE_SCREEN, [self._NOT_QUEUED, nextPosition])

            # SELECT ARMY
            else:
                if self.current_agent == 1:
                    func = self.select_1
                else:
                    func = self.select_2

        if self.current_agent == 1:
            #agent 1
            self.num_action_1 = action
            self.action_1 = func
        else:
            #agent 2
            self.num_action_2 = action
            self.action_2 = func

    '''
        Return reward
    '''
    def get_reward(self, action):
        reward = 0

        if self.current_agent == 1:
            #agent 1
            if self.end_1: 
                return reward

            actual_ally_totalHP = self._get_group_totalHP(group_type=self.UNIT_AGENT_1)
            last_can_shoot = self.last_can_shoot_1
            # check if we were hurt
            if actual_ally_totalHP < self.ally_totalHP_1:
                reward += -1
                self.end_1 = False
                
        else:
            #agent 2
            if self.end_2: 
                return reward

            actual_ally_totalHP = self._get_group_totalHP(group_type=self.UNIT_AGENT_2)
            last_can_shoot = self.last_can_shoot_2
            # check if we were hurt
            if actual_ally_totalHP < self.ally_totalHP_2:
                reward += -1
                self.end_2 = False

        actual_enemy_totalHP = self._get_group_totalHP(group_type=self.UNIT_ENEMY)
        actual_enemy_onlyHP = self._get_group_onlyHP(group_type=self.UNIT_ENEMY)

        # check if we made some damage and we have shot with this action
        if actual_enemy_totalHP < self.enemy_totalHP and (action == self._ATTACK_CLOSEST or action == self._ATTACK_LOWEST) and last_can_shoot:
            reward += 1

        # check if some zealot was defeated
        zealots = self._get_group(group_type=self.UNIT_ENEMY)

        if self.current_enemies != len(zealots):
            reward += 2
            self.current_enemies = len(zealots)

        #update values
        self.enemy_totalHP = actual_enemy_totalHP
        self.enemy_onlyHP = actual_enemy_onlyHP

        if self.current_agent == 1:
            self.ally_totalHP_1 = actual_ally_totalHP
        else:
            self.ally_totalHP_2 = actual_ally_totalHP

        return reward


    '''
        Return True if we must end this episode
    '''
    def get_end(self):
        ally = self._get_group(group_type=self.UNIT_AGENT_1)
        self.end_1 = not ally

        ally = self._get_group(group_type=self.UNIT_AGENT_2)
        self.end_2 = not ally

        return self.end_1 and self.end_2

    '''
        (Protected method)
        Return specified group
    '''
    def _get_group(self, group_type):
        group = [unit for unit in self.obs.observation['feature_units'] 
                    if unit.unit_type == group_type]
        return group

    '''
        (Protected method)
        Return mean position of a group
    '''
    def _get_meangroup_position(self, group_type):
        group = self._get_group(group_type=group_type)

        unitx = unity = 0
        for unit in group:
            unitx += unit.x
            unity += unit.y
        if len(group) < 0 or len(group) == 0:
                       
            unitx /= 1
            unity /= 1

        else:
            unitx /= len(group)
            unity /= len(group)
        return unitx, unity

    '''
        (Protected method)
        Return totalHP of a group = (unit health plus unit shield)
    '''
    def _get_group_totalHP(self, group_type):
        group = self._get_group(group_type=group_type)
        totalHP = 0
        for unit in group:
            totalHP += unit.health + unit.shield
        return totalHP

    '''
        (Protected method)
        Return totalHP of a group = (unit health plus unit shield)
    '''

    def _get_group_onlyHP(self, group_type):
        group = self._get_group(group_type=group_type)
        totalHP = 0
        for unit in group:
            totalHP += unit.health
        return totalHP

    '''
        (Protected method)
        Return zealot with lowest healt (health + shield)
    '''
    def _get_zealot_closest(self, unit):
        zealots = self._get_group(group_type=self.UNIT_ENEMY)
        # search who has lower hp and lower shield
        target = zealots[0]
        for i in range(1, len(zealots)):
            if self._get_dist([unit[0], unit[1]], [zealots[i].x, zealots[i].y]) < self._get_dist([unit[0], unit[1]], [target.x, target.y]) :
                target = zealots[i]
                
        return target

    '''
        (Protected method)
        Return zealot with lowest healt (health + shield)
    '''
    def _get_zealot_lowest(self):
        zealots = self._get_group(group_type=self.UNIT_ENEMY)
        # search who has lower hp and lower shield
        target = zealots[0]
        for i in range(1, len(zealots)):
            if zealots[i].health < target.health or (zealots[i].health == target.health and zealots[i].shield < target.shield):
                target = zealots[i]
                
        return target

    '''
        (Protected method)
        Return True if group cooldown == 0
    '''
    def _can_shoot(self, group_type):
        group = self._get_group(group_type=group_type)

        for unit in group:
            if unit.weapon_cooldown != 0:
                return False

        return True

    '''
        (Protected method)
        Return dist between A and B
    '''
    def _get_dist(self, A, B):
        newDist = math.sqrt(pow(A[0] - B[0], 2) + pow(A[1] - B[1], 2))
        return newDist

    '''
        (Protected method)
        Return angle formed from two lines
    '''
    def _ang(self, lineA, lineB):
        # Get nicer vector form
        vA = lineA
        vB = lineB
        # Get dot prod
        dot_prod = np.dot(vA, vB)
        # Get magnitudes
        magA = np.dot(vA, vA)**0.5
        magB = np.dot(vB, vB)**0.5
        # Get cosine value
        cos = dot_prod/magA/magB
        # Get angle in radians and then convert to degrees
        angle = math.acos(dot_prod/magB/magA)
        # Basically doing angle <- angle mod 360
        ang_deg = math.degrees(angle)%360


        return ang_deg

    '''
        (Private method)
        Return if current action is available in the environment
    '''
    def __check_action_available(self):
        if self.current_agent == 1:
            # agent 1
            unit = self._get_group(group_type=self.UNIT_AGENT_1)[0]

            if self.possible_actions[self.num_action_1] == self._ATTACK_CLOSEST or self._ATTACK_LOWEST:
                if not (self._ATTACK_SCREEN in self.obs.observation.available_actions):
                    self.action_1 = actions.FUNCTIONS.select_point("select",(unit.x, unit.y))
            else:
                if not (self._MOVE_SCREEN in self.obs.observation.available_actions):
                    self.action_1 = actions.FUNCTIONS.select_point("select",(unit.x, unit.y))
        else:
            # agent 2
            unit = self._get_group(group_type=self.UNIT_AGENT_2)[0]

            if self.possible_actions[self.num_action_2] == self._ATTACK_CLOSEST or self._ATTACK_LOWEST:
                if not (self._ATTACK_SCREEN in self.obs.observation.available_actions):
                    self.action_2 = actions.FUNCTIONS.select_point("select",(unit.x, unit.y))
            else:
                if not (self._MOVE_SCREEN in self.obs.observation.available_actions):
                    self.action_2 = actions.FUNCTIONS.select_point("select",(unit.x, unit.y))
