import magent2


def get_config(map_size):
    gw = magent2.gridworld
    cfg = gw.Config()

    cfg.set({"map_width": map_size, "map_height": map_size})

    predator = cfg.register_agent_type(
        "predator",
        {
            "width": 2,
            "length": 2,
            "hp": 1,
            "speed": 1,
            "view_range": gw.CircleRange(5),
            "attack_range": gw.CircleRange(2),
            "attack_penalty": -0.2,
        },
    )

    prey = cfg.register_agent_type(
        "prey",
        {
            "width": 1,
            "length": 1,
            "hp": 1,
            "speed": 1.5,
            "view_range": gw.CircleRange(4),
            "attack_range": gw.CircleRange(0),
        },
    )

    predator_group = cfg.add_group(predator)
    prey_group = cfg.add_group(prey)

    a = gw.AgentSymbol(predator_group, index="any")
    b = gw.AgentSymbol(prey_group, index="any")

    cfg.add_reward_rule(gw.Event(a, "attack", b), receiver=[a, b], value=[1, -1])

    return cfg
