#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/7/3
file:       base.py
version:
doc:        智能体规范声明.
"""
import abc
import os
from typing import Any, Optional, Type

import gymnasium as gym

from corekit.utils import get_debug_logger
from corekit.utils import get_file_paths
from corekit.utils.annotations import StanardAPI
from corekit.utils.typings import ActionArray, ActionExe, AgentConfig, AgentID, ObsArray

_logger = get_debug_logger()


@StanardAPI
class Agent(abc.ABC):
    """智能体基类"""
    # all_ids: set = set()  # 无法处理多线程/进程
    agent_config: AgentConfig = {}  # 智能体配置字典
    dead: bool = False  # True代表智能体死亡
    uid: AgentID = None  # 智能体ID
    # camp: str = 'red'  # 智能体阵营, 比如 'red'
    is_learn: bool = True  # 是否是学习智能体
    # model: Type = None  # 决策模型
    action_space: Optional[gym.Space] = None  # 动作空间维度
    observation_space: Optional[gym.Space] = None  # 状态空间维度

    # rule_type = None  # 规则策略指向

    def __init__(self, agent_config: AgentConfig = None):
        if agent_config is None:
            agent_config = {}
        self.from_dict(agent_config)

    def from_dict(self, agent_config: AgentConfig):
        self.agent_config.update(agent_config)
        # 避免赋值为 0, '', False 等时更新失败
        self.uid = agent_config['uid'] if 'uid' in agent_config else self.uid
        # self.camp = agent_config['camp'] if 'camp' in agent_config else self.camp
        self.is_learn = agent_config['is_learn'] if 'is_learn' in agent_config else self.is_learn

        # 特殊数据类型, 不会为更新失败
        # self.model = agent_config.get('model', self.model)
        self.action_space = agent_config.get('action_space', self.action_space)
        self.observation_space = agent_config.get('observation_space', self.observation_space)

    def __post_init__(self):
        if not (
                self.observation_space
                and self.action_space
                and isinstance(self.action_space, gym.Space)
                and isinstance(self.observation_space, gym.Space)
        ):
            raise ValueError(
                f'你应该设置智能体的动作空间和观测空间为 gym.Space 类, '
                f'当前为{self.action_space}, {self.observation_space}请检查'
            )

    def act_trans(self, action_array, env_info: dict = None) -> ActionExe:
        """动作指令转换, 决策模型的决策指令数列 -> 引擎的可执行指令格式

        学习策略: 环境信息 -> 学习策略() -> 指令数组 -> 指令转换() -> 可执行指令 \n
        规则策略: 环境信息 -> 规则策略() -> 可执行指令
        :param action_array: 学习策略模型输出的 指令数组
        :param env_info: dict, 环境信息
        :return:
        """
        # 获得本机的控制指令数列
        array = action_array
        if isinstance(action_array, dict):
            array = action_array[self.uid]

        # 转换成 引擎的可执行指令格式
        executable_act = array
        return executable_act

    def obs_trans(self, raw_obs: Any) -> ObsArray:
        """引擎返回的观测信息 -> 决策模型可执行的观测信息数列"""
        obs_array = raw_obs
        return obs_array

    def get_reward(self, *args, **kwargs) -> float:
        return 0.

    def save_model(self, filepath):
        pass

    def step(self, *args):
        return self.observation_space.sample()

    def load_model(self):
        """加载神经网络模型
        :return:
        """
        pass
