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_250326_erase 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
import pickle

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=2)
    name = "d1a3o14_erase_250830_20D_05D_3_63delay"

    # model = PPO.load(os.path.join(parent_dir, "models", "250729", "d1a3o14_erase_250830_20D_05D_2_65delay.zip"), env=vec_env, device=torch.device("cuda:0"))

    model = PPO(
        "MlpPolicy", 
        policy_kwargs=dict(activation_fn=Sin),
        env=vec_env, 
        device=torch.device("cuda:2"),
        # n_steps=3000,
        # batch_size=300,
        verbose=0)

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

    history_data = []

    for i in range(500):
        model.learn(total_timesteps=400)
        test_env = model.get_env()
        test_obs = test_env.reset()
        list_reward = []
        episolde_data = {'actions': [], 'observations': [], 'rewards': []}

        for step in range(200):
            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)
            episolde_data['actions'].append(test_action[0, :])
            episolde_data['observations'].append(np.array(test_obs))
            episolde_data['rewards'].append(test_rewards)

        history_data.append(episolde_data)

        avg_reward = np.mean(list_reward[-100:])
        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", "250729", name + ".zip"))
        # if i % 10 == 0:
        #     model.save(os.path.join(parent_dir, "models", "250329", name + f"_{i}.zip"))

    # with open(os.path.join(parent_dir, "output", name + ".pkl"), 'wb') as f:
    #     pickle.dump(history_data, f)