import json, copy, re, inspect, os
import numpy as np
from UTIL.tensor_ops import my_view, repeat_at
from ...common.base_env import RawObsArray
from ..actset_lookup import digit2act_dictionary, AgentPropertyDefaults
from ..actset_lookup import decode_action_as_string, decode_action_as_string
from ..uhmap_env_wrapper import UhmapEnv, ScenarioConfig
from .UhmapFlagCaptureConf import SubTaskConfig
from .cython_func import tear_num_arr
from .SubtaskCommonFn import UhmapCommonFn




class UhmapFlagCapture(UhmapCommonFn, UhmapEnv):
    def __init__(self, rank) -> None:
        super().__init__(rank)
        self.observation_space = self.make_obs(get_shape=True)
        self.SubTaskConfig = SubTaskConfig
        assert os.path.basename(inspect.getfile(SubTaskConfig)) == type(self).__name__+'Conf.py', \
                ('make sure you have imported the correct SubTaskConfig class')

    def extract_key_gameobj(self, resp):
        keyObjArr = resp['dataGlobal']['keyObjArr']
        return keyObjArr

    def reset(self):
        self.dominate_team_score = [0 for _ in range(self.n_teams)]
        return super().reset()

    def get_distance_of_each_agent_to_center(self, resp):
        return np.linalg.norm(np.array([q['agentLocationArr'][:2] for q in resp['dataArr']]), axis=1)


    def gen_reward_and_win(self, resp):
        # self.simple_render_with_threejs()

        reward = [0]*self.n_teams
        events = resp['dataGlobal']['events']
        WinningResult = None
        flag = [k for k in resp['dataGlobal']['keyObjArr'] if k['className']=='FlagToCapture_C']
        assert len(flag) == 1
        flag = flag[0]
        info = flag['rSVD1']

        if info == 'Captured By No Team':
            pass

        for team_id in range(self.n_teams):
            if info == f'Captured By Team {team_id}':
                reward[team_id] += 0.005
                self.dominate_team_score[team_id] += 1


        for event in events:
            event_parsed = self.parse_event(event)

            if event_parsed['Event'] == 'EndEpisode':
                WinTeam = np.argmax(self.dominate_team_score)
                reward[WinTeam] = 1

                sorted_indices = np.argsort(-np.array(self.dominate_team_score))
                team_ranking = list(range(self.n_teams))
                for i, j in enumerate(sorted_indices):
                    team_ranking[j] = i

                WinningResult = {
                    "team_ranking": team_ranking,
                    "end_reason": "Team{} Win".format(WinTeam)
                }


        return reward, WinningResult

    @staticmethod
    def item_random_mv(src,dst,prob,rand=False):
        assert len(src.shape)==1; assert len(dst.shape)==1
        if rand: np.random.shuffle(src)
        len_src = len(src)
        n_mv = (np.random.rand(len_src) < prob).sum()
        item_mv = src[range(len_src-n_mv,len_src)]
        src = src[range(0,0+len_src-n_mv)]
        dst = np.concatenate((item_mv, dst))
        return src, dst

    @staticmethod
    def get_binary_array(n_int, n_bits=8, dtype=np.float32):
        arr = np.zeros((*n_int.shape, n_bits), dtype=dtype)
        for i in range(n_bits):
            arr[:, i] = (n_int%2==1).astype(int)
            n_int = n_int / 2
            n_int = n_int.astype(np.int8)
        return arr


    def make_obs(self, resp=None, get_shape=False):
        # CORE_DIM = 38
        CORE_DIM = 23
        assert ScenarioConfig.obs_vec_length == CORE_DIM
        if get_shape:
            return CORE_DIM

        # temporary parameters
        OBS_RANGE_PYTHON_SIDE = SubTaskConfig.OBS_RANGE_PYTHON_SIDE
        MAX_NUM_OPP_OBS       = SubTaskConfig.MAX_NUM_OPP_OBS
        MAX_NUM_ALL_OBS       = SubTaskConfig.MAX_NUM_ALL_OBS
        MAX_OBJ_NUM_ACCEPT    = SubTaskConfig.MAX_OBJ_NUM_ACCEPT

        # get and calculate distance array
        pos3d_arr = np.zeros(shape=(self.n_agents, 3), dtype=np.float32)
        for i, agent in enumerate(self.agents): pos3d_arr[i] = agent.pos3d
        # use the distance matrix calculated by unreal engine to accelerate
        # dis_mat = distance_matrix(pos3d_arr)    # dis_mat is a matrix, shape = (n_agent, n_agent)
        dis_mat = resp['dataGlobal']['distanceMat']
        alive_all = np.array([agent.alive for agent in self.agents])
        try:
            dis_mat[~alive_all,:] = +np.inf
            dis_mat[:,~alive_all] = +np.inf
        except:
            pass

        # get team list
        team_belonging = np.array([agent.team for agent in self.agents])

        # gather the obs arr of all known agents
        obs_arr = RawObsArray(key='Agent')

        if not hasattr(self, "uid_binary"):
            self.uid_binary = self.get_binary_array(np.arange(self.n_agents), 10)

        for i, agent in enumerate(self.agents):
            assert agent.location is not None
            assert agent.uid == i

            obs_arr.append(
                self.uid_binary[i]  # 0~9
            )
            obs_arr.append([
                agent.index,    # 10
                agent.team,     # 11
                agent.alive,    # 12
                agent.uid_remote, # 13
            ])
            obs_arr.append( #[14,15,16,17,18,19]
                agent.pos3d
                # tear_num_arr(agent.pos3d, n_digits=6, base=10, mv_left=0)
                # tear_num_arr(agent.pos3d, 6, ScenarioConfig.ObsBreakBase, 0) # 3 -- > 3*6 = 18 , 18-3=15, 23+15 = 38
            )
            obs_arr.append(
                agent.vel3d
            )
            obs_arr.append([
                agent.hp,
                agent.yaw,
                agent.max_speed,
            ])
        obs_ = obs_arr.get()
        new_obs = my_view(obs_, [self.n_agents, -1])

        assert CORE_DIM == new_obs.shape[-1]
        OBS_ALL_AGENTS = np.zeros(shape=(
            self.n_agents,
            MAX_NUM_OPP_OBS+MAX_NUM_ALL_OBS,
            CORE_DIM
            ))

        # now arranging the individual obs
        for i, agent in enumerate(self.agents):
            if not agent.alive:
                OBS_ALL_AGENTS[i, :] = np.nan
                continue

            # if alive
            # scope <all>
            dis2all = dis_mat[i, :]
            is_ally = (team_belonging == agent.team)

            # scope <opp/hostile>
            a2h_dis = dis2all[~is_ally]
            h_alive = alive_all[~is_ally]
            h_feature = new_obs[~is_ally]
            h_iden_sort  = np.argsort(a2h_dis)[:MAX_NUM_OPP_OBS]
            a2h_dis_sorted = a2h_dis[h_iden_sort]
            h_alive_sorted = h_alive[h_iden_sort]
            h_vis_mask = (a2h_dis_sorted <= OBS_RANGE_PYTHON_SIDE) & h_alive_sorted

            # scope <all>
            h_vis_index = h_iden_sort[h_vis_mask]
            h_invis_index = h_iden_sort[~h_vis_mask]
            h_vis_index, h_invis_index = self.item_random_mv(src=h_vis_index, dst=h_invis_index,prob=0, rand=True)
            h_ind = np.concatenate((h_vis_index, h_invis_index))
            h_msk = np.concatenate((h_vis_index<0, h_invis_index>=0)) # "<0" project to False; ">=0" project to True
            a2h_feature_sort = h_feature[h_ind]
            a2h_feature_sort[h_msk] = 0
            if len(a2h_feature_sort)<MAX_NUM_OPP_OBS:
                a2h_feature_sort = np.concatenate((
                    a2h_feature_sort,
                    np.ones(shape=(MAX_NUM_OPP_OBS-len(a2h_feature_sort), CORE_DIM))+np.nan
                ), axis=0)

            # scope <ally/friend>
            a2f_dis = dis2all[is_ally]
            f_alive = alive_all[is_ally]
            f_feature = new_obs[is_ally]
            f_iden_sort  = np.argsort(a2f_dis)[:MAX_NUM_ALL_OBS]
            a2f_dis_sorted = a2f_dis[f_iden_sort]
            f_alive_sorted = f_alive[f_iden_sort]
            f_vis_mask = (a2f_dis_sorted <= OBS_RANGE_PYTHON_SIDE) & f_alive_sorted

            # scope <all>
            f_vis_index = f_iden_sort[f_vis_mask]
            self_vis_index = f_vis_index[:1] # seperate self and ally
            f_vis_index = f_vis_index[1:]    # seperate self and ally
            f_invis_index = f_iden_sort[~f_vis_mask]
            f_vis_index, f_invis_index = self.item_random_mv(src=f_vis_index, dst=f_invis_index,prob=0, rand=True)
            f_ind = np.concatenate((self_vis_index, f_vis_index, f_invis_index))
            f_msk = np.concatenate((self_vis_index<0, f_vis_index<0, f_invis_index>=0)) # "<0" project to False; ">=0" project to True
            self_ally_feature_sort = f_feature[f_ind]
            self_ally_feature_sort[f_msk] = 0
            if len(self_ally_feature_sort)<MAX_NUM_ALL_OBS:
                self_ally_feature_sort = np.concatenate((
                    self_ally_feature_sort,
                    np.ones(shape=(MAX_NUM_ALL_OBS-len(self_ally_feature_sort), CORE_DIM))+np.nan
                ), axis=0)
            OBS_ALL_AGENTS[i,:] = np.concatenate((self_ally_feature_sort, a2h_feature_sort), axis = 0)


        # the last part of observation is the list of core game objects
        self.N_Obj = len(self.key_obj)
        if self.N_Obj > 0:
            OBJ_UID_OFFSET = 32768

            obs_arr = RawObsArray(key = 'GameObj')

            for i, obj in enumerate(self.key_obj):
                assert obj['uId'] - OBJ_UID_OFFSET == i
                obs_arr.append(
                    -self.uid_binary[i] # reverse uid binary, self.uid_binary[i]
                )
                obs_arr.append([
                    obj['uId'] - OBJ_UID_OFFSET,    #agent.index,
                    -1,                             #agent.team,
                    True,                           #agent.alive,
                    obj['uId'] - OBJ_UID_OFFSET,    #agent.uid_remote,
                ])
                # tear_num_arr(agent.pos3d, n_digits=6, base=10, mv_left=0)
                obs_arr.append(
                    [
                        obj['location']['x'], obj['location']['y'], obj['location']['z']  # agent.pos3d
                    ]
                    # tear_num_arr([
                    #     obj['location']['x'], obj['location']['y'], obj['location']['z']  # agent.pos3d
                    # ], 6, ScenarioConfig.ObsBreakBase, 0)
                )

                obs_arr.append([
                    obj['velocity']['x'], obj['velocity']['y'], obj['velocity']['z']  # agent.vel3d
                ]+
                [
                    -1,                         # hp
                    obj['rotation']['yaw'],     # yaw
                    0,                          # max_speed
                ])

            OBS_GameObj = my_view(obs_arr.get(), [len(self.key_obj), -1])
            OBS_GameObj = OBS_GameObj[:MAX_OBJ_NUM_ACCEPT, :]
            OBS_GameObj = repeat_at(OBS_GameObj, insert_dim=0, n_times=self.n_agents)
            OBS_ALL_AGENTS = np.concatenate((OBS_ALL_AGENTS, OBS_GameObj), axis=1)

        return OBS_ALL_AGENTS


    def init_drone(self, agent_info, pos_ro):
        N_COL = 5
        team = agent_info['team']
        n_teams = self.n_teams
        n_team_agent = agent_info['n_team_agent']
        agent_class = agent_info['type']
        tid = agent_info['tid']
        uid = agent_info['uid']

        center_distance=4000
        x_offset_unit=400
        y_offset_unit=400

        if self.n_teams == 2:
            N_COL = int(np.sqrt(n_team_agent))
            center_x = (center_distance / 2) * (-1)**team
            center_y = 0

            x_offset = x_offset_unit * ((tid % N_COL) - N_COL // 2)
            y_offset = y_offset_unit * ((tid // N_COL) - N_COL // 2)
        else:
            N_COL = 5
            team_angle = 2 * np.pi * team / self.n_teams
            center_x = np.cos(team_angle) * center_distance
            center_y = np.sin(team_angle) * center_distance

            x_offset = x_offset_unit * ((tid % N_COL) - N_COL // 2)
            y_offset = y_offset_unit * (tid // N_COL - n_team_agent // (2 * N_COL))

        x = center_x + x_offset
        y = center_y + y_offset

        yaw = np.degrees(2 * np.pi * team / self.n_teams) % 360


        def hsl_to_rgba(h, s, l, a=1):
            import colorsys
            r, g, b = colorsys.hls_to_rgb(h, l, s)
            return f'(R={r:.2f},G={g:.2f},B={b:.2f},A={a})'

        hue = team / n_teams  
        saturation = 0.7  
        lightness = 0.5  
        color = hsl_to_rgba(hue, saturation, lightness)

        team_color = color


        z = 500 # 500 is slightly above the ground
        assert np.abs(x) < 15000.0 and np.abs(y) < 15000.0
        agent_property = copy.deepcopy(AgentPropertyDefaults)
        agent_property.update({

                'DebugAgent': False,
                # max drive/fly speed
                'MaxMoveSpeed':  400,
                # also influence object mass, please change it with causion!
                'AgentScale'  : { 'x': 0.7,  'y': 0.7, 'z': 0.7, },
                # probability of escaping dmg 闪避
                "DodgeProb": 0.0,
                # ms explode dmg
                "ExplodeDmg": 75,
                # team belonging
                'AgentTeam': team,
                # choose ue class to init
                'ClassName': agent_class,
                # Weapon CD
                'WeaponCD': 1,
                # open fire range
                "PerceptionRange":  2000,
                "GuardRange":       1400,
                "FireRange":        1300,
                # debugging
                'RSVD1': f'-CarrierName=T{team}-0 -NumDrone={self.n_team_agent[team]-1}',
                # regular
                'RSVD2': '-InitAct=ActionSet2::Idle;StaticAlert',
                # agent hp
                'AgentHp': 140,
                # the rank of agent inside the team
                'IndexInTeam': tid,
                # the unique identity of this agent in simulation system
                'UID': uid,
                # show color
                'Color':team_color,
                # 'Color':'(R=0,G=1,B=0,A=1)' if team==0 else '(R=0,G=0,B=1,A=1)',
                # initial location
                'InitLocation': { 'x': x,  'y': y, 'z': z, },
                # initial facing direction et.al.
                'InitRotator': { 'pitch': 0,  'roll': 0, 'yaw': yaw, },
        }),
        return agent_property
