#! /usr/bin/env python
# _*_ coding: utf-8 _*_
"""
author:     Chen XuHui
date:       2024/9/4
file:       demo.py
version:
doc:        示例智能体.
"""
import random

import gymnasium as gym
import numpy as np

from corekit.agents.base import Agent
from corekit.dnn.custom_ann import SharedWeightsModel1, TF2SharedWeightsModel


class DemoAgent(Agent):
    """智能体示例"""
    action_space: gym.Space = gym.spaces.Box(low=-1, high=1, shape=(7,))  # 动作空间维度
    observation_space: gym.Space = gym.spaces.Box(low=-np.inf, high=+np.inf, shape=(73,))  # 状态空间维度
    model = SharedWeightsModel1

    def __init__(self, agent_conf: dict = None):
        """配置状态空间、动作空间"""
        if agent_conf is None:
            agent_conf = {}
        super().__init__(agent_conf)

        # 空战智能体神经网络模型
        framework: str = agent_conf.get('framework', 'tf')
        if framework == 'tf':
            self.model = SharedWeightsModel1
        elif framework == 'tf2':
            self.model = TF2SharedWeightsModel
        else:
            raise NotImplemented

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