import os
os.environ['MKL_THREADING_LAYER'] = 'GNU'
os.environ["OMP_NUM_THREADS"] = "8"
os.environ["MKL_NUM_THREADS"] = "8"
import torch
import numpy as np
from torch.nn import Module
import gymnasium as gym
from gym_env import CustomEnv
from stable_baselines3 import PPO
from stable_baselines3.common.vec_env import SubprocVecEnv
from stable_baselines3.common.vec_env import DummyVecEnv
from sb3_contrib import RecurrentPPO
from torch.utils.tensorboard import SummaryWriter

current_dir = os.path.dirname(os.path.abspath("__file__"))
parent_dir = os.path.abspath(os.path.join(current_dir, os.pardir))

class Sin(Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.sin(x)

if __name__ == '__main__':
    
    vec_env = CustomEnv(device_id=1)
    name = "d1a3o12_c1"

    model = PPO.load(os.path.join(parent_dir, "models", "d1a3o12_a0"), env=vec_env, device=torch.device("cuda:1"))

    # model = PPO(
    #     "MlpPolicy", 
    #     policy_kwargs=dict(activation_fn=Sin),
    #     env=vec_env, 
    #     device=torch.device("cuda:1"),
    #     verbose=0)

    writer = SummaryWriter(log_dir=os.path.join(parent_dir, "tensorboard", name))
    max_reward = 0

    for i in range(100):
        model.learn(total_timesteps=480)
        test_env = model.get_env()
        test_obs = test_env.reset()
        list_reward = []
        for step in range(480):
            test_action, _states = model.predict(observation=test_obs)
            test_obs, test_rewards, test_dones, info = test_env.step(test_action)
            list_reward.append(test_rewards)

        avg_reward = np.mean(list_reward[-240:])
        writer.add_scalar('Reward', np.mean(avg_reward), i)
        if avg_reward > max_reward:
            max_reward = avg_reward
            model.save(os.path.join(parent_dir, "models", name + ".zip"))