import os
import sumo_rl
from sumo_rl import env, parallel_env

PATH = os.path.dirname(sumo_rl.__file__)

"""
Environments from RESCO: https://github.com/jault/RESCO, paper https://people.engr.tamu.edu/guni/Papers/NeurIPS-signals.pdf
"""

def grid4x4(parallel=True, **kwargs):
    """
    Number of agents = 16
    Number of actions = 4
    Agents have the same obsevation and action space
    """
    kwargs.update({'net_file': PATH+'/../nets/RESCO/grid4x4/grid4x4.net.xml',
                   'route_file': PATH+'/../nets/RESCO/grid4x4/grid4x4_1.rou.xml',
                   'num_seconds': 3600})
    if parallel:
        return parallel_env(**kwargs)
    else:
        return env(**kwargs)

def arterial4x4(parallel=True, **kwargs):
    """
    Number of agents = 16
    Number of actions = 5
    Agents have the same obsevation and action space
    """
    kwargs.update({'net_file': PATH+'/../nets/RESCO/arterial4x4/arterial4x4.net.xml',
                   'route_file': PATH+'/../nets/RESCO/arterial4x4/arterial4x4_1.rou.xml',
                   'num_seconds': 3600})
    if parallel:
        return parallel_env(**kwargs)
    else:
        return env(**kwargs)

def cologne1(parallel=True, **kwargs):
    """ 
    Number of agents: 1
    Number of actions: 4
    """
    kwargs.update({'net_file': PATH+'/../nets/RESCO/cologne1/cologne1.net.xml',
                   'route_file': PATH+'/../nets/RESCO/cologne1/cologne1.rou.xml',
                   'begin_time': 25200,
                   'num_seconds': 28800})
    if parallel:
        return parallel_env(**kwargs)
    else:
        return env(**kwargs)

def cologne3(parallel=True, **kwargs):
    """ 
    Number of agents: 3
    Number of actions: 2 agents with 4 actions and 1 agent with 3 actions
    2 agents have the same obsevation and action space and 1 has different spaces
    """
    kwargs.update({'net_file': PATH+'/../nets/RESCO/cologne3/cologne3.net.xml',
                   'route_file': PATH+'/../nets/RESCO/cologne3/cologne3.rou.xml',
                   'begin_time': 25200,
                   'num_seconds': 28800})
    if parallel:
        return parallel_env(**kwargs)
    else:
        return env(**kwargs)

def cologne8(parallel=True, **kwargs):
    """ 
    Number of agents: 8
    Number of actions: variable
    """
    kwargs.update({'net_file': PATH+'/../nets/RESCO/cologne8/cologne8.net.xml',
                   'route_file': PATH+'/../nets/RESCO/cologne8/cologne8.rou.xml',
                   'begin_time': 25200,
                   'num_seconds': 28800})
    if parallel:
        return parallel_env(**kwargs)
    else:
        return env(**kwargs)

def ingolstadt1(parallel=True, **kwargs):
    """ 
    Number of agents: 1
    Number of actions: 3
    """
    kwargs.update({'net_file': PATH+'/../nets/RESCO/ingolstadt1/ingolstadt1.net.xml',
                   'route_file': PATH+'/../nets/RESCO/ingolstadt1/ingolstadt1.rou.xml',
                   'begin_time': 57600,
                   'num_seconds': 61200})
    if parallel:
        return parallel_env(**kwargs)
    else:
        return env(**kwargs)

def ingolstadt7(parallel=True, **kwargs):
    """ 
    Number of agents: 7
    Number of actions: variable
    """
    kwargs.update({'net_file': PATH+'/../nets/RESCO/ingolstadt7/ingolstadt7.net.xml',
                   'route_file': PATH+'/../nets/RESCO/ingolstadt7/ingolstadt7.rou.xml',
                   'begin_time': 57600,
                   'num_seconds': 61200})
    if parallel:
        return parallel_env(**kwargs)
    else:
        return env(**kwargs)

def ingolstadt21(parallel=True, **kwargs):
    """ 
    Number of agents: 21
    Number of actions: variable
    """
    kwargs.update({'net_file': PATH+'/../nets/RESCO/ingolstadt21/ingolstadt21.net.xml',
                   'route_file': PATH+'/../nets/RESCO/ingolstadt21/ingolstadt21.rou.xml',
                   'begin_time': 57600,
                   'num_seconds': 61200})
    if parallel:
        return parallel_env(**kwargs)
    else:
        return env(**kwargs)
