
from pcgym.policy_evaluation import policy_eval
# from pcgym import make_env
# import numpy as np
# from stable_baselines3 import PPO

# T = 26
# nsteps = 120
# SP = {
#     'Ca': [0.85 for i in range(int(nsteps/2))] + [0.9 for i in range(int(nsteps/2))],
# }

# disturbance = {'Caf': np.repeat([1, 1.05, 1], [nsteps//3, nsteps//3, nsteps//3])}

# #Continuous box action space
# action_space = {
#     'low': np.array([295]),
#     'high':np.array([302]) 
# }
# #Continuous box observation space
# observation_space = {
#     'low' : np.array([0.7,300,0.8]),
#     'high' : np.array([1,350,0.9])  
# }
# disturbance_space ={
#   'low': np.array([1]),
#   'high': np.array([1.05])
# }
# env_params = {
#     'N': nsteps, # Number of time steps
#     'tsim':T, # Simulation Time
#     'SP':SP, #Setpoint
#     'o_space' : observation_space, #Observation space
#     'a_space' : action_space, # Action space
#     'x0': np.array([0.8,330,0.8]), # Initial conditions (torch.tensor)
#     'model': 'cstr_ode', #Select the model
#     'normalise_a': True, #Normalise the actions
#     'normalise_o':True, #Normalise the states,
#     'r_scale': {'Ca':100},
#     'noise':True, #Add noise to the states
#     'integration_method': 'casadi', #Select the integration method
#     'noise_percentage':0, #Noise percentage
#     # 'disturbance_bounds':disturbance_space,
#     # 'disturbances': disturbance
# }

# env = make_env(env_params)
# env.reset()
# env.step(0.1)
# model = PPO('MlpPolicy', env, verbose=1,learning_rate=1e-3)
# model.learn(total_timesteps=3e4)
# model.save('pse_track_ppo.zip')
# model.load('pse_track_ppo.zip')
# env.plot_rollout({'Random policy':model},1,oracle = True,dist_reward=True,MPC_params={'N':10,'R':0.001})