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_uniflow 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 = "d0a3o12_c0"

    # 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"),
        n_steps=2400,
        batch_size=240,
        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=2400)
        test_env = model.get_env()
        test_obs = test_env.reset()
        list_reward = []
        for step in range(240):
            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[-120:])
        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"))