#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/8/2
file:       hierarchical.py
version:
doc:        分层智能体基类.
"""
import abc
from typing import Any, List, Union

import gymnasium as gym

from corekit.agents.base import Agent
from corekit.utils.annotations import StanardAPI
from corekit.utils.typings import ActionDict, AgentID


@StanardAPI
class HierarchicalAgent(Agent, abc.ABC):
    """适用于分层学习的智能体"""
    observation_space: gym.spaces.Tuple

    level: Union[str, int] = ''  # 分层级别
    superior: "HierarchicalAgent" = None  # 上级智能体
    colleagues: List["HierarchicalAgent"] = []  # 同事智能体
    subordinates: List["HierarchicalAgent"] = []  # 下属智能体
    superior_command = None  # 上级智能体的决策

    def __init__(self, agent_conf: dict = None):
        super().__init__(agent_conf)

        if agent_conf is not None:
            self.level = agent_conf.get('level') or self.level

    def __post_init__(self):
        super().__post_init__()
        if not (
            self.observation_space
            and isinstance(self.observation_space, gym.spaces.Tuple)
        ):
            raise ValueError(
                f'分层智能体的观测空间应是 gym.spaces.Tuple 类型,'
                f'而当前设置的观测空间为 {self.observation_space}, 请检查'
            )

    def issue_commands(self, action_dict: ActionDict) -> None:
        """给下属的智能体下发指令"""
        for agent in self.subordinates:
            aid: AgentID = ''
            if self.uid in action_dict:
                aid = self.uid
            else:
                for key in action_dict:  # key: str, 比如 'task_red-team_red', 即 f'{higher.uid}-{lower.uid}'
                    if agent.uid in key:
                        aid = key
                        break
            if aid:
                agent.superior_command = action_dict[aid]  # 下属智能体更新上级指令

    def get_self_status(self, env_info: Any):
        """获得自身的状态"""
        obs = self.observation_space.sample()
        if not isinstance(obs, tuple):
            return obs

        if self.superior_command is None:
            return obs,
        else:
            return obs[0],  # self.superior_command

    def get_subordinate_status(self, env_info: Any) -> dict:
        """获得下属智能体的状态"""
        obs = {}
        for agent in self.subordinates:
            obs[agent.uid] = agent.get_self_status(env_info)
        return obs

    def get_status(self):
        """没使用过"""
        obs, rewards, terminateds, truncateds, infos = [{}] * 5
        for index, agent in enumerate(self.subordinates):
            aid = f'{self.level}-{index}'
            obs_i = agent.observation_space.sample()
            rew = agent.get_reward()
            terminated = False
            truncated = False
            info = {}

            obs[aid] = obs_i
            rewards[aid] = rew
            terminateds[aid] = terminated
            truncateds[aid] = truncated
            infos[aid] = info

        # return obs, rewards, terminateds, truncateds, infos
        raise NotImplementedError

    def get_reward(self, *args, **kwargs) -> float:
        if self.subordinates:  # 如果本身属于高级智能体, 则奖励为下属智能体奖励的平均值
            rewards = [
                agent.get_reward(*args, **kwargs)
                for agent in self.subordinates
            ]
            return sum(rewards) / len(rewards)  # 下属智能体的平均奖励

        return 0.


@StanardAPI
class TaskAgent(HierarchicalAgent):
    """任务智能体"""
    level = 'task'
    colleagues: List["TaskAgent"] = []  # 同事智能体
    subordinates: List["TeamAgent"] = []  # 下属智能体


@StanardAPI
class TeamAgent(HierarchicalAgent):
    """编队智能体"""
    level = 'team'
    superior: "TaskAgent" = None  # 上级智能体
    colleagues: List["TeamAgent"] = []  # 同事智能体
    subordinates: List["UnitAgent"] = []  # 下属智能体


@StanardAPI
class UnitAgent(HierarchicalAgent):
    """成员智能体"""
    level = 'unit'
    superior: "TeamAgent" = None  # 上级智能体
    colleagues: List["UnitAgent"] = []  # 同事智能体

    def get_self_status(self, env_info: Any):
        return self.obs_trans(env_info),  # , self.superior_command
