import gymnasium as gym
import mo_gymnasium as mo_gym
from stable_baselines3.common.monitor import Monitor


from wrappers.normalize_action_wrapper import check_and_normalize_box_actions

import envs
import numpy as np
import os

# Register all custom envs
envs.register_custom_envs()

def make_env(args, w=None, monitor=True, is_mogym=False, render=False):
    if is_mogym:
        if render:
            env = mo_gym.make(args.env.name, render_mode="human")
        else:
            env = mo_gym.make(args.env.name)
    else:
        if render:
            env = gym.make(args.env.name, render_mode="human")
        else:
            env = gym.make(args.env.name)

    if isinstance(w, np.ndarray):
        env.set_weight(w)

    if monitor:
        env = Monitor(env, "gym")


    # Normalize box actions to [-1, 1]
    env = check_and_normalize_box_actions(env)
    return env