import subprocess
from collections import defaultdict
from pettingzoo.tests.all_modules import all_environments
import math
import os

def generate_cycle_words(agents, is_classic):
    if is_classic:
        words = []
        for agent in agents:
            words.append("env")
            words.append(agent)
    else:
        if len(agents) >= 10:
            types = []
            type_nums = defaultdict(list)
            for name in agents:
                splitname = name.split("_")
                type = splitname[0]
                type_nums[type].append(int(splitname[1]))
                if type not in types:
                    types.append(type)
            words = ["env"]
            for type in types:
                tyrange = list(range(type_nums[type][0],type_nums[type][-1]+1))
                assert tyrange == type_nums[type]
                type_range = f"{type}_[{tyrange[0]}...{tyrange[-1]}]"
                words.append(type_range)
        else:
            words = ["env"] + agents
    return words

def generate_graphviz(words):
    max_chars = max(len(w) for w in words)
    node_width = max_chars * 0.1 + 0.2
    innards = ""#"overlap = false;\n"
    innards += f'node [shape = circle,fixedsize=true,width={node_width},fontname="Segoe UI"];\n'
    for i,word in enumerate(words):
        theta = 2 * math.pi * i / len(words)
        rad_const = 1. if len(words) <= 3 else 0.8
        rad = (len(words)) * rad_const * node_width / math.pi
        xpos = rad * math.sin(theta)
        ypos = rad * math.cos(theta)
        innards += f'a{i} [label="{word}",pos="{xpos},{ypos}!"];\n'
    for i in range(len(words)-1):
        innards += f'a{i} -> a{i+1};\n'
    innards += f'a{len(words)-1} -> a{0};\n'
    return 'digraph G {\n%s\n}'%innards

for name, module in list(all_environments.items()):
    env = module.env()
    agents = env.agents
    words = generate_cycle_words(env.agents, "classic/" in name)
    vis_code = generate_graphviz(words)
    code_path = "graphviz/"+name+".vis"
    os.makedirs(os.path.dirname(code_path),exist_ok=True)
    with open(code_path,'w') as file:
        file.write(vis_code)
    out_path = f"docs/assets/img/aec/{name.replace('/','_')}_aec.svg"
    cmd = ['neato','-Tsvg',"-o",out_path, code_path]
    print(" ".join(cmd))
    subprocess.Popen(cmd)
