# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/agents/benchmark_agents/99_TEMP_PPO_SB3.ipynb.

# %% auto 0
__all__ = ['PPO_TEMP_Agent', 'GymEnv', 'PPOPolicy']

# %% ../../../nbs/agents/benchmark_agents/99_TEMP_PPO_SB3.ipynb 3
# General libraries:
import numpy as np
from scipy.stats import norm
from tqdm import tqdm

# TEMP libraries:
import gym
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# Mushroom libraries
from mushroom_rl.core import Agent

# %% ../../../nbs/agents/benchmark_agents/99_TEMP_PPO_SB3.ipynb 5
class PPO_TEMP_Agent(Agent):

    train_directly=True

    """_summary_
    """

    def __init__(self,
                    mdp_info,
                    mdp,
                    agent_name = None
                    ):

        print("initializing PPO agent")

        mdp._mdp_info.horizon = mdp.demand.shape[0]
        mdp.reset(0)

        policy = PPOPolicy(
            mdp_info=mdp_info,
            mdp=mdp,
                        )
        
        if agent_name is None:
            self.name = 'PPO_SB3_Agent'
        else:
            self.name = agent_name
        
        self.train_directly=True
        self.train_mode="epochs" # try without this first

        super().__init__(mdp_info, policy)
    
    def fit_epoch(self, features=None, demand=None, n_steps=None, n_steps_per_fit=None):

        self.policy.fit(n_steps=n_steps, n_steps_per_fit=n_steps_per_fit)
    
    def train(self):
        self.policy.train()
    
    def eval(self):
        self.policy.eval()


class GymEnv(gym.Env):
    def __init__(self, mdp, observation_space, action_space):

        self.observation_space = gym.spaces.Box(observation_space.low, observation_space.high, observation_space.shape)
        self.action_space = gym.spaces.Box(action_space.low, action_space.high, action_space.shape)
        
        self.render_mode = None
        
        self.mdp = mdp

    def reset(self, seed=None, options=None):
        state = self.mdp.reset()
        return state, None
    
    def step(self, action):

        s, r, absorbing, info = self.mdp.step(action) 
        return s, r, absorbing, False, info # False for gym logic of terminal state

    def render(self):
        pass


class PPOPolicy():
    """
    

    """

    def __init__(self,
                    mdp_info,
                    mdp,
                    ):

        self.mdp = mdp
        self.mdp_info = mdp_info
        self.preprocessors = []
        self._train = True

        self.gym_env = GymEnv(mdp, mdp_info.observation_space, mdp_info.action_space)

        self.agent = PPO('MlpPolicy', self.gym_env, verbose=2, use_sde=True, sde_sample_freq=10)

    def fit(self, n_steps=None, n_steps_per_fit=None):

        # vec_env = make_vec_env("CartPole-v1", n_envs=4)
        # model = PPO("MlpPolicy", vec_env, verbose=1)
        
        # model.n_steps=n_steps_per_fit
        # model.learn(total_timesteps=n_steps)

        self.agent.learn(n_steps)

    def draw_action(self, input):

        """
        Note: only designed for single product case
        """

        for preprocessor in self.preprocessors:
            # input must be a vector containing the inventory and the pipeline vector
            input = preprocessor(input)

        action, _ = self.agent.predict(input, deterministic=True)

        # print(input)
        # breakpoint()
        

        return action

    def train(self):
        self._train = True
    
    def eval(self):
        self._train = False

    def reset(self):
        pass
