#
# (C) 2021 Alex Mayfield
#
# A script to count standard Doom 2 monster counts in a WAD.
# pip install --user omgifol
#

import csv, sys

from pathlib import Path

from omg import WAD, HeaderGroup
from omg.mapedit import MapEditor

MONSTERS: dict[int, str] = {
    3004: "Zombieman",
    9: "Shotgun guy",
    3001: "Imp",
    3002: "Demon",
    58: "Spectre",
    3006: "Lost Soul",
    3005: "Cacodemon",
    3003: "Baron of Hell",
    16: "Cyberdemon",
    7: "Spiderdemon",
 
    65: "Chaingunner",
    69: "Hell knight",
    66: "Revenant",
    67: "Mancubus",
    68: "Arachnotron",
    71: "Pain elemental",
    64: "Arch-vile",
    84: "Wolfenstein SS",
    72: "Commander Keen",
    88: "Romero Head",

    888: "Helper Dog",
}

def main(wadpath: Path):
    csvh = open(wadpath.stem + ".csv", 'w', newline='')
    csvwrite = csv.writer(csvh)

    wad = WAD(str(wadpath))
    if not isinstance(wad.maps, HeaderGroup):
        raise Exception("Maps not found in WAD")

    csvwrite.writerow(["", *MONSTERS.values()])

    for maplump in wad.maps:
        mapthing = dict.fromkeys(MONSTERS.keys(), 0)

        editor = MapEditor(wad.maps[maplump])
        for thing in editor.things:
            if thing.type not in MONSTERS:
                continue
            if not thing.hard:
                continue
            if thing.multiplayer:
                continue

            mapthing[thing.type] += 1

        csvwrite.writerow([maplump, *mapthing.values()])

if __name__ == "__main__":
    main(Path(sys.argv[1]))