from functools import partial
from smac.env import MultiAgentEnv
from smac_plus import StarCraft2Env, Tracker1Env, Join1Env,Room1Env,Room2Env,Room2p5Env,Room3Env,Room4Env,Room4v3rdEnv,Room5Env,Room6Env,Room7Env,RoomEnv,\
	Calc1Env,Calc2Env,Calc3Env,Calc4Env,CalcdbEnv,Radar1Env,Radar2Env,Radar4Env,Bridge1Env,Toy1Env,PPEnv,Point1Env,SASEnv
import sys
import os


def env_fn(env, **kwargs) -> MultiAgentEnv:
	return env(**kwargs)


REGISTRY = {
	"sc2": partial(env_fn, env=StarCraft2Env),
	"tracker1": partial(env_fn, env=Tracker1Env),
	"join1": partial(env_fn, env=Join1Env),
	"room1": partial(env_fn, env=Room1Env),
	"room2": partial(env_fn, env=Room2Env),
	"room2p5": partial(env_fn, env=Room2p5Env),
	"room3": partial(env_fn, env=Room3Env),
	"room4": partial(env_fn, env=Room4Env),
	"room4v3rd": partial(env_fn, env=Room4v3rdEnv),
	"room5": partial(env_fn, env=Room5Env),
	"room6": partial(env_fn, env=Room6Env),
	"room7": partial(env_fn, env=Room7Env),
	"room": partial(env_fn, env=RoomEnv),
	"calc1": partial(env_fn, env=Calc1Env),
	"calc2": partial(env_fn, env=Calc2Env),
	"calc3": partial(env_fn, env=Calc3Env),
	"calc4": partial(env_fn, env=Calc4Env),
	"calcdb": partial(env_fn, env=CalcdbEnv),
	"radar1": partial(env_fn, env=Radar1Env),
	"radar2": partial(env_fn, env=Radar2Env),
	"radar4": partial(env_fn, env=Radar4Env),
	"bridge1": partial(env_fn, env=Bridge1Env),
	"toy1": partial(env_fn, env=Toy1Env),
	"pp": partial(env_fn, env=PPEnv),
	"point1": partial(env_fn, env=Point1Env),
	"sas": partial(env_fn, env=SASEnv),
}

if sys.platform == "linux":
	os.environ.setdefault("SC2PATH",
	                      os.path.join(os.getcwd(), "3rdparty", "StarCraftII"))
